forked from ccf-ai-infra/GPUCodeForces
feat:add high performance rmsnormresidual #32
This commit is contained in:
parent
f989885dde
commit
2d69fa9275
|
|
@ -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]
|
||||
|
|
@ -0,0 +1,135 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
rmsnorm_residual_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
// 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<<<grid, block_size, shared_mem_size>>>(
|
||||
x.data_ptr<float>(),
|
||||
residual.data_ptr<float>(),
|
||||
weight.data_ptr<float>(),
|
||||
y.data_ptr<float>(),
|
||||
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)
|
||||
|
|
@ -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]
|
||||
|
|
@ -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()
|
||||
Loading…
Reference in New Issue