[CPU] Add Gemma3RMSNorm kernel in sgl-kernel and add ut (#9324)

test/srt/cpu/test_norm.py

@@ -36,6 +36,35 @@ class TestNorm(CustomTestCase):
         else:
             return x, residual
 
+    def _norm(self, x, eps):
+        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
+    ):
+        output = self._norm(x.float(), variance_epsilon)
+        output = output * (1.0 + weight.float())
+        return output.type_as(x)
+
+    def _gemma_rmsnorm_native(
+        self,
+        x: torch.Tensor,
+        weight: torch.Tensor,
+        variance_epsilon: float = 1e-6,
+        residual: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        orig_dtype = x.dtype
+        if residual is not None:
+            x = x + residual
+            residual = x
+
+        x = x.float()
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(variance + variance_epsilon)
+        x = x * (1.0 + weight.float())
+        x = x.to(orig_dtype)
+        return x if residual is None else (x, residual)
+
     def _norm_test(self, m, n, dtype):
 
         x = torch.randn([m, n], dtype=dtype)
@@ -78,11 +107,58 @@ 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):
+
+        x = torch.randn([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.gemma_rmsnorm_cpu(x, weight, variance_epsilon)
+        ref_out = self._gemma_rmsnorm_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.gemma_fused_add_rmsnorm_cpu(
+            x, residual, weight, variance_epsilon
+        )
+
+        ref_x, ref_residual = self._gemma_rmsnorm_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)
+
+    def _gemma3_rmsnorm_test(self, m, n, dtype):
+        x_list = [
+            torch.randn([m, n], dtype=dtype),
+            torch.randn([1, m, 2, n], 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, variance_epsilon)
+            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)
+
     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)
                 self._l2norm_test(*params)
+                self._gemma_rmsnorm_test(*params)
+                self._gemma3_rmsnorm_test(*params)
 
 
 class TestFusedRMSNormGated(CustomTestCase):