forked from ccf-ai-infra/GPUCodeForces
103 lines
2.7 KiB
Python
103 lines
2.7 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
rmsnorm_source = """
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
#include <math.h>
|
|
|
|
__global__ void rmsnorm_kernel(
|
|
const float* __restrict__ x,
|
|
const float* __restrict__ weight,
|
|
float* __restrict__ y,
|
|
int batch,
|
|
int features,
|
|
float eps
|
|
) {
|
|
int row = blockIdx.x;
|
|
if (row >= batch) return;
|
|
int tid = threadIdx.x;
|
|
|
|
extern __shared__ float sdata[];
|
|
float sum_sq = 0.0f;
|
|
|
|
for (int i = tid; i < features; i += blockDim.x) {
|
|
float v = x[row * features + i];
|
|
sum_sq += v * v;
|
|
}
|
|
sdata[tid] = sum_sq;
|
|
__syncthreads();
|
|
|
|
for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
|
|
if (tid < offset) {
|
|
sdata[tid] += sdata[tid + offset];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
float rms = rsqrtf(sdata[0] / features + eps);
|
|
__syncthreads(); // 保证 rms 可见
|
|
|
|
for (int i = tid; i < features; i += blockDim.x) {
|
|
float v = x[row * features + i];
|
|
float w = weight[i];
|
|
y[row * features + i] = v * rms * w;
|
|
}
|
|
}
|
|
|
|
torch::Tensor rmsnorm_cuda(torch::Tensor x, torch::Tensor weight, float eps) {
|
|
TORCH_CHECK(x.is_cuda(), "x 必须是 CUDA 张量");
|
|
TORCH_CHECK(weight.is_cuda(), "weight 必须是 CUDA 张量");
|
|
TORCH_CHECK(x.dim() == 2, "当前内核仅支持二维输入张量");
|
|
TORCH_CHECK(weight.dim() == 1, "RMSNorm 权重必须是一维向量");
|
|
TORCH_CHECK(x.size(1) == weight.size(0), "输入最后一维与权重长度不匹配");
|
|
|
|
int batch = x.size(0);
|
|
int features = x.size(1);
|
|
|
|
auto y = torch::empty_like(x);
|
|
|
|
int threads = 256;
|
|
if (features < threads) {
|
|
threads = 1;
|
|
while (threads < features) threads <<= 1;
|
|
if (threads < 32) threads = 32;
|
|
}
|
|
size_t shared = threads * sizeof(float);
|
|
|
|
rmsnorm_kernel<<<batch, threads, shared>>>(
|
|
x.data_ptr<float>(),
|
|
weight.data_ptr<float>(),
|
|
y.data_ptr<float>(),
|
|
batch,
|
|
features,
|
|
eps
|
|
);
|
|
return y;
|
|
}
|
|
"""
|
|
|
|
rmsnorm_cpp_source = """
|
|
torch::Tensor rmsnorm_cuda(torch::Tensor x, torch::Tensor weight, float eps);
|
|
"""
|
|
|
|
rmsnorm = load_inline(
|
|
name="rmsnorm",
|
|
cpp_sources=rmsnorm_cpp_source,
|
|
cuda_sources=rmsnorm_source,
|
|
functions=["rmsnorm_cuda"],
|
|
verbose=True
|
|
)
|
|
|
|
class ModelNew(nn.Module):
|
|
def __init__(self, weight: torch.Tensor, eps: float = 1e-6):
|
|
super().__init__()
|
|
if weight.dim() != 1:
|
|
raise ValueError("RMSNorm 权重必须是一维向量。")
|
|
self.weight = nn.Parameter(weight.clone())
|
|
self.eps = eps
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return rmsnorm.rmsnorm_cuda(x, self.weight, self.eps)
|