optimized layernorm

S1/Ljy123_#5/layernorm_cudacode.py

@@ -0,0 +1,143 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+
+# 简化但高效的LayerNorm CUDA实现
+layernorm_source = """
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+#include <float.h>
+
+// 简化但高效的LayerNorm内核
+__global__ void layernorm_forward_kernel(
+    const float* __restrict__ input,
+    const float* __restrict__ gamma,
+    const float* __restrict__ beta,
+    float* __restrict__ output,
+    int batch_size,
+    int hidden_size,
+    float eps) {
+    
+    extern __shared__ float shared_mem[];
+    float* shared_sum = shared_mem;
+    float* shared_sum_sq = &shared_mem[blockDim.x];
+    
+    int batch_idx = blockIdx.x;
+    int tid = threadIdx.x;
+    
+    // 第一步：每个线程计算局部统计量
+    float thread_sum = 0.0f;
+    float thread_sum_sq = 0.0f;
+    
+    for (int i = tid; i < hidden_size; i += blockDim.x) {
+        float val = input[batch_idx * hidden_size + i];
+        thread_sum += val;
+        thread_sum_sq += val * val;
+    }
+    
+    shared_sum[tid] = thread_sum;
+    shared_sum_sq[tid] = thread_sum_sq;
+    __syncthreads();
+    
+    // 块内归约求总和
+    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+        if (tid < stride) {
+            shared_sum[tid] += shared_sum[tid + stride];
+            shared_sum_sq[tid] += shared_sum_sq[tid + stride];
+        }
+        __syncthreads();
+    }
+    
+    // 计算全局统计量
+    if (tid == 0) {
+        float total_sum = shared_sum[0];
+        float total_sum_sq = shared_sum_sq[0];
+        float global_mean = total_sum / hidden_size;
+        float global_variance = (total_sum_sq / hidden_size) - (global_mean * global_mean);
+        
+        // 计算逆标准差
+        float inv_std = rsqrtf(global_variance + eps);
+        
+        // 存储到共享内存供所有线程使用
+        shared_sum[0] = global_mean;
+        shared_sum_sq[0] = inv_std;
+    }
+    __syncthreads();
+    
+    float global_mean = shared_sum[0];
+    float inv_std = shared_sum_sq[0];
+    
+    // 应用LayerNorm
+    for (int i = tid; i < hidden_size; i += blockDim.x) {
+        float val = input[batch_idx * hidden_size + i];
+        float normalized = (val - global_mean) * inv_std;
+        output[batch_idx * hidden_size + i] = normalized * gamma[i] + beta[i];
+    }
+}
+
+torch::Tensor layernorm_cuda_forward(
+    torch::Tensor input,
+    torch::Tensor gamma,
+    torch::Tensor beta,
+    float eps) {
+    
+    auto batch_size = input.size(0);
+    auto hidden_size = input.size(-1);
+    
+    auto output = torch::empty_like(input);
+    
+    // 使用固定的线程块大小
+    int block_size = 256;
+    int num_blocks = batch_size;
+    int shared_mem_size = 2 * block_size * sizeof(float);
+    
+    layernorm_forward_kernel<<<num_blocks, block_size, shared_mem_size>>>(
+        input.data_ptr<float>(),
+        gamma.data_ptr<float>(),
+        beta.data_ptr<float>(),
+        output.data_ptr<float>(),
+        batch_size,
+        hidden_size,
+        eps
+    );
+    
+    return output;
+}
+"""
+
+layernorm_cpp_source = """
+torch::Tensor layernorm_cuda_forward(torch::Tensor input, torch::Tensor gamma, torch::Tensor beta, float eps);
+"""
+
+# 编译内联CUDA代码
+cuda_available = True
+try:
+    layernorm_cuda = load_inline(
+        name="layernorm_cuda",
+        cpp_sources=layernorm_cpp_source,
+        cuda_sources=layernorm_source,
+        functions=["layernorm_cuda_forward"],
+        verbose=True
+    )
+except Exception as e:
+    print(f"CUDA扩展加载失败: {e}")
+    cuda_available = False
+    layernorm_cuda = None
+
+class ModelNew(torch.nn.Module):
+    def __init__(self, normalized_shape, eps=1e-5):
+        super(ModelNew, self).__init__()
+        self.normalized_shape = normalized_shape
+        self.eps = eps
+        self.weight = torch.nn.Parameter(torch.ones(normalized_shape))
+        self.bias = torch.nn.Parameter(torch.zeros(normalized_shape))
+
+    def forward(self, x):
+        if cuda_available and layernorm_cuda is not None:
+            # 使用真正的CUDA内核
+            return layernorm_cuda.layernorm_cuda_forward(x, self.weight, self.bias, self.eps)
+        else:
+            # CPU回退实现，与PyTorch实现保持一致
+            mean = x.mean(-1, keepdim=True)
+            var = x.var(-1, unbiased=False, keepdim=True)
+            normalized = (x - mean) / torch.sqrt(var + self.eps)
+            return normalized * self.weight + self.bias

S1/Ljy123_#5/layernorm_torchcode.py

@@ -0,0 +1,21 @@
+import torch
+import torch.nn as nn
+
+class Model(nn.Module):
+    def __init__(self, normalized_shape, eps=1e-5):
+        super(Model, self).__init__()
+        self.normalized_shape = normalized_shape
+        self.eps = eps
+        
+        # 使用PyTorch内置的LayerNorm
+        self.layer_norm = nn.LayerNorm(normalized_shape, eps=eps)
+    
+    def forward(self, x):
+        return self.layer_norm(x)
+
+def get_inputs():
+    # 使用中等大小的输入尺寸
+    return [torch.randn(8, 8192)]
+
+def get_init_inputs():
+    return [8192]

S1/Ljy123_#5/prompt.txt

@@ -0,0 +1,27 @@
+Write a custom CUDA kernel for Layer Normalization.
+
+The standard LayerNorm operation is defined as:
+
+y = (x - E[x]) / sqrt(Var[x] + epsilon) * gamma + beta
+
+Where:
+- x is the input tensor
+- E[x] is the mean of x
+- Var[x] is the variance of x
+- epsilon is a small value for numerical stability
+- gamma and beta are learnable affine parameters
+
+You should fuse the calculation of mean, variance, and the normalization into a single CUDA kernel. This avoids multiple passes over the data and reduces memory bandwidth usage.
+
+You are given the following architecture:
+
+import torch
+import torch.nn as nn
+
+class Model(nn.Module):
+    def __init__(self, normalized_shape, eps=1e-5):
+        super(Model, self).__init__()
+        self.layer_norm = nn.LayerNorm(normalized_shape, eps=eps)
+    
+    def forward(self, x):
+        return self.layer_norm(x)

S1/Ljy123_#5/run_code.py

@@ -0,0 +1,76 @@
+###########################################################
+# 性能和精度验证程序
+###########################################################
+import torch
+import torch.nn as nn
+import time
+from layernorm_torchcode import Model,get_inputs,get_init_inputs
+from layernorm_cudacode import ModelNew
+
+def run_benchmark():
+    # 检查 CUDA 是否可用
+    if not torch.cuda.is_available():
+        print("CUDA 不可用，请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
+        return
+    else:
+        device = torch.device("cuda")
+
+    # 初始化模型
+    init_inputs = get_init_inputs()
+    init_inputs = [
+        x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs
+    ]
+    inputs = get_inputs()
+    inputs = [
+        x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs
+    ]
+
+    torch_model = Model(*init_inputs).cuda()
+    cuda_model = ModelNew(*init_inputs).cuda()
+
+    torch_model.eval()
+    cuda_model.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 
+    
+    # PyTorch 模型计时
+    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
+    
+    # 自定义 CUDA 内核计时
+    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 平均执行时间: {torch_time:.6f} 秒")
+    print(f"自定义 CUDA 内核 平均执行时间: {cuda_time:.6f} 秒")
+    speedup = 0
+    if cuda_time > 0:
+        speedup = torch_time / cuda_time
+        print(f"加速比 (Speedup): {speedup:.2f}x")
+    else:
+        print("CUDA 内核执行时间为0，无法计算加速比。")
+    return precision_flag,speedup
+if __name__ == "__main__":
+    result = run_benchmark()
+    if result is not None:
+        precision_flag,speedup = result