From 2f4babe32b329cd8e12bc024ed95fb0ba6523c0d Mon Sep 17 00:00:00 2001 From: Zaili Wang <109502517+ZailiWang@users.noreply.github.com> Date: Thu, 19 Mar 2026 13:15:24 +0800 Subject: [PATCH] [CPU] support LayerNorm with 3D shape (#15075) Co-authored-by: Ma Mingfei --- python/sglang/srt/layers/layernorm.py | 3 +- sgl-kernel/csrc/cpu/norm.cpp | 122 +++++++++++++++----- sgl-kernel/csrc/cpu/torch_extension_cpu.cpp | 16 ++- test/srt/cpu/test_norm.py | 109 +++++++++++++---- 4 files changed, 196 insertions(+), 54 deletions(-) diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 401b5f131..8d81125d3 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -421,8 +421,9 @@ class LayerNorm(MultiPlatformOp): x: torch.Tensor, ) -> torch.Tensor: if _is_cpu_amx_available: + bias_data = self.bias.data if self.use_bias else None return torch.ops.sgl_kernel.layernorm_cpu( - x, self.weight.data, self.variance_epsilon + x, self.weight.data, bias_data, self.variance_epsilon ) else: return self.forward_native(x) diff --git a/sgl-kernel/csrc/cpu/norm.cpp b/sgl-kernel/csrc/cpu/norm.cpp index d822c0d44..031797ec5 100644 --- a/sgl-kernel/csrc/cpu/norm.cpp +++ b/sgl-kernel/csrc/cpu/norm.cpp @@ -388,26 +388,32 @@ void fused_rmsnorm_gated_kernel_impl( template void fused_add_layernorm_kernel_impl( - scalar_t* __restrict__ input, + scalar_t* __restrict__ output, + const scalar_t* __restrict__ input, scalar_t* __restrict__ residual, const scalar_t* __restrict__ weight, + const scalar_t* __restrict__ bias, float* __restrict__ buffer, int64_t batch_size, + int64_t seq_len, int64_t hidden_size, int64_t input_strideN, 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) { - int tid = at::get_thread_num(); - float* __restrict__ buffer_ptr = buffer + tid * hidden_size; + + const bool has_residual{residual != nullptr}; + const bool has_bias{bias != nullptr}; + const int64_t parallel_size{batch_size * seq_len}; + at::parallel_for(0, parallel_size, 0, [&](int64_t begin, int64_t end) { + float* __restrict__ buffer_ptr = buffer + at::get_thread_num() * hidden_size; for (int64_t i = begin; i < end; ++i) { - scalar_t* __restrict__ input_ptr = input + i * input_strideN; + scalar_t* __restrict__ out_ptr = output + i * hidden_size; + const scalar_t* __restrict__ input_ptr = input + i * input_strideN; scalar_t* __restrict__ residual_ptr{(scalar_t*)nullptr}; - if (residual != nullptr) { + if (has_residual) { residual_ptr = residual + i * hidden_size; } @@ -422,7 +428,7 @@ void fused_add_layernorm_kernel_impl( fVec x_fvec0, x_fvec1; std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); - if (residual_ptr != nullptr) { + if (has_residual) { bVec r_bvec = bVec::loadu(residual_ptr + d); fVec r_fvec0, r_fvec1; std::tie(r_fvec0, r_fvec1) = at::vec::convert_to_float(r_bvec); @@ -445,7 +451,7 @@ void fused_add_layernorm_kernel_impl( #pragma GCC unroll 4 for (; d < hidden_size; ++d) { float x_val = static_cast(input_ptr[d]); - if (residual_ptr != nullptr) { + if (has_residual) { float r_val = static_cast(residual_ptr[d]); x_val += r_val; residual_ptr[d] = static_cast(x_val); @@ -475,7 +481,6 @@ void fused_add_layernorm_kernel_impl( for (d = 0; d <= hidden_size - kVecSize; d += kVecSize) { fVec x_fvec0 = fVec::loadu(buffer_ptr + d); fVec x_fvec1 = fVec::loadu(buffer_ptr + d + fVec::size()); - bVec w_bvec = bVec::loadu(weight + d); fVec w_fvec0, w_fvec1; std::tie(w_fvec0, w_fvec1) = at::vec::convert_to_float(w_bvec); @@ -483,14 +488,25 @@ void fused_add_layernorm_kernel_impl( x_fvec0 = (x_fvec0 - mean_fvec) * scale_fvec * w_fvec0; x_fvec1 = (x_fvec1 - mean_fvec) * scale_fvec * w_fvec1; - bVec x_bvec = convert_from_float_ext(x_fvec0, x_fvec1); - x_bvec.store(input_ptr + d); + if (has_bias) { + bVec b_bvec = bVec::loadu(bias + d); + fVec b_fvec0, b_fvec1; + std::tie(b_fvec0, b_fvec1) = at::vec::convert_to_float(b_bvec); + x_fvec0 += b_fvec0; + x_fvec1 += b_fvec1; + } + + bVec o_bvec = convert_from_float_ext(x_fvec0, x_fvec1); + o_bvec.store(out_ptr + d); } #pragma GCC unroll 4 for (; d < hidden_size; ++d) { float normalized = (buffer_ptr[d] - mean) * rsqrt_var; float x_val = normalized * static_cast(weight[d]); - input_ptr[d] = static_cast(x_val); + if (has_bias) { + x_val += static_cast(bias[d]); + } + out_ptr[d] = static_cast(x_val); } } }); @@ -543,33 +559,52 @@ at::Tensor rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) { return output; } -// input : {batch_size, hidden_size} +// input : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size} // weight: {hidden_size} -void layernorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) { +// bias : {hidden_size} +at::Tensor +layernorm_cpu(const at::Tensor& input, const at::Tensor& weight, const std::optional& bias, double eps) { RECORD_FUNCTION("sgl-kernel::layernorm_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)); - int64_t batch_size = input.size(0); - int64_t hidden_size = input.size(1); - int64_t input_strideN = input.stride(0); + if (bias.has_value()) { + CHECK_DIM(1, bias.value()); + CHECK_EQ(bias.value().size(0), weight.size(0)); + } + int64_t batch_size{input.size(0)}, seq_len{1}, hidden_size{input.size(1)}, input_strideN{input.stride(0)}; + if (inp_dim == 3) { + CHECK_EQ(input.size(2), weight.size(0)); + seq_len = input.size(1); + hidden_size = input.size(2); + input_strideN = input.stride(1); + } else { + CHECK_EQ(input.size(1), weight.size(0)); + } + + at::Tensor output = at::empty_like(input); int64_t num_threads = at::get_num_threads(); at::Tensor buffer = at::empty({num_threads, hidden_size}, input.options().dtype(at::kFloat)); + AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "layernorm_kernel", [&] { fused_add_layernorm_kernel_impl( + output.data_ptr(), input.data_ptr(), nullptr, weight.data_ptr(), + conditional_data_ptr(bias), buffer.data_ptr(), batch_size, + seq_len, hidden_size, input_strideN, eps); }); + return output; } at::Tensor gemma_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) { @@ -776,36 +811,65 @@ void gemma_fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Te }); } -// 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_layernorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps) { +// bias : {hidden_size} +at::Tensor fused_add_layernorm_cpu( + const at::Tensor& input, + at::Tensor& residual, + const at::Tensor& weight, + const std::optional& bias, + double eps) { RECORD_FUNCTION("sgl-kernel::fused_add_layernorm_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); + TORCH_CHECK(res_dim == 2 || res_dim == 3, "Expected residual dim to be 2 or 3, but got ", res_dim); + CHECK_DIM(1, weight); + if (bias.has_value()) { + CHECK_DIM(1, bias.value()); + CHECK_EQ(bias.value().size(0), weight.size(0)); + } CHECK_EQ(input.size(0), residual.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); + if (inp_dim == 3) { + CHECK_EQ(input.size(2), residual.size(2)); + CHECK_EQ(input.size(2), weight.size(0)); + } else { + CHECK_EQ(input.size(1), weight.size(0)); + } + int64_t batch_size{input.size(0)}, seq_len{1}, hidden_size{input.size(1)}, input_strideN{input.stride(0)}; + if (inp_dim == 3) { + seq_len = input.size(1); + hidden_size = input.size(2); + input_strideN = input.stride(1); + } + at::Tensor output = at::empty_like(input); + + // Allocate temp buffer to store x in float32 per thread + // It is necessary to store FP32 precision of residual-add results to pass UT acc test int64_t num_threads = at::get_num_threads(); at::Tensor buffer = at::empty({num_threads, hidden_size}, input.options().dtype(at::kFloat)); AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "fused_add_layernorm_kernel", [&] { fused_add_layernorm_kernel_impl( + output.data_ptr(), input.data_ptr(), residual.data_ptr(), weight.data_ptr(), + conditional_data_ptr(bias), buffer.data_ptr(), batch_size, + seq_len, hidden_size, input_strideN, eps); }); + return output; } diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index cc487c3af..563a7d561 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -36,7 +36,8 @@ 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); // layernorm -void layernorm_cpu(at::Tensor& input, at::Tensor& weight, double eps); +at::Tensor +layernorm_cpu(const at::Tensor& input, const at::Tensor& weight, const std::optional& bias, double eps); // qwen3_next_rmsnorm_gated at::Tensor fused_rmsnorm_gated_cpu(at::Tensor& input, at::Tensor& weight, at::Tensor& gate, double eps); @@ -46,7 +47,12 @@ void fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& void gemma_fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps); // fused_add_layernorm -void fused_add_layernorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps); +at::Tensor fused_add_layernorm_cpu( + const at::Tensor& input, + at::Tensor& residual, + const at::Tensor& weight, + const std::optional& bias, + double eps); // topk std::tuple @@ -364,7 +370,7 @@ 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("layernorm_cpu(Tensor(a!) input, Tensor weight, float eps) -> ()"); + 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"); m.impl("l2norm_cpu", torch::kCPU, &l2norm_cpu); @@ -374,7 +380,9 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.impl("fused_add_rmsnorm_cpu", torch::kCPU, &fused_add_rmsnorm_cpu); m.def("gemma_fused_add_rmsnorm_cpu(Tensor(a!) input, Tensor(a!) residual, Tensor weight, float eps) -> ()"); m.impl("gemma_fused_add_rmsnorm_cpu", torch::kCPU, &gemma_fused_add_rmsnorm_cpu); - m.def("fused_add_layernorm_cpu(Tensor(a!) input, Tensor(a!) residual, Tensor weight, float eps) -> ()"); + m.def( + "fused_add_layernorm_cpu(Tensor input, Tensor residual, Tensor weight, Tensor? bias, float eps) -> " + "Tensor"); m.impl("fused_add_layernorm_cpu", torch::kCPU, &fused_add_layernorm_cpu); // topk diff --git a/test/srt/cpu/test_norm.py b/test/srt/cpu/test_norm.py index 434736a82..923e96a07 100644 --- a/test/srt/cpu/test_norm.py +++ b/test/srt/cpu/test_norm.py @@ -215,6 +215,7 @@ class TestLayerNorm(CustomTestCase): weight: torch.Tensor, variance_epsilon: float, residual: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: orig_dtype = x.dtype x = x.to(torch.float32) @@ -224,45 +225,113 @@ class TestLayerNorm(CustomTestCase): variance, mean = torch.var_mean(x, dim=-1, keepdim=True, correction=0) x = (x - mean) * torch.rsqrt(variance + variance_epsilon) - x = x.to(orig_dtype) * weight - if residual is None: - return x - else: - return x, residual + x = x * weight.to(torch.float32) + if bias is not None: + x = x + bias.to(torch.float32) + x = x.to(orig_dtype) + return x if residual is None else (x, residual) @parametrize( m=[4096, 1024], n=[4096, 4109], dtype=[torch.float16, torch.bfloat16], ) - def test_norm(self, m: int, n: int, dtype: torch.dtype) -> None: - x_ln = torch.randn([m, n], dtype=dtype) - x_ln = make_non_contiguous(x_ln) - ref_x_ln = x_ln.clone() - hidden_size = x_ln.size(-1) + def test_norm_input_2d(self, m: int, n: int, dtype: torch.dtype) -> None: + x = torch.randn([m, n], dtype=dtype) + x = make_non_contiguous(x) + hidden_size = x.size(-1) weight = torch.randn(hidden_size, dtype=dtype) + bias = torch.randn(hidden_size, dtype=dtype) variance_epsilon = 1e-6 - torch.ops.sgl_kernel.layernorm_cpu(x_ln, weight, variance_epsilon) - ref_ln_out = self._forward_native(ref_x_ln, weight, variance_epsilon) + ln_out = torch.ops.sgl_kernel.layernorm_cpu(x, weight, None, variance_epsilon) + ref_ln_out = self._forward_native(x, weight, variance_epsilon) atol = rtol = precision[ref_ln_out.dtype] - torch.testing.assert_close(x_ln, ref_ln_out, atol=atol, rtol=rtol) + torch.testing.assert_close(ln_out, ref_ln_out, atol=atol, rtol=rtol) + + ln_out = torch.ops.sgl_kernel.layernorm_cpu(x, weight, bias, variance_epsilon) + ref_ln_out = self._forward_native( + x, weight, variance_epsilon, residual=None, bias=bias + ) + torch.testing.assert_close(ln_out, ref_ln_out, atol=atol, rtol=rtol) - x_add_ln = torch.randn([m, n], dtype=dtype) - x_add_ln = make_non_contiguous(x_add_ln) - ref_x_add_ln = x_add_ln.clone() residual = torch.randn([m, hidden_size], dtype=dtype) ref_residual = residual.clone() - torch.ops.sgl_kernel.fused_add_layernorm_cpu( - x_add_ln, residual, weight, variance_epsilon + add_ln_out = torch.ops.sgl_kernel.fused_add_layernorm_cpu( + x, residual, weight, None, variance_epsilon ) ref_add_ln_out, ref_residual = self._forward_native( - ref_x_add_ln, weight, variance_epsilon, ref_residual + x, weight, variance_epsilon, residual=ref_residual ) - torch.testing.assert_close(x_add_ln, ref_add_ln_out, atol=atol, rtol=rtol) + torch.testing.assert_close(add_ln_out, ref_add_ln_out, atol=atol, rtol=rtol) + torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol) + + residual = torch.randn([m, hidden_size], dtype=dtype) + ref_residual = residual.clone() + + add_ln_out = torch.ops.sgl_kernel.fused_add_layernorm_cpu( + x, residual, weight, bias, variance_epsilon + ) + ref_add_ln_out, ref_residual = self._forward_native( + x, weight, variance_epsilon, residual=ref_residual, bias=bias + ) + + torch.testing.assert_close(add_ln_out, ref_add_ln_out, atol=atol, rtol=rtol) + torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol) + + @parametrize( + l=[4096, 1024], + m=[1, 4], + n=[4096, 4109, 2304], + dtype=[torch.float16, torch.bfloat16], + ) + def test_norm_input_3d(self, l: int, m: int, n: int, dtype: torch.dtype) -> None: + 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) + bias = torch.randn(hidden_size, dtype=dtype) + variance_epsilon = 1e-6 + + ln_out = torch.ops.sgl_kernel.layernorm_cpu(x, weight, None, variance_epsilon) + ref_ln_out = self._forward_native(x, weight, variance_epsilon) + + atol = rtol = precision[ref_ln_out.dtype] + torch.testing.assert_close(ln_out, ref_ln_out, atol=atol, rtol=rtol) + + ln_out = torch.ops.sgl_kernel.layernorm_cpu(x, weight, bias, variance_epsilon) + ref_ln_out = self._forward_native( + x, weight, variance_epsilon, residual=None, bias=bias + ) + torch.testing.assert_close(ln_out, ref_ln_out, atol=atol, rtol=rtol) + + residual = torch.randn([l, m, hidden_size], dtype=dtype) + ref_residual = residual.clone() + + add_ln_out = torch.ops.sgl_kernel.fused_add_layernorm_cpu( + x, residual, weight, None, variance_epsilon + ) + ref_add_ln_out, ref_residual = self._forward_native( + x, weight, variance_epsilon, ref_residual + ) + + torch.testing.assert_close(add_ln_out, ref_add_ln_out, atol=atol, rtol=rtol) + torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol) + + residual = torch.randn([l, m, hidden_size], dtype=dtype) + ref_residual = residual.clone() + + add_ln_out = torch.ops.sgl_kernel.fused_add_layernorm_cpu( + x, residual, weight, bias, variance_epsilon + ) + ref_add_ln_out, ref_residual = self._forward_native( + x, weight, variance_epsilon, residual=ref_residual, bias=bias + ) + + torch.testing.assert_close(add_ln_out, ref_add_ln_out, atol=atol, rtol=rtol) torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol)