[CPU] add fused_qk_gemma_norm and refactor norm kernel implementation (#30216)

python/sglang/srt/models/qwen3_5.py

@@ -143,6 +143,10 @@ if _is_cuda:
 
 if _is_cpu:
     fused_sigmoid_mul = torch.ops.sgl_kernel.fused_sigmoid_mul_cpu
+    fused_qk_gemma_rmsnorm = torch.ops.sgl_kernel.fused_qk_gemma_rmsnorm_cpu
+    fused_qk_gemma_rmsnorm_with_gate = (
+        torch.ops.sgl_kernel.fused_qk_gemma_rmsnorm_with_gate_cpu
+    )
 
 if _is_npu:
     from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import (
@@ -876,7 +880,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
                 k_by_head = k.reshape(-1, self.head_dim)
                 k_by_head = self.k_norm(k_by_head)
             current_stream.wait_stream(self.alt_stream)
-        elif _is_hip or _is_xpu:
+        elif _is_hip or _is_xpu or _is_cpu:
             q_by_head, k_by_head = fused_qk_gemma_rmsnorm(
                 q,
                 k,
@@ -1001,7 +1005,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
                 positions=positions,
                 hidden_states=hidden_states,
             )
-        elif (_is_hip or _is_xpu) and self.attn_output_gate:
+        elif (_is_hip or _is_xpu or _is_cpu) and self.attn_output_gate:
             q, k, v, gate = self.forward_prepare_fused_gate(
                 positions=positions,
                 hidden_states=hidden_states,

sgl-kernel/csrc/cpu/torch_extension_cpu.cpp

@@ -58,6 +58,23 @@ at::Tensor fused_add_layernorm_cpu(
     const std::optional<at::Tensor>& bias,
     double eps);
 
+// fused_qk_gemma_rmsnorm
+std::tuple<at::Tensor, at::Tensor> fused_qk_gemma_rmsnorm_cpu(
+    const at::Tensor& q,
+    const at::Tensor& k,
+    const at::Tensor& q_weight,
+    const at::Tensor& k_weight,
+    double eps,
+    int64_t head_dim);
+std::tuple<at::Tensor, at::Tensor, at::Tensor> fused_qk_gemma_rmsnorm_with_gate_cpu(
+    const at::Tensor& q_gate,
+    const at::Tensor& k,
+    const at::Tensor& q_weight,
+    const at::Tensor& k_weight,
+    double eps,
+    int64_t head_dim,
+    int64_t num_head);
+
 // topk
 std::tuple<at::Tensor, at::Tensor>
 topk_sigmoid_cpu(at::Tensor& hidden_states, at::Tensor& gating_output, int64_t topk, bool renormalize);
@@ -468,6 +485,15 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       "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);
+  m.def(
+      "fused_qk_gemma_rmsnorm_cpu(Tensor q, Tensor k, Tensor q_weight, Tensor k_weight, float eps, int head_dim) -> "
+      "(Tensor, Tensor)");
+  m.impl("fused_qk_gemma_rmsnorm_cpu", torch::kCPU, &fused_qk_gemma_rmsnorm_cpu);
+  m.def(
+      "fused_qk_gemma_rmsnorm_with_gate_cpu(Tensor q_gate, Tensor k, Tensor q_weight, Tensor k_weight, float eps, int "
+      "head_dim, int num_head) -> "
+      "(Tensor, Tensor, Tensor)");
+  m.impl("fused_qk_gemma_rmsnorm_with_gate_cpu", torch::kCPU, &fused_qk_gemma_rmsnorm_with_gate_cpu);
 
   // topk
   m.def("topk_sigmoid_cpu(Tensor hidden_states, Tensor gating_output, int topk, bool renormalize) -> (Tensor, Tensor)");

sgl-kernel/csrc/cpu/vec.h

@@ -451,6 +451,62 @@ inline std::tuple<__m512i, __m512i> transpose_2x32_16bit(__m512i r0, __m512i r1)
 }
 #pragma GCC diagnostic pop
 
+// Note: mapped from aten exp_u20
+inline __attribute__((always_inline)) __m512 _mm512_exp_u20_ps(const __m512 values) {
+  const __m512 vec_factorial_1 = _mm512_set1_ps(0.999999701f);
+  const __m512 vec_factorial_2 = _mm512_set1_ps(0.499991506f);
+  const __m512 vec_factorial_3 = _mm512_set1_ps(0.166676521f);
+  const __m512 vec_factorial_4 = _mm512_set1_ps(0.0418978221f);
+  const __m512 vec_factorial_5 = _mm512_set1_ps(0.00828929059f);
+  const __m512 vec_exp_log2ef = _mm512_castsi512_ps(_mm512_set1_epi32(0x3fb8aa3b));  // log2(e)
+  const __m512 vec_half = _mm512_set1_ps(0.5f);
+  const __m512 vec_one = _mm512_set1_ps(1.f);
+  const __m512 vec_zero = _mm512_set1_ps(0.f);
+  const __m512 vec_two = _mm512_set1_ps(2.f);
+  const __m512 vec_ln2f = _mm512_castsi512_ps(_mm512_set1_epi32(0x3f317218));
+  const __m512 vec_ln_flt_min = _mm512_castsi512_ps(_mm512_set1_epi32(0xc2aeac50));
+  const __m512 vec_ln_flt_max = _mm512_castsi512_ps(_mm512_set1_epi32(0x42b17218));
+  const __m512i vec_127 = _mm512_set1_epi32(0x0000007f);
+  const int n_mantissa_bits = 23;
+
+  // exp(x) =
+  // = exp(n * ln(2) + r) // divide x by ln(2) and get quot and rem
+  // = 2^n * exp(r) // simplify the exp(n*ln(2)) expression
+
+  auto less_ln_flt_min_mask = _mm512_cmp_ps_mask(values, vec_ln_flt_min, 1 /*_CMP_LT_OS*/);
+  auto vec_src = _mm512_min_ps(values, vec_ln_flt_max);
+  vec_src = _mm512_max_ps(vec_src, vec_ln_flt_min);
+
+  // fx = floorf(x * log2ef + 0.5)
+  auto vec_fx = _mm512_fmadd_ps(vec_src, vec_exp_log2ef, vec_half);
+  auto vec_fx_i = _mm512_cvt_roundps_epi32(vec_fx, _MM_FROUND_TO_NEG_INF | _MM_FROUND_NO_EXC);
+  vec_fx = _mm512_cvtepi32_ps(vec_fx_i);
+
+  // x = x - fx * ln2
+  auto vec_exp_poly = _mm512_fnmadd_ps(vec_fx, vec_ln2f, vec_src);
+
+  // compute polynomial
+  auto vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_factorial_5, vec_factorial_4);
+  vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_3);
+  vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_2);
+  vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_1);
+  vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_one);
+
+  // compute 2^(n-1)
+  auto vec_exp_number = _mm512_sub_ps(vec_fx, vec_one);
+  auto vec_exp_number_i = _mm512_cvtps_epi32(vec_exp_number);
+  auto vec_two_pow_n_i = _mm512_add_epi32(vec_exp_number_i, vec_127);
+  vec_two_pow_n_i = _mm512_slli_epi32(vec_two_pow_n_i, n_mantissa_bits);
+  auto vec_two_pow_n = _mm512_castsi512_ps(vec_two_pow_n_i);
+  vec_two_pow_n = _mm512_mask_blend_ps(less_ln_flt_min_mask, vec_two_pow_n, vec_zero);
+
+  // y = y * 2^n
+  vec_res = _mm512_mul_ps(vec_res, vec_two_pow_n);
+  vec_res = _mm512_mul_ps(vec_res, vec_two);
+  return vec_res;
+}
+
+// Note: mapped from aten fexp_u20
 inline __attribute__((always_inline)) __m512 _mm512_fexp_u20_ps(const __m512 values) {
   const __m512 vec_c0 = _mm512_set1_ps(0.00010703434948458272f);
   const __m512 vec_c1 = _mm512_set1_ps(0.30354260500649682f);

test/registered/cpu/test_norm.py

@@ -1,26 +1,29 @@
-import itertools
+import sys
-import unittest
 from typing import Optional, Tuple, Union
 
+import pytest
 import torch
-from utils import make_non_contiguous, parametrize, precision
+from utils import make_non_contiguous, precision
 
 from sglang.test.ci.ci_register import register_cpu_ci
-from sglang.test.test_utils import CustomTestCase
 
 register_cpu_ci(est_time=10, suite="base-b-test-cpu")
 register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64")
 
 torch.manual_seed(1234)
 
+DTYPES = [torch.float16, torch.bfloat16]
+DTYPE_IDS = ["float16", "bfloat16"]
+eps = 1e-6
 
-class TestNorm(CustomTestCase):
+
+class TestNorm:
 
     def _forward_native(
         self,
         x: torch.Tensor,
         weight: torch.Tensor,
-        variance_epsilon: float = 1e-6,
+        variance_epsilon: float = eps,
         residual: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         orig_dtype = x.dtype
@@ -41,7 +44,7 @@ class TestNorm(CustomTestCase):
         return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
 
     def _gemma3_rmsnorm_native(
-        self, x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float = 1e-6
+        self, x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float = eps
     ):
         output = self._norm(x.float(), variance_epsilon)
         output = output * (1.0 + weight.float())
@@ -51,7 +54,7 @@ class TestNorm(CustomTestCase):
         self,
         x: torch.Tensor,
         weight: torch.Tensor,
-        variance_epsilon: float = 1e-6,
+        variance_epsilon: float = eps,
         residual: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         orig_dtype = x.dtype
@@ -66,144 +69,91 @@ class TestNorm(CustomTestCase):
         x = x.to(orig_dtype)
         return x if residual is None else (x, residual)
 
-    @parametrize(
+    @pytest.mark.parametrize("dtype", DTYPES, ids=DTYPE_IDS)
-        m=[4096, 1024],
+    @pytest.mark.parametrize("hidden_size", [2048, 512])
-        n=[4096, 4109],
+    @pytest.mark.parametrize("batch_size", [32, 121])
-        dtype=[torch.float16, torch.bfloat16],
+    def test_l2norm(self, batch_size, hidden_size, dtype):
-    )
-    def test_norm(self, m, n, dtype):
 
-        x = torch.randn([m, n], dtype=dtype)
+        x = torch.randn([batch_size, hidden_size], 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([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(
-        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)
         fake_ones_weight = torch.ones(hidden_size, dtype=dtype)
-        variance_epsilon = 1e-6
 
-        out = torch.ops.sgl_kernel.l2norm_cpu(x, variance_epsilon)
+        out = torch.ops.sgl_kernel.l2norm_cpu(x, eps)
-        ref_out = self._forward_native(x, fake_ones_weight, variance_epsilon)
+        ref_out = self._forward_native(x, fake_ones_weight, eps)
 
         atol = rtol = precision[ref_out.dtype]
         torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
 
-    @parametrize(
+    @pytest.mark.parametrize("dtype", DTYPES, ids=DTYPE_IDS)
-        m=[4096, 1024],
+    @pytest.mark.parametrize("hidden_size", [2048, 512])
-        n=[4096, 4109],
+    @pytest.mark.parametrize("batch_size", [32, 121])
-        dtype=[torch.float16, torch.bfloat16],
+    @pytest.mark.parametrize("seq_len", [None, 2], ids=["2d", "3d"])
-    )
+    def test_rmsnorm(self, seq_len, batch_size, hidden_size, dtype):
-    def test_gemma_rmsnorm(self, m, n, dtype):
 
-        x = torch.randn([m, n], dtype=dtype)
+        if seq_len is None:
+            x = torch.randn([batch_size, hidden_size], dtype=dtype)
+        else:
+            x = torch.randn([batch_size, seq_len, hidden_size], dtype=dtype)
         x = make_non_contiguous(x)
-        hidden_size = x.size(-1)
+        residual = torch.randn(x.shape, dtype=dtype)
         weight = torch.randn(hidden_size, dtype=dtype)
-        variance_epsilon = 1e-6
 
-        out = torch.ops.sgl_kernel.gemma_rmsnorm_cpu(x, weight, variance_epsilon)
+        out = torch.ops.sgl_kernel.rmsnorm_cpu(x, weight, eps)
-        ref_out = self._gemma_rmsnorm_native(x, weight, variance_epsilon)
+        ref_out = self._forward_native(x, weight, eps)
 
         atol = rtol = precision[ref_out.dtype]
         torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
 
         ref_x = x.clone()
-        residual = torch.randn([m, hidden_size], dtype=dtype)
         ref_residual = residual.clone()
 
-        torch.ops.sgl_kernel.gemma_fused_add_rmsnorm_cpu(
+        torch.ops.sgl_kernel.fused_add_rmsnorm_cpu(x, residual, weight, eps)
-            x, residual, weight, variance_epsilon
+
-        )
+        ref_x, ref_residual = self._forward_native(ref_x, weight, eps, ref_residual)
+
+        torch.testing.assert_close(x, ref_x, atol=atol, rtol=rtol)
+        torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol)
+
+    @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bfloat16"])
+    @pytest.mark.parametrize("hidden_size", [2048, 256, 33])
+    @pytest.mark.parametrize("batch_size", [32, 121])
+    def test_gemma_rmsnorm(self, batch_size, hidden_size, dtype):
+
+        x = torch.randn([batch_size, hidden_size], dtype=dtype)
+        x = make_non_contiguous(x)
+        weight = torch.randn(hidden_size, dtype=dtype)
+
+        out = torch.ops.sgl_kernel.gemma_rmsnorm_cpu(x, weight, eps)
+        ref_out = self._gemma_rmsnorm_native(x, weight, eps)
+
+        atol = rtol = precision[ref_out.dtype]
+        torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
+
+        ref_x = x.clone()
+        residual = torch.randn([batch_size, hidden_size], dtype=dtype)
+        ref_residual = residual.clone()
+
+        torch.ops.sgl_kernel.gemma_fused_add_rmsnorm_cpu(x, residual, weight, eps)
 
         ref_x, ref_residual = self._gemma_rmsnorm_native(
-            ref_x, weight, variance_epsilon, ref_residual
+            ref_x, weight, eps, 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(
+    @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bfloat16"])
-        m=[4096, 1024],
+    @pytest.mark.parametrize("hidden_size", [128, 256])
-        n=[4096, 4109],
+    @pytest.mark.parametrize("batch_size", [32, 121])
-        dtype=[torch.float16, torch.bfloat16],
+    def test_gemma3_rmsnorm(self, batch_size, hidden_size, dtype):
-    )
-    def test_gemma3_rmsnorm(self, m, n, dtype):
         x_list = [
-            torch.randn([m, n], dtype=dtype),
+            torch.randn([batch_size, hidden_size], dtype=dtype),
-            torch.randn([1, m, 2, n], dtype=dtype),
+            torch.randn([batch_size, 16, 2, hidden_size], dtype=dtype),
         ]
         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.gemma3_rmsnorm_cpu(x, weight, eps)
-            out = torch.ops.sgl_kernel.gemma3_rmsnorm_cpu(x, weight, variance_epsilon)
+            ref_out = self._gemma3_rmsnorm_native(x, weight, eps)
-            ref_out = self._gemma3_rmsnorm_native(x, weight, variance_epsilon)
 
             atol = rtol = precision[ref_out.dtype]
             torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
@@ -212,7 +162,7 @@ class TestNorm(CustomTestCase):
         self,
         x: torch.Tensor,
         weight: torch.Tensor,
-        variance_epsilon: float = 1e-6,
+        variance_epsilon: float = eps,
         scale_shift: float = 0.0,
         with_scale: bool = True,
     ):
@@ -221,53 +171,41 @@ class TestNorm(CustomTestCase):
             output = output * (weight.float() + scale_shift)
         return output.type_as(x)
 
-    @parametrize(
+    @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bfloat16"])
-        m=[4096, 1024],
+    @pytest.mark.parametrize("hidden_size", [128, 2048])
-        n=[4096, 4109],
+    @pytest.mark.parametrize("batch_size", [32, 121])
-        dtype=[torch.float16, torch.bfloat16],
+    @pytest.mark.parametrize("scale_shift", [0.0, 1.0], ids=["shift0.0", "shift1.0"])
-    )
+    @pytest.mark.parametrize("with_scale", [True, False], ids=["scale", "no-scale"])
-    def test_gemma4_rmsnorm(self, m, n, dtype):
+    def test_gemma4_rmsnorm(
-        for scale_shift, with_scale in [
+        self, batch_size, hidden_size, dtype, scale_shift, with_scale
-            (0.0, True),
+    ):
-            (1.0, True),
-            (0.0, False),
-            (1.0, False),
-        ]:
         x_list = [
-                torch.randn([m, n], dtype=dtype),
+            torch.randn([batch_size, hidden_size], dtype=dtype),
-                torch.randn([4, m, n], dtype=dtype),
+            torch.randn([batch_size, 4, hidden_size], 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
+                x, weight, eps, scale_shift, with_scale
             )
             ref_out = self._gemma4_rmsnorm_native(
-                    x, weight, variance_epsilon, scale_shift, with_scale
+                x, weight, eps, scale_shift, with_scale
             )
 
             atol = rtol = precision[ref_out.dtype]
             torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
 
 
-class TestFusedRMSNormGated(CustomTestCase):
+class TestFusedRMSNormGated:
-    M = [4096, 1024]
-    N = [4096, 4096 + 13]
-    dtype = [torch.float16, torch.bfloat16]
 
     def _forward_native(
         self,
         hidden_states: torch.Tensor,
         weight: torch.Tensor,
-        variance_epsilon: float = 1e-6,
+        variance_epsilon: float = eps,
         gate: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         input_dtype = hidden_states.dtype
@@ -280,37 +218,29 @@ class TestFusedRMSNormGated(CustomTestCase):
 
         return hidden_states.to(input_dtype)
 
-    def _norm_test(self, m, n, dtype):
+    @pytest.mark.parametrize("dtype", DTYPES, ids=DTYPE_IDS)
-
+    @pytest.mark.parametrize("hidden_size", [64, 1024 + 13])
-        x = torch.randn([m, n], dtype=dtype)
+    @pytest.mark.parametrize("batch_size", [32, 121])
+    def test_fused_rmsnorm_gated(self, batch_size, hidden_size, dtype):
+        x = torch.randn([batch_size, hidden_size], dtype=dtype)
         x = make_non_contiguous(x)
-        batch_size = x.size(0)
-        hidden_size = x.size(-1)
         weight = torch.randn(hidden_size, dtype=dtype)
-        variance_epsilon = 1e-6
         gate = torch.randn([batch_size, hidden_size], dtype=dtype)
 
-        out = torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu(
+        out = torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu(x, weight, gate, eps)
-            x, weight, gate, variance_epsilon
+        ref_out = self._forward_native(x, weight, eps, gate)
-        )
-        ref_out = self._forward_native(x, weight, variance_epsilon, gate)
 
-        atol = rtol = precision[ref_out.dtype] * 2
+        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)
 
-
+class TestLayerNorm:
-class TestLayerNorm(CustomTestCase):
 
     def _forward_native(
         self,
         x: torch.Tensor,
         weight: torch.Tensor,
-        variance_epsilon: float,
+        variance_epsilon: float = eps,
         residual: Optional[torch.Tensor] = None,
         bias: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
@@ -328,109 +258,141 @@ class TestLayerNorm(CustomTestCase):
         x = x.to(orig_dtype)
         return x if residual is None else (x, residual)
 
-    @parametrize(
+    @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bfloat16"])
-        m=[4096, 1024],
+    @pytest.mark.parametrize("batch_size", [32, 121])
-        n=[4096, 4109],
+    @pytest.mark.parametrize("hidden_size", [128, 4096, 533])
-        dtype=[torch.float16, torch.bfloat16],
+    @pytest.mark.parametrize("has_bias", [False, True], ids=["no-bias", "bias"])
-    )
+    def test_layernorm(
-    def test_norm_input_2d(self, m: int, n: int, dtype: torch.dtype) -> None:
+        self,
-        x = torch.randn([m, n], dtype=dtype)
+        batch_size: int,
-        x = make_non_contiguous(x)
+        hidden_size: int,
-        hidden_size = x.size(-1)
+        has_bias: bool,
-        weight = torch.randn(hidden_size, dtype=dtype)
+        dtype: torch.dtype,
-        bias = torch.randn(hidden_size, dtype=dtype)
+    ) -> None:
-        variance_epsilon = 1e-6
+        x_list = [
+            torch.randn([batch_size, hidden_size], dtype=dtype),
+            torch.randn([batch_size, 3, hidden_size], dtype=dtype),
+        ]
 
-        ln_out = torch.ops.sgl_kernel.layernorm_cpu(x, weight, None, variance_epsilon)
+        for x in x_list:
-        ref_ln_out = self._forward_native(x, weight, variance_epsilon)
+            x = make_non_contiguous(x)
+            weight = torch.randn(hidden_size, dtype=dtype)
+            bias = torch.randn(hidden_size, dtype=dtype) if has_bias else None
+
+            ln_out = torch.ops.sgl_kernel.layernorm_cpu(x, weight, bias, eps)
+            ref_ln_out = self._forward_native(x, weight, eps, residual=None, bias=bias)
 
             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)
+            residual = torch.randn(x.shape, dtype=dtype)
-        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([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
+                x, residual, weight, bias, eps
             )
             ref_add_ln_out, ref_residual = self._forward_native(
-            x, weight, variance_epsilon, residual=ref_residual
+                x, weight, eps, 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)
 
-        residual = torch.randn([m, hidden_size], dtype=dtype)
-        ref_residual = residual.clone()
 
-        add_ln_out = torch.ops.sgl_kernel.fused_add_layernorm_cpu(
+class TestFusedQKGemmaRMSNorm:
-            x, residual, weight, bias, variance_epsilon
+
+    def _gemma_rmsnorm_per_head_native(
+        self,
+        x: torch.Tensor,
+        weight: torch.Tensor,
+        head_dim: int,
+        variance_epsilon: float = eps,
+    ) -> torch.Tensor:
+        orig_dtype = x.dtype
+        x_f = x.to(torch.float32).reshape(-1, head_dim)
+        variance = x_f.pow(2).mean(dim=-1, keepdim=True)
+        x_f = x_f * torch.rsqrt(variance + variance_epsilon)
+        x_f = x_f * (1.0 + weight.to(torch.float32))
+        return x_f.to(orig_dtype).reshape_as(x)
+
+    @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bfloat16"])
+    @pytest.mark.parametrize(
+        "batch_size,num_head,num_head_kv,head_dim",
+        [
+            (8, 4, 2, 128),
+            (17, 8, 2, 64),
+            (5, 3, 1, 96),
+        ],
     )
-        ref_add_ln_out, ref_residual = self._forward_native(
+    def test_fused_qk_gemma_rmsnorm(
-            x, weight, variance_epsilon, residual=ref_residual, bias=bias
+        self, batch_size: int, num_head: int, num_head_kv: int, head_dim: int, dtype
+    ):
+        q = torch.randn([batch_size, num_head * head_dim], dtype=dtype)
+        k = torch.randn([batch_size, num_head_kv * head_dim], dtype=dtype)
+
+        # Keep last dim contiguous but make base storage non-contiguous to stress stride handling.
+        q = make_non_contiguous(q)
+        k = make_non_contiguous(k)
+
+        q_weight = torch.randn(head_dim, dtype=dtype)
+        k_weight = torch.randn(head_dim, dtype=dtype)
+
+        q_out, k_out = torch.ops.sgl_kernel.fused_qk_gemma_rmsnorm_cpu(
+            q, k, q_weight, k_weight, eps, head_dim
         )
 
-        torch.testing.assert_close(add_ln_out, ref_add_ln_out, atol=atol, rtol=rtol)
+        ref_q_out = self._gemma_rmsnorm_per_head_native(q, q_weight, head_dim, eps)
-        torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol)
+        ref_k_out = self._gemma_rmsnorm_per_head_native(k, k_weight, head_dim, eps)
 
-    @parametrize(
+        atol = rtol = precision[ref_q_out.dtype]
-        l=[4096, 1024],
+        torch.testing.assert_close(q_out, ref_q_out, atol=atol, rtol=rtol)
-        m=[1, 4],
+        torch.testing.assert_close(k_out, ref_k_out, atol=atol, rtol=rtol)
-        n=[4096, 4109, 2304],
+
-        dtype=[torch.float16, torch.bfloat16],
+    @pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bfloat16"])
+    @pytest.mark.parametrize(
+        "batch_size,num_head,num_head_kv,head_dim",
+        [
+            (8, 4, 2, 128),
+            (17, 8, 2, 64),
+            (5, 3, 1, 96),
+        ],
     )
-    def test_norm_input_3d(self, l: int, m: int, n: int, dtype: torch.dtype) -> None:
+    def test_fused_qk_gemma_rmsnorm_with_gate(
-        x = torch.randn([l, m, n], dtype=dtype)
+        self, batch_size: int, num_head: int, num_head_kv: int, head_dim: int, dtype
-        x = make_non_contiguous(x)
+    ):
-        hidden_size = x.size(-1)
+        q = torch.randn([batch_size, num_head, head_dim], dtype=dtype)
-        weight = torch.randn(hidden_size, dtype=dtype)
+        gate = torch.randn([batch_size, num_head, head_dim], dtype=dtype)
-        bias = torch.randn(hidden_size, dtype=dtype)
+        q_gate = torch.cat((q, gate), dim=-1).reshape(
-        variance_epsilon = 1e-6
+            batch_size, num_head * head_dim * 2
-
-        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)
+        k = torch.randn([batch_size, num_head_kv * head_dim], dtype=dtype)
 
-        residual = torch.randn([l, m, hidden_size], dtype=dtype)
+        q_gate = make_non_contiguous(q_gate)
-        ref_residual = residual.clone()
+        k = make_non_contiguous(k)
 
-        add_ln_out = torch.ops.sgl_kernel.fused_add_layernorm_cpu(
+        q_weight = torch.randn(head_dim, dtype=dtype)
-            x, residual, weight, None, variance_epsilon
+        k_weight = torch.randn(head_dim, dtype=dtype)
+
+        q_out, k_out, gate_out = (
+            torch.ops.sgl_kernel.fused_qk_gemma_rmsnorm_with_gate_cpu(
+                q_gate, k, q_weight, k_weight, eps, head_dim, num_head
             )
-        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)
+        ref_q_out = self._gemma_rmsnorm_per_head_native(q, q_weight, head_dim, eps)
-        torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol)
+        ref_k_out = self._gemma_rmsnorm_per_head_native(k, k_weight, head_dim, eps)
 
-        residual = torch.randn([l, m, hidden_size], dtype=dtype)
+        atol = rtol = precision[ref_q_out.dtype]
-        ref_residual = residual.clone()
+        torch.testing.assert_close(
-
+            q_out, ref_q_out.reshape(-1, head_dim), atol=atol, rtol=rtol
-        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(
+        torch.testing.assert_close(
-            x, weight, variance_epsilon, residual=ref_residual, bias=bias
+            k_out, ref_k_out.reshape(-1, head_dim), atol=atol, rtol=rtol
+        )
+        torch.testing.assert_close(
+            gate_out, gate.reshape(-1, head_dim), 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)
 
 
 if __name__ == "__main__":
-    unittest.main()
+    sys.exit(pytest.main([__file__]))

test/registered/cpu/utils.py

@@ -511,11 +511,14 @@ class MXFP4QuantizeUtil:
 
 
 def make_non_contiguous(x: torch.Tensor) -> torch.Tensor:
-    """
+    # Make a tensor non-contiguous without changing shape.
-    Make a tensor non-contiguous by slicing it via last dimension.
+    if not x.is_contiguous():
-    """
+        return x
+
     last_dim = x.shape[-1]
-    return x[..., : last_dim // 2] if x.is_contiguous() else x
+    expanded = torch.empty(*x.shape[:-1], last_dim + 32, dtype=x.dtype, device=x.device)
+    expanded[..., :last_dim].copy_(x)
+    return expanded.narrow(-1, 0, last_dim)
 
 
 def awq_reverse_reorder_int_tensor(int_tensor, bits: int):