minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CPU] Add gemma4_rmsnorm_cpu kernel (#22842)

Browse Source 
Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com>
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
blzheng
2026-04-17 13:03:16 +08:00
committed by GitHub
co-authored by Copilot Ma Mingfei
parent 6c89214584
commit 0dcfae5553
5 changed files with 280 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/norm.cpp
+162 -26
View File
@@ -10,17 +10,24 @@ void l2norm_kernel_impl(
scalar_t* __restrict__ output,
const scalar_t* __restrict__ input,
int64_t batch_size,
int64_t seq_len,
int64_t hidden_size,
int64_t input_strideB,
int64_t input_strideS,
int64_t output_strideB,
int64_t output_strideS,
float eps = 1e-5) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
at::parallel_for(0, batch_size, 0, [&](int64_t begin, int64_t end) {
at::parallel_for(0, batch_size * seq_len, 0, [&](int64_t begin, int64_t end) {
int64_t bi{0}, si{0};
data_index_init(begin, bi, batch_size, si, seq_len);
for (int64_t i = begin; i < end; ++i) {
// local ptrs
scalar_t* __restrict__ out_ptr = output + i * hidden_size;
const scalar_t* __restrict__ input_ptr = input + i * hidden_size;
scalar_t* __restrict__ out_ptr = output + bi * output_strideB + si * output_strideS;
const scalar_t* __restrict__ input_ptr = input + bi * input_strideB + si * input_strideS;
fVec sum_fvec = fVec(float(0));
float sum_val = float(0);
@@ -62,17 +69,24 @@ void l2norm_kernel_impl(
float x_val = static_cast<float>(input_ptr[d]);
out_ptr[d] = static_cast<scalar_t>(x_val * rsqrt_var);
}
// move to the next index
data_index_step(bi, batch_size, si, seq_len);
}
});
}
template <typename scalar_t, typename func_t, typename vec_func_t>
void rmsnorm_kernel_impl(
scalar_t* __restrict__ output,
const scalar_t* __restrict__ input,
const scalar_t* __restrict__ weight,
int64_t batch_size,
int64_t seq_len,
int64_t hidden_size,
int64_t input_strideN,
int64_t input_strideB,
int64_t input_strideS,
int64_t output_strideB,
int64_t output_strideS,
const func_t& f,
const vec_func_t& vf,
float eps = 1e-5) {
@@ -80,11 +94,13 @@ void rmsnorm_kernel_impl(
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
at::parallel_for(0, batch_size, 0, [&](int64_t begin, int64_t end) {
at::parallel_for(0, batch_size * seq_len, 0, [&](int64_t begin, int64_t end) {
int64_t bi{0}, si{0};
data_index_init(begin, bi, batch_size, si, seq_len);
for (int64_t i = begin; i < end; ++i) {
// local ptrs
scalar_t* __restrict__ out_ptr = output + i * hidden_size;
const scalar_t* __restrict__ input_ptr = input + i * input_strideN;
scalar_t* __restrict__ out_ptr = output + bi * output_strideB + si * output_strideS;
const scalar_t* __restrict__ input_ptr = input + bi * input_strideB + si * input_strideS;
fVec sum_fvec = fVec(float(0));
float sum_val = float(0);
@@ -131,6 +147,8 @@ void rmsnorm_kernel_impl(
float w_val = static_cast<float>(weight[d]);
out_ptr[d] = static_cast<scalar_t>(x_val * rsqrt_var * f(w_val));
}
// move to the next index
data_index_step(bi, batch_size, si, seq_len);
}
});
}
@@ -222,8 +240,10 @@ void fused_add_rmsnorm_kernel_impl(
const scalar_t* __restrict__ weight,
float* __restrict__ buffer,
int64_t batch_size,
int64_t seq_len,
int64_t hidden_size,
int64_t input_strideN,
int64_t input_strideB,
int64_t input_strideS,
const func_t& f,
const vec_func_t& vf,
float eps = 1e-5) {
@@ -231,13 +251,15 @@ void fused_add_rmsnorm_kernel_impl(
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
at::parallel_for(0, batch_size, 0, [&](int64_t begin, int64_t end) {
at::parallel_for(0, batch_size * seq_len, 0, [&](int64_t begin, int64_t end) {
int64_t bi{0}, si{0};
data_index_init(begin, bi, batch_size, si, seq_len);
int tid = at::get_thread_num();
float* __restrict__ buffer_ptr = buffer + tid * hidden_size;
for (int64_t i = begin; i < end; ++i) {
// local ptrs
scalar_t* __restrict__ input_ptr = input + i * input_strideN;
scalar_t* __restrict__ input_ptr = input + bi * input_strideB + si * input_strideS;
scalar_t* __restrict__ residual_ptr = residual + i * hidden_size;
fVec sum_fvec = fVec(float(0));
@@ -301,6 +323,8 @@ void fused_add_rmsnorm_kernel_impl(
float x_val = buffer_ptr[d] * rsqrt_var * static_cast<float>(f(weight[d]));
input_ptr[d] = x_val;
}
// move to the next index
data_index_step(bi, batch_size, si, seq_len);
}
});
}
@@ -523,25 +547,46 @@ at::Tensor l2norm_cpu(at::Tensor& input, double eps) {
at::Tensor output = at::empty_like(input);
AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "l2norm_kernel", [&] {
l2norm_kernel_impl<scalar_t>(output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), batch_size, hidden_size, eps);
l2norm_kernel_impl<scalar_t>(
output.data_ptr<scalar_t>(),
input.data_ptr<scalar_t>(),
batch_size,
1,
hidden_size,
hidden_size,
0,
hidden_size,
0,
eps);
});
return output;
}
// input : {batch_size, hidden_size}
// input : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// weight: {hidden_size}
at::Tensor rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) {
RECORD_FUNCTION("sgl-kernel::rmsnorm_cpu", std::vector<c10::IValue>({input, weight}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
CHECK_INPUT(weight);
CHECK_DIM(2, input);
int64_t inp_dim{input.dim()};
TORCH_CHECK(inp_dim == 2 || inp_dim == 3, "Expected input dim to be 2 or 3, but got ", inp_dim);
CHECK_DIM(1, weight);
CHECK_EQ(input.size(1), weight.size(0));
CHECK_EQ(input.size(-1), weight.size(0));
int64_t batch_size = input.size(0);
int64_t hidden_size = input.size(1);
int64_t seq_len = 1;
int64_t hidden_size = input.size(-1);
int64_t input_strideB = input.stride(0);
int64_t input_strideS = 0;
at::Tensor output = at::empty_like(input);
int64_t input_strideN = input.stride(0);
int64_t output_strideB = output.stride(0);
int64_t output_strideS = 0;
if (inp_dim == 3) {
seq_len = input.size(1);
input_strideS = input.stride(1);
output_strideS = output.stride(1);
}
AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "rmsnorm_kernel", [&] {
using Vec = at::vec::Vectorized<float>;
@@ -550,8 +595,12 @@ at::Tensor rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) {
input.data_ptr<scalar_t>(),
weight.data_ptr<scalar_t>(),
batch_size,
seq_len,
hidden_size,
input_strideN,
input_strideB,
input_strideS,
output_strideB,
output_strideS,
[](float x) { return x; },
[](Vec x) { return x; },
eps);
@@ -619,6 +668,7 @@ at::Tensor gemma_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps)
int64_t hidden_size = input.size(1);
at::Tensor output = at::empty_like(input);
int64_t input_strideN = input.stride(0);
int64_t output_strideN = output.stride(0);
AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "gemma_rmsnorm_kernel", [&] {
using Vec = at::vec::Vectorized<float>;
@@ -628,8 +678,12 @@ at::Tensor gemma_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps)
input.data_ptr<scalar_t>(),
weight.data_ptr<scalar_t>(),
batch_size,
1,
hidden_size,
input_strideN,
0,
output_strideN,
0,
[](float x) { return x + 1; },
[one_vec](Vec x) { return x + one_vec; },
eps);
@@ -653,6 +707,7 @@ at::Tensor gemma3_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps)
at::Tensor output = at::empty_like(input);
if (input.dim() == 2) {
int64_t input_strideN = input.stride(0);
int64_t output_strideN = output.stride(0);
AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "gemma3_rmsnorm_kernel", [&] {
using Vec = at::vec::Vectorized<float>;
@@ -662,8 +717,12 @@ at::Tensor gemma3_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps)
input.data_ptr<scalar_t>(),
weight.data_ptr<scalar_t>(),
batch_size,
1,
hidden_size,
input_strideN,
0,
output_strideN,
0,
[](float x) { return x + 1; },
[one_vec](Vec x) { return x + one_vec; },
eps);
@@ -698,6 +757,71 @@ at::Tensor gemma3_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps)
return output;
}
// Gemma4RMSNorm: with_scale ? norm(x) * (weight + scale_shift) : norm(x)
// input : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// weight: {hidden_size}
at::Tensor gemma4_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps, double scale_shift, bool with_scale) {
RECORD_FUNCTION("sgl-kernel::gemma4_rmsnorm_cpu", std::vector<c10::IValue>({input, weight}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
CHECK_INPUT(weight);
int64_t inp_dim{input.dim()};
TORCH_CHECK(inp_dim == 2 || inp_dim == 3, "gemma4_rmsnorm_cpu: expected input dim 2 or 3, got ", inp_dim);
CHECK_DIM(1, weight);
CHECK_EQ(input.size(-1), weight.size(0));
int64_t hidden_size = input.size(-1);
at::Tensor output = at::empty_like(input);
int64_t batch_size = input.size(0);
int64_t seq_len = 1;
int64_t input_strideB = input.stride(0);
int64_t input_strideS = 0;
int64_t output_strideB = output.stride(0);
int64_t output_strideS = 0;
if (inp_dim == 3) {
seq_len = input.size(1);
input_strideS = input.stride(1);
output_strideS = output.stride(1);
}
if (with_scale) {
float shift = static_cast<float>(scale_shift);
AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "gemma4_rmsnorm_kernel", [&] {
using Vec = at::vec::Vectorized<float>;
Vec shift_vec = Vec(shift);
rmsnorm_kernel_impl<scalar_t>(
output.data_ptr<scalar_t>(),
input.data_ptr<scalar_t>(),
weight.data_ptr<scalar_t>(),
batch_size,
seq_len,
hidden_size,
input_strideB,
input_strideS,
output_strideB,
output_strideS,
[shift](float x) { return x + shift; },
[shift_vec](Vec x) { return x + shift_vec; },
eps);
});
} else {
AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "gemma4_rmsnorm_kernel", [&] {
l2norm_kernel_impl<scalar_t>(
output.data_ptr<scalar_t>(),
input.data_ptr<scalar_t>(),
batch_size,
seq_len,
hidden_size,
input_strideB,
input_strideS,
output_strideB,
output_strideS,
eps);
});
}
return output;
}
// input : {batch_size, hidden_size}
// weight: {hidden_size}
// gate: {batch_size, hidden_size}
@@ -732,23 +856,31 @@ at::Tensor fused_rmsnorm_gated_cpu(at::Tensor& input, at::Tensor& weight, at::Te
return output;
}
// input : {batch_size, hidden_size}
// residual: {batch_size, hidden_size}
// input : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// residual: {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
// weight : {hidden_size}
void fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps) {
RECORD_FUNCTION("sgl-kernel::fused_add_rmsnorm_cpu", std::vector<c10::IValue>({input, residual, weight}));
CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
CHECK_INPUT(residual);
CHECK_INPUT(weight);
CHECK_DIM(2, input);
CHECK_DIM(2, residual);
int64_t inp_dim{input.dim()}, res_dim{residual.dim()};
CHECK_EQ(inp_dim, res_dim);
TORCH_CHECK(inp_dim == 2 || inp_dim == 3, "Expected input dim to be 2 or 3, but got ", inp_dim);
CHECK_DIM(1, weight);
CHECK_EQ(input.size(0), residual.size(0));
CHECK_EQ(input.size(1), residual.size(1));
CHECK_EQ(input.size(1), weight.size(0));
CHECK_EQ(input.size(-1), residual.size(-1));
CHECK_EQ(input.size(-1), weight.size(0));
int64_t batch_size = input.size(0);
int64_t hidden_size = input.size(1);
int64_t input_strideN = input.stride(0);
int64_t seq_len = 1;
int64_t hidden_size = input.size(-1);
int64_t input_strideB = input.stride(0);
int64_t input_strideS = 0;
if (inp_dim == 3) {
seq_len = input.size(1);
input_strideS = input.stride(1);
}
// allocate temp buffer to store x in float32 per thread
// TODO: implement a singleton for context
@@ -763,8 +895,10 @@ void fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor&
weight.data_ptr<scalar_t>(),
buffer.data_ptr<float>(),
batch_size,
seq_len,
hidden_size,
input_strideN,
input_strideB,
input_strideS,
[](float x) { return x; },
[](Vec x) { return x; },
eps);
@@ -803,8 +937,10 @@ void gemma_fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Te
weight.data_ptr<scalar_t>(),
buffer.data_ptr<float>(),
batch_size,
1,
hidden_size,
input_strideN,
0,
[](float x) { return x + 1; },
[one_vec](Vec x) { return x + one_vec; },
eps);
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+3
View File
@@ -34,6 +34,7 @@ at::Tensor l2norm_cpu(at::Tensor& input, double eps);
at::Tensor rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps);
at::Tensor gemma_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps);
at::Tensor gemma3_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps);
at::Tensor gemma4_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps, double scale_shift, bool with_scale);
// layernorm
at::Tensor
@@ -408,6 +409,8 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
m.impl("gemma_rmsnorm_cpu", torch::kCPU, &gemma_rmsnorm_cpu);
m.def("gemma3_rmsnorm_cpu(Tensor input, Tensor weight, float eps) -> Tensor");
m.impl("gemma3_rmsnorm_cpu", torch::kCPU, &gemma3_rmsnorm_cpu);
m.def("gemma4_rmsnorm_cpu(Tensor input, Tensor weight, float eps, float scale_shift, bool with_scale) -> Tensor");
m.impl("gemma4_rmsnorm_cpu", torch::kCPU, &gemma4_rmsnorm_cpu);
m.def("layernorm_cpu(Tensor input, Tensor weight, Tensor? bias, float eps) -> Tensor");
m.impl("layernorm_cpu", torch::kCPU, &layernorm_cpu);
m.def("l2norm_cpu(Tensor input, float eps) -> Tensor");