finish layernorm-silu #64

S1/Ljy123_#64/cudacode.py

@@ -0,0 +1,130 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+
+source = """
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+#include <algorithm>
+
+__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__(256) void layernorm_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 s_mean;
+    __shared__ float s_invstd;
+    // Welford per-thread
+    float count = 0.0f;
+    float mean = 0.0f;
+    float M2 = 0.0f;
+    for(int i = tid; i < D; i += blockDim.x){
+        float xk = xr[i];
+        count += 1.0f;
+        float delta = xk - mean;
+        mean += delta / count;
+        float delta2 = xk - mean;
+        M2 += delta * delta2;
+    }
+    // Warp reduce Welford
+    for(int offset=16; offset>0; offset/=2){
+        float mean_other = __shfl_down_sync(0xffffffff, mean, offset);
+        float M2_other = __shfl_down_sync(0xffffffff, M2, offset);
+        float cnt_other = __shfl_down_sync(0xffffffff, count, offset);
+        if(cnt_other > 0.0f){
+            float delta = mean_other - mean;
+            float tot = count + cnt_other;
+            M2 += M2_other + delta * delta * (count * cnt_other) / tot;
+            mean += (mean_other - mean) * (cnt_other / tot);
+            count = tot;
+        }
+    }
+    __shared__ float s_mean_warp[32];
+    __shared__ float s_M2_warp[32];
+    __shared__ float s_cnt_warp[32];
+    if(lane == 0){
+        s_mean_warp[warp_id] = mean;
+        s_M2_warp[warp_id] = M2;
+        s_cnt_warp[warp_id] = count;
+    }
+    __syncthreads();
+    if(warp_id == 0){
+        float mean0 = (lane < warps) ? s_mean_warp[lane] : 0.0f;
+        float M20 = (lane < warps) ? s_M2_warp[lane] : 0.0f;
+        float cnt0 = (lane < warps) ? s_cnt_warp[lane] : 0.0f;
+        // reduce across warps
+        for(int offset=16; offset>0; offset/=2){
+            float mean_other = __shfl_down_sync(0xffffffff, mean0, offset);
+            float M2_other = __shfl_down_sync(0xffffffff, M20, offset);
+            float cnt_other = __shfl_down_sync(0xffffffff, cnt0, offset);
+            if(cnt_other > 0.0f){
+                float delta = mean_other - mean0;
+                float tot = cnt0 + cnt_other;
+                M20 += M2_other + delta * delta * (cnt0 * cnt_other) / tot;
+                mean0 += (mean_other - mean0) * (cnt_other / tot);
+                cnt0 = tot;
+            }
+        }
+        if(lane == 0){
+            float var = M20 / fmaxf(cnt0, 1.0f);
+            s_mean = mean0;
+            s_invstd = rsqrtf(fmaxf(var, 0.0f) + eps);
+        }
+    }
+    __syncthreads();
+    for(int i = tid; i < D; i += blockDim.x){
+        float v = (__ldg(xr + i) - s_mean) * s_invstd;
+        float u = v * __ldg(gamma + i) + __ldg(beta + i);
+        float g = __fdividef(1.0f, 1.0f + __expf(-u));
+        yr[i] = u * g;
+    }
+}
+
+torch::Tensor layernorm_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 = 256;
+    dim3 grid(B);
+    layernorm_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 layernorm_silu_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor eps);
+"""
+
+ops = load_inline(
+    name="layernorm_silu",
+    cpp_sources=cpp_source,
+    cuda_sources=source,
+    functions=["layernorm_silu_cuda"],
+    extra_cflags=["-O3","-std=c++17"],
+    extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17"],
+    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.layernorm_silu_cuda(x, self.gamma, self.beta, self.eps)

S1/Ljy123_#64/prompt.txt

@@ -0,0 +1,15 @@
+Objective: Fused LayerNorm + SiLU CUDA kernel with rtol=1e-3 accuracy and ≥1.3x speedup.
+
+Computation:
+- Row layer normalization (mean/var over dim) with epsilon, then affine transform, then SiLU: `y = u * sigmoid(u)`.
+
+Numerical Method:
+- Use Welford’s online algorithm for mean/variance to match PyTorch precision and avoid catastrophic cancellation.
+- Perform warp-level reductions and shared-memory aggregation across warps.
+
+Performance:
+- 512 threads per block; coalesced loads and parallel write-back.
+- Fuse normalization, affine, and activation to minimize memory traffic.
+
+Benchmark:
+- Batch size: 16; Dim: 16384; 100 iterations; target speedup ≥1.3x.

S1/Ljy123_#64/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 LayerNorm-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_#64/torchcode.py

@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+
+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:
+        mean = x.mean(dim=1, keepdim=True)
+        var = x.var(dim=1, unbiased=False, keepdim=True)
+        xn = (x - mean) / torch.sqrt(var + self.eps)
+        u = xn * self.gamma + self.beta
+        g = torch.sigmoid(u)
+        return u * g
+
+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]
+