finish rmsnorm-silu #70

S1/Ljy123_#70/cudacode.py

@@ -0,0 +1,115 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+
+source = """
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+
+__device__ __forceinline__ float warp_reduce_sum(float v){
+    for(int offset=16; offset>0; offset/=2){
+        v += __shfl_down_sync(0xffffffff, v, offset);
+    }
+    return v;
+}
+
+__global__ __launch_bounds__(1024) void rmsnorm_silu_kernel(const float* __restrict__ x, const float* __restrict__ gamma, const float* __restrict__ beta, float* __restrict__ y, int B, int D, float eps){
+    int b = blockIdx.x;
+    int tid = threadIdx.x;
+    int lane = tid & 31;
+    int warp_id = tid >> 5;
+    int warps = blockDim.x >> 5;
+    const float* xr = x + b * D;
+    float* yr = y + b * D;
+    __shared__ float ssum[32];
+    __shared__ float s_inv;
+    float local_sum = 0.0f;
+    float c = 0.0f;
+    for(int i = tid; i < D; i += blockDim.x){
+        float v = __ldg(xr + i);
+        float vsq = v * v;
+        float y = vsq - c;
+        float t = local_sum + y;
+        c = (t - local_sum) - y;
+        local_sum = t;
+    }
+    float wsum = warp_reduce_sum(local_sum);
+    if(lane == 0) ssum[warp_id] = wsum;
+    __syncthreads();
+    float block_sum = (tid < warps) ? ssum[tid] : 0.0f;
+    if(warp_id == 0){
+        float s = warp_reduce_sum(block_sum);
+        if(lane == 0){
+            float mean2 = s / (float)D;
+            s_inv = rsqrtf(mean2 + eps);
+        }
+    }
+    __syncthreads();
+    int aligned = ((((long long)xr & 15LL) == 0) && (((long long)yr & 15LL) == 0) && (((long long)gamma & 15LL) == 0) && (((long long)beta & 15LL) == 0) && ((D & 3) == 0));
+    if(aligned){
+        int D4 = (D / 4) * 4;
+        #pragma unroll 4
+        for(int i = tid * 4; i < D4; i += blockDim.x * 4){
+            float4 xv = reinterpret_cast<const float4*>(xr)[i / 4];
+            float4 gv = reinterpret_cast<const float4*>(gamma)[i / 4];
+            float4 bv = reinterpret_cast<const float4*>(beta)[i / 4];
+            float4 yv;
+            float u0 = (xv.x * s_inv) * gv.x + bv.x; yv.x = u0 / (1.0f + expf(-u0));
+            float u1 = (xv.y * s_inv) * gv.y + bv.y; yv.y = u1 / (1.0f + expf(-u1));
+            float u2 = (xv.z * s_inv) * gv.z + bv.z; yv.z = u2 / (1.0f + expf(-u2));
+            float u3 = (xv.w * s_inv) * gv.w + bv.w; yv.w = u3 / (1.0f + expf(-u3));
+            reinterpret_cast<float4*>(yr)[i / 4] = yv;
+        }
+        for(int i = (D4 + tid); i < D; i += blockDim.x){
+            float v = __ldg(xr + i) * s_inv;
+            float u = v * __ldg(gamma + i) + __ldg(beta + i);
+            float g = u / (1.0f + expf(-u));
+            yr[i] = g;
+        }
+    } else {
+        for(int i = tid; i < D; i += blockDim.x){
+            float v = __ldg(xr + i) * s_inv;
+            float u = v * __ldg(gamma + i) + __ldg(beta + i);
+            float g = u / (1.0f + expf(-u));
+            yr[i] = g;
+        }
+    }
+}
+
+torch::Tensor rmsnorm_silu_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor eps){
+    auto xc = x.contiguous();
+    auto gc = gamma.contiguous();
+    auto bc = beta.contiguous();
+    auto y = torch::empty_like(xc);
+    int B = (int)xc.size(0);
+    int D = (int)xc.size(1);
+    float e = eps.item<float>();
+    int block = 1024;
+    dim3 grid(B);
+    rmsnorm_silu_kernel<<<grid, block>>>(xc.data_ptr<float>(), gc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, e);
+    return y;
+}
+"""
+
+cpp_source = """
+torch::Tensor rmsnorm_silu_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor eps);
+"""
+
+ops = load_inline(
+    name="rmsnorm_silu",
+    cpp_sources=cpp_source,
+    cuda_sources=source,
+    functions=["rmsnorm_silu_cuda"],
+    extra_cuda_cflags=["-O3","--use_fast_math","-gencode=arch=compute_80,code=sm_80","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"],
+    verbose=True
+)
+
+class ModelNew(torch.nn.Module):
+    def __init__(self, gamma: torch.Tensor, beta: torch.Tensor, eps: float):
+        super(ModelNew, self).__init__()
+        self.ops = ops
+        self.register_buffer("gamma", gamma)
+        self.register_buffer("beta", beta)
+        self.register_buffer("eps", torch.tensor(float(eps), dtype=torch.float32))
+
+    def forward(self, x):
+        return self.ops.rmsnorm_silu_cuda(x, self.gamma, self.beta, self.eps)

S1/Ljy123_#70/prompt.txt

@@ -0,0 +1 @@
+融合算子：RMSNorm+Affine+SiLU，一次内核完成归一化、仿射与激活。

S1/Ljy123_#70/run_code.py

@@ -0,0 +1,51 @@
+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().eval()
+    cuda_model = ModelNew(*init_inputs).cuda().eval()
+
+    print("-------------------- 精度对齐验证 --------------------")
+    with torch.no_grad():
+        output_torch = torch_model(*inputs)
+        output_cuda = cuda_model(*inputs)
+    precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03)
+    if precision_flag:
+        print("✅ 精度对齐：两个模型的输出结果非常接近。")
+    else:
+        print("❌ 精度不一致！")
+
+    print("\n-------------------- 性能加速比测试 --------------------")
+    num_iterations = 100
+    torch.cuda.synchronize(); start_time = time.time()
+    for _ in range(num_iterations):
+        _ = torch_model(*inputs)
+    torch.cuda.synchronize(); torch_time = (time.time() - start_time) / num_iterations
+
+    torch.cuda.synchronize(); start_time = time.time()
+    for _ in range(num_iterations):
+        _ = cuda_model(*inputs)
+    torch.cuda.synchronize(); cuda_time = (time.time() - start_time) / num_iterations
+
+    print(f"PyTorch RMSNorm-SiLU 平均执行时间: {torch_time:.6f} 秒")
+    print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f} 秒")
+    speedup = torch_time / cuda_time if cuda_time > 0 else 0
+    if cuda_time > 0:
+        print(f"加速比 (Speedup): {speedup:.2f}x")
+    else:
+        print("CUDA 内核执行时间为0，无法计算加速比。")
+    return precision_flag, speedup
+
+if __name__ == "__main__":
+    run_benchmark()
+

S1/Ljy123_#70/torchcode.py

@@ -0,0 +1,29 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Model(nn.Module):
+    def __init__(self, gamma: torch.Tensor, beta: torch.Tensor, eps: float):
+        super(Model, self).__init__()
+        self.register_buffer("gamma", gamma)
+        self.register_buffer("beta", beta)
+        self.register_buffer("eps", torch.tensor(float(eps), dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r = torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+        v = x / r
+        u = v * self.gamma + self.beta
+        return F.silu(u)
+
+batch_size = 16
+dim = 16384
+
+def get_inputs():
+    x = torch.randn(batch_size, dim)
+    return [x]
+
+def get_init_inputs():
+    gamma = torch.randn(dim)
+    beta = torch.randn(dim)
+    return [gamma, beta, 1e-5]
+