finish rms-gate-norm #32

S1/Ljy123_#32/cudacode.py

@@ -0,0 +1,94 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+
+source = """
+#include <torch/extension.h>
+
+__inline__ __device__ float warpReduceSum(float val) {
+    for (int offset = 16; offset > 0; offset >>= 1) {
+        val += __shfl_down_sync(0xffffffff, val, offset);
+    }
+    return val;
+}
+
+__global__ void rms_gate_kernel(const float* x, float* y, int B, int D, float gamma, float beta) {
+    int b = blockIdx.x;
+    int tid = threadIdx.x;
+    int lane = tid & 31;
+    int warpId = tid >> 5;
+    extern __shared__ float sh[];
+    float sumsq = 0.0f;
+    for (int i = tid * 4; i < D; i += blockDim.x * 4) {
+        int base = b * D + i;
+        if (i + 3 < D) {
+            float4 v4 = reinterpret_cast<const float4*>(x)[(b * D + i) / 4];
+            sumsq += v4.x * v4.x + v4.y * v4.y + v4.z * v4.z + v4.w * v4.w;
+        } else {
+            for (int k = 0; k < 4 && i + k < D; ++k) {
+                float v = x[base + k];
+                sumsq += v * v;
+            }
+        }
+    }
+    sumsq = warpReduceSum(sumsq);
+    if (lane == 0) sh[warpId] = sumsq;
+    __syncthreads();
+    float total_sum = 0.0f;
+    if (warpId == 0) {
+        float v = (lane < ((blockDim.x + 31) >> 5)) ? sh[lane] : 0.0f;
+        v = warpReduceSum(v);
+        if (lane == 0) sh[0] = v;
+    }
+    __syncthreads();
+    float rms = sqrtf(sh[0] / (float)D);
+    float g = 1.0f / (1.0f + expf(-(gamma * rms + beta)));
+    for (int i = tid * 4; i < D; i += blockDim.x * 4) {
+        int base = b * D + i;
+        if (i + 3 < D) {
+            float4 v4 = reinterpret_cast<const float4*>(x)[(b * D + i) / 4];
+            v4.x *= g; v4.y *= g; v4.z *= g; v4.w *= g;
+            reinterpret_cast<float4*>(y)[(b * D + i) / 4] = v4;
+        } else {
+            for (int k = 0; k < 4 && i + k < D; ++k) {
+                y[base + k] = x[base + k] * g;
+            }
+        }
+    }
+}
+
+torch::Tensor rms_gate_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta) {
+    auto xc = x.contiguous();
+    auto y = torch::empty_like(xc);
+    int B = (int)xc.size(0);
+    int D = (int)xc.size(1);
+    float g = gamma.item<float>();
+    float b = beta.item<float>();
+    int block = 256;
+    int grid = B;
+    size_t shmem = (size_t)((block + 31) >> 5) * sizeof(float);
+    rms_gate_kernel<<<grid, block, shmem>>>(xc.data_ptr<float>(), y.data_ptr<float>(), B, D, g, b);
+    return y;
+}
+"""
+
+cpp_source = """
+torch::Tensor rms_gate_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta);
+"""
+
+ops = load_inline(
+    name="rms_gate_norm",
+    cpp_sources=cpp_source,
+    cuda_sources=source,
+    functions=["rms_gate_cuda"],
+    verbose=True
+)
+
+class ModelNew(torch.nn.Module):
+    def __init__(self, gamma: float, beta: float):
+        super(ModelNew, self).__init__()
+        self.ops = ops
+        self.register_buffer("gamma", torch.tensor(float(gamma), dtype=torch.float32))
+        self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor):
+        return self.ops.rms_gate_cuda(x, self.gamma, self.beta)

S1/Ljy123_#32/prompt.txt

@@ -0,0 +1,2 @@
+You write custom CUDA kernels to replace PyTorch operators for speedups.
+Implement RMS-Gated Normalization: Given a 2D tensor x of shape [B, D], compute the per-row RMS r = sqrt(mean(x^2, dim=1)), then a scalar gate g = sigmoid(gamma * r + beta) per row, and output y = x * g. The CUDA kernel must perform a block-level reduction to obtain RMS for each row (one block per row), then apply the gate to all elements in that row. Use contiguous memory layout, dynamic shared memory for the reduction, and avoid multiple kernel launches. Provide a PyTorch reference module that computes the same operation using nn.Parameters for gamma and beta. Ensure numerical stability and accuracy matching torch reference within rtol=1e-3.

S1/Ljy123_#32/run_code.py

@@ -0,0 +1,52 @@
+import torch
+import time
+from torchcode import Model, get_inputs, get_init_inputs
+from cudacode import ModelNew
+
+def run_benchmark():
+    if not torch.cuda.is_available():
+        print("CUDA 不可用，请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
+        return
+    device = torch.device("cuda")
+
+    init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
+    inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
+
+    torch_model = Model(*init_inputs).cuda()
+    cuda_model = ModelNew(*init_inputs).cuda()
+    torch_model.eval(); cuda_model.eval()
+
+    print("-------------------- 精度对齐验证 --------------------")
+    with torch.no_grad():
+        out_torch = torch_model(*inputs)
+        out_cuda = cuda_model(*inputs)
+    flag = torch.allclose(out_torch, out_cuda, rtol=1e-03)
+    if flag:
+        print("✅ 精度对齐：两个模型的输出结果非常接近。")
+    else:
+        print("❌ 精度不一致！")
+        print(f"最大绝对误差: {(out_torch - out_cuda).abs().max().item()}" )
+
+    print("\n-------------------- 性能加速比测试 --------------------")
+    iters = 100
+    torch.cuda.synchronize(); t0 = time.time()
+    for _ in range(iters):
+        _ = torch_model(*inputs)
+    torch.cuda.synchronize(); t_torch = (time.time() - t0) / iters
+
+    torch.cuda.synchronize(); t0 = time.time()
+    for _ in range(iters):
+        _ = cuda_model(*inputs)
+    torch.cuda.synchronize(); t_cuda = (time.time() - t0) / iters
+
+    print(f"PyTorch RMS-Gated-Norm 平均执行时间: {t_torch:.6f} 秒")
+    print(f"自定义 CUDA 融合内核 平均执行时间: {t_cuda:.6f} 秒")
+    sp = t_torch / t_cuda if t_cuda > 0 else 0
+    if t_cuda > 0:
+        print(f"加速比 (Speedup): {sp:.2f}x")
+    else:
+        print("CUDA 内核执行时间为0，无法计算加速比。")
+    return flag, sp
+
+if __name__ == "__main__":
+    run_benchmark()

S1/Ljy123_#32/torchcode.py

@@ -0,0 +1,23 @@
+import torch
+import torch.nn as nn
+
+class Model(nn.Module):
+    def __init__(self, gamma: float, beta: float):
+        super(Model, self).__init__()
+        self.gamma = nn.Parameter(torch.tensor(float(gamma), dtype=torch.float32))
+        self.beta = nn.Parameter(torch.tensor(float(beta), dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        rms = torch.sqrt(torch.mean(x * x, dim=1, keepdim=True))
+        gate = torch.sigmoid(self.gamma * rms + self.beta)
+        return x * gate
+
+batch_size = 32
+dim = 8192
+
+def get_inputs():
+    x = torch.randn(batch_size, dim)
+    return [x]
+
+def get_init_inputs():
+    return [1.0, 0.0]