[CPU] support LayerNorm with 3D shape (#15075)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

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)

sgl-kernel/csrc/cpu/norm.cpp

@@ -388,26 +388,32 @@ void fused_rmsnorm_gated_kernel_impl(
 
 template <typename scalar_t>
 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<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) {
+
-    int tid = at::get_thread_num();
+  const bool has_residual{residual != nullptr};
-    float* __restrict__ buffer_ptr = buffer + tid * hidden_size;
+  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<float>(input_ptr[d]);
-        if (residual_ptr != nullptr) {
+        if (has_residual) {
           float r_val = static_cast<float>(residual_ptr[d]);
           x_val += r_val;
           residual_ptr[d] = static_cast<scalar_t>(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<scalar_t>(x_fvec0, x_fvec1);
+        if (has_bias) {
-        x_bvec.store(input_ptr + d);
+          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<scalar_t>(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<float>(weight[d]);
-        input_ptr[d] = static_cast<scalar_t>(x_val);
+        if (has_bias) {
+          x_val += static_cast<float>(bias[d]);
+        }
+        out_ptr[d] = static_cast<scalar_t>(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<at::Tensor>& bias, double eps) {
   RECORD_FUNCTION("sgl-kernel::layernorm_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));
+  if (bias.has_value()) {
-  int64_t batch_size = input.size(0);
+    CHECK_DIM(1, bias.value());
-  int64_t hidden_size = input.size(1);
+    CHECK_EQ(bias.value().size(0), weight.size(0));
-  int64_t input_strideN = input.stride(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<scalar_t>(
+        output.data_ptr<scalar_t>(),
         input.data_ptr<scalar_t>(),
         nullptr,
         weight.data_ptr<scalar_t>(),
+        conditional_data_ptr<scalar_t>(bias),
         buffer.data_ptr<float>(),
         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}
+// input   : {batch_size, hidden_size} or {batch_size, seq_len, hidden_size}
-// residual: {batch_size, 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<at::Tensor>& bias,
+    double eps) {
   RECORD_FUNCTION("sgl-kernel::fused_add_layernorm_cpu", std::vector<c10::IValue>({input, residual, weight}));
   CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
   CHECK_INPUT(residual);
   CHECK_INPUT(weight);
-  CHECK_DIM(2, input);
+  int64_t inp_dim{input.dim()}, res_dim{residual.dim()};
-  CHECK_DIM(2, residual);
+  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));
+  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);
+  }
-  int64_t hidden_size = input.size(1);
-  int64_t input_strideN = input.stride(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<scalar_t>(
+        output.data_ptr<scalar_t>(),
         input.data_ptr<scalar_t>(),
         residual.data_ptr<scalar_t>(),
         weight.data_ptr<scalar_t>(),
+        conditional_data_ptr<scalar_t>(bias),
         buffer.data_ptr<float>(),
         batch_size,
+        seq_len,
         hidden_size,
         input_strideN,
         eps);
   });
+  return output;
 }

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<at::Tensor>& 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<at::Tensor>& bias,
+    double eps);
 
 // topk
 std::tuple<at::Tensor, at::Tensor>
@@ -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

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
+        x = x * weight.to(torch.float32)
-        if residual is None:
+        if bias is not None:
-            return x
+            x = x + bias.to(torch.float32)
-        else:
+        x = x.to(orig_dtype)
-            return x, residual
+        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:
+    def test_norm_input_2d(self, m: int, n: int, dtype: torch.dtype) -> None:
-        x_ln = torch.randn([m, n], dtype=dtype)
+        x = torch.randn([m, n], dtype=dtype)
-        x_ln = make_non_contiguous(x_ln)
+        x = make_non_contiguous(x)
-        ref_x_ln = x_ln.clone()
+        hidden_size = x.size(-1)
-        hidden_size = x_ln.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)
+        ln_out = torch.ops.sgl_kernel.layernorm_cpu(x, weight, None, variance_epsilon)
-        ref_ln_out = self._forward_native(ref_x_ln, weight, 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(
+        add_ln_out = torch.ops.sgl_kernel.fused_add_layernorm_cpu(
-            x_add_ln, residual, weight, variance_epsilon
+            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)