rwkv-ops 0.2.2__py3-none-any.whl → 0.3.1__py3-none-any.whl

rwkv_ops/rwkv7_kernel/jax_cuda_kernel/wkv7_jax.py

@@ -0,0 +1,237 @@
+"""
+JAX 版 RWKV7 wkv kernel + generalized_delta_rule
+延迟编译 CUDA 扩展，接口与 Torch 版本 1:1 对齐
+"""
+
+from __future__ import annotations
+import pathlib
+import subprocess
+import ctypes
+import jax
+import jax.numpy as jnp
+from typing import Optional, Tuple, Union
+from jax.ad_checkpoint import checkpoint_policies as cp
+CHUNK_LEN = 16  # 这是一个常数
+# ---------- 延迟编译（改到当前目录） ----------
+_CURRENT_DIR = pathlib.Path(
+    __file__
+).parent.absolute()  # rwkv_ops/rwkv7_kernel/jax_cuda_kernel
+
+
+def get_jax_generalized_delta_rule(HEAD_SIZE=64):
+    _BUILD_DIR = _CURRENT_DIR / f"build_{HEAD_SIZE}"
+    _SO_PATH = _CURRENT_DIR / f"build_{HEAD_SIZE}/wkv7.so"
+
+    def _ensure_compiled() -> pathlib.Path:
+        """首次调用时编译 CUDA 扩展，产出放在当前源码目录"""
+        if _SO_PATH.exists():
+            return _SO_PATH
+
+        print("[rwkv7_jax] First use – compiling CUDA kernel…")
+        src_dir = _CURRENT_DIR
+        build_dir = _BUILD_DIR
+        build_dir.mkdir(exist_ok=True)
+
+        # ---------- 关键：拿到 JAX 的 XLA 头文件路径 ----------
+        xla_include_dir = jax.ffi.include_dir()  # 方案 3 核心 API
+        if not xla_include_dir:
+            raise RuntimeError("jax.ffi.include_dir() 返回空，请检查 JAX >= 0.4.31")
+
+        # ---------- 关键：把数值稳定性 flag 写死 ----------
+        cuda_flags = [
+            "-ftz=true",  # flush sub-normal to zero
+            "-prec-div=false",  # 更快除法，避免特殊路径
+            "-prec-sqrt=false",  # 更快开方
+            "--use_fast_math",  # 统一 fast math
+            "-O3",
+            "-Xptxas=-O3",
+            "-res-usage",
+            "--extra-device-vectorization",
+            "-D_C_=64",
+            f"-D_C_={HEAD_SIZE}",
+            f"-D_CHUNK_LEN_={CHUNK_LEN}",
+        ]
+
+        # 1. 配置
+        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}",  # 传给 CMake
+            f"-DCMAKE_CUDA_FLAGS={' '.join(cuda_flags)}",
+        ]
+        subprocess.check_call(cmake_args)
+
+        # 2. 构建
+        subprocess.check_call(["cmake", "--build", str(build_dir), "-j"])
+
+        # 3. 安装（把 .so 拷贝到当前目录）
+        subprocess.check_call(["cmake", "--install", str(build_dir)])
+
+        if not _SO_PATH.exists():
+            raise RuntimeError("Compilation failed – wkv7.so not found.")
+
+        print("[rwkv7_jax] Compilation finished – output at", _SO_PATH)
+        return _SO_PATH
+
+    # 注册 FFI 符号
+    _lib = ctypes.CDLL(_ensure_compiled())
+    jax.ffi.register_ffi_target(
+        "wkv7_fwd", jax.ffi.pycapsule(_lib.Wkv7Fwd), platform="CUDA"
+    )
+    jax.ffi.register_ffi_target(
+        "wkv7_bwd", jax.ffi.pycapsule(_lib.Wkv7Bwd), platform="CUDA"
+    )
+
+    # ---------- 工具 ----------
+    def _transpose_head(x: jnp.ndarray, head_first: bool) -> jnp.ndarray:
+        """(B, T, H, K) <-> (B, H, T, K)"""
+        x = jnp.asarray(x, dtype=jnp.bfloat16)
+        if head_first:
+            return jnp.transpose(x, (0, 2, 1, 3))
+        return x
+
+    # ---------- 前向 + 反向 kernel ----------
+
+    def _wkv7_kernel(
+        w: jnp.ndarray,
+        q: jnp.ndarray,
+        k: jnp.ndarray,
+        v: jnp.ndarray,
+        a: jnp.ndarray,
+        b: jnp.ndarray,
+        h0: jnp.ndarray,
+    ):
+        """
+        内部 kernel 接口
+        参数顺序与 wkv7_ffi.cc 声明完全一致：
+        w,q,k,v,z,a,b  -> y,s,sa
+        """
+        B, T, H, K = q.shape
+        dtype = q.dtype
+        chunk_num = int(T // CHUNK_LEN)
+        out_type = jax.ShapeDtypeStruct((B, T, H, K), dtype)
+        s_type = jax.ShapeDtypeStruct((B, H, chunk_num, K, K), jnp.float32)
+        sa_type = jax.ShapeDtypeStruct((B, T, H, K), jnp.float32)
+
+        y, s, sa = jax.ffi.ffi_call(
+            "wkv7_fwd", (out_type, s_type, sa_type), vmap_method="broadcast_all"
+        )(w, q, k, v, a, b, h0)
+
+        return y, s, sa
+
+    @jax.custom_vjp
+    def wk7_kernel(
+        w: jnp.ndarray,
+        q: jnp.ndarray,
+        k: jnp.ndarray,
+        v: jnp.ndarray,
+        a: jnp.ndarray,
+        b: jnp.ndarray,
+        h0: jnp.ndarray,
+    ):
+        y, s, sa = _wkv7_kernel(w, q, k, v, a, b, h0)
+        finnal_state = s[:, :, -1]
+        return (y, jnp.transpose(finnal_state, [0, 1, 3, 2]))
+
+    # 前向定义
+    def _fwd(
+        w: jnp.ndarray,
+        q: jnp.ndarray,
+        k: jnp.ndarray,
+        v: jnp.ndarray,
+        a: jnp.ndarray,
+        b: jnp.ndarray,
+        h0: jnp.ndarray,
+    ):
+        y, s, sa = _wkv7_kernel(w, q, k, v, a, b, h0)
+        finnal_state = s[:, :, -1]
+        return (y, jnp.transpose(finnal_state, [0, 1, 3, 2])), (w, q, k, v, a, b, s, sa)
+
+    def _wkv7_bwd_kernel(w, q, k, v, a, b, dy, s, sa, dht):
+        dh0_type = jax.ShapeDtypeStruct(dht.shape, dht.dtype)
+        dw_type = jax.ShapeDtypeStruct(w.shape, w.dtype)
+        dq_type = jax.ShapeDtypeStruct(q.shape, q.dtype)
+        dk_type = jax.ShapeDtypeStruct(k.shape, k.dtype)
+        dv_type = jax.ShapeDtypeStruct(v.shape, v.dtype)
+        da_type = jax.ShapeDtypeStruct(a.shape, a.dtype)
+        db_type = jax.ShapeDtypeStruct(b.shape, b.dtype)
+
+        dh0, dw, dq, dk, dv, da, db = jax.ffi.ffi_call(
+            "wkv7_bwd",
+            (dh0_type, dw_type, dq_type, dk_type, dv_type, da_type, db_type),
+            vmap_method="broadcast_all",
+        )(w, q, k, v, a, b, dy, s, sa, dht)
+
+        return dw, dq, dk, dv, da, db, dh0
+
+    # 反向定义
+    def _bwd(res, grads):
+        w, q, k, v, a, b, s, sa = res
+        dy, dht = grads
+        dy = jnp.asarray(dy, jnp.bfloat16)
+        # 调用反向 kernel
+        return _wkv7_bwd_kernel(w, q, k, v, a, b, dy, s, sa, dht)
+
+    wk7_kernel.defvjp(_fwd, _bwd)
+
+    def generalized_delta_rule(
+        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 规则，接口与 Torch 实现完全一致
+        参数:
+            r,w,k,v,a,b: 输入张量，形状 (B, T, H, K) 或 (B, H, T, K) 当 head_first=True
+            initial_state: 可选 (B, H, K, K) 初始状态，None 则零初始化
+            output_final_state: 是否同时返回最后状态
+            head_first: 是否将 head 维提前
+            chunk_len: 必须整除 T，默认 16
+        返回:
+            out: (B, T, H, K)  与输入 dtype 一致
+            last_state: (B, H, K, K) 当 output_final_state=True
+        """
+        # 统一转 (B, T, H, K)
+        dtype = r.dtype
+        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 % CHUNK_LEN:
+            raise ValueError(
+                f"Sequence length T={T} must be divisible by chunk_len={CHUNK_LEN}"
+            )
+
+        # 处理初始状态
+        if initial_state is None:
+            h0 = jnp.zeros((B, H, K, K), jnp.float32)
+        else:
+            h0 = jnp.asarray(initial_state, jnp.float32)
+
+        # 调用 kernel
+
+        out, last_state = jax.checkpoint(
+            wk7_kernel, policy=cp.save_anything_except_these_names(())
+        )(w, r, k, v, a, b, h0)
+        out = jnp.asarray(out, dtype)  # 保证输出 dtype 与输入一致
+
+        if output_final_state:
+            return out, last_state
+        return out
+
+    return generalized_delta_rule, _wkv7_kernel, _wkv7_bwd_kernel

rwkv_ops/rwkv7_kernel/jax_op.py

@@ -14,9 +14,11 @@ from .jax_kernel.chunk_o_fwd import chunk_dplr_fwd_o
 from .jax_kernel.wy_fast_bwd import chunk_dplr_bwd_wy
 from .jax_kernel.wy_fast_fwd import prepare_wy_repr_fwd
 from .jax_kernel.cumsum import chunk_rwkv6_fwd_cumsum
-from jax.ad_checkpoint import checkpoint_policies 
+from jax.ad_checkpoint import checkpoint_policies as cp
+
 CHUNKSIZE = 16
 
+
 def chunk_dplr_fwd(
     q: jax.Array,
     k: jax.Array,
@@ -307,7 +309,6 @@ def transpose_head(x, head_first):
         return x
 
 
-# @partial(jax.jit, static_argnames=['initial_state',"output_final_state","head_first","use_chunk"])
 def generalized_delta_rule(
     r: jax.Array,
     w: jax.Array,
@@ -365,7 +366,9 @@ def generalized_delta_rule(
     else:
         assert log_w is not None, "Either w or log_w must be provided!"
     log_w = transpose_head(log_w, head_first)
-    o, final_state = chunk_dplr(
+    o, final_state = jax.checkpoint(
+        chunk_dplr, policy=cp.save_anything_except_these_names(())
+    )(
         r=r,
         k=k,
         v=v,
@@ -377,5 +380,3 @@ def generalized_delta_rule(
     if output_final_state:
         return jnp.asarray(o, DTYPE), final_state
     return jnp.asarray(o, DTYPE)
-
-

rwkv_ops/rwkv7_kernel/native_keras_op.py

@@ -1,4 +1,3 @@
-import keras
 from keras import ops
 
 
@@ -62,8 +61,9 @@ def generalized_delta_rule(
         if ops.shape(state)[0] == 1:
             state = ops.broadcast_to(state, (B, H, N, N))
     else:
-        state = ops.zeros((B, H, N, N), dtype="float32")
-    out = ops.zeros((B, T, H, N), dtype=r.dtype)
+        state = ops.zeros((B, H, N, N))
+    state = ops.cast(state, "float32")
+    out = ops.zeros((B, T, H, N), DTYPE)
 
     def step(t, inputs):
         """
@@ -83,9 +83,8 @@ def generalized_delta_rule(
         aa = ops.reshape(a[:, t, :], (B, H, N, 1))
         bb = ops.reshape(b[:, t, :], (B, H, 1, N))
         state = state * w[:, t, :, None, :] + state @ aa @ bb + vv @ kk
-        out = ops.slice_update(
-            out, [0, t, 0, 0], ops.reshape((state @ rr), (B, 1, H, N))
-        )
+        o = ops.cast((state @ rr), out.dtype)
+        out = ops.slice_update(out, [0, t, 0, 0], ops.reshape(o, (B, 1, H, N)))
         return [state, out]
 
     state, out = ops.fori_loop(0, T, step, [state, out])

rwkv_ops/rwkv7_kernel/tf_eager_kernel.py

@@ -0,0 +1,123 @@
+"""
+TensorFlow 版 generalized_delta_rule
+前向用 tf.py_function 调 JAX CUDA 内核，反向同样走 JAX。
+可 @tf.function 编译，可 tf.GradientTape 训练。
+"""
+
+import tensorflow as tf
+from typing import Optional, Tuple
+import jax.numpy as jnp
+from .jax_cuda_kernel.wkv7_jax import get_jax_generalized_delta_rule
+
+
+def transpose_head(x, head_first: bool):
+    """(B, T, H, K) <-> (B, H, T, K)"""
+    x = tf.cast(x, dtype=tf.float32)
+    if head_first:
+        return tf.transpose(x, (0, 2, 1, 3))
+    return x
+
+
+def get_tf_generalized_delta_rule(HEAD_SIZE=64):
+    _, _wkv7_kernel, _wkv7_bwd_kernel = get_jax_generalized_delta_rule(HEAD_SIZE)
+
+    # ---------- 底层 kernel 包装 ----------
+    @tf.py_function(Tout=[tf.bfloat16, tf.float32, tf.float32])
+    def _tf_wkv7_fwd(w, q, k, v, a, b, h0):
+        """tf.py_function 包装 JAX 前向"""
+        y, s, sa = _wkv7_kernel(
+            jnp.asarray(w, jnp.bfloat16),
+            jnp.asarray(q, jnp.bfloat16),
+            jnp.asarray(k, jnp.bfloat16),
+            jnp.asarray(v, jnp.bfloat16),
+            jnp.asarray(a, jnp.bfloat16),
+            jnp.asarray(b, jnp.bfloat16),
+            jnp.asarray(h0, jnp.float32),
+        )
+        return (
+            tf.convert_to_tensor(y, tf.bfloat16),
+            tf.convert_to_tensor(s, tf.float32),
+            tf.convert_to_tensor(sa, tf.float32),
+        )
+
+    @tf.py_function(Tout=[tf.bfloat16] * 6 + [tf.float32])
+    def _tf_wkv7_bwd(w, q, k, v, a, b, dy, s, sa, dht):
+        """tf.py_function 包装 JAX 反向"""
+        dw, dq, dk, dv, da, db, dh0 = _wkv7_bwd_kernel(
+            jnp.asarray(w, jnp.bfloat16),
+            jnp.asarray(q, jnp.bfloat16),
+            jnp.asarray(k, jnp.bfloat16),
+            jnp.asarray(v, jnp.bfloat16),
+            jnp.asarray(a, jnp.bfloat16),
+            jnp.asarray(b, jnp.bfloat16),
+            jnp.asarray(dy, jnp.bfloat16),
+            jnp.asarray(s, jnp.float32),
+            jnp.asarray(sa, jnp.float32),
+            jnp.asarray(dht, jnp.bfloat16),
+        )
+        return tuple(
+            tf.convert_to_tensor(g, dtype)
+            for g, dtype in zip((dw, dq, dk, dv, da, db), [tf.bfloat16] * 6)
+        ) + (tf.convert_to_tensor(dh0, tf.float32),)
+
+    # ---------- 带梯度的前向 ----------
+    @tf.custom_gradient
+    def _wk7_tf(w, q, k, v, a, b, h0):
+        y, s, sa = _tf_wkv7_fwd(w, q, k, v, a, b, h0)
+
+        def grad(dy, dht):
+            # dy 上层传来的 loss 对 y 的梯度
+            # dht 对最后状态的梯度（没有就传 0）
+            if dht is None:
+                dht = tf.zeros_like(h0)
+            grads = _tf_wkv7_bwd(w, q, k, v, a, b, dy, s, sa, dht)
+            return grads  # (dw, dq, dk, dv, da, db, dh0)
+
+        final_state = s[:, :, -1]  # (B, H, K, K)
+        final_state = tf.transpose(final_state, [0, 1, 3, 2])  # 与 JAX 对齐
+        return (y, final_state), grad
+
+    # ---------- 用户接口 ----------
+    def generalized_delta_rule(
+        r: tf.Tensor,  # (B, T, H, K) 或 (B, H, T, K)
+        w: tf.Tensor,
+        k: tf.Tensor,
+        v: tf.Tensor,
+        a: tf.Tensor,
+        b: tf.Tensor,
+        initial_state: Optional[tf.Tensor] = None,
+        output_final_state: bool = True,
+        head_first: bool = False,
+        chunk_len: int = 16,
+    ) -> Tuple[tf.Tensor, tf.Tensor]:
+        """
+        与 JAX 版接口 1:1 对齐，返回 (out, last_state)
+        可 @tf.function  compile，可 tf.GradientTape 训练
+        """
+        dtype = r.dtype
+
+        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 = tf.unstack(tf.shape(r), num=4)
+        if T % chunk_len != 0:
+            raise ValueError(f"T={T} must be divisible by chunk_len={chunk_len}")
+
+        if initial_state is None:
+            h0 = tf.zeros([B, H, K, K], dtype=tf.float32)
+        else:
+            h0 = tf.cast(initial_state, tf.float32)
+
+        # 带梯度前向
+        out, last_state = _wk7_tf(w, r, k, v, a, b, h0)
+
+        # 转回用户期望 dtype
+        out = tf.cast(out, dtype)
+
+        return (out, last_state) if output_final_state else out
+
+    return generalized_delta_rule, _tf_wkv7_fwd, _tf_wkv7_bwd

rwkv_ops/rwkv7_kernel/torch_cuda_kernel/wkv7_cuda.cu

@@ -0,0 +1,165 @@
+#include <cuda_bf16.h>
+#include <assert.h>
+
+using bf = __nv_bfloat16;
+
+__device__ inline float to_float(const bf & u) {
+    return __bfloat162float(u);
+}
+
+__device__ inline bf to_bf(const float & u) {
+    return __float2bfloat16_rn(u);
+}
+typedef bf * __restrict__ F_;
+
+__global__ void forward_kernel(int T, int H,
+     F_ w_, F_ q_, F_ k_, F_ v_, F_ a_, F_ b_,
+      bf* y_, float* s_, float* sa_, float* h0_) {
+    constexpr int C = _C_;
+    int bb = blockIdx.y, hh = blockIdx.x, i = threadIdx.x;
+    float state[C] =  {0};
+    __shared__ float q[C], k[C], w[C], a[C], b[C];
+    int h0_base =( (bb*H + hh)*C + i)*C;
+#pragma unroll
+        for (int j = 0; j < C; j++) {
+            state[j] = h0_[h0_base + j];
+        }
+    for (int t = 0; t < T; t++) {
+        int ind = bb*T*H*C + t*H*C + hh * C + i;
+        __syncthreads();
+        q[i] = to_float(q_[ind]);
+        w[i] = __expf(-__expf(to_float(w_[ind])));
+        k[i] = to_float(k_[ind]);
+        a[i] = to_float(a_[ind]);
+        b[i] = to_float(b_[ind]);
+        __syncthreads();
+        float sa = 0;
+#pragma unroll
+        for (int j = 0; j < C; j++) {
+            sa += a[j] * state[j];
+        }
+        sa_[ind] = sa;
+        float v = to_float(v_[ind]);
+        float y = 0;
+#pragma unroll
+        for (int j = 0; j < C; j++) {
+            float& s = state[j];
+            s = s * w[j] + sa * b[j] + k[j] * v;
+            y += s * q[j];
+        }
+        y_[ind] = to_bf(y);
+
+        if ((t+1)%_CHUNK_LEN_ == 0) {
+            int base = (bb*H+hh)*(T/_CHUNK_LEN_)*C*C + (t/_CHUNK_LEN_)*C*C + i;
+#pragma unroll
+            for (int j = 0; j < C; j++) {
+                s_[base + j*C] = state[j];
+            }
+        }
+    }
+}
+__global__ void backward_kernel(int T, int H, 
+    F_ w_, F_ q_, F_ k_, F_ v_, F_ a_, F_ b_, F_ dy_,
+float * __restrict__ s_, float * __restrict__ sa_,
+float * __restrict__ dht_,float * __restrict__ dh0_,
+bf* dw_, bf* dq_, bf* dk_, bf* dv_, bf* da_, bf* db_) {
+    constexpr int C = _C_;
+    int bb = blockIdx.y, hh = blockIdx.x, i = threadIdx.x;
+    float stateT[C] =  {
+        0
+    }
+    , dstate[C] =  {
+        0
+    }
+    , dstateT[C] =  {
+        0
+    }
+    ;
+    int dht_base =( (bb*H + hh)*C + i)*C;
+#pragma unroll
+        for (int j = 0; j < C; j++) {
+            dstate[j] = dht_[dht_base + j];
+            dstateT[j] = dht_[dht_base + j];
+        }
+    __shared__ float w[C], q[C], k[C], v[C], a[C], b[C], dy[C], sa[C], dSb_shared[C];
+    float qi, wi, ki, ai, bi, dyi;
+
+    for (int t = T-1; t >= 0; t--) {
+        int ind = bb*T*H*C + t*H*C + hh * C + i;
+        __syncthreads();
+        q[i] = qi = to_float(q_[ind]);
+        float wi_fac = -__expf(to_float(w_[ind]));
+        w[i] = wi = __expf(wi_fac);
+        k[i] = ki = to_float(k_[ind]);
+        a[i] = ai = to_float(a_[ind]);
+        b[i] = bi = to_float(b_[ind]);
+        v[i] = to_float(v_[ind]);
+        dy[i] = dyi = to_float(dy_[ind]);
+        sa[i] = sa_[ind];
+        __syncthreads();
+
+        if ((t+1)%_CHUNK_LEN_ == 0) {
+            int base = (bb*H+hh)*(T/_CHUNK_LEN_)*C*C + (t/_CHUNK_LEN_)*C*C + i*C;
+#pragma unroll
+            for (int j = 0; j < C; j++) {
+                stateT[j] = s_[base + j];
+            }
+        }
+        float dq = 0;
+#pragma unroll
+        for (int j = 0; j < C; j++) {
+            dq += stateT[j]*dy[j];
+        }
+        dq_[ind] = to_bf(dq);
+        float iwi = 1.0f/(wi+0.000001f);
+#pragma unroll
+        for (int j = 0; j < C; j++) {
+            stateT[j] = (stateT[j] - ki*v[j] - bi*sa[j]) * iwi;
+            dstate[j] += dyi * q[j];
+            dstateT[j] += qi * dy[j];
+        }
+        float dw = 0, dk = 0, dv = 0, db = 0, dSb = 0;
+#pragma unroll
+        for (int j = 0; j < C; j++) {
+            dw += dstateT[j]*stateT[j];
+            dk += dstateT[j]*v[j];
+            dv += dstate[j]*k[j];
+            dSb += dstate[j]*b[j];
+            db += dstateT[j]*sa[j];
+        }
+        dw_[ind] = to_bf(dw * wi * wi_fac);
+        dk_[ind] = to_bf(dk);
+        dv_[ind] = to_bf(dv);
+        db_[ind] = to_bf(db);
+        __syncthreads();
+        dSb_shared[i] = dSb;
+        __syncthreads();
+        float da = 0;
+#pragma unroll
+        for (int j = 0; j < C; j++) {
+            da += stateT[j]*dSb_shared[j];
+        }
+        da_[ind] = to_bf(da);
+#pragma unroll
+        for (int j = 0; j < C; j++) {
+            dstate[j] = dstate[j]*w[j] + dSb * a[j];
+            dstateT[j] = dstateT[j]*wi + ai * dSb_shared[j];
+            if (t==0){
+                dh0_[dht_base + j] = dstate[j];
+            }
+        }
+    }
+}
+
+void cuda_forward(int B, int T, int H, bf*w, bf*q, bf*k, bf*v, bf*z, bf*a, bf*y, float*s, float*sa, float* h0) {
+    forward_kernel<<<dim3(H,B), dim3(_C_)>>>(T,H,w,q,k,v,z,a,y,s,sa,h0);
+}
+
+void cuda_backward(int B, int T, int H,
+     bf*w, bf*q, bf*k, bf*v, bf*z, bf*a, bf*dy,
+    float*s, float*sa,float*dht,float*dh0,
+    bf*dw, bf*dq, bf*dk, bf*dv, bf*dz, bf*da
+    ) {
+    assert(T%_CHUNK_LEN_ == 0);
+    backward_kernel<<<dim3(H,B), dim3(_C_)>>>(T,H,w,q,k,v,z,a,dy,s,sa,dht,dh0,dw,dq,dk,dv,dz,da);
+}

rwkv_ops/rwkv7_kernel/torch_cuda_kernel/wkv7_op.cpp

@@ -0,0 +1,35 @@
+#include <torch/extension.h>
+#include <cuda_bf16.h>
+
+using bf = __nv_bfloat16;
+void cuda_forward(int B, int T, int H, bf*w, bf*q, bf*k, bf*v, bf*z, bf*a, bf*y, float*s, float*sa,float*h0);
+
+void forward(torch::Tensor &w, torch::Tensor &q, torch::Tensor &k, torch::Tensor &v, 
+    torch::Tensor &z, torch::Tensor &a,
+    torch::Tensor &y, 
+    torch::Tensor &s, torch::Tensor &sa,torch::Tensor &h0) {
+    int B = w.sizes()[0], T = w.sizes()[1], H = w.sizes()[2];
+    cuda_forward(B, T, H, 
+        (bf*)w.data_ptr(), (bf*)q.data_ptr(), (bf*)k.data_ptr(), (bf*)v.data_ptr(), (bf*)z.data_ptr(), (bf*)a.data_ptr(), (bf*)y.data_ptr(), 
+        (float*)s.data_ptr(), (float*)sa.data_ptr(),(float*)h0.data_ptr());
+}
+void cuda_backward(int B, int T, int H, bf*w, bf*q, bf*k, bf*v, bf*z, bf*a, bf*dy, float*s, float*sa,float*dht,float*dh0, bf*dw, bf*dq, bf*dk, bf*dv, bf*dz, bf*da);
+void backward(torch::Tensor &w, torch::Tensor &q, torch::Tensor &k, torch::Tensor &v, torch::Tensor &z, torch::Tensor &a, torch::Tensor &dy,
+torch::Tensor &s, torch::Tensor &sa,torch::Tensor &dht,torch::Tensor &dh0,
+ torch::Tensor &dw, torch::Tensor &dq, torch::Tensor &dk, torch::Tensor &dv, torch::Tensor &dz, torch::Tensor &da) {
+    int B = w.sizes()[0], T = w.sizes()[1], H = w.sizes()[2];
+    cuda_backward(B, T, H, (bf*)w.data_ptr(), (bf*)q.data_ptr(), (bf*)k.data_ptr(), (bf*)v.data_ptr(), (bf*)z.data_ptr(), (bf*)a.data_ptr(),
+     (bf*)dy.data_ptr(),
+    (float*)s.data_ptr(), (float*)sa.data_ptr(),(float*)dht.data_ptr(),(float*)dh0.data_ptr(),
+     (bf*)dw.data_ptr(), (bf*)dq.data_ptr(), (bf*)dk.data_ptr(), (bf*)dv.data_ptr(), (bf*)dz.data_ptr(), (bf*)da.data_ptr());
+}
+
+TORCH_LIBRARY(wind_backstepping, m) {
+    m.def("forward(Tensor w, Tensor q, Tensor k, Tensor v, Tensor z, Tensor a, Tensor(a!) y, Tensor(b!) s, Tensor(c!) sa, Tensor(f!) h0) -> ()");
+    m.def("backward(Tensor w, Tensor q, Tensor k, Tensor v, Tensor z, Tensor a, Tensor dy, Tensor s, Tensor sa,Tensor dht,Tensor(a!) dh0, Tensor(b!) dw, Tensor(c!) dq, Tensor(d!) dk, Tensor(e!) dv, Tensor(f!) dz, Tensor(g!) da) -> ()");
+}
+
+TORCH_LIBRARY_IMPL(wind_backstepping, CUDA, m) {
+    m.impl("forward", &forward);
+    m.impl("backward", &backward);
+}