From ac05feb06c23adebd267912eb579944545d97732 Mon Sep 17 00:00:00 2001 From: gsd <2396583337@qq.com> Date: Fri, 14 Nov 2025 10:54:10 +0800 Subject: [PATCH] fix layernorm #3 --- S1/gsd123_#3/layernorm_cuda.py | 205 ++++++++++++++++++++++++++++++++ S1/gsd123_#3/layernorm_torch.py | 52 ++++++++ S1/gsd123_#3/prompt.txt | 82 +++++++++++++ S1/gsd123_#3/run_code.py | 78 ++++++++++++ 4 files changed, 417 insertions(+) create mode 100644 S1/gsd123_#3/layernorm_cuda.py create mode 100644 S1/gsd123_#3/layernorm_torch.py create mode 100644 S1/gsd123_#3/prompt.txt create mode 100644 S1/gsd123_#3/run_code.py diff --git a/S1/gsd123_#3/layernorm_cuda.py b/S1/gsd123_#3/layernorm_cuda.py new file mode 100644 index 0000000..da9d57c --- /dev/null +++ b/S1/gsd123_#3/layernorm_cuda.py @@ -0,0 +1,205 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline + +# LayerNorm CUDA 实现 - 增强优化版本 +layernorm_source = """ +#include +#include +#include + +#define WARP_SIZE 32 + +// Warp级归约函数 +__device__ __forceinline__ float warp_reduce_sum(float val) { + for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + return val; +} + +__global__ void layernorm_kernel_optimized( + const float* __restrict__ x, + const float* __restrict__ weight, + const float* __restrict__ bias, + float* __restrict__ y, + int batch, + int features, + float eps +) { + int row = blockIdx.x; + if (row >= batch) return; + + int tid = threadIdx.x; + int warp_id = tid / WARP_SIZE; + int lane_id = tid % WARP_SIZE; + int num_warps = (blockDim.x + WARP_SIZE - 1) / WARP_SIZE; + + __shared__ float s_mean; + __shared__ float s_inv_std; + __shared__ float s_warp_sums[32]; // 支持最多1024个线程 + __shared__ float s_warp_sum_sqs[32]; + + const float* x_row = x + row * features; + float* y_row = y + row * features; + + // 第一步：并行计算均值和方差 + float thread_sum = 0.0f; + float thread_sum_sq = 0.0f; + + // 使用向量化加载（如果特征数是4的倍数） + if (features % 4 == 0) { + for (int i = tid * 4; i < features; i += blockDim.x * 4) { + float4 vec = *reinterpret_cast(x_row + i); + thread_sum += vec.x + vec.y + vec.z + vec.w; + thread_sum_sq += vec.x * vec.x + vec.y * vec.y + vec.z * vec.z + vec.w * vec.w; + } + } else { + // 标量版本 + for (int i = tid; i < features; i += blockDim.x) { + float v = x_row[i]; + thread_sum += v; + thread_sum_sq += v * v; + } + } + + // Warp级归约 + float warp_sum = warp_reduce_sum(thread_sum); + float warp_sum_sq = warp_reduce_sum(thread_sum_sq); + + // 将warp结果写入共享内存 + if (lane_id == 0) { + s_warp_sums[warp_id] = warp_sum; + s_warp_sum_sqs[warp_id] = warp_sum_sq; + } + __syncthreads(); + + // Block级归约（在第一个warp中完成） + if (warp_id == 0) { + float block_sum = (lane_id < num_warps) ? s_warp_sums[lane_id] : 0.0f; + float block_sum_sq = (lane_id < num_warps) ? s_warp_sum_sqs[lane_id] : 0.0f; + + block_sum = warp_reduce_sum(block_sum); + block_sum_sq = warp_reduce_sum(block_sum_sq); + + if (lane_id == 0) { + float mean = block_sum / features; + float var = (block_sum_sq / features) - (mean * mean); + s_mean = mean; + s_inv_std = rsqrtf(fmaxf(var, 0.0f) + eps); + } + } + __syncthreads(); + + float mean = s_mean; + float inv_std = s_inv_std; + + // 第二步：应用归一化（向量化存储） + if (features % 4 == 0) { + for (int i = tid * 4; i < features; i += blockDim.x * 4) { + float4 vec = *reinterpret_cast(x_row + i); + float4 w_vec = *reinterpret_cast(weight + i); + float4 b_vec = *reinterpret_cast(bias + i); + + vec.x = (vec.x - mean) * inv_std * w_vec.x + b_vec.x; + vec.y = (vec.y - mean) * inv_std * w_vec.y + b_vec.y; + vec.z = (vec.z - mean) * inv_std * w_vec.z + b_vec.z; + vec.w = (vec.w - mean) * inv_std * w_vec.w + b_vec.w; + + *reinterpret_cast(y_row + i) = vec; + } + } else { + // 标量版本 + for (int i = tid; i < features; i += blockDim.x) { + float v = x_row[i]; + float w = weight[i]; + float b = bias[i]; + y_row[i] = (v - mean) * inv_std * w + b; + } + } +} + +torch::Tensor layernorm_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, float eps) { + TORCH_CHECK(x.is_cuda(), "x 必须是 CUDA 张量"); + TORCH_CHECK(weight.is_cuda(), "weight 必须是 CUDA 张量"); + TORCH_CHECK(bias.is_cuda(), "bias 必须是 CUDA 张量"); + TORCH_CHECK(x.dim() == 2, "当前内核仅支持二维输入张量"); + TORCH_CHECK(weight.dim() == 1, "LayerNorm 权重必须是一维向量"); + TORCH_CHECK(bias.dim() == 1, "LayerNorm 偏置必须是一维向量"); + TORCH_CHECK(x.size(1) == weight.size(0), "输入最后一维与权重长度不匹配"); + TORCH_CHECK(weight.size(0) == bias.size(0), "权重和偏置长度必须相同"); + + int batch = x.size(0); + int features = x.size(1); + + auto y = torch::empty_like(x); + + // 智能线程配置 + int threads; + if (features <= 64) { + threads = 64; + } else if (features <= 256) { + threads = 128; + } else if (features <= 1024) { + threads = 256; + } else { + threads = 512; + } + + // 确保线程数是warp大小的倍数 + threads = (threads + WARP_SIZE - 1) / WARP_SIZE * WARP_SIZE; + threads = min(threads, features); + + // 计算共享内存大小 + size_t shared_mem = 2 * ((threads + WARP_SIZE - 1) / WARP_SIZE) * sizeof(float) + 2 * sizeof(float); + + layernorm_kernel_optimized<<>>( + x.data_ptr(), + weight.data_ptr(), + bias.data_ptr(), + y.data_ptr(), + batch, + features, + eps + ); + + return y; +} +""" + +layernorm_cpp_source = """ +torch::Tensor layernorm_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, float eps); +""" + +# 编译 CUDA 代码 +layernorm = load_inline( + name="layernorm", + cpp_sources=layernorm_cpp_source, + cuda_sources=layernorm_source, + functions=["layernorm_cuda"], + extra_cuda_cflags=["-O3", "--use_fast_math"], + verbose=True +) + + +class ModelNew(nn.Module): + def __init__(self, eps: float = 1e-5): + super(ModelNew, self).__init__() + self.eps = eps + # 在forward中动态确定特征维度 + self.weight = None + self.bias = None + self.layernorm = layernorm + self._initialized = False + + def forward(self, x): + # 动态初始化权重和偏置（只初始化一次） + if self.weight is None: + feature_dim = x.size(1) + self.weight = nn.Parameter(torch.ones(feature_dim, device=x.device)) + self.bias = nn.Parameter(torch.zeros(feature_dim, device=x.device)) + self._initialized = True + + + + return self.layernorm.layernorm_cuda(x, self.weight, self.bias, self.eps) \ No newline at end of file diff --git a/S1/gsd123_#3/layernorm_torch.py b/S1/gsd123_#3/layernorm_torch.py new file mode 100644 index 0000000..a053e6f --- /dev/null +++ b/S1/gsd123_#3/layernorm_torch.py @@ -0,0 +1,52 @@ +import torch +import torch.nn as nn + +class Model(nn.Module): + """ + Simple model that performs LayerNorm normalization using PyTorch's built-in nn.LayerNorm. + """ + + def __init__(self, normalized_shape=None, eps=1e-5, elementwise_affine=True): + super(Model, self).__init__() + # 如果未指定normalized_shape，将在forward中动态设置 + self.normalized_shape = normalized_shape + self.eps = eps + self.elementwise_affine = elementwise_affine + self.layernorm = None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Applies LayerNorm to the input tensor. + + Args: + x (torch.Tensor): Input tensor of any shape. + + Returns: + torch.Tensor: Output tensor with LayerNorm applied, same shape as input. + """ + # 如果layernorm未初始化，根据输入形状动态创建 + if self.layernorm is None: + if self.normalized_shape is None: + # 默认对最后一个维度进行归一化 + self.normalized_shape = x.shape[1:] + self.layernorm = nn.LayerNorm( + normalized_shape=self.normalized_shape, + eps=self.eps, + elementwise_affine=self.elementwise_affine + ).to(x.device) + + return self.layernorm(x) + + +batch_size = 16 +dim = 16384 + + +def get_inputs(): + x = torch.randn(batch_size, dim) + return [x] + + +def get_init_inputs(): + # 可以传入归一化形状、eps等参数，保持向后兼容 + return [] # 使用默认参数 \ No newline at end of file diff --git a/S1/gsd123_#3/prompt.txt b/S1/gsd123_#3/prompt.txt new file mode 100644 index 0000000..016892c --- /dev/null +++ b/S1/gsd123_#3/prompt.txt @@ -0,0 +1,82 @@ +LayerNorm CUDA Implementation - Enhanced Optimized Version + +Key optimization techniques used in this implementation: + +1.Warp-Level Parallel Reduction: Implements efficient warp-level reduction for mean and variance calculations using warp shuffle operations +2.Vectorized Memory Access: Utilizes float4 vector loads/stores for coalesced memory access when feature dimension is divisible by 4 +3.Shared Memory Hierarchy: Employs multi-level shared memory for intermediate results between warp and block levels +4.Dynamic Thread Configuration: Automatically adjusts thread block size based on feature dimension for optimal occupancy +5.Numerical Stability: Maintains numerical precision with robust variance calculation and epsilon handling + +Bank Conflict Avoidance: Carefully structures shared memory access patterns to minimize bank conflicts + +The custom kernel eliminates multiple memory passes by computing mean, variance, and normalization in a single fused operation with optimized memory hierarchy usage across warp, shared, and global memory levels. + +Technical Features: + +1.Warp Reduction: Efficient 32-thread warp reduction using __shfl_down_sync +2.Vectorization: Automatic fallback between vectorized (float4) and scalar operations +3.Smart Block Sizing: Adaptive thread configuration (64-512 threads) based on feature dimension +4.Memory Coalescing: Organized memory access patterns for maximum bandwidth utilization +5.Fused Operations: Combines statistics computation and normalization in one kernel + +Performance Benefits: + +1.Reduces global memory traffic by processing entire LayerNorm operation in-place +2.Eliminates intermediate tensor allocations between mean/variance calculations +3.Optimizes for various feature dimensions through adaptive thread configuration +4.Leverages CUDA memory hierarchy for maximum data reuse + +Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: +import torch +import torch.nn as nn + +class Model(nn.Module): + """ + Simple model that performs LayerNorm normalization using PyTorch's built-in nn.LayerNorm. + """ + + def __init__(self, normalized_shape=None, eps=1e-5, elementwise_affine=True): + super(Model, self).__init__() + # 如果未指定normalized_shape，将在forward中动态设置 + self.normalized_shape = normalized_shape + self.eps = eps + self.elementwise_affine = elementwise_affine + self.layernorm = None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Applies LayerNorm to the input tensor. + + Args: + x (torch.Tensor): Input tensor of any shape. + + Returns: + torch.Tensor: Output tensor with LayerNorm applied, same shape as input. + """ + # 如果layernorm未初始化，根据输入形状动态创建 + if self.layernorm is None: + if self.normalized_shape is None: + # 默认对最后一个维度进行归一化 + self.normalized_shape = x.shape[1:] + self.layernorm = nn.LayerNorm( + normalized_shape=self.normalized_shape, + eps=self.eps, + elementwise_affine=self.elementwise_affine + ).to(x.device) + + return self.layernorm(x) + + +batch_size = 16 +dim = 16384 + + +def get_inputs(): + x = torch.randn(batch_size, dim) + return [x] + + +def get_init_inputs(): + # 可以传入归一化形状、eps等参数，保持向后兼容 + return [] # 使用默认参数 \ No newline at end of file diff --git a/S1/gsd123_#3/run_code.py b/S1/gsd123_#3/run_code.py new file mode 100644 index 0000000..831c751 --- /dev/null +++ b/S1/gsd123_#3/run_code.py @@ -0,0 +1,78 @@ +import torch +import time +from layernorm_torch import Model, get_inputs, get_init_inputs +from layernorm_cuda import ModelNew + +def run_benchmark(): + if not torch.cuda.is_available(): + print("CUDA 不可用") + return + + device = torch.device("cuda") + + # 准备输入数据 + inputs = [x.cuda(device=device) for x in get_inputs()] + init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()] + + # 初始化模型 + torch_model = Model(*init_inputs).cuda() + cuda_model = ModelNew(*init_inputs).cuda() + + torch_model.eval() + cuda_model.eval() + + print("-------------------- 精度对齐验证 --------------------") + with torch.no_grad(): + # 预热GPU + _ = torch_model(*inputs) + _ = cuda_model(*inputs) + + # 正式测试 + output_torch = torch_model(*inputs) + output_cuda = cuda_model(*inputs) + + # 精度验证 + abs_diff = torch.abs(output_torch - output_cuda) + max_diff = torch.max(abs_diff).item() + mean_diff = torch.mean(abs_diff).item() + + if max_diff < 1e-4 and mean_diff < 1e-5: + print(f"✅ 精度对齐：最大误差 {max_diff:.6f}，平均误差 {mean_diff:.6f}") + precision_flag = True + else: + print(f"❌ 精度不一致：最大误差 {max_diff:.6f}，平均误差 {mean_diff:.6f}") + precision_flag = False + + print("\n-------------------- 性能加速比测试 --------------------") + num_iterations = 100 + + # 预热GPU + for _ in range(10): + _ = torch_model(*inputs) + _ = cuda_model(*inputs) + + # 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内置Swish平均执行时间: {torch_time:.6f}秒") + print(f"自定义CUDA Swish平均执行时间: {cuda_time:.6f}秒") + speedup = torch_time / cuda_time if cuda_time > 0 else 0 + print(f"加速比 (Speedup): {speedup:.2f}x") + + return precision_flag, speedup + +if __name__ == "__main__": + precision_flag, speedup = run_benchmark() \ No newline at end of file