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
python/sglang/srt/layers/layernorm.py
+7
View File
@@ -700,6 +700,13 @@ class Gemma4RMSNorm(MultiPlatformOp):
normed_output = normed_output * (self.weight.float() + self.scale_shift)
return normed_output.type_as(x)
def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
if _is_cpu_amx_available:
return torch.ops.sgl_kernel.gemma4_rmsnorm_cpu(
x, self.weight.data, self.eps, self.scale_shift, self.with_scale
)
return self.forward_native(x)
def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
if x.numel() == 0:
return x
python/sglang/srt/model_executor/cpu_graph_runner.py
+1
View File
@@ -140,6 +140,7 @@ def register_fake_ops():
"causal_conv1d_fwd_cpu",
"gemma_rmsnorm_cpu",
"gemma3_rmsnorm_cpu",
"gemma4_rmsnorm_cpu",
]:
@torch.library.register_fake(f"sgl_kernel::{op}")
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");
test/srt/cpu/test_norm.py
+107 -14
View File
@@ -11,9 +11,6 @@ torch.manual_seed(1234)
class TestNorm(CustomTestCase):
M = [4096, 1024]
N = [4096, 4096 + 13]
dtype = [torch.float16, torch.bfloat16]
def _forward_native(
self,
@@ -65,7 +62,12 @@ class TestNorm(CustomTestCase):
x = x.to(orig_dtype)
return x if residual is None else (x, residual)
def _norm_test(self, m, n, dtype):
@parametrize(
m=[4096, 1024],
n=[4096, 4109],
dtype=[torch.float16, torch.bfloat16],
)
def test_norm(self, m, n, dtype):
x = torch.randn([m, n], dtype=dtype)
x = make_non_contiguous(x)
@@ -94,7 +96,47 @@ class TestNorm(CustomTestCase):
torch.testing.assert_close(x, ref_x, atol=atol, rtol=rtol)
torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol)
def _l2norm_test(self, m, n, dtype):
@parametrize(
l=[1, 2],
m=[4096, 1024],
n=[4096, 4109],
dtype=[torch.float16, torch.bfloat16],
)
def test_norm_3d(self, l, m, n, dtype):
x = torch.randn([l, m, n], dtype=dtype)
x = make_non_contiguous(x)
hidden_size = x.size(-1)
weight = torch.randn(hidden_size, dtype=dtype)
variance_epsilon = 1e-6
out = torch.ops.sgl_kernel.rmsnorm_cpu(x, weight, variance_epsilon)
ref_out = self._forward_native(x, weight, variance_epsilon)
atol = rtol = precision[ref_out.dtype]
torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
ref_x = x.clone()
residual = torch.randn([l, m, hidden_size], dtype=dtype)
ref_residual = residual.clone()
torch.ops.sgl_kernel.fused_add_rmsnorm_cpu(
x, residual, weight, variance_epsilon
)
ref_x, ref_residual = self._forward_native(
ref_x, weight, variance_epsilon, ref_residual
)
torch.testing.assert_close(x, ref_x, atol=atol, rtol=rtol)
torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol)
@parametrize(
m=[4096, 1024],
n=[4096, 4109],
dtype=[torch.float16, torch.bfloat16],
)
def test_l2norm(self, m, n, dtype):
x = torch.randn([m, n], dtype=dtype)
hidden_size = x.size(-1)
@@ -107,7 +149,12 @@ class TestNorm(CustomTestCase):
atol = rtol = precision[ref_out.dtype]
torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
def _gemma_rmsnorm_test(self, m, n, dtype):
@parametrize(
m=[4096, 1024],
n=[4096, 4109],
dtype=[torch.float16, torch.bfloat16],
)
def test_gemma_rmsnorm(self, m, n, dtype):
x = torch.randn([m, n], dtype=dtype)
x = make_non_contiguous(x)
@@ -136,7 +183,12 @@ class TestNorm(CustomTestCase):
torch.testing.assert_close(x, ref_x, atol=atol, rtol=rtol)
torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol)
def _gemma3_rmsnorm_test(self, m, n, dtype):
@parametrize(
m=[4096, 1024],
n=[4096, 4109],
dtype=[torch.float16, torch.bfloat16],
)
def test_gemma3_rmsnorm(self, m, n, dtype):
x_list = [
torch.randn([m, n], dtype=dtype),
torch.randn([1, m, 2, n], dtype=dtype),
@@ -152,13 +204,54 @@ class TestNorm(CustomTestCase):
atol = rtol = precision[ref_out.dtype]
torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
def test_norm(self):
for params in itertools.product(self.M, self.N, self.dtype):
with self.subTest(m=params[0], n=params[1], dtype=params[2]):
self._norm_test(*params)
self._l2norm_test(*params)
self._gemma_rmsnorm_test(*params)
self._gemma3_rmsnorm_test(*params)
def _gemma4_rmsnorm_native(
self,
x: torch.Tensor,
weight: torch.Tensor,
variance_epsilon: float = 1e-6,
scale_shift: float = 0.0,
with_scale: bool = True,
):
output = self._norm(x.float(), variance_epsilon)
if with_scale:
output = output * (weight.float() + scale_shift)
return output.type_as(x)
@parametrize(
m=[4096, 1024],
n=[4096, 4109],
dtype=[torch.float16, torch.bfloat16],
)
def test_gemma4_rmsnorm(self, m, n, dtype):
for scale_shift, with_scale in [
(0.0, True),
(1.0, True),
(0.0, False),
(1.0, False),
]:
x_list = [
torch.randn([m, n], dtype=dtype),
torch.randn([4, m, n], dtype=dtype),
]
# Add non-block-contiguous 3D input
base = torch.randn([4, 2 * m, n], dtype=dtype)
x_list.append(base[:, :m, :])
for x in x_list:
x = make_non_contiguous(x)
hidden_size = x.size(-1)
weight = torch.randn(hidden_size, dtype=dtype)
variance_epsilon = 1e-6
out = torch.ops.sgl_kernel.gemma4_rmsnorm_cpu(
x, weight, variance_epsilon, scale_shift, with_scale
)
ref_out = self._gemma4_rmsnorm_native(
x, weight, variance_epsilon, scale_shift, with_scale
)
atol = rtol = precision[ref_out.dtype]
torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
class TestFusedRMSNormGated(CustomTestCase):