From 6d03e9b65dd29ffa052e7e39a7944fe35d69e0cb Mon Sep 17 00:00:00 2001 From: wut0n <3455534242@qq.com> Date: Wed, 10 Dec 2025 18:00:16 +0800 Subject: [PATCH] Feat:add layernorm+residual #81 --- S1/wut0n_#81/layernorm_residual_cudacode.py | 199 ++++++++++++++++ S1/wut0n_#81/layernorm_residual_torchcode.py | 23 ++ S1/wut0n_#81/prompt.txt | 229 +++++++++++++++++++ S1/wut0n_#81/run_code.py | 74 ++++++ 4 files changed, 525 insertions(+) create mode 100644 S1/wut0n_#81/layernorm_residual_cudacode.py create mode 100644 S1/wut0n_#81/layernorm_residual_torchcode.py create mode 100644 S1/wut0n_#81/prompt.txt create mode 100644 S1/wut0n_#81/run_code.py diff --git a/S1/wut0n_#81/layernorm_residual_cudacode.py b/S1/wut0n_#81/layernorm_residual_cudacode.py new file mode 100644 index 00000000..fce6a75f --- /dev/null +++ b/S1/wut0n_#81/layernorm_residual_cudacode.py @@ -0,0 +1,199 @@ +import torch +from torch.utils.cpp_extension import load_inline + +layernorm_residual_source = """ +#include +#include +#include + +struct Float4 { + float x, y, z, w; + __device__ __forceinline__ Float4() {} + __device__ __forceinline__ Float4(float f) : x(f), y(f), z(f), w(f) {} + __device__ __forceinline__ Float4(float x_, float y_, float z_, float w_) : x(x_), y(y_), z(z_), w(w_) {} +}; + +__device__ __forceinline__ Float4 load_float4(const float* addr) { + return *reinterpret_cast(addr); +} + +__device__ __forceinline__ void store_float4(float* addr, Float4 val) { + *reinterpret_cast(addr) = val; +} + +// 高效的warp reduction +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + val += __shfl_xor_sync(0xffffffff, val, mask); + } + return val; +} + + +__global__ void layernorm_residual_kernel_vectorized( + const float* __restrict__ x, + const float* __restrict__ residual, + const float* __restrict__ weight, + const float* __restrict__ bias, + float* __restrict__ output, + int N, int D, float eps +) { + // 每个warp处理一行，使用向量化访问 + int row = blockIdx.x * (blockDim.x / 32) + (threadIdx.x / 32); + int lane_id = threadIdx.x % 32; + + if (row >= N) return; + + const float* x_row = x + row * D; + const float* residual_row = residual + row * D; + float* output_row = output + row * D; + + float sum = 0.0f; + int vec_elems = D / 4; + int remaining = D % 4; + + for (int i = lane_id; i < vec_elems; i += 32) { + Float4 x_vec = load_float4(x_row + i * 4); + Float4 residual_vec = load_float4(residual_row + i * 4); + + sum += x_vec.x + residual_vec.x; + sum += x_vec.y + residual_vec.y; + sum += x_vec.z + residual_vec.z; + sum += x_vec.w + residual_vec.w; + } + + for (int i = vec_elems * 4 + lane_id; i < D; i += 32) { + sum += x_row[i] + residual_row[i]; + } + + sum = warp_reduce_sum(sum); + float mean = sum / D; + + float var_sum = 0.0f; + + for (int i = lane_id; i < vec_elems; i += 32) { + Float4 x_vec = load_float4(x_row + i * 4); + Float4 residual_vec = load_float4(residual_row + i * 4); + + float diff_x = x_vec.x + residual_vec.x - mean; + float diff_y = x_vec.y + residual_vec.y - mean; + float diff_z = x_vec.z + residual_vec.z - mean; + float diff_w = x_vec.w + residual_vec.w - mean; + + var_sum += diff_x * diff_x + diff_y * diff_y + diff_z * diff_z + diff_w * diff_w; + } + + for (int i = vec_elems * 4 + lane_id; i < D; i += 32) { + float diff = x_row[i] + residual_row[i] - mean; + var_sum += diff * diff; + } + + var_sum = warp_reduce_sum(var_sum); + float inv_std = rsqrtf(var_sum / D + eps); + + for (int i = lane_id; i < vec_elems; i += 32) { + Float4 x_vec = load_float4(x_row + i * 4); + Float4 residual_vec = load_float4(residual_row + i * 4); + Float4 weight_vec = load_float4(weight + i * 4); + Float4 bias_vec = load_float4(bias + i * 4); + + Float4 result; + float val_x = x_vec.x + residual_vec.x; + float val_y = x_vec.y + residual_vec.y; + float val_z = x_vec.z + residual_vec.z; + float val_w = x_vec.w + residual_vec.w; + + result.x = (val_x - mean) * inv_std * weight_vec.x + bias_vec.x; + result.y = (val_y - mean) * inv_std * weight_vec.y + bias_vec.y; + result.z = (val_z - mean) * inv_std * weight_vec.z + bias_vec.z; + result.w = (val_w - mean) * inv_std * weight_vec.w + bias_vec.w; + + store_float4(output_row + i * 4, result); + } + + for (int i = vec_elems * 4 + lane_id; i < D; i += 32) { + float val = x_row[i] + residual_row[i]; + output_row[i] = (val - mean) * inv_std * weight[i] + bias[i]; + } +} + +torch::Tensor layernorm_residual_cuda( + torch::Tensor x, + torch::Tensor residual, + torch::Tensor weight, + torch::Tensor bias, + float eps +) { + TORCH_CHECK(x.is_cuda() && x.is_contiguous()); + TORCH_CHECK(residual.is_cuda() && residual.is_contiguous()); + TORCH_CHECK(weight.is_cuda() && weight.is_contiguous()); + TORCH_CHECK(bias.is_cuda() && bias.is_contiguous()); + TORCH_CHECK(x.dim() == 2 && residual.dim() == 2); + TORCH_CHECK(x.sizes() == residual.sizes()); + TORCH_CHECK(weight.numel() == x.size(1)); + TORCH_CHECK(bias.numel() == x.size(1)); + + int N = x.size(0); + int D = x.size(1); + + auto output = torch::empty_like(x); + + int warps_per_block = 8; // 256 threads per block + int block_size = warps_per_block * 32; + int grid_size = (N + warps_per_block - 1) / warps_per_block; + + layernorm_residual_kernel_vectorized<<>>( + x.data_ptr(), + residual.data_ptr(), + weight.data_ptr(), + bias.data_ptr(), + output.data_ptr(), + N, D, eps + ); + + return output; +} +""" + +layernorm_residual_cpp_source = """ +torch::Tensor layernorm_residual_cuda( + torch::Tensor x, + torch::Tensor residual, + torch::Tensor weight, + torch::Tensor bias, + float eps +); +""" + +layernorm_residual_module = load_inline( + name="layernorm_residual", + cpp_sources=layernorm_residual_cpp_source, + cuda_sources=layernorm_residual_source, + functions=["layernorm_residual_cuda"], + verbose=True, + extra_cuda_cflags=["-O3", "--use_fast_math", "-maxrregcount=64"] +) + +class ModelNew(torch.nn.Module): + def __init__(self, normalized_shape=512, eps=1e-5): + super(ModelNew, self).__init__() + self.normalized_shape = normalized_shape + self.eps = eps + if normalized_shape % 4 != 0: + raise ValueError("normalized_shape must be multiple of 4 for vectorization") + self.weight = torch.nn.Parameter(torch.ones(normalized_shape)) + self.bias = torch.nn.Parameter(torch.zeros(normalized_shape)) + self.layernorm_residual = layernorm_residual_module + + def forward(self, x, residual): + if not x.is_contiguous(): x = x.contiguous() + if not residual.is_contiguous(): residual = residual.contiguous() + if not self.weight.is_contiguous(): + self.weight.data = self.weight.data.contiguous() + if not self.bias.is_contiguous(): + self.bias.data = self.bias.data.contiguous() + + return self.layernorm_residual.layernorm_residual_cuda( + x, residual, self.weight, self.bias, self.eps + ) diff --git a/S1/wut0n_#81/layernorm_residual_torchcode.py b/S1/wut0n_#81/layernorm_residual_torchcode.py new file mode 100644 index 00000000..05ab7a04 --- /dev/null +++ b/S1/wut0n_#81/layernorm_residual_torchcode.py @@ -0,0 +1,23 @@ +import torch +import torch.nn as nn + +class Model(nn.Module): + def __init__(self, normalized_shape=512, eps=1e-5): + super(Model, self).__init__() + self.layer_norm = nn.LayerNorm(normalized_shape, eps=eps) + + def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + add_result = x + residual + output = self.layer_norm(add_result) + return output + +N = 1024 +D = 512 + +def get_inputs(): + x = torch.randn(N, D) + residual = torch.randn(N, D) + return [x, residual] + +def get_init_inputs(): + return [512, 1e-5] diff --git a/S1/wut0n_#81/prompt.txt b/S1/wut0n_#81/prompt.txt new file mode 100644 index 00000000..35757206 --- /dev/null +++ b/S1/wut0n_#81/prompt.txt @@ -0,0 +1,229 @@ +You write custom CUDA kernels to replace the PyTorch operators to get speedups. +You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace the combined broadcasting subtraction and L2 norm operators with a custom CUDA kernel or adjust algorithms for better performance. You are only limited by your imagination. + +Here's an example to show you the syntax of inline embedding custom CUDA operators in PyTorch (for reference of code structure, not functional alignment): +The example given architecture (sample structure): +import torch +import torch.nn as nn + +class Model(nn.Module): + def __init__(self, normalized_shape=512, eps=1e-5): + super(Model, self).__init__() + self.layer_norm = nn.LayerNorm(normalized_shape, eps=eps) + + def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + add_result = x + residual + output = self.layer_norm(add_result) + return output + +N = 1024 +D = 512 + +def get_inputs(): + x = torch.randn(N, D) + residual = torch.randn(N, D) + return [x, residual] + +def get_init_inputs(): + return [512, 1e-5] + +The example new arch with custom CUDA kernels (sample structure): +import torch +from torch.utils.cpp_extension import load_inline + +layernorm_residual_source = """ +#include +#include +#include + +struct Float4 { + float x, y, z, w; + __device__ __forceinline__ Float4() {} + __device__ __forceinline__ Float4(float f) : x(f), y(f), z(f), w(f) {} + __device__ __forceinline__ Float4(float x_, float y_, float z_, float w_) : x(x_), y(y_), z(z_), w(w_) {} +}; + +__device__ __forceinline__ Float4 load_float4(const float* addr) { + return *reinterpret_cast(addr); +} + +__device__ __forceinline__ void store_float4(float* addr, Float4 val) { + *reinterpret_cast(addr) = val; +} + +// 高效的warp reduction +__device__ __forceinline__ float warp_reduce_sum(float val) { + #pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + val += __shfl_xor_sync(0xffffffff, val, mask); + } + return val; +} + + +__global__ void layernorm_residual_kernel_vectorized( + const float* __restrict__ x, + const float* __restrict__ residual, + const float* __restrict__ weight, + const float* __restrict__ bias, + float* __restrict__ output, + int N, int D, float eps +) { + // 每个warp处理一行，使用向量化访问 + int row = blockIdx.x * (blockDim.x / 32) + (threadIdx.x / 32); + int lane_id = threadIdx.x % 32; + + if (row >= N) return; + + const float* x_row = x + row * D; + const float* residual_row = residual + row * D; + float* output_row = output + row * D; + + float sum = 0.0f; + int vec_elems = D / 4; + int remaining = D % 4; + + for (int i = lane_id; i < vec_elems; i += 32) { + Float4 x_vec = load_float4(x_row + i * 4); + Float4 residual_vec = load_float4(residual_row + i * 4); + + sum += x_vec.x + residual_vec.x; + sum += x_vec.y + residual_vec.y; + sum += x_vec.z + residual_vec.z; + sum += x_vec.w + residual_vec.w; + } + + for (int i = vec_elems * 4 + lane_id; i < D; i += 32) { + sum += x_row[i] + residual_row[i]; + } + + sum = warp_reduce_sum(sum); + float mean = sum / D; + + float var_sum = 0.0f; + + for (int i = lane_id; i < vec_elems; i += 32) { + Float4 x_vec = load_float4(x_row + i * 4); + Float4 residual_vec = load_float4(residual_row + i * 4); + + float diff_x = x_vec.x + residual_vec.x - mean; + float diff_y = x_vec.y + residual_vec.y - mean; + float diff_z = x_vec.z + residual_vec.z - mean; + float diff_w = x_vec.w + residual_vec.w - mean; + + var_sum += diff_x * diff_x + diff_y * diff_y + diff_z * diff_z + diff_w * diff_w; + } + + for (int i = vec_elems * 4 + lane_id; i < D; i += 32) { + float diff = x_row[i] + residual_row[i] - mean; + var_sum += diff * diff; + } + + var_sum = warp_reduce_sum(var_sum); + float inv_std = rsqrtf(var_sum / D + eps); + + for (int i = lane_id; i < vec_elems; i += 32) { + Float4 x_vec = load_float4(x_row + i * 4); + Float4 residual_vec = load_float4(residual_row + i * 4); + Float4 weight_vec = load_float4(weight + i * 4); + Float4 bias_vec = load_float4(bias + i * 4); + + Float4 result; + float val_x = x_vec.x + residual_vec.x; + float val_y = x_vec.y + residual_vec.y; + float val_z = x_vec.z + residual_vec.z; + float val_w = x_vec.w + residual_vec.w; + + result.x = (val_x - mean) * inv_std * weight_vec.x + bias_vec.x; + result.y = (val_y - mean) * inv_std * weight_vec.y + bias_vec.y; + result.z = (val_z - mean) * inv_std * weight_vec.z + bias_vec.z; + result.w = (val_w - mean) * inv_std * weight_vec.w + bias_vec.w; + + store_float4(output_row + i * 4, result); + } + + for (int i = vec_elems * 4 + lane_id; i < D; i += 32) { + float val = x_row[i] + residual_row[i]; + output_row[i] = (val - mean) * inv_std * weight[i] + bias[i]; + } +} + +torch::Tensor layernorm_residual_cuda( + torch::Tensor x, + torch::Tensor residual, + torch::Tensor weight, + torch::Tensor bias, + float eps +) { + TORCH_CHECK(x.is_cuda() && x.is_contiguous()); + TORCH_CHECK(residual.is_cuda() && residual.is_contiguous()); + TORCH_CHECK(weight.is_cuda() && weight.is_contiguous()); + TORCH_CHECK(bias.is_cuda() && bias.is_contiguous()); + TORCH_CHECK(x.dim() == 2 && residual.dim() == 2); + TORCH_CHECK(x.sizes() == residual.sizes()); + TORCH_CHECK(weight.numel() == x.size(1)); + TORCH_CHECK(bias.numel() == x.size(1)); + + int N = x.size(0); + int D = x.size(1); + + auto output = torch::empty_like(x); + + int warps_per_block = 8; // 256 threads per block + int block_size = warps_per_block * 32; + int grid_size = (N + warps_per_block - 1) / warps_per_block; + + layernorm_residual_kernel_vectorized<<>>( + x.data_ptr(), + residual.data_ptr(), + weight.data_ptr(), + bias.data_ptr(), + output.data_ptr(), + N, D, eps + ); + + return output; +} +""" + +layernorm_residual_cpp_source = """ +torch::Tensor layernorm_residual_cuda( + torch::Tensor x, + torch::Tensor residual, + torch::Tensor weight, + torch::Tensor bias, + float eps +); +""" + +layernorm_residual_module = load_inline( + name="layernorm_residual", + cpp_sources=layernorm_residual_cpp_source, + cuda_sources=layernorm_residual_source, + functions=["layernorm_residual_cuda"], + verbose=True, + extra_cuda_cflags=["-O3", "--use_fast_math", "-maxrregcount=64"] +) + +class ModelNew(torch.nn.Module): + def __init__(self, normalized_shape=512, eps=1e-5): + super(ModelNew, self).__init__() + self.normalized_shape = normalized_shape + self.eps = eps + if normalized_shape % 4 != 0: + raise ValueError("normalized_shape must be multiple of 4 for vectorization") + self.weight = torch.nn.Parameter(torch.ones(normalized_shape)) + self.bias = torch.nn.Parameter(torch.zeros(normalized_shape)) + self.layernorm_residual = layernorm_residual_module + + def forward(self, x, residual): + if not x.is_contiguous(): x = x.contiguous() + if not residual.is_contiguous(): residual = residual.contiguous() + if not self.weight.is_contiguous(): + self.weight.data = self.weight.data.contiguous() + if not self.bias.is_contiguous(): + self.bias.data = self.bias.data.contiguous() + + return self.layernorm_residual.layernorm_residual_cuda( + x, residual, self.weight, self.bias, self.eps + ) diff --git a/S1/wut0n_#81/run_code.py b/S1/wut0n_#81/run_code.py new file mode 100644 index 00000000..2596fcf0 --- /dev/null +++ b/S1/wut0n_#81/run_code.py @@ -0,0 +1,74 @@ +########################################################### +# 性能和精度验证程序 +########################################################### +import torch +import torch.nn as nn +import time +from layernorm_residual_torchcode import Model,get_inputs,get_init_inputs +from layernorm_residual_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 matmul_relu 平均执行时间: {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