rwkv-ops 0.6.1__py3-none-any.whl

rwkv_ops/mhc_kernel/native_keras_op.py

@@ -0,0 +1,193 @@
+from keras import ops
+
+# --- 辅助函数：确保在 fp32 下计算以保证数值稳定性 ---
+
+
+def fp32_sigmoid(x):
+    dtype = x.dtype
+    return ops.cast(ops.nn.sigmoid(ops.cast(x, "float32")), dtype)
+
+
+# --- 核心 MHC 算子 ---
+
+
+def sinkhorn_knopp(inp, num_iters=20, eps=1e-8):
+    """
+    将输入矩阵投影为双拟随机矩阵 (Doubly Stochastic Matrix)。
+    通常 inp 是 log 域的矩阵 (H_res_raw)。
+    """
+    dtype = inp.dtype
+    # 转换到 fp32 并应用 exp (论文 Eq. 9 之前的步骤)
+    x = ops.cast(inp, "float32")
+    # 防溢出技巧：减去最大值
+    x = x - ops.max(x, axis=(-1, -2), keepdims=True)
+    P = ops.exp(x)
+
+    for _ in range(num_iters):
+        # 行归一化
+        P = P / (ops.sum(P, axis=-1, keepdims=True) + eps)
+        # 列归一化
+        P = P / (ops.sum(P, axis=-2, keepdims=True) + eps)
+
+    return ops.cast(P, dtype)
+
+
+def rmsnorm(inp, eps=1e-5):
+    """
+    标准 RMSNorm 算子。
+    inp: [..., C], weight: [C]
+    """
+    dtype = inp.dtype
+    x = ops.cast(inp, "float32")
+    # 计算均方根
+    rms = ops.sqrt(ops.mean(ops.square(x), axis=-1, keepdims=True) + eps)
+    x_normed = x / rms
+    # 应用权重
+    return ops.cast(x_normed, dtype)
+
+
+def stream_aggregate(inp, H_pre):
+    # 1. 转换为 float32 进行高精度计算
+    inp_f32 = ops.cast(inp, "float32")
+    H_f32 = ops.cast(H_pre, "float32")
+
+    # 2. 在 float32 空间完成乘法和累加
+    out_f32 = ops.sum(inp_f32 * ops.expand_dims(H_f32, -1), axis=-2)
+
+    # 3. 最后转回原先的格式 (如 bf16)
+    return ops.cast(out_f32, inp.dtype)
+
+
+def stream_distribute(inp, H_post):
+    """
+    Distribute (1 -> n): 将单流输出分发回多流。
+    对齐精度版：强制在 FP32 下进行广播乘法。
+
+    inp: [B, T, C] (BF16)
+    H_post: [B, T, n] (FP32)
+    """
+    # 1. 记录原始类型
+    original_dtype = inp.dtype
+
+    # 2. 提升到 FP32 进行运算 (对齐 CUDA 内核内部的 to_float 逻辑)
+    # [B, T, 1, C]
+    x_fp32 = ops.cast(ops.expand_dims(inp, -2), "float32")
+
+    # [B, T, n, 1]
+    w_fp32 = ops.cast(ops.expand_dims(H_post, -1), "float32")
+
+    # 3. 执行广播乘法
+    # 结果为 [B, T, n, C]
+    res_fp32 = x_fp32 * w_fp32
+
+    # 4. 转回原始类型 (对齐 CUDA 内核末尾的 to_bf 逻辑)
+    return ops.cast(res_fp32, original_dtype)
+
+
+def stream_mix(inp, M):
+    """
+    Mix (n -> n): 残差流之间的线性交互。
+    inp: [B, T, n, C]
+    M: [B, T, n, n] 或 [n, n] (由 sinkhorn_knopp 生成的方阵)
+    """
+    # 使用 einsum 表达矩阵乘法：M @ inp
+    # i,j 是流索引，k 是通道索引
+    dtype = inp.dtype
+    inp = ops.cast(inp, M.dtype)
+    if len(ops.shape(M)) == 2:
+        out = ops.einsum("ij,btjk->btik", M, inp)
+    else:
+        out = ops.einsum("btij,btjk->btik", M, inp)
+    return ops.cast(out, dtype)
+
+
+def stream_mix_fp32(x_expanded, H_res):
+    """内部强制使用 FP32 计算的流混合"""
+    # x_expanded: [B, T, n, C], H_res: [B, T, n, n]
+    x_f32 = ops.cast(x_expanded, "float32")
+    h_f32 = ops.cast(H_res, "float32")
+    # 执行矩阵乘法: [B, T, n, n] @ [B, T, n, C] -> [B, T, n, C]
+    return ops.matmul(h_f32, x_f32)
+
+
+def stream_distribute_fp32(layer_out, H_post):
+    """内部强制使用 FP32 计算的分发"""
+    # layer_out: [B, T, C], H_post: [B, T, n]
+    l_f32 = ops.cast(layer_out, "float32")
+    h_f32 = ops.cast(H_post, "float32")
+
+    # [B, T, 1, C] * [B, T, n, 1] -> [B, T, n, C]
+    return ops.expand_dims(l_f32, -2) * ops.expand_dims(h_f32, -1)
+
+
+def mhc_post_op(layer_out, x_expanded, H_post, H_res):
+    """
+    mHC 后处理融合算子
+    输入:
+        layer_out: [B, T, C] - 核心层 (Attention/FFN) 处理后的输出
+        x_expanded: [B, T, n, C] - 之前的扩展残差流 (Pre-Op 之前的状态)
+        H_post: [B, T, n] - 分发权重 (来自 Pre-Op)，2*sigmoid后的数值
+        H_res: [B, T, n, n] - 流混合矩阵 (来自 Pre-Op)
+    返回:
+        x_next: [B, T, n, C] - 更新后的扩展残差流
+    """
+    # 1. 在 FP32 下计算混合路径
+    x_mixed_f32 = stream_mix_fp32(x_expanded, H_res)
+
+    # 2. 在 FP32 下计算增量路径
+    x_delta_f32 = stream_distribute_fp32(layer_out, H_post)
+
+    # 3. 在 FP32 下完成最后的残差加法
+    x_next_f32 = x_mixed_f32 + x_delta_f32
+
+    # 4. 只在最后输出时进行一次 BF16 转换
+    # 这一步对应 CUDA 内核中最后的 to_bf()
+    return ops.cast(x_next_f32, x_expanded.dtype)
+
+
+def mhc_pre_op(x_expanded, h_pre_raw, h_post_raw, h_res_raw, num_iters=20, eps=1e-8):
+    """
+    mHC 前处理融合算子
+    输入:
+        x_expanded: [B, T, n, C] - 当前的扩展残差流
+        h_pre_raw, h_post_raw: [B, T, n] - 线性投影后的原始激活值
+        h_res_raw: [B, T, n*n] - 用于生成 Sinkhorn 矩阵的原始值
+    返回:
+        x_layer_in: [B, T, C] - 聚合后准备进入 Layer (Attention/FFN) 的输入
+        H_post: [B, T, n] - 激活后的分发权重
+        H_res: [B, T, n, n] - 经过流形约束后的混合矩阵
+    """
+    original_dtype = x_expanded.dtype
+    B, T, n, C = ops.shape(x_expanded)
+
+    # --- 0. 提升精度 ---
+    x_exp_f32 = ops.cast(x_expanded, "float32")
+    h_pre_f32 = ops.cast(h_pre_raw, "float32")
+    h_post_f32 = ops.cast(h_post_raw, "float32")
+    h_res_f32 = ops.cast(h_res_raw, "float32")
+
+    # --- 1. Stream Aggregate (n -> 1) ---
+    # H_pre = sigmoid(h_pre)
+    H_pre_f32 = ops.nn.sigmoid(h_pre_f32)
+
+    # x_layer_in = sum(H_pre_i * x_expanded_i)
+    # 使用 expand_dims 确保广播正确：[B, T, n, 1] * [B, T, n, C] -> [B, T, n, C] -> sum -> [B, T, C]
+    x_layer_in_f32 = ops.sum(ops.expand_dims(H_pre_f32, -1) * x_exp_f32, axis=-2)
+
+    # --- 2. H_post 计算 (带 2.0 缩放) ---
+    # 根据论文，这里使用 2.0 * sigmoid 确保恒等映射的初始化稳定性
+    H_post_f32 = 2.0 * ops.nn.sigmoid(h_post_f32)
+
+    # --- 3. Sinkhorn-Knopp (n x n 投影) ---
+    h_res_reshaped_f32 = ops.reshape(h_res_f32, (B, T, n, n))
+
+    # 内部执行 sinkhorn 迭代（确保 sinkhorn_knopp 内部也是 float32）
+    H_res_f32 = sinkhorn_knopp(h_res_reshaped_f32, num_iters=num_iters, eps=eps)
+
+    # --- 4. 转换回原始格式 ---
+    # 模拟 CUDA Kernel 最后写回显存时的 cast 操作
+    return (
+        ops.cast(x_layer_in_f32, original_dtype),
+        ops.cast(H_post_f32, "float32"),  # H 权重通常在模型中保持 FP32 精度
+        ops.cast(H_res_f32, "float32"),
+    )

rwkv_ops/mhc_kernel/torch_kernel/mhc_cuda.cu

@@ -0,0 +1,207 @@
+#include <cuda_runtime.h>
+#include <cuda_bf16.h>
+#include "../common_kernel/include/mhc_types.h"
+#include "../common_kernel/kernels/sinkhorn_knopp.cuh"
+#include "../common_kernel/kernels/rmsnorm.cuh"
+#include "../common_kernel/kernels/stream_mix.cuh"
+#include "../common_kernel/kernels/stream_aggregate.cuh"
+#include "../common_kernel/kernels/stream_distribute.cuh"
+#include "../common_kernel/kernels/mhc_post_op.cuh"
+#include "../common_kernel/kernels/mhc_pre_op.cuh"
+
+namespace mhc {
+// --- Post-Op 融合前向 ---
+void cuda_mhc_post_op_fwd(
+    nv_bfloat16* out,               // [B, T, n, C]
+    const nv_bfloat16* layer_out,   // [B, T, C]
+    const nv_bfloat16* x_expanded,  // [B, T, n, C]
+    const float* H_post,            // [B, T, n]
+    const float* H_res,             // [B, T, n, n]
+    int64_t B, int64_t T, int n, int64_t C, 
+    cudaStream_t stream) 
+{
+    // 调用 .cuh 中的 inline 包装函数
+    mhc::mhc_post_op_forward(
+        reinterpret_cast<mhc::floatX*>(out),
+        reinterpret_cast<const mhc::floatX*>(layer_out),
+        reinterpret_cast<const mhc::floatX*>(x_expanded),
+        H_post,
+        H_res,
+        B, T, n, C,
+        stream
+    );
+}
+
+// --- 反向传播包装 (融合版) ---
+void cuda_mhc_post_op_bwd(
+    nv_bfloat16* d_layer_out,      // [B, T, C]
+    nv_bfloat16* d_x_expanded,     // [B, T, n, C]
+    float* d_H_post,               // [B, T, n]
+    float* d_H_res,                // [B, T, n, n]
+    const nv_bfloat16* grad_next,  // [B, T, n, C]
+    const nv_bfloat16* layer_out,  // [B, T, C]
+    const nv_bfloat16* x_expanded, // [B, T, n, C]
+    const float* H_post,
+    const float* H_res,
+    int64_t B, int64_t T, int n, int64_t C, 
+    cudaStream_t stream) 
+{
+    // 调用 .cuh 中的全量反向融合内核
+    // 该内核内部会处理 dx, dl, dH_post, dH_res 的全部逻辑
+    mhc::mhc_post_op_backward_full(
+        reinterpret_cast<mhc::floatX*>(d_layer_out),
+        reinterpret_cast<mhc::floatX*>(d_x_expanded),
+        d_H_post,
+        d_H_res,
+        reinterpret_cast<const mhc::floatX*>(grad_next),
+        reinterpret_cast<const mhc::floatX*>(layer_out),
+        reinterpret_cast<const mhc::floatX*>(x_expanded),
+        H_post,
+        H_res,
+        B, T, n, C,
+        stream
+    );
+}
+
+// --- Sinkhorn 包装 ---
+void cuda_sinkhorn_fwd(float* out, const float* inp, int64_t B, int64_t M, int64_t N, int iters, float eps, cudaStream_t stream) {
+    for (int64_t b = 0; b < B; b++) {
+        mhc::sinkhorn_knopp_forward(out + b * M * N, inp + b * M * N, (int)M, (int)N, iters, eps, stream);
+    }
+}
+
+void cuda_sinkhorn_bwd(float* d_inp, const float* grad, const float* M_out, const float* M_inp, int64_t B, int64_t N, int iters, float eps, cudaStream_t stream) {
+    for (int64_t b = 0; b < B; b++) {
+        mhc::sinkhorn_knopp_backward(d_inp + b * N * N, grad + b * N * N, M_out + b * N * N, M_inp + b * N * N, (int)N, iters, eps, stream);
+    }
+}
+
+// --- RMSNorm 包装 ---
+void cuda_rmsnorm_fwd(nv_bfloat16* out, const nv_bfloat16* inp, int64_t N, int64_t C, float eps, cudaStream_t stream) {
+    mhc::rmsnorm_forward(reinterpret_cast<mhc::floatX*>(out), reinterpret_cast<const mhc::floatX*>(inp), N, C, eps, stream);
+}
+
+void cuda_rmsnorm_bwd(nv_bfloat16* dx, const nv_bfloat16* grad, const nv_bfloat16* x, int64_t N, int64_t C, float eps, cudaStream_t stream) {
+    mhc::rmsnorm_backward(reinterpret_cast<mhc::floatX*>(dx), reinterpret_cast<const mhc::floatX*>(grad), reinterpret_cast<const mhc::floatX*>(x), N, C, eps, stream);
+}
+
+// --- Stream Mix 包装 ---
+void cuda_stream_mix_fwd(nv_bfloat16* out, const nv_bfloat16* inp, const float* M, int64_t B, int64_t T, int n, int64_t C, cudaStream_t stream) {
+    mhc::stream_mix_forward(reinterpret_cast<mhc::floatX*>(out), reinterpret_cast<const mhc::floatX*>(inp), M, B, T, n, C, stream);
+}
+
+void cuda_stream_mix_bwd(nv_bfloat16* d_inp, float* d_M, const float* grad, const nv_bfloat16* inp, const float* M, int64_t B, int64_t T, int n, int64_t C, cudaStream_t stream) {
+    mhc::stream_mix_backward(reinterpret_cast<mhc::floatX*>(d_inp), d_M, grad, reinterpret_cast<const mhc::floatX*>(inp), M, B, T, n, C, stream);
+}
+
+// --- 新增：Stream Aggregate 包装 ---
+void cuda_stream_aggregate_fwd(nv_bfloat16* out, const nv_bfloat16* inp, const float* H_pre, int64_t B, int64_t T, int n, int64_t C, bool per_token, cudaStream_t stream) {
+    mhc::stream_aggregate_forward(reinterpret_cast<mhc::floatX*>(out), reinterpret_cast<const mhc::floatX*>(inp), H_pre, B * T, n, C, per_token, stream);
+}
+
+void cuda_stream_aggregate_bwd(nv_bfloat16* d_inp, float* d_H_pre, const float* grad, const nv_bfloat16* inp, const float* H_pre, int64_t B, int64_t T, int n, int64_t C, bool per_token, cudaStream_t stream) {
+    mhc::stream_aggregate_backward(reinterpret_cast<mhc::floatX*>(d_inp), d_H_pre, grad, reinterpret_cast<const mhc::floatX*>(inp), H_pre, B * T, n, C, per_token, stream);
+}
+
+void cuda_stream_distribute_fwd(
+    nv_bfloat16* out, 
+    const nv_bfloat16* inp, 
+    const float* H, 
+    int64_t B, int64_t T, int n, int64_t C, 
+    cudaStream_t stream) 
+{
+    // 计算总线程数（针对单流输入 B*T*C）
+    int64_t total_btc = B * T * C;
+    dim3 threads(256);
+    // x轴覆盖所有的元素，y轴覆盖流索引 n
+    dim3 blocks((total_btc + 255) / 256, (unsigned int)n);
+
+    mhc::stream_distribute_fwd_kernel<<<blocks, threads, 0, stream>>>(
+        reinterpret_cast<mhc::floatX*>(out),
+        reinterpret_cast<const mhc::floatX*>(inp),
+        H, B, T, n, C
+    );
+}
+
+void cuda_stream_distribute_bwd(
+    nv_bfloat16* d_inp, 
+    float* d_H, 
+    const nv_bfloat16* grad, 
+    const nv_bfloat16* inp, 
+    const float* H, 
+    int64_t B, int64_t T, int n, int64_t C, 
+    cudaStream_t stream) 
+{
+    // 1. 计算 dx [B, T, C]
+    int64_t total_btc = B * T * C;
+    dim3 threads_dx(256);
+    dim3 blocks_dx((total_btc + 255) / 256);
+    
+    mhc::stream_distribute_bwd_dx_kernel<<<blocks_dx, threads_dx, 0, stream>>>(
+        reinterpret_cast<mhc::floatX*>(d_inp),
+        reinterpret_cast<const mhc::floatX*>(grad),
+        H, B, T, n, C
+    );
+
+    // 2. 计算 dH [B, T, n]
+    // 每个 (bt, i) 分配一个 block 进行通道 C 维度的规约
+    dim3 threads_dh(256);
+    dim3 blocks_dh((unsigned int)(B * T), (unsigned int)n);
+
+    mhc::stream_distribute_bwd_dh_kernel<256><<<blocks_dh, threads_dh, 0, stream>>>(
+        d_H,
+        reinterpret_cast<const mhc::floatX*>(grad),
+        reinterpret_cast<const mhc::floatX*>(inp),
+        B, T, n, C
+    );
+}
+void cuda_mhc_pre_op_fwd(
+    nv_bfloat16* x_layer_in, 
+    float* H_pre, 
+    float* H_post, 
+    float* H_res,
+    const nv_bfloat16* 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) 
+{
+    // 调用 .cuh 中的融合接口
+    mhc::mhc_pre_op_forward(
+        reinterpret_cast<mhc::floatX*>(x_layer_in),
+        H_pre, H_post, H_res,
+        reinterpret_cast<const mhc::floatX*>(x_expanded),
+        h_pre_raw, h_post_raw, h_res_raw,
+        B, T, n, C, sinkhorn_iters, eps, stream
+    );
+}
+
+void cuda_mhc_pre_op_bwd(
+    nv_bfloat16* d_x_expanded, 
+    float* d_h_pre_raw, 
+    float* d_h_post_raw, 
+    float* d_h_res_raw,
+    const nv_bfloat16* grad_layer_in, 
+    const float* grad_H_post, 
+    const float* grad_H_res,
+    const nv_bfloat16* 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) 
+{
+    mhc::mhc_pre_op_backward(
+        reinterpret_cast<mhc::floatX*>(d_x_expanded),
+        d_h_pre_raw, d_h_post_raw, d_h_res_raw,
+        reinterpret_cast<const mhc::floatX*>(grad_layer_in),
+        grad_H_post, grad_H_res,
+        reinterpret_cast<const mhc::floatX*>(x_expanded),
+        H_pre, H_post, H_res_out, H_res_in_raw,
+        B, T, n, C, sinkhorn_iters, eps, stream
+    );
+}
+
+} // namespace mhc