Merge pull request 'finish layernorm #1' (#107) from Ljy123/GPUCodeForces:layer into main

S1/Ljy123_#1/layernorm_cudacode.py

@@ -0,0 +1,190 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+
+# LayerNorm的CUDA实现
+layernorm_source = """
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cub/cub.cuh>
+
+// 使用高度优化的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;
+    
+    // 使用向量化加载，每个线程处理4个元素
+    float4 thread_sum = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
+    float4 thread_sum_sq = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
+    
+    // 向量化处理，提高内存带宽利用率
+    for (int i = tid * 4; i < hidden_size; i += blockDim.x * 4) {
+        if (i + 3 < hidden_size) {
+            float4 vals = *reinterpret_cast<const float4*>(&input[batch_idx * hidden_size + i]);
+            thread_sum.x += vals.x; thread_sum_sq.x += vals.x * vals.x;
+            thread_sum.y += vals.y; thread_sum_sq.y += vals.y * vals.y;
+            thread_sum.z += vals.z; thread_sum_sq.z += vals.z * vals.z;
+            thread_sum.w += vals.w; thread_sum_sq.w += vals.w * vals.w;
+        } else {
+            // 处理剩余元素
+            for (int j = 0; j < 4 && i + j < hidden_size; j++) {
+                float val = input[batch_idx * hidden_size + i + j];
+                thread_sum.x += val; thread_sum_sq.x += val * val;
+            }
+        }
+    }
+    
+    // 归约线程内的4个分量
+    float thread_total_sum = thread_sum.x + thread_sum.y + thread_sum.z + thread_sum.w;
+    float thread_total_sum_sq = thread_sum_sq.x + thread_sum_sq.y + thread_sum_sq.z + thread_sum_sq.w;
+    
+    shared_sum[tid] = thread_total_sum;
+    shared_sum_sq[tid] = thread_total_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 * 4; i < hidden_size; i += blockDim.x * 4) {
+        if (i + 3 < hidden_size) {
+            float4 vals = *reinterpret_cast<const float4*>(&input[batch_idx * hidden_size + i]);
+            float4 normalized;
+            normalized.x = (vals.x - global_mean) * inv_std;
+            normalized.y = (vals.y - global_mean) * inv_std;
+            normalized.z = (vals.z - global_mean) * inv_std;
+            normalized.w = (vals.w - global_mean) * inv_std;
+            
+            float4 result;
+            result.x = normalized.x * gamma[i] + beta[i];
+            result.y = normalized.y * gamma[i+1] + beta[i+1];
+            result.z = normalized.z * gamma[i+2] + beta[i+2];
+            result.w = normalized.w * gamma[i+3] + beta[i+3];
+            
+            *reinterpret_cast<float4*>(&output[batch_idx * hidden_size + i]) = result;
+        } else {
+            // 处理剩余元素
+            for (int j = 0; j < 4 && i + j < hidden_size; j++) {
+                float val = input[batch_idx * hidden_size + i + j];
+                float normalized = (val - global_mean) * inv_std;
+                output[batch_idx * hidden_size + i + j] = normalized * gamma[i + j] + beta[i + j];
+            }
+        }
+    }
+}
+
+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);
+    
+    // 优化线程块大小，根据hidden_size动态调整，使用更激进的优化
+    int block_size = 256;  // 固定使用256线程，适合大多数GPU架构
+    if (hidden_size <= 512) {
+        block_size = 128;
+    } else if (hidden_size <= 1024) {
+        block_size = 256;
+    } else {
+        block_size = 512;
+    }
+    
+    // 确保block_size不超过硬件限制
+    block_size = min(1024, max(32, block_size));
+    
+    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_#1/layernorm_torchcode.py

@@ -0,0 +1,23 @@
+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):
+        # 使用PyTorch内置LayerNorm，这是标准的实现方式
+        return self.layer_norm(x)
+
+def get_inputs():
+    # 使用更大的输入尺寸以获得更好的性能对比
+    # 增加batch size和hidden size，模拟真实场景
+    return [torch.randn(16, 16384)]
+
+def get_init_inputs():
+    return [16384]

S1/Ljy123_#1/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_#1/run_code.py

@@ -0,0 +1,74 @@
+###########################################################
+# 性能和精度验证程序
+###########################################################
+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__":
+    precision_flag,speedup = run_benchmark()