forked from ccf-ai-infra/GPUCodeForces
Feat:add layernorm+residual #81
This commit is contained in:
parent
f989885dde
commit
6d03e9b65d
|
|
@ -0,0 +1,199 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
layernorm_residual_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
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<const Float4*>(addr);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void store_float4(float* addr, Float4 val) {
|
||||
*reinterpret_cast<Float4*>(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<<<grid_size, block_size>>>(
|
||||
x.data_ptr<float>(),
|
||||
residual.data_ptr<float>(),
|
||||
weight.data_ptr<float>(),
|
||||
bias.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
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
|
||||
)
|
||||
|
|
@ -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]
|
||||
|
|
@ -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 <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
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<const Float4*>(addr);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void store_float4(float* addr, Float4 val) {
|
||||
*reinterpret_cast<Float4*>(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<<<grid_size, block_size>>>(
|
||||
x.data_ptr<float>(),
|
||||
residual.data_ptr<float>(),
|
||||
weight.data_ptr<float>(),
|
||||
bias.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
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
|
||||
)
|
||||
|
|
@ -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()
|
||||
Loading…
Reference in New Issue