rwkv-ops 0.6.1__py3-none-any.whl

rwkv_ops/__init__.py

@@ -0,0 +1,45 @@
+__version__ = "0.6.0"
+import os
+
+KERNEL_TYPE = os.environ.get("KERNEL_TYPE", "cuda").lower()
+KERAS_BACKEND = os.environ.get("KERAS_BACKEND")
+BACKEND = os.environ.get("KERNEL_BACKEND")
+
+
+if KERAS_BACKEND is not None:
+    BACKEND = KERAS_BACKEND.lower()
+elif BACKEND is not None:
+    os.environ["KERAS_BACKEND"] = BACKEND.lower()
+else:
+    import keras
+
+    BACKEND = "torch"
+    os.environ["KERAS_BACKEND"] = BACKEND
+    keras.config.set_backend("torch")
+assert KERNEL_TYPE in ["triton", "cuda", "native"]
+assert BACKEND in ["torch", "jax", "numpy", "tensorflow"]
+from .rwkv7_kernel import get_generalized_delta_rule, get_rnn_generalized_delta_rule
+from .rwkv6_kernel import get_rwkv6_kernel
+from .mhc_kernel import get_mhu_kernel
+
+(
+    mhc_sinkhorn_knopp,
+    mhc_rmsnorm,
+    mhc_stream_aggregate,
+    mhc_stream_distribute,
+    mhc_stream_mix,
+    mhc_post_op,
+    mhc_pre_op,
+) = get_mhu_kernel(KERNEL_TYPE)
+
+generalized_delta_rule, generalized_delta_rule_inference, RWKV7_USE_TRITON_KERNEL = (
+    get_generalized_delta_rule(KERNEL_TYPE=KERNEL_TYPE)
+)
+rwkv7_op = generalized_delta_rule
+rwkv7_op_inference = generalized_delta_rule_inference
+
+rnn_generalized_delta_rule = get_rnn_generalized_delta_rule(KERNEL_TYPE=KERNEL_TYPE)
+rwkv7_op_rnn = rnn_generalized_delta_rule
+
+
+RWKV6_OP = get_rwkv6_kernel(KERNEL_TYPE=KERNEL_TYPE)

rwkv_ops/mhc_kernel/__init__.py

@@ -0,0 +1,50 @@
+import keras
+
+
+def get_mhu_kernel(KERNEL_TYPE):
+    from .native_keras_op import (
+        sinkhorn_knopp,
+        rmsnorm,
+        stream_aggregate,
+        stream_distribute,
+        stream_mix,
+        mhc_post_op,
+        mhc_pre_op,
+    )
+
+    if KERNEL_TYPE == "cuda":
+        if keras.config.backend() == "torch":
+            import torch
+
+            if torch.cuda.is_available():
+                from .torch_kernel.mhc_torch import (
+                    sinkhorn_knopp,
+                    rmsnorm,
+                    stream_aggregate,
+                    stream_distribute,
+                    stream_mix,
+                    mhc_post_op,
+                    mhc_pre_op,
+                )
+        elif keras.config.backend() == "jax":
+            import jax
+
+            if jax.devices()[0].platform == "gpu":
+                from .jax_kernel.mhu_jax import (
+                    sinkhorn_knopp,
+                    rmsnorm,
+                    stream_aggregate,
+                    stream_distribute,
+                    stream_mix,
+                    mhc_post_op,
+                    mhc_pre_op,
+                )
+    return (
+        sinkhorn_knopp,
+        rmsnorm,
+        stream_aggregate,
+        stream_distribute,
+        stream_mix,
+        mhc_post_op,
+        mhc_pre_op,
+    )

rwkv_ops/mhc_kernel/common_kernel/include/mhc_types.h

@@ -0,0 +1,66 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cublasLt.h>
+#include <assert.h>
+#include <cstdint>
+
+namespace mhc {
+
+using floatX = __nv_bfloat16;
+using floatN = float;
+
+// 定义统一的转换工具，供所有 .cuh 和 .cu 使用
+__device__ inline float to_float(const floatX& u) {
+    return __bfloat162float(u);
+}
+
+__device__ inline floatX to_bf(const float& u) {
+    #if __CUDA_ARCH__ >= 800
+        return __float2bfloat16(u);
+    #else
+        // 兼容旧架构或强制舍入
+        return __float2bfloat16_rn(u);
+    #endif
+}
+
+
+struct MHCConfig {
+    int sinkhorn_iters;
+    int nC;
+    float eps;
+    bool use_pdl;
+};
+
+struct RMSNormParams {
+    int n;
+    float eps;
+};
+
+inline void check_cuda(cudaError_t err, const char* file, int line) {
+    if (err != cudaSuccess) {
+        fprintf(stderr, "CUDA error at %s:%d: %s\n", file, line, cudaGetErrorString(err));
+        exit(EXIT_FAILURE);
+    }
+}
+
+inline void check_cublas(cublasStatus_t status, const char* file, int line) {
+    if (status != CUBLAS_STATUS_SUCCESS) {
+        fprintf(stderr, "cuBLAS error at %s:%d: %d\n", file, line, (int)status);
+        exit(EXIT_FAILURE);
+    }
+}
+// 错误检查宏
+#define CHECK_CUDA(call) \
+do { \
+    cudaError_t err = call; \
+    if (err != cudaSuccess) { \
+        printf("CUDA Error at %s:%d - %s\n", __FILE__, __LINE__, cudaGetErrorString(err)); \
+    } \
+} while (0)
+
+#define CHECK_CUDA(call) mhc::check_cuda((call), __FILE__, __LINE__)
+#define CHECK_CUBLAS(call) mhc::check_cublas((call), __FILE__, __LINE__)
+} // namespace mhc

rwkv_ops/mhc_kernel/common_kernel/kernels/mhc_post_op.cuh

@@ -0,0 +1,197 @@
+#ifndef MHC_POST_OP_CUH
+#define MHC_POST_OP_CUH
+
+#include <cuda_runtime.h>
+#include <cuda_bf16.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+#include "../include/mhc_types.h"
+#include "type_conversions.cuh"
+
+namespace cg = cooperative_groups;
+
+namespace mhc {
+
+/**
+ * 1. Fused Forward Kernel
+ * 公式: x_next[b,t,i,c] = sum_j(H_res[b,t,i,j] * x_expanded[b,t,j,c]) + layer_out[b,t,c] * H_post[b,t,i]
+ * 精度策略：所有中间累加使用 FP32
+ */
+template<int MAX_N = 8>
+__global__ void mhc_post_op_fwd_kernel(
+    floatX* __restrict__ x_next,           // [B, T, n, C]
+    const floatX* __restrict__ layer_out,  // [B, T, C]
+    const floatX* __restrict__ x_expanded, // [B, T, n, C]
+    const float* __restrict__ H_post,      // [B, T, n]
+    const float* __restrict__ H_res,       // [B, T, n, n]
+    int64_t B, int64_t T, int n, int64_t C) 
+{
+    int64_t btc = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    int i = blockIdx.y; // 输出流索引
+
+    if (btc < B * T * C && i < n) {
+        int64_t bt = btc / C;
+        int64_t c = btc % C;
+        // [修改]: 增加 (int64_t)n 强转
+        int64_t bt_offset_n_c = bt * (int64_t)n * C;
+
+        // --- Stream Mix 部分 ---
+        float mixed_val = 0.0f;
+        // [修改]: 增加 (int64_t)n 和 (int64_t)i 强转，防止 i*n 在 32 位下溢出
+        int64_t res_base = bt * (int64_t)n * n + (int64_t)i * n;
+        #pragma unroll
+        for (int j = 0; j < MAX_N; j++) {
+            if (j < n) {
+                float w_res = H_res[res_base + j];
+                float val_x = to_float(x_expanded[bt_offset_n_c + (int64_t)j * C + c]);
+                mixed_val += w_res * val_x;
+            }
+        }
+
+        // --- Stream Distribute 部分 ---
+        float l_val = to_float(layer_out[btc]);
+        // [修改]: 增加 (int64_t)n 强转
+        float w_post = H_post[bt * (int64_t)n + i];
+        float dist_val = l_val * w_post;
+
+        x_next[bt_offset_n_c + (int64_t)i * C + c] = to_bf(mixed_val + dist_val);
+    }
+}
+
+/**
+ * 2. Fused Backward Full Kernel
+ * 计算: 
+ * dl = sum_i(grad_next_i * H_post_i)
+ * dx_j = sum_i(grad_next_i * H_res_ij)
+ * dH_post_i = sum_c(grad_next_i * layer_out)
+ * dH_res_ij = sum_c(grad_next_i * x_expanded_j)
+ */
+template<int BLOCK_SIZE, int MAX_N = 8>
+__global__ void mhc_post_op_bwd_full_kernel(
+    floatX* __restrict__ d_layer_out,      // [B, T, C]
+    floatX* __restrict__ d_x_expanded,     // [B, T, n, C]
+    float* __restrict__ d_H_post,          // [B, T, n]
+    float* __restrict__ d_H_res,           // [B, T, n, n]
+    const floatX* __restrict__ grad_next,  // [B, T, n, C]
+    const floatX* __restrict__ layer_out,  // [B, T, C]
+    const floatX* __restrict__ x_expanded, // [B, T, n, C]
+    const float* __restrict__ H_post,
+    const float* __restrict__ H_res,
+    int64_t B, int64_t T, int n, int64_t C) 
+{
+    cg::thread_block block = cg::this_thread_block();
+    cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
+
+    int64_t bt = blockIdx.x;
+    if (bt >= B * T) return;
+
+    // [修改]: 增加 (int64_t)n 强转
+    int64_t bt_offset_n_c = bt * (int64_t)n * C;
+    int64_t bt_offset_c = bt * C;
+
+    // 局部存储用于参数梯度的 Reduce
+    // dw_post: [n], dw_res: [n, n]
+    float thread_dw_post[MAX_N] = {0.0f};
+    float thread_dw_res[MAX_N][MAX_N] = {0.0f};
+
+    // --- 1. 计算数据梯度 (dl, dx) ---
+    // 每个线程处理一个通道 c
+    for (int64_t c = threadIdx.x; c < C; c += BLOCK_SIZE) {
+        float l_val = to_float(layer_out[bt_offset_c + c]);
+        
+        // 先读取所有流在该通道的梯度
+        float g_vals[MAX_N];
+        #pragma unroll
+        for(int i=0; i<MAX_N; i++) {
+            if(i < n) g_vals[i] = to_float(grad_next[bt_offset_n_c + (int64_t)i * C + c]);
+        }
+
+        // 计算 dl (数据梯度)
+        float dl_sum = 0.0f;
+        #pragma unroll
+        for(int i=0; i<MAX_N; i++) {
+            // [修改]: 增加 (int64_t)n 强转
+            if(i < n) dl_sum += g_vals[i] * H_post[bt * (int64_t)n + i];
+        }
+        d_layer_out[bt_offset_c + c] = to_bf(dl_sum);
+
+        // 计算 dx (数据梯度) 和 累加参数梯度局部和
+        #pragma unroll
+        for(int j=0; j<MAX_N; j++) {
+            if (j < n) {
+                float dx_j = 0.0f;
+                float xj_val = to_float(x_expanded[bt_offset_n_c + (int64_t)j * C + c]);
+                
+                #pragma unroll
+                for(int i=0; i<MAX_N; i++) {
+                    if (i < n) {
+                        // [修改]: 增加 (int64_t)n 和 (int64_t)i 强转
+                        dx_j += g_vals[i] * H_res[bt * (int64_t)n * n + (int64_t)i * n + j];
+                        // 顺便计算 dH_res 的线程局部部分
+                        thread_dw_res[i][j] += g_vals[i] * xj_val;
+                    }
+                }
+                d_x_expanded[bt_offset_n_c + (int64_t)j * C + c] = to_bf(dx_j);
+            }
+        }
+
+        // 计算 dH_post 的线程局部部分
+        #pragma unroll
+        for(int i=0; i<MAX_N; i++) {
+            if(i < n) thread_dw_post[i] += g_vals[i] * l_val;
+        }
+    }
+
+    // --- 2. 参数梯度规约 (C 维度的 Reduction) ---
+    // 使用 Warp Shuffle 规约并写回
+    #pragma unroll
+    for(int i=0; i<MAX_N; i++) {
+        if(i < n) {
+            float sum_p = cg::reduce(warp, thread_dw_post[i], cg::plus<float>());
+            // [修改]: 增加 (int64_t)n 强转
+            if (warp.thread_rank() == 0) atomicAdd(&d_H_post[bt * (int64_t)n + i], sum_p);
+            
+            #pragma unroll
+            for(int j=0; j<MAX_N; j++) {
+                if(j < n) {
+                    float sum_r = cg::reduce(warp, thread_dw_res[i][j], cg::plus<float>());
+                    // [修改]: 增加 (int64_t)n 和 (int64_t)i 强转
+                    if (warp.thread_rank() == 0) atomicAdd(&d_H_res[bt * (int64_t)n * n + (int64_t)i * n + j], sum_r);
+                }
+            }
+        }
+    }
+}
+
+/* -------------------- API 包装函数 -------------------- */
+
+inline void mhc_post_op_forward(
+    floatX* x_next, const floatX* layer_out, const floatX* x_expanded,
+    const float* H_post, const float* H_res,
+    int64_t B, int64_t T, int n, int64_t C, cudaStream_t stream) 
+{
+    int64_t total_btc = B * T * C;
+    dim3 threads(256);
+    dim3 blocks((total_btc + 255) / 256, n);
+    mhc_post_op_fwd_kernel<8><<<blocks, threads, 0, stream>>>(
+        x_next, layer_out, x_expanded, H_post, H_res, B, T, n, C);
+}
+
+inline void mhc_post_op_backward_full(
+    floatX* d_layer_out, floatX* d_x_expanded, float* d_H_post, float* d_H_res,
+    const floatX* grad_next, const floatX* layer_out, const floatX* x_expanded,
+    const float* H_post, const float* H_res,
+    int64_t B, int64_t T, int n, int64_t C, cudaStream_t stream) 
+{
+    // 每个 Block 负责一个 Token (bt) 的所有通道规约
+    const int BLOCK_SIZE = 256;
+    dim3 threads(BLOCK_SIZE);
+    dim3 blocks(B * T);
+    mhc_post_op_bwd_full_kernel<BLOCK_SIZE, 8><<<blocks, threads, 0, stream>>>(
+        d_layer_out, d_x_expanded, d_H_post, d_H_res,
+        grad_next, layer_out, x_expanded, H_post, H_res, B, T, n, C);
+}
+
+} // namespace mhc
+
+#endif

rwkv_ops/mhc_kernel/common_kernel/kernels/mhc_pre_op.cuh

@@ -0,0 +1,212 @@
+#ifndef MHC_PRE_OP_CUH
+#define MHC_PRE_OP_CUH
+
+#include <cuda_runtime.h>
+#include <cuda_bf16.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+#include "../include/mhc_types.h"
+#include "type_conversions.cuh"
+#include "sinkhorn_knopp.cuh"
+
+namespace cg = cooperative_groups;
+
+namespace mhc {
+
+/**
+ * 1. Fused Pre-Op Forward Kernel
+ * 修复：确保所有索引步进均使用 int64_t 避免在大模型/长序列下溢出
+ */
+template<int MAX_N = 8>
+__global__ void mhc_pre_op_fwd_kernel(
+    floatX* __restrict__ x_layer_in,
+    float* __restrict__ H_pre_out,
+    float* __restrict__ H_post_out,
+    const floatX* __restrict__ x_expanded,
+    const float* __restrict__ h_pre_raw,
+    const float* __restrict__ h_post_raw,
+    int64_t B, int64_t T, int n, int64_t C) 
+{
+    int64_t btc = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    if (btc >= B * T * C) return;
+
+    int64_t bt = btc / C;
+    int64_t c = btc % C;
+    int64_t bt_offset_n = bt * (int64_t)n;
+
+    float H_pre[MAX_N];
+    #pragma unroll
+    for (int i = 0; i < MAX_N; i++) {
+        if (i < n) {
+            float val_pre = h_pre_raw[bt_offset_n + (int64_t)i];
+            H_pre[i] = 1.0f / (1.0f + __expf(-val_pre));
+            
+            if (c == 0) {
+                H_pre_out[bt_offset_n + (int64_t)i] = H_pre[i];
+                float val_post = h_post_raw[bt_offset_n + (int64_t)i];
+                // 2.0 * sigmoid 逻辑保持不变
+                H_post_out[bt_offset_n + (int64_t)i] = 2.0f * (1.0f / (1.0f + __expf(-val_post)));
+            }
+        }
+    }
+
+    float sum_val = 0.0f;
+    int64_t bt_offset_n_c = bt * (int64_t)n * C;
+    #pragma unroll
+    for (int i = 0; i < MAX_N; i++) {
+        if (i < n) {
+            float xi = to_float(x_expanded[bt_offset_n_c + (int64_t)i * C + c]);
+            sum_val += H_pre[i] * xi;
+        }
+    }
+
+    x_layer_in[btc] = to_bf(sum_val);
+}
+
+/**
+ * 2. Fused Pre-Op Backward Kernel
+ * 修复重点：
+ * 1. 强化规约逻辑：使用 block 级规约确保 sum_grad_x 的准确性。
+ * 2. 检查发现之前的 atomicAdd 虽然逻辑正确，但若输入 tensor 未在 python/cpp 层清零会导致 Fail。
+ * 3. 这里的 MAX_N 限制了能够并行处理的流数量。
+ */
+template<int BLOCK_SIZE, int MAX_N = 8>
+__global__ void mhc_pre_op_bwd_kernel(
+    floatX* __restrict__ d_x_expanded,
+    float* __restrict__ d_h_pre_raw,
+    float* __restrict__ d_h_post_raw,
+    const floatX* __restrict__ grad_layer_in,
+    const float* __restrict__ grad_H_post,
+    const floatX* __restrict__ x_expanded,
+    const float* __restrict__ H_pre,
+    const float* __restrict__ H_post,
+    int64_t B, int64_t T, int n, int64_t C) 
+{
+    // 定义共享内存用于 Block 级规约 (大小为 BLOCK_SIZE * n)
+    // 假设 MAX_N 很小 (如 8)，256 * 8 * 4 bytes = 8KB，远小于显卡限制
+    __shared__ float s_reduce[BLOCK_SIZE][MAX_N];
+
+    int64_t bt = blockIdx.x; 
+    if (bt >= B * T) return;
+
+    int tid = threadIdx.x;
+    int64_t bt_offset_n = bt * (int64_t)n;
+    int64_t bt_offset_c = bt * C;
+    int64_t bt_offset_n_c = bt * (int64_t)n * C;
+
+    // 初始化局部累加器
+    float thread_dh_pre_sum[MAX_N];
+    #pragma unroll
+    for(int i=0; i<MAX_N; ++i) thread_dh_pre_sum[i] = 0.0f;
+
+    // 1. 计算 dx 并收集局部和
+    for (int64_t c = (int64_t)tid; c < C; c += (int64_t)BLOCK_SIZE) {
+        float g_in = to_float(grad_layer_in[bt_offset_c + c]);
+
+        #pragma unroll
+        for (int i = 0; i < MAX_N; i++) {
+            if (i < n) {
+                float h_pre_i = H_pre[bt_offset_n + (int64_t)i];
+                d_x_expanded[bt_offset_n_c + (int64_t)i * C + c] = to_bf(g_in * h_pre_i);
+                
+                float xi = to_float(x_expanded[bt_offset_n_c + (int64_t)i * C + c]);
+                thread_dh_pre_sum[i] += g_in * xi;
+            }
+        }
+    }
+
+    // 2. 将结果存入共享内存准备规约
+    #pragma unroll
+    for (int i = 0; i < MAX_N; i++) {
+        s_reduce[tid][i] = thread_dh_pre_sum[i];
+    }
+    __syncthreads();
+
+    // 3. 树状规约 (Tree Reduction)
+    for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) {
+        if (tid < stride) {
+            #pragma unroll
+            for (int i = 0; i < MAX_N; i++) {
+                if (i < n) {
+                    s_reduce[tid][i] += s_reduce[tid + stride][i];
+                }
+            }
+        }
+        __syncthreads();
+    }
+
+    // 4. 写回结果
+    if (tid == 0) {
+        #pragma unroll
+        for (int i = 0; i < MAX_N; i++) {
+            if (i < n) {
+                int64_t idx = bt_offset_n + (int64_t)i;
+                float sum_grad_x = s_reduce[0][i];
+                
+                // d_h_pre_raw 梯度逻辑
+                float s_pre = H_pre[idx];
+                d_h_pre_raw[idx] = sum_grad_x * (s_pre * (1.0f - s_pre));
+                
+                // d_h_post_raw 梯度逻辑
+                float s_post = H_post[idx] * 0.5f; 
+                d_h_post_raw[idx] = grad_H_post[idx] * 2.0f * (s_post * (1.0f - s_post));
+            }
+        }
+    }
+}
+
+/* -------------------- API 封装 -------------------- */
+
+inline void mhc_pre_op_forward(
+    floatX* x_layer_in, float* H_pre, float* H_post, float* H_res,
+    const floatX* x_expanded, const float* h_pre_raw, const float* h_post_raw, const float* h_res_raw,
+    int64_t B, int64_t T, int n, int64_t C, int sinkhorn_iters, float eps, cudaStream_t stream) 
+{
+    int64_t total_elements = B * T * C;
+    dim3 threads(256);
+    dim3 blocks((unsigned int)((total_elements + 255) / 256));
+    
+    mhc_pre_op_fwd_kernel<8><<<blocks, threads, 0, stream>>>(
+        x_layer_in, H_pre, H_post, x_expanded, h_pre_raw, h_post_raw, B, T, n, C);
+
+    // 处理 Sinkhorn 投影
+    for (int64_t i = 0; i < B * T; i++) {
+        sinkhorn_knopp_forward(
+            H_res + i * (int64_t)n * n, 
+            h_res_raw + i * (int64_t)n * n, 
+            n, n, sinkhorn_iters, eps, stream
+        );
+    }
+}
+
+inline void mhc_pre_op_backward(
+    floatX* d_x_expanded, float* d_h_pre_raw, float* d_h_post_raw, float* d_h_res_raw,
+    const floatX* grad_layer_in, const float* grad_H_post, const float* grad_H_res,
+    const floatX* x_expanded, const float* H_pre, const float* H_post, 
+    const float* H_res_out, const float* H_res_in_raw,
+    int64_t B, int64_t T, int n, int64_t C, int sinkhorn_iters, float eps, cudaStream_t stream) 
+{
+    const int BLOCK_SIZE = 256;
+    dim3 threads(BLOCK_SIZE);
+    dim3 blocks((unsigned int)(B * T));
+    
+    // 调用反向内核
+    mhc_pre_op_bwd_kernel<BLOCK_SIZE, 8><<<blocks, threads, 0, stream>>>(
+        d_x_expanded, d_h_pre_raw, d_h_post_raw,
+        grad_layer_in, grad_H_post, x_expanded, H_pre, H_post, B, T, n, C);
+
+    // 处理 Sinkhorn 梯度
+    for (int64_t i = 0; i < B * T; i++) {
+        sinkhorn_knopp_backward(
+            d_h_res_raw + i * (int64_t)n * n,
+            grad_H_res + i * (int64_t)n * n,
+            H_res_out + i * (int64_t)n * n,
+            H_res_in_raw + i * (int64_t)n * n,
+            n, sinkhorn_iters, eps, stream
+        );
+    }
+}
+
+} // namespace mhc
+
+#endif