forked from ccf-ai-infra/GPUCodeForces
finish layernorm #1
This commit is contained in:
parent
10eed82956
commit
2f8387c7f5
|
|
@ -0,0 +1,190 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
# LayerNorm的CUDA实现
|
||||
layernorm_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
// 使用高度优化的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<const float4*>(&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<const float4*>(&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<float4*>(&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<<<num_blocks, block_size, shared_mem_size>>>(
|
||||
input.data_ptr<float>(),
|
||||
gamma.data_ptr<float>(),
|
||||
beta.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
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
|
||||
|
|
@ -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]
|
||||
|
|
@ -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)
|
||||
|
|
@ -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()
|
||||
Loading…
Reference in New Issue