rwkv-ops 0.6.1__py3-none-any.whl

rwkv_ops/mhc_kernel/common_kernel/kernels/rmsnorm.cuh

@@ -0,0 +1,152 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cuda_bf16.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+#include "../include/mhc_types.h"
+
+namespace cg = cooperative_groups;
+
+namespace mhc {
+
+/* -------------------- 前向传播 (修正索引溢出) -------------------- */
+template<int BLOCK_SIZE>
+__global__ void rmsnorm_fwd_kernel(
+    floatX* __restrict__ out, 
+    const floatX* __restrict__ inp, 
+    int N, int C, float eps) {
+    
+    cg::thread_block block = cg::this_thread_block();
+    cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
+
+    // 使用 int64_t 接收 blockIdx.x
+    int64_t row_idx = blockIdx.x;
+    if (row_idx >= N) return;
+
+    // 关键修正：强制 size_t 运算，防止 N * C 溢出
+    size_t offset = (size_t)row_idx * C;
+    const floatX* x_ptr = inp + offset;
+    floatX* o_ptr = out + offset;
+
+    // 共享内存用于存储每个 Warp 的局部和
+    extern __shared__ float s_reduce[]; 
+
+    // 1. 计算平方和 (使用 FP32 累加)
+    float thread_sum_sq = 0.0f;
+    for (int i = threadIdx.x; i < C; i += BLOCK_SIZE) {
+        float val = mhc::to_float(x_ptr[i]); 
+        thread_sum_sq += val * val;
+    }
+
+    // 2. Warp 级规约
+    float warp_sum = cg::reduce(warp, thread_sum_sq, cg::plus<float>());
+    
+    int warp_id = threadIdx.x / 32;
+    int lane_id = threadIdx.x % 32;
+    if (lane_id == 0) s_reduce[warp_id] = warp_sum;
+    block.sync();
+
+    // 3. Block 级规约
+    if (warp_id == 0) {
+        float b_sum = (lane_id < (BLOCK_SIZE / 32)) ? s_reduce[lane_id] : 0.0f;
+        b_sum = cg::reduce(warp, b_sum, cg::plus<float>());
+        if (lane_id == 0) s_reduce[0] = b_sum;
+    }
+    block.sync();
+
+    // 4. 计算 RMS 逆
+    float rms_inv = rsqrtf((s_reduce[0] / (float)C) + eps);
+
+    // 5. 写回结果
+    for (int i = threadIdx.x; i < C; i += BLOCK_SIZE) {
+        float val = mhc::to_float(x_ptr[i]);
+        o_ptr[i] = mhc::to_bf(val * rms_inv);
+    }
+}
+
+/* -------------------- 反向传播 (修正索引溢出) -------------------- */
+template<int BLOCK_SIZE>
+__global__ void rmsnorm_bwd_kernel(
+    floatX* __restrict__ dx,
+    const floatX* __restrict__ grad,
+    const floatX* __restrict__ x,
+    int N, int C, float eps) {
+    
+    cg::thread_block block = cg::this_thread_block();
+    cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
+
+    int64_t row_idx = blockIdx.x;
+    if (row_idx >= N) return;
+
+    // 关键修正：强制 size_t 运算
+    size_t offset = (size_t)row_idx * C;
+    const floatX* g_ptr = grad + offset;
+    const floatX* x_ptr = x + offset;
+    floatX* dx_ptr = dx + offset;
+
+    extern __shared__ float s_mem[]; 
+    int num_warps = BLOCK_SIZE / 32;
+    float* s_sum_sq = s_mem;
+    float* s_dot = s_mem + num_warps;
+
+    // 1. 局部累加
+    float t_sum_sq = 0.0f;
+    float t_dot = 0.0f;
+    for (int i = threadIdx.x; i < C; i += BLOCK_SIZE) {
+        float xv = mhc::to_float(x_ptr[i]);
+        float gv = mhc::to_float(g_ptr[i]);
+        t_sum_sq += xv * xv;
+        t_dot += gv * xv;
+    }
+
+    // 2. Warp 级规约
+    float w_sum = cg::reduce(warp, t_sum_sq, cg::plus<float>());
+    float w_dot = cg::reduce(warp, t_dot, cg::plus<float>());
+
+    int warp_id = threadIdx.x / 32;
+    int lane_id = threadIdx.x % 32;
+    if (lane_id == 0) {
+        s_sum_sq[warp_id] = w_sum;
+        s_dot[warp_id] = w_dot;
+    }
+    block.sync();
+
+    // 3. Block 级规约
+    if (warp_id == 0) {
+        float v1 = (lane_id < num_warps) ? s_sum_sq[lane_id] : 0.0f;
+        float v2 = (lane_id < num_warps) ? s_dot[lane_id] : 0.0f;
+        s_sum_sq[0] = cg::reduce(warp, v1, cg::plus<float>());
+        s_dot[0] = cg::reduce(warp, v2, cg::plus<float>());
+    }
+    block.sync();
+
+    // 4. 计算中间项
+    float r2_inv = 1.0f / (s_sum_sq[0] / (float)C + eps); 
+    float rms_inv = sqrtf(r2_inv);
+    float projection = s_dot[0] * (r2_inv * rms_inv) / (float)C;
+
+    // 5. 应用公式并写回
+    for (int i = threadIdx.x; i < C; i += BLOCK_SIZE) {
+        float xv = mhc::to_float(x_ptr[i]);
+        float gv = mhc::to_float(g_ptr[i]);
+        dx_ptr[i] = mhc::to_bf(gv * rms_inv - xv * projection);
+    }
+}
+
+/* -------------------- 包装函数 -------------------- */
+
+inline void rmsnorm_forward(floatX* out, const floatX* inp, int N, int C, float eps, cudaStream_t stream) {
+    const int BLOCK_SIZE = 256;
+    size_t smem = (BLOCK_SIZE / 32) * sizeof(float);
+    // Grid size 使用 int64 兼容的 N
+    rmsnorm_fwd_kernel<BLOCK_SIZE><<<N, BLOCK_SIZE, smem, stream>>>(out, inp, N, C, eps);
+}
+
+inline void rmsnorm_backward(floatX* dx, const floatX* grad, const floatX* x, int N, int C, float eps, cudaStream_t stream) {
+    const int BLOCK_SIZE = 256;
+    size_t smem = (BLOCK_SIZE / 32) * 2 * sizeof(float);
+    rmsnorm_bwd_kernel<BLOCK_SIZE><<<N, BLOCK_SIZE, smem, stream>>>(dx, grad, x, N, C, eps);
+}
+
+} // namespace mhc

rwkv_ops/mhc_kernel/common_kernel/kernels/sinkhorn_knopp.cuh

@@ -0,0 +1,158 @@
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cuda_bf16.h>
+#include <cooperative_groups.h>
+#include "../include/mhc_types.h"
+#include "type_conversions.cuh"
+
+namespace mhc {
+
+/**
+ * 1. 前向传播 Kernel (保持不变，但确保 eps 使用一致)
+ */
+template<int BLOCK_SIZE>
+__global__ void sinkhorn_knopp_fwd_kernel(
+    float* __restrict__ out,
+    const float* __restrict__ inp,
+    int M, int N, int num_iters, float eps) {
+    
+    extern __shared__ float smem[];
+    float* tile = smem; 
+    float* row_sums = smem + M * N; 
+    float* col_sums = row_sums + M;
+
+    int tid = threadIdx.x;
+    int64_t total = (int64_t)M * N;
+
+    for (int64_t i = tid; i < total; i += BLOCK_SIZE) {
+        tile[i] = __expf(inp[i]);
+    }
+    __syncthreads();
+
+    for (int iter = 0; iter < num_iters; iter++) {
+        // 行归一化
+        for (int r = tid; r < M; r += BLOCK_SIZE) {
+            float sum = 0.0f;
+            for (int c = 0; c < N; c++) sum += tile[(int64_t)r * N + c];
+            row_sums[r] = 1.0f / (sum + eps); 
+        }
+        __syncthreads();
+
+        for (int64_t i = tid; i < total; i += BLOCK_SIZE) {
+            tile[i] *= row_sums[i / N];
+        }
+        __syncthreads();
+
+        // 列归一化
+        for (int c = tid; c < N; c += BLOCK_SIZE) {
+            float sum = 0.0f;
+            for (int r = 0; r < M; r++) sum += tile[(int64_t)r * N + c];
+            col_sums[c] = 1.0f / (sum + eps);
+        }
+        __syncthreads();
+
+        for (int64_t i = tid; i < total; i += BLOCK_SIZE) {
+            tile[i] *= col_sums[i % N];
+        }
+        __syncthreads();
+    }
+
+    for (int64_t i = tid; i < total; i += BLOCK_SIZE) out[i] = tile[i];
+}
+
+/**
+ * 2. 反向传播 Kernel (修正：匹配自动微分迭代)
+ */
+template<int BLOCK_SIZE>
+__global__ void sinkhorn_knopp_bwd_kernel(
+    float* __restrict__ d_inp,
+    const float* __restrict__ grad,
+    const float* __restrict__ out_fwd,
+    int M, int N, int num_iters, float eps) {
+    
+    extern __shared__ float smem[];
+    float* P = smem;               // 前向输出 P [M*N]
+    float* dP = smem + M * N;      // 传入梯度 G [M*N]
+    float* alpha = smem + 2 * M * N; // 辅助变量 alpha [M]
+    float* beta = alpha + M;       // 辅助变量 beta [N]
+
+    int tid = threadIdx.x;
+    int64_t total = (int64_t)M * N;
+
+    // 加载数据
+    for (int64_t i = tid; i < total; i += BLOCK_SIZE) {
+        P[i] = out_fwd[i];
+        dP[i] = grad[i];
+    }
+    __syncthreads();
+
+    // 核心：Sinkhorn 梯度的迭代解法 (对应 PyTorch 的循环展开)
+    // 我们需要求解 alpha 和 beta 满足：
+    // alpha_i = sum_j (P_ij * (dP_ij - beta_j))
+    // beta_j = sum_i (P_ij * (dP_ij - alpha_i))
+    
+    // 初始化 alpha, beta 为 0
+    for(int i = tid; i < M; i += BLOCK_SIZE) alpha[i] = 0.0f;
+    for(int j = tid; j < N; j += BLOCK_SIZE) beta[j] = 0.0f;
+    __syncthreads();
+
+    for (int iter = 0; iter < num_iters; iter++) {
+        // 更新 beta
+        for (int j = tid; j < N; j += BLOCK_SIZE) {
+            float sum = 0.0f;
+            for (int i = 0; i < M; i++) {
+                sum += P[(int64_t)i * N + j] * (dP[(int64_t)i * N + j] - alpha[i]);
+            }
+            beta[j] = sum;
+        }
+        __syncthreads();
+
+        // 更新 alpha
+        for (int i = tid; i < M; i += BLOCK_SIZE) {
+            float sum = 0.0f;
+            for (int j = 0; j < N; j++) {
+                sum += P[(int64_t)i * N + j] * (dP[(int64_t)i * N + j] - beta[j]);
+            }
+            alpha[i] = sum;
+        }
+        __syncthreads();
+    }
+
+    // 最终梯度公式: dL/dA = P * (G - alpha - beta)
+    for (int64_t i = tid; i < total; i += BLOCK_SIZE) {
+        int64_t r = i / N;
+        int64_t c = i % N;
+        d_inp[i] = P[i] * (dP[i] - alpha[r] - beta[c]);
+    }
+}
+
+/* -------------------- API 接口 -------------------- */
+
+inline void sinkhorn_knopp_forward(
+    float* out, const float* inp, 
+    int M, int N, int num_iters, float eps, 
+    cudaStream_t stream = nullptr) {
+    
+    const int BLOCK_SIZE = 256;
+    size_t smem_size = ((size_t)M * N + M + N) * sizeof(float);
+    sinkhorn_knopp_fwd_kernel<BLOCK_SIZE><<<1, BLOCK_SIZE, smem_size, stream>>>(
+        out, inp, M, N, num_iters, eps
+    );
+}
+
+inline void sinkhorn_knopp_backward(
+    float* d_inp, const float* grad, const float* M_out, const float* M_inp, 
+    int N, int num_iters, float eps, 
+    cudaStream_t stream = nullptr) {
+    
+    const int BLOCK_SIZE = 256;
+    // 显存占用: P(N*N) + dP(N*N) + alpha(N) + beta(N)
+    size_t smem_size = (2 * (size_t)N * N + 2 * N) * sizeof(float);
+
+    sinkhorn_knopp_bwd_kernel<BLOCK_SIZE><<<1, BLOCK_SIZE, smem_size, stream>>>(
+        d_inp, grad, M_out, N, N, num_iters, eps
+    );
+}
+
+} // namespace mhc

rwkv_ops/mhc_kernel/common_kernel/kernels/stream_aggregate.cuh

@@ -0,0 +1,141 @@
+#ifndef MHC_STREAM_AGGREGATE_CUH
+#define MHC_STREAM_AGGREGATE_CUH
+
+#include <cuda_runtime.h>
+#include <cuda_bf16.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+#include "../include/mhc_types.h"
+
+namespace cg = cooperative_groups;
+
+namespace mhc {
+
+/**
+ * 前向传播：高精度版
+ * Out = sum(H_i * X_i)
+ */
+template<bool PER_TOKEN_H>
+__global__ void stream_aggregate_fwd_kernel(
+    floatX* __restrict__ out,
+    const floatX* __restrict__ inp,
+    const float* __restrict__ H_pre,
+    int64_t BT, int n, int64_t C) {
+    
+    // blockIdx.x 强转为 int64_t，防止 32 位溢出
+    int64_t btc = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    if (btc >= BT * C) return;
+
+    int64_t bt = btc / C;
+    int64_t c = btc % C;
+
+    // 核心：使用 FP32 寄存器进行所有乘加运算
+    float sum = 0.0f;
+
+    #pragma unroll
+    for (int i = 0; i < 8; i++) { // 假设 n <= 8，由 mHC 论文设定
+        if (i < n) {
+            // [修改]: bt * n 可能超过 32 位，且 n 为 int，显式强转 n 为 int64_t
+            float h_val = PER_TOKEN_H ? H_pre[bt * (int64_t)n + i] : H_pre[i];
+            // 关键：读取 bf16 后立即转为 fp32 参与运算
+            // [修改]: 显式强转 n 为 int64_t，确保 bt * n * C 全程为 64 位运算
+            float x_val = to_float(inp[bt * (int64_t)n * C + (int64_t)i * C + c]);
+            sum += h_val * x_val;
+        }
+    }
+    // 最后一次性转回 bf16
+    out[btc] = to_bf(sum);
+}
+
+/**
+ * 反向传播 dx: d_inp = d_out * H_pre
+ */
+template<bool PER_TOKEN_H>
+__global__ void stream_aggregate_bwd_dx_kernel(
+    floatX* __restrict__ d_inp,
+    const float* __restrict__ d_out, 
+    const float* __restrict__ H_pre,
+    int64_t BT, int n, int64_t C) {
+    
+    // blockIdx.x 强转为 int64_t
+    int64_t btc = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    int i = blockIdx.y; 
+
+    if (btc < BT * C && i < n) {
+        int64_t bt = btc / C;
+        // [修改]: bt * n 增加 (int64_t)n 强转
+        float h_val = PER_TOKEN_H ? H_pre[bt * (int64_t)n + i] : H_pre[i];
+        float grad_val = d_out[btc]; // 接收 FP32 梯度
+        
+        // [修改]: bt * n * C 增加 (int64_t)n 强转
+        d_inp[bt * (int64_t)n * C + (int64_t)i * C + (btc % C)] = to_bf(grad_val * h_val);
+    }
+}
+
+/**
+ * 反向传播 dH: d_H = sum_over_C(d_out * inp)
+ * 采用并行规约（Parallel Reduction）以保证高精度
+ */
+template<int BLOCK_SIZE>
+__global__ void stream_aggregate_bwd_dh_kernel(
+    float* __restrict__ d_H_pre,
+    const float* __restrict__ d_out,
+    const floatX* __restrict__ inp,
+    int64_t BT, int n, int64_t C) {
+    
+    cg::thread_block block = cg::this_thread_block();
+    cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
+
+    // blockIdx.x 在 C++ 侧对应 BT，若 BT 很大这里 int64_t 是必要的
+    int64_t bt = blockIdx.x; 
+    int i = blockIdx.y;
+    if (bt >= BT || i >= n) return;
+
+    extern __shared__ float s_reduce[];
+
+    float thread_sum = 0.0f;
+    for (int64_t c = threadIdx.x; c < C; c += BLOCK_SIZE) {
+        float g_val = d_out[bt * C + c]; 
+        // [修改]: bt * n * C 增加 (int64_t)n 强转
+        float x_val = to_float(inp[bt * (int64_t)n * C + (int64_t)i * C + c]);
+        thread_sum += g_val * x_val;
+    }
+
+    float warp_sum = cg::reduce(warp, thread_sum, cg::plus<float>());
+    int warp_id = threadIdx.x / 32;
+    int lane_id = threadIdx.x % 32;
+
+    if (lane_id == 0) s_reduce[warp_id] = warp_sum;
+    block.sync();
+
+    if (warp_id == 0) {
+        float val = (lane_id < (BLOCK_SIZE / 32)) ? s_reduce[lane_id] : 0.0f;
+        float block_sum = cg::reduce(warp, val, cg::plus<float>());
+        // [修改]: bt * n 增加 (int64_t)n 强转
+        if (lane_id == 0) d_H_pre[bt * (int64_t)n + i] = block_sum;
+    }
+}
+
+inline void stream_aggregate_forward(floatX* out, const floatX* inp, const float* H_pre, int64_t BT, int n, int64_t C, bool per_token, cudaStream_t stream) {
+    dim3 threads(256);
+    // 注意: BT * C 如果极大，超过 uint32 范围，dim3.x 会截断，这是 CUDA 硬件限制。
+    // 但计算过程本身 BT*C 是 int64 不会溢出。
+    dim3 blocks((BT * C + 255) / 256);
+    if (per_token) stream_aggregate_fwd_kernel<true><<<blocks, threads, 0, stream>>>(out, inp, H_pre, BT, n, C);
+    else stream_aggregate_fwd_kernel<false><<<blocks, threads, 0, stream>>>(out, inp, H_pre, BT, n, C);
+}
+
+inline void stream_aggregate_backward(floatX* d_inp, float* d_H_pre, const float* d_out, const floatX* inp, const float* H_pre, int64_t BT, int n, int64_t C, bool per_token, cudaStream_t stream) {
+    dim3 threads_dx(256);
+    dim3 blocks_dx((BT * C + 255) / 256, n);
+    if (per_token) stream_aggregate_bwd_dx_kernel<true><<<blocks_dx, threads_dx, 0, stream>>>(d_inp, d_out, H_pre, BT, n, C);
+    else stream_aggregate_bwd_dx_kernel<false><<<blocks_dx, threads_dx, 0, stream>>>(d_inp, d_out, H_pre, BT, n, C);
+
+    constexpr int DH_BLOCK_SIZE = 256;
+    dim3 grid_dh(BT, n);
+    size_t smem_size = (DH_BLOCK_SIZE / 32) * sizeof(float);
+    stream_aggregate_bwd_dh_kernel<DH_BLOCK_SIZE><<<grid_dh, DH_BLOCK_SIZE, smem_size, stream>>>(d_H_pre, d_out, inp, BT, n, C);
+}
+
+} // namespace mhc
+#endif

rwkv_ops/mhc_kernel/common_kernel/kernels/stream_distribute.cuh

@@ -0,0 +1,111 @@
+#ifndef MHC_STREAM_DISTRIBUTE_CUH
+#define MHC_STREAM_DISTRIBUTE_CUH
+
+#include <cuda_runtime.h>
+#include <cuda_bf16.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+#include "../include/mhc_types.h"
+
+namespace cg = cooperative_groups;
+
+namespace mhc {
+
+/**
+ * Forward: Out = Inp * H_post
+ * Shape: Inp [B, T, C], H_post [B, T, n] -> Out [B, T, n, C]
+ */
+__global__ void stream_distribute_fwd_kernel(
+    floatX* __restrict__ out,
+    const floatX* __restrict__ inp,
+    const float* __restrict__ H_post,
+    int64_t B, int64_t T, int n, int64_t C) {
+    
+    // 索引加固：强制使用 int64_t 防止 (B*T*C) 超过 21 亿
+    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;
+
+        float val = to_float(inp[btc]);
+        float weight = H_post[bt * n + i];
+        
+        // 计算 64 位偏移量
+        int64_t target_idx = bt * n * C + (int64_t)i * C + c;
+        out[target_idx] = to_bf(val * weight);
+    }
+}
+
+/**
+ * Backward dx: dx = sum_i(grad_i * H_post_i)
+ * 精度：FP32 累加
+ */
+__global__ void stream_distribute_bwd_dx_kernel(
+    floatX* __restrict__ dx,
+    const floatX* __restrict__ grad,
+    const float* __restrict__ H_post,
+    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;
+
+    float sum = 0.0f;
+    for (int i = 0; i < n; i++) {
+        // 强制 64 位偏移计算
+        float g = to_float(grad[bt * n * C + (int64_t)i * C + c]);
+        float w = H_post[bt * n + i];
+        sum += g * w; 
+    }
+    dx[btc] = to_bf(sum);
+}
+
+/**
+ * Backward dH: dH = sum_c(grad * inp)
+ * 精度：利用 Warp Shuffle 在通道维度 C 上进行全精度规约
+ */
+template<int BLOCK_SIZE>
+__global__ void stream_distribute_bwd_dh_kernel(
+    float* __restrict__ d_H_post,
+    const floatX* __restrict__ grad,
+    const floatX* __restrict__ inp,
+    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;
+    int i = blockIdx.y;
+    if (bt >= B * T || i >= n) return;
+
+    int64_t base_grad = bt * n * C + (int64_t)i * C;
+    int64_t base_inp = bt * C;
+
+    float thread_sum = 0.0f;
+    for (int64_t c = threadIdx.x; c < C; c += BLOCK_SIZE) {
+        float g = to_float(grad[base_grad + c]);
+        float x = to_float(inp[base_inp + c]);
+        thread_sum += g * x;
+    }
+
+    // Block 级规约
+    float sum = cg::reduce(warp, thread_sum, cg::plus<float>());
+    static __shared__ float s_reduce[32];
+    int warp_id = threadIdx.x / 32;
+    int lane_id = threadIdx.x % 32;
+    if (lane_id == 0) s_reduce[warp_id] = sum;
+    block.sync();
+
+    if (warp_id == 0) {
+        float val = (lane_id < (BLOCK_SIZE / 32)) ? s_reduce[lane_id] : 0.0f;
+        float block_sum = cg::reduce(warp, val, cg::plus<float>());
+        if (lane_id == 0) d_H_post[bt * n + i] = block_sum;
+    }
+}
+
+} // namespace mhc
+#endif