From 0dcfae55531f7e9b0e640ebd759b4944eeaca72f Mon Sep 17 00:00:00 2001 From: blzheng Date: Fri, 17 Apr 2026 13:03:16 +0800 Subject: [PATCH] [CPU] Add gemma4_rmsnorm_cpu kernel (#22842) Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> Co-authored-by: Ma Mingfei --- python/sglang/srt/layers/layernorm.py | 7 + .../srt/model_executor/cpu_graph_runner.py | 1 + sgl-kernel/csrc/cpu/norm.cpp | 188 +++++++++++++++--- sgl-kernel/csrc/cpu/torch_extension_cpu.cpp | 3 + test/srt/cpu/test_norm.py | 121 +++++++++-- 5 files changed, 280 insertions(+), 40 deletions(-) diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 55d880643..3d914813a 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -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 diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index 28b808fd0..d945c4837 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/python/sglang/srt/model_executor/cpu_graph_runner.py @@ -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}") diff --git a/sgl-kernel/csrc/cpu/norm.cpp b/sgl-kernel/csrc/cpu/norm.cpp index 031797ec5..6641e81a3 100644 --- a/sgl-kernel/csrc/cpu/norm.cpp +++ b/sgl-kernel/csrc/cpu/norm.cpp @@ -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; using fVec = at::vec::Vectorized; 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(input_ptr[d]); out_ptr[d] = static_cast(x_val * rsqrt_var); } + // move to the next index + data_index_step(bi, batch_size, si, seq_len); } }); } + template 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; 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(weight[d]); out_ptr[d] = static_cast(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; 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(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(output.data_ptr(), input.data_ptr(), batch_size, hidden_size, eps); + l2norm_kernel_impl( + output.data_ptr(), + input.data_ptr(), + 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({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; @@ -550,8 +595,12 @@ at::Tensor rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) { input.data_ptr(), weight.data_ptr(), 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; @@ -628,8 +678,12 @@ at::Tensor gemma_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) input.data_ptr(), weight.data_ptr(), 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; @@ -662,8 +717,12 @@ at::Tensor gemma3_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) input.data_ptr(), weight.data_ptr(), 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({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(scale_shift); + AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "gemma4_rmsnorm_kernel", [&] { + using Vec = at::vec::Vectorized; + Vec shift_vec = Vec(shift); + rmsnorm_kernel_impl( + output.data_ptr(), + input.data_ptr(), + weight.data_ptr(), + 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( + output.data_ptr(), + input.data_ptr(), + 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({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(), buffer.data_ptr(), 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(), buffer.data_ptr(), batch_size, + 1, hidden_size, input_strideN, + 0, [](float x) { return x + 1; }, [one_vec](Vec x) { return x + one_vec; }, eps); diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index 31a6d95fd..e5f279866 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -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"); diff --git a/test/srt/cpu/test_norm.py b/test/srt/cpu/test_norm.py index 923e96a07..e123f3e35 100644 --- a/test/srt/cpu/test_norm.py +++ b/test/srt/cpu/test_norm.py @@ -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):