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

rwkv_ops/rwkv7_kernel/__init__.py

@@ -12,18 +12,17 @@ def transpose_head(x, head_first):
 
 
 def get_generalized_delta_rule(HEAD_SIZE=64, KERNEL_TYPE="native"):
-    USE_KERNEL = False
+    USE_TRITON_KERNEL = False
     if keras.config.backend() == "torch":
         import torch
 
         if KERNEL_TYPE.lower() == "triton":
             from .torch_op import generalized_delta_rule
 
-            USE_KERNEL = True
+            USE_TRITON_KERNEL = True
 
         elif KERNEL_TYPE.lower() == "cuda":
             CHUNK_LEN = 16
-            USE_KERNEL = True
             from torch.utils.cpp_extension import load
             import os
 
@@ -44,8 +43,8 @@ def get_generalized_delta_rule(HEAD_SIZE=64, KERNEL_TYPE="native"):
             load(
                 name="wind_backstepping",
                 sources=[
-                    os.path.join(current_dir_path, "cuda_kernel/wkv7_cuda.cu"),
-                    os.path.join(current_dir_path, "cuda_kernel/wkv7_op.cpp"),
+                    os.path.join(current_dir_path, "torch_cuda_kernel/wkv7_cuda.cu"),
+                    os.path.join(current_dir_path, "torch_cuda_kernel/wkv7_op.cpp"),
                 ],
                 is_python_module=False,
                 verbose=True,
@@ -54,8 +53,8 @@ def get_generalized_delta_rule(HEAD_SIZE=64, KERNEL_TYPE="native"):
 
             class WindBackstepping(torch.autograd.Function):
                 @staticmethod
-                def forward(ctx, w, q, k, v, z, b):
-                    B, T, H, C = w.shape
+                def forward(ctx, w, q, k, v, z, b, h0):
+                    B, T, H, N = w.shape
                     DTYPE = q.dtype
                     q = ops.cast(q, "bfloat16")
                     k = ops.cast(k, "bfloat16")
@@ -63,30 +62,42 @@ def get_generalized_delta_rule(HEAD_SIZE=64, KERNEL_TYPE="native"):
                     z = ops.cast(z, "bfloat16")
                     b = ops.cast(b, "bfloat16")
                     w = ops.cast(w, "bfloat16")
-                    assert T % CHUNK_LEN == 0
+                    if T % CHUNK_LEN != 0:
+                        raise ValueError(
+                            "RWKV输入的序列长度必须可以被16整除"
+                            "Please make sure the sequence length is divisible by 16"
+                        )
                     assert all(i.is_contiguous() for i in [w, q, k, v, z, b])
                     y = torch.empty_like(v)
                     s = torch.empty(
-                        B, H, T // CHUNK_LEN, C, C, dtype=torch.float32, device=w.device
+                        B, H, T // CHUNK_LEN, N, N, dtype=torch.float32, device=w.device
                     )
-                    sa = torch.empty(B, T, H, C, dtype=torch.float32, device=w.device)
-                    torch.ops.wind_backstepping.forward(w, q, k, v, z, b, y, s, sa)
+                    sa = torch.empty(B, T, H, N, dtype=torch.float32, device=w.device)
+                    torch.ops.wind_backstepping.forward(w, q, k, v, z, b, y, s, sa, h0)
                     ctx.save_for_backward(w, q, k, v, z, b, s, sa)
-                    return ops.cast(y, DTYPE)
+                    last_state = torch.empty_like(h0)
+                    last_state.copy_(ops.transpose(s[:, :, -1], [0, 1, 3, 2]))
+
+                    return ops.cast(y, DTYPE), last_state
 
                 @staticmethod
-                def backward(ctx, dy):
+                def backward(ctx, dy, dht):
                     DTYPE = dy.dtype
                     dy = ops.cast(dy, torch.bfloat16)
                     dy = dy.contiguous()
-                    assert all(i.dtype == torch.bfloat16 for i in [dy])
-                    assert all(i.is_contiguous() for i in [dy])
+
                     w, q, k, v, z, b, s, sa = ctx.saved_tensors
+                    dht = ops.cast(dht, "float32")
+                    dht = dht.contiguous()
+                    assert all(i.dtype == torch.bfloat16 for i in [dy])
+                    assert all(i.is_contiguous() for i in [dy, dht])
+                    dh0 = torch.empty(dht.shape, dtype=dht.dtype, device=dht.device)
                     dw, dq, dk, dv, dz, db = [
                         torch.empty_like(x) for x in [w, q, k, v, z, b]
                     ]
+
                     torch.ops.wind_backstepping.backward(
-                        w, q, k, v, z, b, dy, s, sa, dw, dq, dk, dv, dz, db
+                        w, q, k, v, z, b, dy, s, sa, dht, dh0, dw, dq, dk, dv, dz, db
                     )
                     return (
                         ops.cast(dw, DTYPE),
@@ -95,9 +106,10 @@ def get_generalized_delta_rule(HEAD_SIZE=64, KERNEL_TYPE="native"):
                         ops.cast(dv, DTYPE),
                         ops.cast(dz, DTYPE),
                         ops.cast(db, DTYPE),
+                        dh0,
                     )
 
-            def RUN_CUDA_RWKV7g(q, w, k, v, a, b):
+            def RUN_CUDA_RWKV7g(q, w, k, v, a, b, h0):
                 B, T, H, C = q.shape
                 q = q.contiguous()
                 w = w.contiguous()
@@ -105,7 +117,8 @@ def get_generalized_delta_rule(HEAD_SIZE=64, KERNEL_TYPE="native"):
                 v = v.contiguous()
                 a = a.contiguous()
                 b = b.contiguous()
-                return WindBackstepping.apply(w, q, k, v, a, b).view(B, T, H * C)
+                out, state = WindBackstepping.apply(w, q, k, v, a, b, h0)
+                return out, state
 
             def generalized_delta_rule(
                 r: torch.Tensor,
@@ -125,26 +138,61 @@ def get_generalized_delta_rule(HEAD_SIZE=64, KERNEL_TYPE="native"):
                 a = transpose_head(a, head_first)
                 b = transpose_head(b, head_first)
                 w = transpose_head(w, head_first)
-                return RUN_CUDA_RWKV7g(r, w, k, v, a, b), None
+                B, T, H, N = w.shape
+                if initial_state is None:
+                    initial_state = ops.zeros((B, H, N, N), "float32")
+                else:
+                    initial_state = ops.cast(initial_state, "float32")
+                out, state = RUN_CUDA_RWKV7g(r, w, k, v, a, b, initial_state)
+                if output_final_state:
+                    return out, state
+                return out
         else:
             from .native_keras_op import generalized_delta_rule
 
-            USE_KERNEL = False
+            USE_TRITON_KERNEL = False
     elif keras.config.backend() == "jax":
         from jax.lib import xla_bridge
         import os
 
-        if (
-            xla_bridge.get_backend().platform == "gpu"
-            and KERNEL_TYPE.lower() == "triton"
-        ):
-            os.environ["JAX_LOG_COMPUTATION"] = "0"
-            from .jax_op import generalized_delta_rule
+        if xla_bridge.get_backend().platform == "gpu":
+            if KERNEL_TYPE.lower() == "triton":
+                os.environ["JAX_LOG_COMPUTATION"] = "0"
+                from .jax_op import generalized_delta_rule
 
-            USE_KERNEL = True
+                USE_TRITON_KERNEL = True
+            elif KERNEL_TYPE.lower() == "cuda":
+                from .jax_cuda_kernel.wkv7_jax import get_jax_generalized_delta_rule
+
+                generalized_delta_rule = get_jax_generalized_delta_rule(HEAD_SIZE)[0]
+            else:
+                from .native_keras_op import generalized_delta_rule
+        else:
+            from .native_keras_op import generalized_delta_rule
+    elif keras.config.backend() == "tensorflow":
+        import tensorflow as tf
+
+        if len(tf.config.list_physical_devices("GPU")) > 0:
+            if KERNEL_TYPE.lower() == "cuda" and HEAD_SIZE == 64:
+                try:
+                    from jax.lib import xla_bridge
+
+                    assert xla_bridge.get_backend().platform == "gpu"
+                except:
+                    raise (
+                        "The operation of the TensorFlow kernel depends on the JAX kernel."
+                        "Therefore, it is necessary to ensure that it can be used in JAX, so that TensorFlow can be used."
+                    )
+                print("🎉" * 10)
+                print("Tensorflow CUDA kernel onlt support Forward,not get graident")
+                print("🎉" * 10)
+                from .tf_eager_kernel import get_tf_generalized_delta_rule
+
+                generalized_delta_rule = get_tf_generalized_delta_rule(HEAD_SIZE)[0]
+            else:
+                from .native_keras_op import generalized_delta_rule
         else:
             from .native_keras_op import generalized_delta_rule
-
     else:
         from .native_keras_op import generalized_delta_rule
-    return generalized_delta_rule, USE_KERNEL
+    return generalized_delta_rule, USE_TRITON_KERNEL

rwkv_ops/rwkv7_kernel/jax_cuda_kernel/CMakeLists.txt

@@ -0,0 +1,42 @@
+cmake_minimum_required(VERSION 3.18)
+project(wkv7 LANGUAGES CXX CUDA)
+
+find_package(CUDAToolkit REQUIRED)
+
+# ---------- 1. 找到 Python ----------
+find_package(Python3 REQUIRED COMPONENTS Interpreter)
+
+# ---------- 2. 取 XLA 头文件路径 ----------
+execute_process(
+  COMMAND "${Python3_EXECUTABLE}" -c "from jax import ffi; print(ffi.include_dir())"
+  OUTPUT_VARIABLE XLA_INCLUDE_DIR
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+)
+if(NOT XLA_INCLUDE_DIR)
+  message(FATAL_ERROR "Cannot get XLA include dir from jax.ffi")
+endif()
+message(STATUS "XLA include directory: ${XLA_INCLUDE_DIR}")
+
+# ---------- 3. 生成共享库 ----------
+add_library(wkv7 SHARED wkv7_ffi.cu)
+
+# 3-1. 头文件搜索路径
+target_include_directories(wkv7 PRIVATE ${XLA_INCLUDE_DIR})
+
+# 3-2. 链接 CUDA 运行时
+target_link_libraries(wkv7 PRIVATE CUDA::cudart)
+
+# 3-3. 关键：C++17 / CUDA17 标准
+target_compile_features(wkv7 PUBLIC cxx_std_17)
+set_target_properties(wkv7 PROPERTIES
+    CUDA_STANDARD          17
+    CUDA_SEPARABLE_COMPILATION ON
+    POSITION_INDEPENDENT_CODE ON
+    PREFIX                 ""        # 去掉默认的 "lib" 前缀
+)
+
+# ---------- 4. 安装 ----------
+# 把 .so 直接装到源码目录（与 wkv7_jax.py 同一级），方便 ctypes.CDLL 加载
+install(TARGETS wkv7
+        LIBRARY DESTINATION "${CMAKE_SOURCE_DIR}"
+        RUNTIME DESTINATION "${CMAKE_SOURCE_DIR}")   # Windows 用 RUNTIME

rwkv_ops/rwkv7_kernel/jax_cuda_kernel/wkv7_ffi.cu

@@ -0,0 +1,279 @@
+/*
+ *  wkv7_ffi_bf16.cu
+ *  BF16 版本，外部接口 BF16，内部 kernel 保持原样
+ */
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <xla/ffi/api/ffi.h>
+#include <vector>
+
+namespace ffi = xla::ffi;
+
+/* -------------------- 类型别名 -------------------- */
+using bf = __nv_bfloat16;
+
+/* -------------------- 设备端辅助（kernel 里用） -------------------- */
+__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_;
+
+/* -------------------- 你的 kernel（禁止修改） -------------------- */
+__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.f;
+#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.f;
+#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 *s_, float *sa_, float *dht_, float *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.f;
+#pragma unroll
+        for (int j = 0; j < C; ++j) dq += stateT[j] * dy[j];
+        dq_[ind] = to_bf(dq);
+
+        float iwi = 1.f / (wi + 1e-6f);
+#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.f, dk = 0.f, dv = 0.f, db = 0.f, dSb = 0.f;
+#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.f;
+#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];
+        }
+    }
+}
+
+/* -------------------- 宿主函数 -------------------- */
+static ffi::Error WKV7FwdHost(
+    cudaStream_t stream,
+    ffi::Buffer<ffi::BF16> w,
+    ffi::Buffer<ffi::BF16> q,
+    ffi::Buffer<ffi::BF16> k,
+    ffi::Buffer<ffi::BF16> v,
+    ffi::Buffer<ffi::BF16> z,
+    ffi::Buffer<ffi::BF16> a,
+    ffi::Buffer<ffi::F32>  h0,   // 保持 float
+    ffi::ResultBuffer<ffi::BF16> y,
+    ffi::ResultBuffer<ffi::F32>  s,
+    ffi::ResultBuffer<ffi::F32>  sa)
+{
+    constexpr int C = _C_;
+    auto dims = w.dimensions();
+    int B = dims[0], T = dims[1], H = dims[2];
+    dim3 block(C);
+    dim3 grid(H, B);
+
+    forward_kernel<<<grid, block, 0, stream>>>(
+        T, H,
+        reinterpret_cast<bf *>(w.typed_data()),
+        reinterpret_cast<bf *>(q.typed_data()),
+        reinterpret_cast<bf *>(k.typed_data()),
+        reinterpret_cast<bf *>(v.typed_data()),
+        reinterpret_cast<bf *>(z.typed_data()),
+        reinterpret_cast<bf *>(a.typed_data()),
+        reinterpret_cast<bf *>(y->typed_data()),
+        s->typed_data(),
+        sa->typed_data(),
+        h0.typed_data());
+
+    cudaError_t err = cudaGetLastError();
+    if (err != cudaSuccess)
+        return ffi::Error::Internal(
+            std::string("CUDA forward_kernel error: ") + cudaGetErrorString(err));
+    return ffi::Error::Success();
+}
+
+static ffi::Error WKV7BwdHost(
+    cudaStream_t stream,
+    ffi::Buffer<ffi::BF16> w,
+    ffi::Buffer<ffi::BF16> q,
+    ffi::Buffer<ffi::BF16> k,
+    ffi::Buffer<ffi::BF16> v,
+    ffi::Buffer<ffi::BF16> z,
+    ffi::Buffer<ffi::BF16> a,
+    ffi::Buffer<ffi::BF16> dy,
+    ffi::Buffer<ffi::F32>  s,
+    ffi::Buffer<ffi::F32>  sa,
+    ffi::Buffer<ffi::F32>  dht,
+    ffi::ResultBuffer<ffi::F32> dh0,
+    ffi::ResultBuffer<ffi::BF16> dw,
+    ffi::ResultBuffer<ffi::BF16> dq,
+    ffi::ResultBuffer<ffi::BF16> dk,
+    ffi::ResultBuffer<ffi::BF16> dv,
+    ffi::ResultBuffer<ffi::BF16> da,
+    ffi::ResultBuffer<ffi::BF16> db)
+{
+    auto dims = w.dimensions();
+    int B = dims[0], T = dims[1], H = dims[2];
+    constexpr int C = _C_;
+    dim3 block(C);
+    dim3 grid(H, B);
+
+    backward_kernel<<<grid, block, 0, stream>>>(
+        T, H,
+        reinterpret_cast<bf *>(w.typed_data()),
+        reinterpret_cast<bf *>(q.typed_data()),
+        reinterpret_cast<bf *>(k.typed_data()),
+        reinterpret_cast<bf *>(v.typed_data()),
+        reinterpret_cast<bf *>(z.typed_data()),
+        reinterpret_cast<bf *>(a.typed_data()),
+        reinterpret_cast<bf *>(dy.typed_data()),
+        s.typed_data(),
+        sa.typed_data(),
+        dht.typed_data(),
+        dh0->typed_data(),
+        reinterpret_cast<bf *>(dw->typed_data()),
+        reinterpret_cast<bf *>(dq->typed_data()),
+        reinterpret_cast<bf *>(dk->typed_data()),
+        reinterpret_cast<bf *>(dv->typed_data()),
+        reinterpret_cast<bf *>(da->typed_data()),
+        reinterpret_cast<bf *>(db->typed_data()));
+
+    cudaError_t err = cudaGetLastError();
+    if (err != cudaSuccess)
+        return ffi::Error::Internal(
+            std::string("CUDA backward_kernel error: ") + cudaGetErrorString(err));
+    return ffi::Error::Success();
+}
+
+/* -------------------- 注册符号 -------------------- */
+XLA_FFI_DEFINE_HANDLER_SYMBOL(
+    Wkv7Fwd, WKV7FwdHost,
+    ffi::Ffi::Bind()
+        .Ctx<ffi::PlatformStream<cudaStream_t>>()
+        .Arg<ffi::Buffer<ffi::BF16>>()   // w
+        .Arg<ffi::Buffer<ffi::BF16>>()   // q
+        .Arg<ffi::Buffer<ffi::BF16>>()   // k
+        .Arg<ffi::Buffer<ffi::BF16>>()   // v
+        .Arg<ffi::Buffer<ffi::BF16>>()   // z
+        .Arg<ffi::Buffer<ffi::BF16>>()   // a
+        .Arg<ffi::Buffer<ffi::F32>>()    // h0  (float)
+        .Ret<ffi::Buffer<ffi::BF16>>()   // y
+        .Ret<ffi::Buffer<ffi::F32>>()    // s
+        .Ret<ffi::Buffer<ffi::F32>>()    // sa
+, {ffi::Traits::kCmdBufferCompatible});
+
+XLA_FFI_DEFINE_HANDLER_SYMBOL(
+    Wkv7Bwd, WKV7BwdHost,
+    ffi::Ffi::Bind()
+        .Ctx<ffi::PlatformStream<cudaStream_t>>()
+        .Arg<ffi::Buffer<ffi::BF16>>()   // w
+        .Arg<ffi::Buffer<ffi::BF16>>()   // q
+        .Arg<ffi::Buffer<ffi::BF16>>()   // k
+        .Arg<ffi::Buffer<ffi::BF16>>()   // v
+        .Arg<ffi::Buffer<ffi::BF16>>()   // z
+        .Arg<ffi::Buffer<ffi::BF16>>()   // a
+        .Arg<ffi::Buffer<ffi::BF16>>()   // dy
+        .Arg<ffi::Buffer<ffi::F32>>()    // s
+        .Arg<ffi::Buffer<ffi::F32>>()    // sa
+        .Arg<ffi::Buffer<ffi::F32>>()    // dht
+        .Ret<ffi::Buffer<ffi::F32>>()   // dh0
+        .Ret<ffi::Buffer<ffi::BF16>>()   // dw
+        .Ret<ffi::Buffer<ffi::BF16>>()   // dq
+        .Ret<ffi::Buffer<ffi::BF16>>()   // dk
+        .Ret<ffi::Buffer<ffi::BF16>>()   // dv
+        .Ret<ffi::Buffer<ffi::BF16>>()   // da
+        .Ret<ffi::Buffer<ffi::BF16>>()   // db
+, {ffi::Traits::kCmdBufferCompatible});