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