From d6bd2d1126632aff68c8045409b78235ca487ab6 Mon Sep 17 00:00:00 2001 From: Zaili Wang <109502517+ZailiWang@users.noreply.github.com> Date: Fri, 12 Dec 2025 08:58:23 +0800 Subject: [PATCH] [CPU] layernorm & fused add-layernorm kernels (#14074) --- .github/CI_PERMISSIONS.json | 17 +- python/sglang/srt/layers/layernorm.py | 7 +- sgl-kernel/csrc/cpu/norm.cpp | 173 ++++++++++++++++++++ sgl-kernel/csrc/cpu/torch_extension_cpu.cpp | 10 ++ test/srt/cpu/test_norm.py | 61 ++++++- 5 files changed, 258 insertions(+), 10 deletions(-) diff --git a/.github/CI_PERMISSIONS.json b/.github/CI_PERMISSIONS.json index 0659f41c4..095457c0d 100644 --- a/.github/CI_PERMISSIONS.json +++ b/.github/CI_PERMISSIONS.json @@ -6,13 +6,6 @@ "reason": "custom override", "can_rerun_stage": true }, - "MingxuZh": { - "can_tag_run_ci_label": true, - "can_rerun_failed_ci": true, - "cooldown_interval_minutes": 0, - "reason": "custom override", - "can_rerun_stage": true - }, "Alcanderian": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, @@ -167,6 +160,13 @@ "reason": "custom override", "can_rerun_stage": true }, + "MingxuZh": { + "can_tag_run_ci_label": true, + "can_rerun_failed_ci": true, + "cooldown_interval_minutes": 0, + "reason": "custom override", + "can_rerun_stage": true + }, "Oasis-Git": { "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, @@ -227,7 +227,8 @@ "can_tag_run_ci_label": true, "can_rerun_failed_ci": true, "cooldown_interval_minutes": 0, - "reason": "custom override" + "reason": "custom override", + "can_rerun_stage": true }, "XiaotongJiang": { "can_tag_run_ci_label": true, diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 3293a8a59..bf9098fc0 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -361,7 +361,12 @@ class LayerNorm(CustomOp): self, x: torch.Tensor, ) -> torch.Tensor: - return self.forward_native(x) + if _is_cpu_amx_available: + return torch.ops.sgl_kernel.layernorm_cpu( + x, self.weight.data, self.variance_epsilon + ) + else: + return self.forward_native(x) class GemmaRMSNorm(CustomOp): diff --git a/sgl-kernel/csrc/cpu/norm.cpp b/sgl-kernel/csrc/cpu/norm.cpp index 239ee1b22..580c51fe6 100644 --- a/sgl-kernel/csrc/cpu/norm.cpp +++ b/sgl-kernel/csrc/cpu/norm.cpp @@ -302,6 +302,116 @@ void fused_rmsnorm_gated_kernel_impl( } // anonymous namespace +template +void fused_add_layernorm_kernel_impl( + scalar_t* __restrict__ input, + scalar_t* __restrict__ residual, + const scalar_t* __restrict__ weight, + float* __restrict__ buffer, + int64_t batch_size, + 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; + + for (int64_t i = begin; i < end; ++i) { + scalar_t* __restrict__ input_ptr = input + i * input_strideN; + scalar_t* __restrict__ residual_ptr{(scalar_t*)nullptr}; + if (residual != nullptr) { + residual_ptr = residual + i * hidden_size; + } + + // First pass: compute mean and var + fVec sum_fvec{fVec(0.0)}, sum_sq_fvec{fVec(0.0)}; + float sum_val{0.0}, sum_sq_val{0.0}; + int64_t d{0}; + +#pragma GCC unroll 4 + for (; d <= hidden_size - kVecSize; d += kVecSize) { + bVec x_bvec = bVec::loadu(input_ptr + d); + fVec x_fvec0, x_fvec1; + std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); + + if (residual_ptr != nullptr) { + 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); + + x_fvec0 += r_fvec0; + x_fvec1 += r_fvec1; + + bVec out_bvec = convert_from_float_ext(x_fvec0, x_fvec1); + out_bvec.store(residual_ptr + d); + } + + sum_fvec += x_fvec0; + sum_fvec += x_fvec1; + sum_sq_fvec += x_fvec0 * x_fvec0; + sum_sq_fvec += x_fvec1 * x_fvec1; + + x_fvec0.store(buffer_ptr + d); + x_fvec1.store(buffer_ptr + d + fVec::size()); + } +#pragma GCC unroll 4 + for (; d < hidden_size; ++d) { + float x_val = static_cast(input_ptr[d]); + if (residual_ptr != nullptr) { + float r_val = static_cast(residual_ptr[d]); + x_val += r_val; + residual_ptr[d] = static_cast(x_val); + } + + sum_val += x_val; + sum_sq_val += x_val * x_val; + buffer_ptr[d] = x_val; + } + + // Var(X) = E(X^2) - (E(X))^2 + // Refer to FlashInfer impl: + // https://github.com/flashinfer-ai/flashinfer/blob/6bb01d19c2d9ab3b6a3a5e9e97448891a5ed2844/include/flashinfer/norm.cuh#L554 + sum_val += vec_reduce_sum(sum_fvec); + sum_sq_val += vec_reduce_sum(sum_sq_fvec); + + float mean{sum_val / hidden_size}; + float mean_sq{sum_sq_val / hidden_size}; + float variance{mean_sq - (mean * mean)}; + float rsqrt_var{float(1) / std::sqrt(variance + eps)}; + + const fVec mean_fvec = fVec(mean); + const fVec scale_fvec = fVec(rsqrt_var); + + // Second pass: apply normalization +#pragma GCC unroll 4 + 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); + + 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); + } +#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); + } + } + }); +} // anonymous namespace + // input : {batch_size, hidden_size} at::Tensor l2norm_cpu(at::Tensor& input, double eps) { RECORD_FUNCTION("sgl-kernel::l2norm_cpu", std::vector({input})); @@ -346,6 +456,35 @@ at::Tensor rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) { return output; } +// input : {batch_size, hidden_size} +// weight: {hidden_size} +void layernorm_cpu(at::Tensor& input, at::Tensor& weight, 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); + 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); + + 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( + input.data_ptr(), + nullptr, + weight.data_ptr(), + buffer.data_ptr(), + batch_size, + hidden_size, + input_strideN, + eps); + }); +} + // input : {batch_size, hidden_size} // weight: {hidden_size} // gate: {batch_size, hidden_size} @@ -415,3 +554,37 @@ void fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& eps); }); } + +// input : {batch_size, hidden_size} +// residual: {batch_size, hidden_size} +// weight : {hidden_size} +void fused_add_layernorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, 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); + 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)); + int64_t batch_size = input.size(0); + int64_t hidden_size = input.size(1); + int64_t input_strideN = input.stride(0); + + 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( + input.data_ptr(), + residual.data_ptr(), + weight.data_ptr(), + buffer.data_ptr(), + batch_size, + hidden_size, + input_strideN, + eps); + }); +} diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index aede839fa..bc8b8fcdd 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -33,12 +33,18 @@ at::Tensor l2norm_cpu(at::Tensor& input, double eps); // rmsnorm at::Tensor rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps); +// layernorm +void layernorm_cpu(at::Tensor& input, at::Tensor& weight, double eps); + // qwen3_next_rmsnorm_gated at::Tensor fused_rmsnorm_gated_cpu(at::Tensor& input, at::Tensor& weight, at::Tensor& gate, double eps); // fused_add_rmsnorm void 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); + // topk std::tuple topk_sigmoid_cpu(at::Tensor& hidden_states, at::Tensor& gating_output, int64_t topk, bool renormalize); @@ -324,12 +330,16 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { // norm m.def("rmsnorm_cpu(Tensor input, Tensor weight, float eps) -> Tensor"); m.impl("rmsnorm_cpu", torch::kCPU, &rmsnorm_cpu); + m.def("layernorm_cpu(Tensor(a!) input, Tensor weight, float eps) -> ()"); + 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); m.def("fused_rmsnorm_gated_cpu(Tensor input, Tensor weight, Tensor gate, float eps) -> Tensor"); m.impl("fused_rmsnorm_gated_cpu", torch::kCPU, &fused_rmsnorm_gated_cpu); m.def("fused_add_rmsnorm_cpu(Tensor(a!) input, Tensor residual, Tensor weight, float eps) -> ()"); m.impl("fused_add_rmsnorm_cpu", torch::kCPU, &fused_add_rmsnorm_cpu); + m.def("fused_add_layernorm_cpu(Tensor(a!) input, Tensor residual, Tensor weight, float eps) -> ()"); + m.impl("fused_add_layernorm_cpu", torch::kCPU, &fused_add_layernorm_cpu); // topk m.def("topk_sigmoid_cpu(Tensor hidden_states, Tensor gating_output, int topk, bool renormalize) -> (Tensor, Tensor)"); diff --git a/test/srt/cpu/test_norm.py b/test/srt/cpu/test_norm.py index 8bb2eba02..ade9a8c22 100644 --- a/test/srt/cpu/test_norm.py +++ b/test/srt/cpu/test_norm.py @@ -3,7 +3,7 @@ import unittest from typing import Optional, Tuple, Union import torch -from utils import make_non_contiguous, precision +from utils import make_non_contiguous, parametrize, precision from sglang.test.test_utils import CustomTestCase @@ -131,5 +131,64 @@ class TestFusedRMSNormGated(CustomTestCase): self._norm_test(*params) +class TestLayerNorm(CustomTestCase): + + def _forward_native( + self, + x: torch.Tensor, + weight: torch.Tensor, + variance_epsilon: float, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + orig_dtype = x.dtype + x = x.to(torch.float32) + if residual is not None: + x = x + residual.to(torch.float32) + residual = x.to(orig_dtype) + + (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 + + @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) + weight = 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) + + atol = rtol = precision[ref_ln_out.dtype] + torch.testing.assert_close(x_ln, 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 + ) + ref_add_ln_out, ref_residual = self._forward_native( + ref_x_add_ln, weight, variance_epsilon, ref_residual + ) + + torch.testing.assert_close(x_add_ln, ref_add_ln_out, atol=atol, rtol=rtol) + torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol) + + if __name__ == "__main__": unittest.main()