finish layernormgelu #69

S1/Ljy123_#69/cudacode.py

@@ -0,0 +1,156 @@
+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;
+}
+
+__device__ __forceinline__ float gelu(float x){
+    float x3 = x * x * x;
+    float t = 0.7978845608028654f * (x + 0.044715f * x3);
+    float h = tanhf(t);
+    return 0.5f * x * (1.0f + h);
+}
+
+__global__ __launch_bounds__(1024) void layernorm_gelu_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;
+    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;
+    }
+    // Reduction
+    float f_mean = mean;
+    float f_M2 = M2;
+    for(int offset=16; offset>0; offset/=2){
+        float mean_other = __shfl_down_sync(0xffffffff, f_mean, offset);
+        float M2_other = __shfl_down_sync(0xffffffff, f_M2, offset);
+        float cnt_other = __shfl_down_sync(0xffffffff, count, offset);
+        if(cnt_other > 0.0f){
+            float delta = mean_other - f_mean;
+            float tot = count + cnt_other;
+            f_M2 += M2_other + delta * delta * (count * cnt_other) / tot;
+            f_mean += (mean_other - f_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] = f_mean;
+        s_M2_warp[warp_id] = f_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;
+        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();
+    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 v0 = (xv.x - s_mean) * s_invstd; float u0 = v0 * gv.x + bv.x; yv.x = gelu(u0);
+            float v1 = (xv.y - s_mean) * s_invstd; float u1 = v1 * gv.y + bv.y; yv.y = gelu(u1);
+            float v2 = (xv.z - s_mean) * s_invstd; float u2 = v2 * gv.z + bv.z; yv.z = gelu(u2);
+            float v3 = (xv.w - s_mean) * s_invstd; float u3 = v3 * gv.w + bv.w; yv.w = gelu(u3);
+            reinterpret_cast<float4*>(yr)[i / 4] = yv;
+        }
+        for(int i = (D4 + tid); i < D; i += blockDim.x){
+            float v = (__ldg(xr + i) - s_mean) * s_invstd;
+            float u = v * __ldg(gamma + i) + __ldg(beta + i);
+            yr[i] = gelu(u);
+        }
+    } else {
+        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);
+            yr[i] = gelu(u);
+        }
+    }
+}
+
+torch::Tensor layernorm_gelu_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);
+    layernorm_gelu_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_gelu_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor eps);
+"""
+
+ops = load_inline(
+    name="layernorm_gelu",
+    cpp_sources=cpp_source,
+    cuda_sources=source,
+    functions=["layernorm_gelu_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.layernorm_gelu_cuda(x, self.gamma, self.beta, self.eps)

S1/Ljy123_#69/prompt.txt

@@ -0,0 +1 @@
+融合算子：LayerNorm+Affine+GELU（tanh 近似），一次内核完成归一化、仿射与激活。

S1/Ljy123_#69/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-GELU 平均执行时间: {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_#69/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:
+        m = x.mean(dim=-1, keepdim=True)
+        v = ((x - m) * (x - m)).mean(dim=-1, keepdim=True)
+        u = ((x - m) / torch.sqrt(v + self.eps)) * self.gamma + self.beta
+        return F.gelu(u, approximate='tanh')
+
+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]
+