rwkv-ops 0.6.1__py3-none-any.whl

rwkv_ops/rwkv7_kernel/jax_cuda_kernel_single/wkv7_single_step_jax.py

@@ -0,0 +1,190 @@
+"""
+JAX 版 RWKV7 单步 wkv kernel（仅前向传播）
+延迟编译 CUDA 扩展，专为 T=1 场景优化
+"""
+
+from __future__ import annotations
+import pathlib
+import subprocess
+import ctypes
+import jax
+import jax.numpy as jnp
+from typing import Optional, Tuple, Union
+
+# ---------- 延迟编译（改到当前目录） ----------
+_CURRENT_DIR = pathlib.Path(__file__).parent.absolute()
+
+
+def get_jax_generalized_delta_rule_single_step(HEAD_SIZE=64):
+    _BUILD_DIR = _CURRENT_DIR / f"build_single_step_{HEAD_SIZE}"
+    _SO_PATH = _CURRENT_DIR / f"build_single_step_{HEAD_SIZE}/wkv7_single_step.so"
+
+    def _ensure_compiled() -> pathlib.Path:
+        """首次调用时编译 CUDA 扩展，产出放在当前源码目录"""
+        if _SO_PATH.exists():
+            return _SO_PATH
+
+        print("[rwkv7_single_step_jax] First use – compiling CUDA kernel…")
+        src_dir = _CURRENT_DIR
+        build_dir = _BUILD_DIR
+        build_dir.mkdir(exist_ok=True)
+
+        # 获取 XLA 头文件路径
+        xla_include_dir = jax.ffi.include_dir()
+        if not xla_include_dir:
+            raise RuntimeError("jax.ffi.include_dir() 返回空，请检查 JAX >= 0.4.31")
+
+        # CUDA 编译 flags（移除 CHUNK_LEN 定义）
+        cuda_flags = [
+            "-ftz=true",
+            "-prec-div=false",
+            "-prec-sqrt=false",
+            "--use_fast_math",
+            "-O3",
+            "-Xptxas=-O3",
+            "-res-usage",
+            "--extra-device-vectorization",
+            f"-D_C_={HEAD_SIZE}",
+        ]
+
+        # CMake 配置
+        cmake_args = [
+            "cmake",
+            "-S",
+            str(src_dir),
+            "-B",
+            str(build_dir),
+            "-DCMAKE_BUILD_TYPE=Release",
+            f"-DCMAKE_INSTALL_PREFIX={_CURRENT_DIR}",
+            f"-DXLA_INCLUDE_DIR={xla_include_dir}",
+            f"-DCMAKE_CUDA_FLAGS={' '.join(cuda_flags)}",
+        ]
+        subprocess.check_call(cmake_args)
+
+        # 构建与安装
+        subprocess.check_call(["cmake", "--build", str(build_dir), "-j"])
+        subprocess.check_call(["cmake", "--install", str(build_dir)])
+
+        if not _SO_PATH.exists():
+            raise RuntimeError("Compilation failed – wkv7_single_step.so not found.")
+
+        print("[rwkv7_single_step_jax] Compilation finished – output at", _SO_PATH)
+        return _SO_PATH
+
+    # 注册 FFI 符号（仅前向）
+    _lib = ctypes.CDLL(_ensure_compiled())
+    jax.ffi.register_ffi_target(
+        "wkv7_single_step_fwd",
+        jax.ffi.pycapsule(_lib.Wkv7SingleStepFwd),
+        platform="CUDA",
+    )
+
+    # ---------- 工具 ----------
+    def _transpose_head(x: jnp.ndarray, head_first: bool) -> jnp.ndarray:
+        """(B, 1, H, K) <-> (B, H, 1, K)"""
+        x = jnp.asarray(x, dtype=jnp.bfloat16)
+        if head_first:
+            return jnp.transpose(x, (0, 2, 1, 3))
+        return x
+
+    # ---------- 前向 kernel ----------
+    def _wkv7_single_step_kernel(
+        w: jnp.ndarray,
+        q: jnp.ndarray,
+        k: jnp.ndarray,
+        v: jnp.ndarray,
+        a: jnp.ndarray,
+        b: jnp.ndarray,
+        h0: jnp.ndarray,
+    ):
+        """
+        内部 kernel 接口
+        参数: w,q,k,v,a,b,h0 -> y,s
+        """
+        B, H, K = q.shape
+        dtype = q.dtype
+
+        out_type = jax.ShapeDtypeStruct((B, H, K), dtype)
+        s_type = jax.ShapeDtypeStruct((B, H, K, K), jnp.float32)
+
+        y, s = jax.ffi.ffi_call(
+            "wkv7_single_step_fwd", (out_type, s_type), vmap_method="broadcast_all"
+        )(w, q, k, v, a, b, h0)
+
+        return y, s
+
+    def wk7_single_step_kernel(
+        w: jnp.ndarray,
+        q: jnp.ndarray,
+        k: jnp.ndarray,
+        v: jnp.ndarray,
+        a: jnp.ndarray,
+        b: jnp.ndarray,
+        h0: jnp.ndarray,
+    ):
+        """前向计算函数"""
+        y, s = _wkv7_single_step_kernel(w, q, k, v, a, b, h0)
+        final_state = s  # 单步后直接返回状态
+        return (y, final_state)
+
+    # ---------- 主接口 ----------
+    def generalized_delta_rule_single_step(
+        r: jnp.ndarray,
+        w: jnp.ndarray,
+        k: jnp.ndarray,
+        v: jnp.ndarray,
+        a: jnp.ndarray,
+        b: jnp.ndarray,
+        initial_state: Optional[jnp.ndarray] = None,
+        output_final_state: bool = True,
+        head_first: bool = False,
+    ) -> Union[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]]:
+        """
+        单步广义 delta 规则（仅前向）
+        参数:
+            r,w,k,v,a,b: 输入张量，形状必须为 (B, 1, H, K) 或 (B, H, 1, K)
+            initial_state: 可选 (B, H, K, K) 初始状态，None 则零初始化
+            output_final_state: 是否同时返回最后状态
+            head_first: 是否将 head 维提前
+        返回:
+            out: (B, 1, H, K)  与输入 dtype 一致
+            last_state: (B, H, K, K) 当 output_final_state=True
+        """
+        # 统一转 (B, 1, H, K) 并验证 T=1
+        r = _transpose_head(r, head_first)
+        w = _transpose_head(w, head_first)
+        k = _transpose_head(k, head_first)
+        v = _transpose_head(v, head_first)
+        a = _transpose_head(a, head_first)
+        b = _transpose_head(b, head_first)
+
+        B, T, H, K = r.shape
+        if T != 1:
+            raise ValueError(f"Single-step kernel requires T=1, but got T={T}.")
+
+        # 处理初始状态
+        if initial_state is None:
+            h0 = jnp.zeros((B, H, K, K), jnp.float32)
+        else:
+            h0 = jnp.asarray(initial_state, jnp.float32)
+
+        # 移除 T 维度后调用 kernel
+        r = r[:, 0, :, :]  # (B, H, K)
+        w = w[:, 0, :, :]
+        k = k[:, 0, :, :]
+        v = v[:, 0, :, :]
+        a = a[:, 0, :, :]
+        b = b[:, 0, :, :]
+
+        # 调用前向 kernel
+        out, last_state = wk7_single_step_kernel(w, r, k, v, a, b, h0)
+
+        # 恢复 T 维度
+        out = jnp.expand_dims(out, axis=1)  # (B, 1, H, K)
+        out = jnp.asarray(out, r.dtype)
+
+        if output_final_state:
+            return out, last_state
+        return out
+
+    return generalized_delta_rule_single_step

rwkv_ops/rwkv7_kernel/jax_kernel/__init__.py

@@ -0,0 +1,9 @@
+from ..jax_kernel.chunk_A_fwd import *
+from ..jax_kernel.chunk_A_bwd import *
+from ..jax_kernel.chunk_h_fwd import *
+from ..jax_kernel.chunk_h_bwd import *
+from ..jax_kernel.chunk_o_fwd import *
+from ..jax_kernel.chunk_o_bwd import *
+from ..jax_kernel.cumsum import *
+from ..jax_kernel.wy_fast_fwd import *
+from ..jax_kernel.wy_fast_bwd import *

rwkv_ops/rwkv7_kernel/jax_kernel/chunk_A_bwd.py

@@ -0,0 +1,95 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025,Qingwen Lin
+
+import jax_triton as jt
+import jax
+import triton
+from ..triton_kernel.chunk_A_bwd import *
+from ..triton_kernel.utils import is_gather_supported
+from ..get_torch_devices_info import check_shared_mem
+
+
+def chunk_dplr_bwd_dqk_intra(
+    q: jax.Array,
+    k: jax.Array,
+    a: jax.Array,
+    b: jax.Array,
+    gi: jax.Array,
+    ge: jax.Array,
+    dAqk: jax.Array,
+    dAqb: jax.Array,
+    dAak: jax.Array,
+    dAab: jax.Array,
+    dqg: jax.Array,
+    dkg: jax.Array,
+    dag: jax.Array,
+    dbg: jax.Array,
+    dgk_last: jax.Array,
+    scale: float = 1.0,
+    chunk_size: int = 16,
+):
+    B, T, H, K = q.shape
+    BT = min(chunk_size, max(16, triton.next_power_of_2(T)))
+    BK = (
+        min(64, triton.next_power_of_2(K))
+        if check_shared_mem()
+        else min(32, triton.next_power_of_2(K))
+    )
+
+    NT = triton.cdiv(T, BT)
+    NK = triton.cdiv(K, BK)
+    grid = (NK, NT, B * H)
+
+    out_shapes = [
+        jax.ShapeDtypeStruct(q.shape, q.dtype),
+        jax.ShapeDtypeStruct(k.shape, k.dtype),
+        jax.ShapeDtypeStruct(a.shape, a.dtype),
+        jax.ShapeDtypeStruct(b.shape, b.dtype),
+        jax.ShapeDtypeStruct(gi.shape, "float32"),
+        jax.ShapeDtypeStruct(gi.shape, "float32"),
+    ]
+
+    dq, dk, da, db, dgk, dgk_offset = jt.triton_call(
+        q,
+        k,
+        a,
+        b,
+        gi,
+        ge,
+        dAqk,
+        dAqb,
+        dAak,
+        dAab,
+        dqg,
+        dkg,
+        dag,
+        dbg,
+        T,
+        scale=scale,
+        H=H,
+        K=K,
+        BT=BT,
+        BC=BT,
+        BK=BK,
+        GATHER_SUPPORTED=is_gather_supported,
+        kernel=chunk_dplr_bwd_kernel_intra,
+        out_shape=out_shapes,
+        grid=grid,
+    )
+
+    def grid(meta):
+        return (NT, triton.cdiv(K, meta["BK"]), B * H)
+
+    dgk_output = jt.triton_call(
+        dgk,
+        dgk_offset,
+        dgk_last,
+        T,
+        H=H,
+        K=K,
+        BT=BT,
+        kernel=chunk_dplr_bwd_dgk_kernel,
+        out_shape=[jax.ShapeDtypeStruct(dgk.shape, "float32")],
+        grid=grid,
+    )[0]
+    return dq, dk, da, db, dgk_output

rwkv_ops/rwkv7_kernel/jax_kernel/chunk_A_fwd.py

@@ -0,0 +1,60 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025,Qingwen Lin
+
+
+import jax_triton as jt
+import jax
+import triton
+
+from ..triton_kernel.utils import is_gather_supported
+
+from ..triton_kernel.chunk_A_fwd import *
+
+
+def chunk_dplr_fwd_intra(
+    q: jax.Array,
+    k: jax.Array,
+    a: jax.Array,
+    b: jax.Array,
+    gi: jax.Array,
+    ge: jax.Array,
+    scale: float,
+    chunk_size: int,
+):
+    B, T, H, K = k.shape
+    BT = min(chunk_size, max(16, triton.next_power_of_2(T)))
+
+    NT = triton.cdiv(T, BT)
+    shape = [B, T, H, BT]
+    out_shapes = [
+        jax.ShapeDtypeStruct(q.shape, q.dtype),
+        jax.ShapeDtypeStruct(k.shape, q.dtype),
+        jax.ShapeDtypeStruct(a.shape, q.dtype),
+        jax.ShapeDtypeStruct(b.shape, q.dtype),
+        jax.ShapeDtypeStruct(shape, q.dtype),
+        jax.ShapeDtypeStruct(shape, q.dtype),
+        jax.ShapeDtypeStruct(shape, "float32"),
+        jax.ShapeDtypeStruct(shape, "float32"),
+    ]
+    grid = (NT, B, H)
+    BK = triton.next_power_of_2(K)
+    qg, kg, ag, bg, Aqk, Aqb, Aab, Aak = jt.triton_call(
+        q,
+        k,
+        a,
+        b,
+        gi,
+        ge,
+        T,
+        scale=scale,
+        H=H,
+        K=K,
+        BT=BT,
+        BC=BT,
+        BK=BK,
+        GATHER_SUPPORTED=is_gather_supported,
+        kernel=chunk_dplr_fwd_A_kernel_intra_sub_intra,
+        out_shape=out_shapes,
+        grid=grid,
+    )
+    return Aab, Aqk, Aak, Aqb, qg, kg, ag, bg

rwkv_ops/rwkv7_kernel/jax_kernel/chunk_h_bwd.py

@@ -0,0 +1,78 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025,Qingwen Lin
+
+from typing import Optional, Tuple
+
+import jax_triton as jt
+import jax
+import triton
+
+from ..get_jax_devices_info import check_shared_mem
+from ..triton_kernel.chunk_h_bwd import *
+
+
+def chunk_dplr_bwd_dhu(
+    qg: jax.Array,
+    bg: jax.Array,
+    w: jax.Array,
+    gk: jax.Array,
+    h0: jax.Array,
+    dht: Optional[jax.Array],
+    do: jax.Array,
+    dv: jax.Array,
+    chunk_size: int = 64,
+) -> Tuple[jax.Array, jax.Array, jax.Array]:
+    B, T, H, K, V = *qg.shape, do.shape[-1]
+    BT = min(chunk_size, max(triton.next_power_of_2(T), 16))
+    BK = triton.next_power_of_2(K)
+    assert BK <= 256, (
+        "current kernel does not support head dimension being larger than 256."
+    )
+    # H100
+    if check_shared_mem("hopper"):
+        BV = 64
+        BC = 64 if K <= 128 else 32
+    elif check_shared_mem("ampere"):  # A100
+        BV = 32
+        BC = 32
+    else:  # Etc: 4090
+        BV = 16
+        BC = 16
+
+    N, NT = B, triton.cdiv(T, BT)
+    BC = min(BT, BC)
+    NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
+    assert NK == 1, (
+        "NK > 1 is not supported because it involves time-consuming synchronization"
+    )
+    dh_shape = (B, NT, H, K, V)
+    out_shapes = [
+        jax.ShapeDtypeStruct(dh_shape, dv.dtype),
+        jax.ShapeDtypeStruct((B, H, K, V), "float32"),
+        jax.ShapeDtypeStruct(dv.shape, dv.dtype),
+    ]
+
+    grid = (NK, NV, N * H)
+    dh, dh0, dv2 = jt.triton_call(
+        qg,
+        bg,
+        w,
+        gk,
+        dht,
+        dv,
+        do,
+        T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BC=BC,
+        BK=BK,
+        BV=BV,
+        kernel=chunk_dplr_bwd_kernel_dhu.fn,
+        out_shape=out_shapes,
+        grid=grid,
+        USE_FINAL_STATE_GRADIENT=dht is not None,
+        USE_INITIAL_STATE=h0 is not None,
+    )
+    return dh, dh0, dv2

rwkv_ops/rwkv7_kernel/jax_kernel/chunk_h_fwd.py

@@ -0,0 +1,80 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025,Qingwen Lin
+
+from typing import Optional, Tuple
+
+import jax_triton as jt
+import jax
+import triton
+
+from ..get_jax_devices_info import check_shared_mem
+from ..triton_kernel.chunk_h_fwd import *
+
+
+def chunk_dplr_fwd_h(
+    kg: jax.Array,
+    v: jax.Array,
+    w: jax.Array,
+    u: jax.Array,
+    bg: jax.Array,
+    gk: jax.Array,
+    initial_state: Optional[jax.Array] = None,
+    output_final_state: bool = False,
+    chunk_size: int = 64,
+) -> Tuple[jax.Array, jax.Array]:
+    B, T, H, K, V = *kg.shape, u.shape[-1]
+    BT = min(chunk_size, max(triton.next_power_of_2(T), 16))
+
+    N, NT, chunk_offsets = B, triton.cdiv(T, BT), None
+    BK = triton.next_power_of_2(K)
+    assert BK <= 256, "current kernel does not support head dimension larger than 256."
+    # H100 can have larger block size
+
+    if check_shared_mem("hopper"):
+        BV = 64
+        BC = 64 if K <= 128 else 32
+    elif check_shared_mem("ampere"):  # A100
+        BV = 32
+        BC = 32
+    else:
+        BV = 16
+        BC = 16
+
+    BC = min(BT, BC)
+    NK = triton.cdiv(K, BK)
+    NV = triton.cdiv(V, BV)
+    assert NK == 1, (
+        "NK > 1 is not supported because it involves time-consuming synchronization"
+    )
+
+    out_shapes = [
+        jax.ShapeDtypeStruct((B, NT, H, K, V), kg.dtype),
+        jax.ShapeDtypeStruct([N, H, K, V], "float32"),
+        jax.ShapeDtypeStruct(u.shape, u.dtype),
+    ]
+    grid = (NK, NV, N * H)
+    if initial_state is None:
+        initial_state = jax.numpy.zeros([N, H, K, V], "float32")
+    h, final_state, v_new = jt.triton_call(
+        kg,
+        v,
+        w,
+        bg,
+        u,
+        gk,
+        initial_state,
+        T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BC=BC,
+        BK=BK,
+        BV=BV,
+        kernel=chunk_dplr_fwd_kernel_h.fn,
+        out_shape=out_shapes,
+        grid=grid,
+        STORE_FINAL_STATE=True,
+        USE_INITIAL_STATE=True,
+    )
+    return h, v_new, final_state

rwkv_ops/rwkv7_kernel/jax_kernel/chunk_o_bwd.py

@@ -0,0 +1,150 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from typing import Tuple
+
+import jax_triton as jt
+import jax
+import triton
+
+from ..get_jax_devices_info import check_shared_mem
+from ..triton_kernel.chunk_o_bwd import *
+
+
+def chunk_dplr_bwd_dv(
+    A_qk: jax.Array,
+    kg: jax.Array,
+    do: jax.Array,
+    dh: jax.Array,
+    chunk_size: int = 64,
+) -> jax.Array:
+    B, T, H, K, V = *kg.shape, do.shape[-1]
+    BT = min(chunk_size, max(16, triton.next_power_of_2(T)))
+
+    NT = triton.cdiv(T, BT)
+
+    def grid(meta):
+        return (triton.cdiv(V, meta["BV"]), NT, B * H)
+
+    dv = jt.triton_call(
+        A_qk,
+        kg,
+        do,
+        dh,
+        T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        kernel=chunk_dplr_bwd_kernel_dv,
+        out_shape=jax.ShapeDtypeStruct(do.shape, do.dtype),
+        grid=grid,
+    )
+    return dv
+
+
+def chunk_dplr_bwd_o(
+    k: jax.Array,
+    b: jax.Array,
+    v: jax.Array,
+    v_new: jax.Array,
+    gk: jax.Array,
+    do: jax.Array,
+    h: jax.Array,
+    dh: jax.Array,
+    dv: jax.Array,
+    w: jax.Array,
+    chunk_size: int = 64,
+    scale: float = 1.0,
+) -> Tuple[jax.Array, jax.Array, jax.Array]:
+    B, T, H, K, V = *w.shape, v.shape[-1]
+
+    BT = min(chunk_size, max(16, triton.next_power_of_2(T)))
+    NT = triton.cdiv(T, BT)
+
+    BK = (
+        min(triton.next_power_of_2(K), 64)
+        if check_shared_mem()
+        else min(triton.next_power_of_2(K), 32)
+    )
+    BV = (
+        min(triton.next_power_of_2(V), 64)
+        if check_shared_mem()
+        else min(triton.next_power_of_2(K), 32)
+    )
+    NK = triton.cdiv(K, BK)
+    grid = (NK, NT, B * H)
+
+    out_shapes = [
+        jax.ShapeDtypeStruct(k.shape, k.dtype),
+        jax.ShapeDtypeStruct(k.shape, k.dtype),
+        jax.ShapeDtypeStruct(w.shape, w.dtype),
+        jax.ShapeDtypeStruct(b.shape, b.dtype),
+        jax.ShapeDtypeStruct([B, NT, H, K], "float32"),
+    ]
+    dq, dk, dw, db, dgk_last = jt.triton_call(
+        v,
+        v_new,
+        h,
+        do,
+        dh,
+        w,
+        dv,
+        gk,
+        k,
+        b,
+        T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=BK,
+        BV=BV,
+        kernel=chunk_dplr_bwd_o_kernel,
+        out_shape=out_shapes,
+        grid=grid,
+    )
+    return (dq, dk, dw, db, dgk_last)
+
+
+def chunk_dplr_bwd_dAu(
+    v: jax.Array,
+    v_new: jax.Array,
+    do: jax.Array,
+    A_qb: jax.Array,
+    scale: float,
+    chunk_size: int = 64,
+) -> jax.Array:
+    B, T, H, V = v.shape
+    BT = min(chunk_size, max(16, triton.next_power_of_2(T)))
+    NT = triton.cdiv(T, BT)
+
+    if check_shared_mem("ampere"):  # A100
+        BV = min(triton.next_power_of_2(V), 128)
+    elif check_shared_mem("ada"):  # 4090
+        BV = min(triton.next_power_of_2(V), 64)
+    else:
+        BV = min(triton.next_power_of_2(V), 32)
+
+    grid = (NT, B * H)
+    out_shapes = [
+        jax.ShapeDtypeStruct([B, T, H, BT], "float32"),
+        jax.ShapeDtypeStruct([B, T, H, BT], "float32"),
+        jax.ShapeDtypeStruct(v_new.shape, v_new.dtype),
+    ]
+    dA_qk, dA_qb, dv_new = jt.triton_call(
+        v,
+        do,
+        v_new,
+        A_qb,
+        T,
+        scale=scale,
+        H=H,
+        V=V,
+        BT=BT,
+        BV=BV,
+        grid=grid,
+        out_shape=out_shapes,
+        kernel=chunk_dplr_bwd_kernel_dAu,
+    )
+    return dv_new, dA_qk, dA_qb

rwkv_ops/rwkv7_kernel/jax_kernel/chunk_o_fwd.py

@@ -0,0 +1,45 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025,Qingwen Lin
+
+
+import jax_triton as jt
+import jax
+import triton
+
+from ..triton_kernel.chunk_o_fwd import *
+
+
+def chunk_dplr_fwd_o(
+    qg: jax.Array,
+    v: jax.Array,
+    v_new: jax.Array,
+    A_qk: jax.Array,
+    A_qb: jax.Array,
+    h: jax.Array,
+    chunk_size: int = 64,
+) -> jax.Array:
+    B, T, H, K, V = *qg.shape, v.shape[-1]
+    BT = min(chunk_size, max(16, triton.next_power_of_2(T)))
+
+    NT = triton.cdiv(T, BT)
+
+    def grid(meta):
+        return (triton.cdiv(V, meta["BV"]), NT, B * H)
+
+    o = jt.triton_call(
+        qg,
+        v,
+        v_new,
+        A_qk,
+        A_qb,
+        h,
+        T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        kernel=chunk_dplr_fwd_kernel_o,
+        out_shape=jax.ShapeDtypeStruct(v.shape, v.dtype),
+        grid=grid,
+    )
+    return o