rwkv-ops 0.6.1__py3-none-any.whl

rwkv_ops/rwkv7_kernel/jax_cuda_kernel/wkv7_jax.py

@@ -0,0 +1,311 @@
+"""
+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"
+    )
+    jax.ffi.register_ffi_target(
+        "wkv7_inference", jax.ffi.pycapsule(_lib.Wkv7Inference), 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
+
+    def _wkv7_inference_kernel(
+        w: jnp.ndarray,
+        q: jnp.ndarray,
+        k: jnp.ndarray,
+        v: jnp.ndarray,
+        a: jnp.ndarray,
+        b: jnp.ndarray,
+        h0: jnp.ndarray,
+    ):
+        """
+        推理专用 kernel，不保存 sa 和中间 s
+        返回: y (B, T, H, K), final_state (B, H, K, K)
+        """
+        B, T, H, K = q.shape
+        dtype = q.dtype
+        out_type = jax.ShapeDtypeStruct((B, T, H, K), dtype)
+        # **关键：仅返回最终状态，非 chunk 历史**
+        s_type = jax.ShapeDtypeStruct((B, H, K, K), jnp.float32)
+
+        y, s = jax.ffi.ffi_call(
+            "wkv7_inference", (out_type, s_type), vmap_method="broadcast_all"
+        )(w, q, k, v, a, b, h0)  # z 参数自动忽略
+
+        return y, s
+
+    # -------------------- 公共推理 API --------------------
+    def generalized_delta_rule_inference(
+        r: jnp.ndarray,
+        w: jnp.ndarray,
+        k: jnp.ndarray,
+        v: jnp.ndarray,
+        a: jnp.ndarray,
+        b: jnp.ndarray,
+        output_final_state: bool = True,
+        initial_state: Optional[jnp.ndarray] = None,
+        head_first: bool = False,
+    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
+        """
+        纯推理版本的广义 delta 规则
+
+        参数:
+            r,w,k,v,a,b: 输入张量，形状 (B, T, H, K) 或 (B, H, T, K)
+            initial_state: (B, H, K, K) 初始状态，None 则零初始化
+            head_first: 是否将 head 维提前
+        返回:
+            out: (B, T, H, K) 输出，dtype 与输入一致
+            final_state: (B, H, K, 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 initial_state is None:
+            h0 = jnp.zeros((B, H, K, K), jnp.float32)
+        else:
+            h0 = jnp.asarray(initial_state, jnp.float32)
+
+        # **无需 checkpoint，推理不保存中间值**
+        out, final_state = _wkv7_inference_kernel(w, r, k, v, a, b, h0)
+        out = jnp.asarray(out, dtype)
+        return out, final_state if output_final_state else out
+
+    # 返回两个函数，用户按需选择
+    return [generalized_delta_rule, generalized_delta_rule_inference]

rwkv_ops/rwkv7_kernel/jax_cuda_kernel_single/CMakeLists.txt

@@ -0,0 +1,42 @@
+cmake_minimum_required(VERSION 3.18)
+project(wkv7_single_step LANGUAGES CXX CUDA)
+
+find_package(CUDAToolkit REQUIRED)
+
+# ---------- 1. 找到 Python ----------
+find_package(Python3 REQUIRED COMPONENTS Interpreter)
+
+# ---------- 2. 取 XLA 头文件路径 ----------
+execute_process(
+  COMMAND "${Python3_EXECUTABLE}" -c "import jax; print(jax.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_single_step SHARED wkv7_single_step_ffi.cu)
+
+# 3-1. 头文件搜索路径
+target_include_directories(wkv7_single_step PRIVATE ${XLA_INCLUDE_DIR})
+
+# 3-2. 链接 CUDA 运行时
+target_link_libraries(wkv7_single_step PRIVATE CUDA::cudart)
+
+# 3-3. 关键：C++17 / CUDA17 标准
+target_compile_features(wkv7_single_step PUBLIC cxx_std_17)
+set_target_properties(wkv7_single_step PROPERTIES
+    CUDA_STANDARD          17
+    CUDA_SEPARABLE_COMPILATION ON
+    POSITION_INDEPENDENT_CODE ON
+    PREFIX                 ""        # 去掉默认的 "lib" 前缀
+)
+
+# ---------- 4. 安装 ----------
+# 把 .so 直接装到源码目录，方便 ctypes.CDLL 加载
+install(TARGETS wkv7_single_step
+        LIBRARY DESTINATION "${CMAKE_SOURCE_DIR}"
+        RUNTIME DESTINATION "${CMAKE_SOURCE_DIR}")   # Windows 用 RUNTIME

rwkv_ops/rwkv7_kernel/jax_cuda_kernel_single/wkv7_single_step_ffi.cu

@@ -0,0 +1,172 @@
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <xla/ffi/api/ffi.h>
+#include <vector>
+#include <cstdint>
+
+namespace ffi = xla::ffi;
+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_;
+
+/* -------------------- 前向 Kernel（修复） -------------------- */
+template<int C> 
+__launch_bounds__(C, 2)
+__global__ void forward_kernel_single_step(
+    int B, int H,
+    F_ w_, F_ q_, F_ k_, F_ v_, F_ a_, F_ b_,
+    bf *y_, float *s_, float *h0_)
+{
+    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];
+    
+    // 加载初始状态 (B, H, C, C)
+    int64_t h0_base = ((int64_t)bb * H + hh) * C * C + i * C;
+#pragma unroll
+    for (int j = 0; j < C; ++j) state[j] = h0_[h0_base + j];
+
+    // 单步索引: (B, H, C)
+    int64_t ind = (int64_t)bb * 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];
+
+    float v_val = 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_val;
+        y += s * q[j];
+    }
+    y_[ind] = to_bf(y);
+
+    // 写入最终状态
+    int64_t s_base = ((int64_t)bb * H + hh) * C * C + i * C;
+#pragma unroll
+    for (int j = 0; j < C; ++j) s_[s_base + j] = state[j];
+}
+
+/* -------------------- 反向 Kernel（补充） -------------------- */
+template<int C> 
+__launch_bounds__(C, 2) 
+__global__ void backward_kernel_single_step(
+    int B, int H,
+    F_ w_, F_ q_, F_ k_, F_ v_, F_ dy_,
+    float *s_, float *dht_, bf *dw_, bf *dq_, bf *dk_, bf *dv_, bf *da_, bf *db_)
+{
+    int bb = blockIdx.y, hh = blockIdx.x, i = threadIdx.x;
+    float stateT[C] = {0}, dstate[C] = {0};
+    
+    int64_t dht_base = ((int64_t)bb * H + hh) * C * C + i * C;
+#pragma unroll
+    for (int j = 0; j < C; ++j) dstate[j] = dht_[dht_base + j];
+
+    __shared__ float w[C], q[C], k[C], v[C], dy[C];
+    int64_t ind = (int64_t)bb * H * C + hh * C + i;
+    
+    __syncthreads();
+    q[i] = to_float(q_[ind]);
+    float wi_fac = -__expf(to_float(w_[ind]));
+    w[i] = __expf(wi_fac);
+    k[i] = to_float(k_[ind]);
+    v[i] = to_float(v_[ind]);
+    dy[i] = to_float(dy_[ind]);
+    __syncthreads();
+
+    // 从 s_ 加载 stateT（float4 优化可在此处添加）
+    int64_t s_base = ((int64_t)bb * H + hh) * C * C + i * C;
+#pragma unroll
+    for (int j = 0; j < C; ++j) stateT[j] = s_[s_base + j];
+
+    float dq_val = 0.f, dw_val = 0.f, dk_val = 0.f, dv_val = 0.f, da_val = 0.f, db_val = 0.f;
+    float iwi = 1.0f / (w[i] + 1e-6f);
+    
+#pragma unroll
+    for (int j = 0; j < C; ++j) {
+        stateT[j] = (stateT[j] - k[i] * v[j]) * iwi;
+        dstate[j] += dy[i] * q[j];
+        
+        dq_val += stateT[j] * dy[j];
+        dw_val += dstate[j] * stateT[j];
+        dk_val += dstate[j] * v[j];
+        dv_val += dstate[j] * k[j];
+    }
+    
+    dq_[ind] = to_bf(dq_val);
+    dw_[ind] = to_bf(dw_val * w[i] * wi_fac);
+    dk_[ind] = to_bf(dk_val);
+    dv_[ind] = to_bf(dv_val);
+}
+
+/* -------------------- Host 函数（修复调用） -------------------- */
+static ffi::Error WKV7SingleStepFwdHost(
+    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> a,
+    ffi::Buffer<ffi::BF16> b,
+    ffi::Buffer<ffi::F32>  h0,
+    ffi::ResultBuffer<ffi::BF16> y,
+    ffi::ResultBuffer<ffi::F32>  s)
+{
+    auto dims = w.dimensions();
+    int B = dims[0], H = dims[1];
+    constexpr int C = _C_;  // 从编译选项获取
+    dim3 block(C);
+    dim3 grid(H, B);
+
+    // ✅ 修复：显式指定模板参数 <_C_>
+    forward_kernel_single_step<_C_><<<grid, block, 0, stream>>>(
+        B, 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 *>(a.typed_data()),
+        reinterpret_cast<bf *>(b.typed_data()),
+        reinterpret_cast<bf *>(y->typed_data()),
+        s->typed_data(),
+        h0.typed_data());
+
+    cudaError_t err = cudaGetLastError();
+    if (err != cudaSuccess)
+        return ffi::Error::Internal(
+            std::string("CUDA forward_kernel_single_step error: ") + cudaGetErrorString(err));
+    return ffi::Error::Success();
+}
+
+/* -------------------- FFI 符号注册 -------------------- */
+XLA_FFI_DEFINE_HANDLER_SYMBOL(
+    Wkv7SingleStepFwd, WKV7SingleStepFwdHost,
+    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>>()   // a
+        .Arg<ffi::Buffer<ffi::BF16>>()   // b
+        .Arg<ffi::Buffer<ffi::F32>>()    // h0
+        .Ret<ffi::Buffer<ffi::BF16>>()   // y
+        .Ret<ffi::Buffer<ffi::F32>>()    // s
+, {ffi::Traits::kCmdBufferCompatible});