diff --git a/S1/wut0n_#32/prompt.txt b/S1/wut0n_#32/prompt.txt new file mode 100644 index 00000000..4929d649 --- /dev/null +++ b/S1/wut0n_#32/prompt.txt @@ -0,0 +1,95 @@ +You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups. + +You have complete freedom to choose the set of operators you want to replace. You may make a decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination. + +Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: + +python +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): +def init(self) -> None: +super().init() + +def forward(self, a, b): + return a + b +def get_inputs(): +# randomly generate input tensors based on the model architecture +a = torch.randn(1, 128).cuda() +b = torch.randn(1, 128).cuda() +return [a, b] + +def get_init_inputs(): +# randomly generate tensors required for initialization based on the model architecture +return [] + + + +The example new arch with custom CUDA kernels looks like this: +python +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): +def init(self) -> None: +super().init() + +def forward(self, a, b): + return a + b +def get_inputs(): +# randomly generate input tensors based on the model architecture +a = torch.randn(1, 128).cuda() +b = torch.randn(1, 128).cuda() +return [a, b] + +def get_init_inputs(): +# randomly generate tensors required for initialization based on the model architecture +return [] + +You are given the following architecture: + +python +import torch +import torch.nn as nn + +class Model(nn.Module): +""" +Simple model that performs RMSNorm + Residual addition. +""" +def init(self, dim, eps=1e-6): +super(Model, self).init() +self.eps = eps +self.weight = torch.nn.Parameter(torch.ones(dim)) + +def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + """ + Applies RMSNorm activation to the input tensor and adds residual connection. + + Args: + x (torch.Tensor): Input tensor of shape (B, L, D). + residual (torch.Tensor): Residual tensor of shape (B, L, D). + + Returns: + torch.Tensor: Output tensor with RMSNorm applied + residual, same shape as input. + """ + dtype = x.dtype + x = x.float() + variance = x.pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + self.eps) + rmsnorm_output = (x * self.weight).to(dtype) + return rmsnorm_output + residual + +batch_size = 16 +seq_len = 512 +dim = 4096 + +def get_inputs(): +x = torch.randn(batch_size, seq_len, dim) +residual = torch.randn(batch_size, seq_len, dim) +return [x, residual] + +def get_init_inputs(): +return [dim] \ No newline at end of file diff --git a/S1/wut0n_#32/rmsnorm_residual_cudacode.py b/S1/wut0n_#32/rmsnorm_residual_cudacode.py new file mode 100644 index 00000000..22ba936d --- /dev/null +++ b/S1/wut0n_#32/rmsnorm_residual_cudacode.py @@ -0,0 +1,135 @@ +import torch +from torch.utils.cpp_extension import load_inline + +rmsnorm_residual_source = """ +#include +#include + +// RMSNorm + Residual融合kernel +__global__ void rmsnorm_residual_kernel( + const float* __restrict__ x, + const float* __restrict__ residual, + const float* __restrict__ weight, + float* __restrict__ y, + int batch_size, + int seq_len, + int dim, + float eps +) { + // 每个block处理一个序列位置 + int seq_idx = blockIdx.x; + int batch_idx = blockIdx.y; + + if (seq_idx >= seq_len || batch_idx >= batch_size) return; + + extern __shared__ float sdata[]; + + // 计算当前序列位置的起始索引 + int base_idx = batch_idx * seq_len * dim + seq_idx * dim; + + // 第一步:计算平方和 + float sum_sq = 0.0f; + for (int i = threadIdx.x; i < dim; i += blockDim.x) { + float val = x[base_idx + i]; + sum_sq += val * val; + } + + sdata[threadIdx.x] = sum_sq; + __syncthreads(); + + // 第二步:规约求和 + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + sdata[threadIdx.x] += sdata[threadIdx.x + stride]; + } + __syncthreads(); + } + + // 第三步:计算RMS + float mean_sq = sdata[0] / dim; + float rms = rsqrtf(mean_sq + eps); + __syncthreads(); + + // 第四步:应用RMSNorm并加上residual + for (int i = threadIdx.x; i < dim; i += blockDim.x) { + int idx = base_idx + i; + // RMSNorm(x) + residual + y[idx] = (x[idx] * rms * weight[i]) + residual[idx]; + } +} + +torch::Tensor rmsnorm_residual_cuda( + torch::Tensor x, + torch::Tensor residual, + torch::Tensor weight, + float eps +) { + auto batch_size = x.size(0); + auto seq_len = x.size(1); + auto dim = x.size(2); + + auto y = torch::empty_like(x); + + const int block_size = 256; + dim3 grid(seq_len, batch_size); + int shared_mem_size = sizeof(float) * block_size; + + rmsnorm_residual_kernel<<>>( + x.data_ptr(), + residual.data_ptr(), + weight.data_ptr(), + y.data_ptr(), + batch_size, + seq_len, + dim, + eps + ); + + return y; +} +""" + +rmsnorm_residual_cpp_source = """ +torch::Tensor rmsnorm_residual_cuda( + torch::Tensor x, + torch::Tensor residual, + torch::Tensor weight, + float eps +); +""" + +# 编译CUDA扩展 +rmsnorm_residual = load_inline( + name="rmsnorm_residual_fused", + cpp_sources=rmsnorm_residual_cpp_source, + cuda_sources=rmsnorm_residual_source, + functions=["rmsnorm_residual_cuda"], + extra_cuda_cflags=[ + "-O3", + "--use_fast_math", + "-std=c++17" + ], + verbose=True +) + +class ModelNew(torch.nn.Module): + """ + RMSNorm + Residual融合模型 + """ + def __init__(self, dim, eps=1e-6): + super(ModelNew, self).__init__() + self.eps = eps + self.weight = torch.nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + """ + 融合实现:RMSNorm + Residual + + Args: + x (torch.Tensor): 输入张量 shape (B, L, D) + residual (torch.Tensor): 残差张量 shape (B, L, D) + + Returns: + torch.Tensor: RMSNorm(x) + residual + """ + return rmsnorm_residual.rmsnorm_residual_cuda(x, residual, self.weight, self.eps) diff --git a/S1/wut0n_#32/rmsnorm_residual_torchcode.py b/S1/wut0n_#32/rmsnorm_residual_torchcode.py new file mode 100644 index 00000000..a6eb7fc2 --- /dev/null +++ b/S1/wut0n_#32/rmsnorm_residual_torchcode.py @@ -0,0 +1,67 @@ +import torch +import torch.nn as nn + +class Model(nn.Module): + """ + 原始模型:RMSNorm + Residual + """ + def __init__(self, dim, eps=1e-6): + super(Model, self).__init__() + self.eps = eps + self.weight = torch.nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + """ + 原始实现:先RMSNorm,再与residual相加 + + Args: + x (torch.Tensor): 输入张量 shape (B, L, D) + residual (torch.Tensor): 残差张量 shape (B, L, D) + + Returns: + torch.Tensor: RMSNorm(x) + residual + """ + dtype = x.dtype + x = x.float() + variance = x.pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + self.eps) + rmsnorm_output = (x * self.weight).to(dtype) + return rmsnorm_output + residual + +class ModelNew(torch.nn.Module): + """ + 融合模型:直接实现RMSNorm + Residual + """ + def __init__(self, dim, eps=1e-6): + super(ModelNew, self).__init__() + self.eps = eps + self.weight = torch.nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + """ + 融合实现:在CUDA kernel中直接完成RMSNorm + Residual + + Args: + x (torch.Tensor): 输入张量 shape (B, L, D) + residual (torch.Tensor): 残差张量 shape (B, L, D) + + Returns: + torch.Tensor: RMSNorm(x) + residual + """ + # 这个将在CUDA中实现 + pass + +# 测试参数 +batch_size = 16 +seq_len = 512 +dim = 4096 + +def get_inputs(): + """生成测试输入""" + x = torch.randn(batch_size, seq_len, dim) + residual = torch.randn(batch_size, seq_len, dim) + return [x, residual] + +def get_init_inputs(): + """获取初始化参数""" + return [dim] diff --git a/S1/wut0n_#32/run_code.py b/S1/wut0n_#32/run_code.py new file mode 100644 index 00000000..47eae315 --- /dev/null +++ b/S1/wut0n_#32/run_code.py @@ -0,0 +1,84 @@ +########################################################### +# 性能和精度验证程序 +########################################################### +import torch +import torch.nn as nn +import time +from rmsnorm_residual_torchcode import Model, get_inputs, get_init_inputs +from rmsnorm_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, atol=1e-05) + max_diff = torch.max(torch.abs(output_torch - output_cuda)).item() + mean_diff = torch.mean(torch.abs(output_torch - output_cuda)).item() + + if precision_flag: + print(f"✅ 精度对齐:两个模型的输出结果非常接近。") + print(f"最大误差: {max_diff:.8f}, 平均误差: {mean_diff:.8f}") + else: + print(f"❌ 精度不一致!最大误差: {max_diff:.8f}, 平均误差: {mean_diff:.8f}") + + 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 rmsnorm_residual 平均执行时间: {torch_time:.6f} 秒") + print(f"自定义 CUDA rmsnorm_residual 平均执行时间: {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()