[CPU] Add gemma4_rmsnorm_cpu kernel (#22842)

Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com>
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

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

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}")

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<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) {
+  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;
+      scalar_t* __restrict__ out_ptr = output + bi * output_strideB + si * output_strideS;
-      const scalar_t* __restrict__ input_ptr = input + i * hidden_size;
+      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<float>(input_ptr[d]);
         out_ptr[d] = static_cast<scalar_t>(x_val * rsqrt_var);
       }
+      // move to the next index
+      data_index_step(bi, batch_size, si, seq_len);
     }
   });
 }
+
 template <typename scalar_t, typename func_t, typename vec_func_t>
 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<float>;
 
   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;
+      scalar_t* __restrict__ out_ptr = output + bi * output_strideB + si * output_strideS;
-      const scalar_t* __restrict__ input_ptr = input + i * input_strideN;
+      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<float>(weight[d]);
         out_ptr[d] = static_cast<scalar_t>(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<float>;
 
   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<float>(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<scalar_t>(output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), batch_size, hidden_size, eps);
+    l2norm_kernel_impl<scalar_t>(
+        output.data_ptr<scalar_t>(),
+        input.data_ptr<scalar_t>(),
+        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<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));
+  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<float>;
@@ -550,8 +595,12 @@ at::Tensor rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps) {
         input.data_ptr<scalar_t>(),
         weight.data_ptr<scalar_t>(),
         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<float>;
@@ -628,8 +678,12 @@ at::Tensor gemma_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps)
         input.data_ptr<scalar_t>(),
         weight.data_ptr<scalar_t>(),
         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<float>;
@@ -662,8 +717,12 @@ at::Tensor gemma3_rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps)
           input.data_ptr<scalar_t>(),
           weight.data_ptr<scalar_t>(),
           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<c10::IValue>({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<float>(scale_shift);
+    AT_DISPATCH_REDUCED_FLOATING_TYPES(input.scalar_type(), "gemma4_rmsnorm_kernel", [&] {
+      using Vec = at::vec::Vectorized<float>;
+      Vec shift_vec = Vec(shift);
+      rmsnorm_kernel_impl<scalar_t>(
+          output.data_ptr<scalar_t>(),
+          input.data_ptr<scalar_t>(),
+          weight.data_ptr<scalar_t>(),
+          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<scalar_t>(
+          output.data_ptr<scalar_t>(),
+          input.data_ptr<scalar_t>(),
+          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}
+// 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_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps) {
   RECORD_FUNCTION("sgl-kernel::fused_add_rmsnorm_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);
   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), residual.size(-1));
-  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 input_strideN = input.stride(0);
+  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<scalar_t>(),
         buffer.data_ptr<float>(),
         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<scalar_t>(),
         buffer.data_ptr<float>(),
         batch_size,
+        1,
         hidden_size,
         input_strideN,
+        0,
         [](float x) { return x + 1; },
         [one_vec](Vec x) { return x + one_vec; },
         eps);

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");

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):
+    def _gemma4_rmsnorm_native(
-        for params in itertools.product(self.M, self.N, self.dtype):
+        self,
-            with self.subTest(m=params[0], n=params[1], dtype=params[2]):
+        x: torch.Tensor,
-                self._norm_test(*params)
+        weight: torch.Tensor,
-                self._l2norm_test(*params)
+        variance_epsilon: float = 1e-6,
-                self._gemma_rmsnorm_test(*params)
+        scale_shift: float = 0.0,
-                self._gemma3_rmsnorm_test(*params)
+        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):