From 2f8387c7f5a99f8e1e06dfa3d106db79b5417dd4 Mon Sep 17 00:00:00 2001 From: Ljy123 Date: Thu, 13 Nov 2025 21:31:19 +0800 Subject: [PATCH] finish layernorm #1 --- S1/Ljy123_#1/layernorm_cudacode.py | 190 ++++++++++++++++++++++++++++ S1/Ljy123_#1/layernorm_torchcode.py | 23 ++++ S1/Ljy123_#1/prompt.txt | 27 ++++ S1/Ljy123_#1/run_code.py | 74 +++++++++++ 4 files changed, 314 insertions(+) create mode 100644 S1/Ljy123_#1/layernorm_cudacode.py create mode 100644 S1/Ljy123_#1/layernorm_torchcode.py create mode 100644 S1/Ljy123_#1/prompt.txt create mode 100644 S1/Ljy123_#1/run_code.py diff --git a/S1/Ljy123_#1/layernorm_cudacode.py b/S1/Ljy123_#1/layernorm_cudacode.py new file mode 100644 index 0000000..4c48f24 --- /dev/null +++ b/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 +#include +#include +#include + +// 使用高度优化的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(&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(&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(&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<<>>( + input.data_ptr(), + gamma.data_ptr(), + beta.data_ptr(), + output.data_ptr(), + 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 \ No newline at end of file diff --git a/S1/Ljy123_#1/layernorm_torchcode.py b/S1/Ljy123_#1/layernorm_torchcode.py new file mode 100644 index 0000000..b40b565 --- /dev/null +++ b/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] \ No newline at end of file diff --git a/S1/Ljy123_#1/prompt.txt b/S1/Ljy123_#1/prompt.txt new file mode 100644 index 0000000..0fad308 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/S1/Ljy123_#1/run_code.py b/S1/Ljy123_#1/run_code.py new file mode 100644 index 0000000..328a1d7 --- /dev/null +++ b/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() \ No newline at end of file