[CPU] layernorm & fused add-layernorm kernels (#14074)

test/srt/cpu/test_norm.py

@@ -3,7 +3,7 @@ import unittest
 from typing import Optional, Tuple, Union
 
 import torch
-from utils import make_non_contiguous, precision
+from utils import make_non_contiguous, parametrize, precision
 
 from sglang.test.test_utils import CustomTestCase
 
@@ -131,5 +131,64 @@ class TestFusedRMSNormGated(CustomTestCase):
                 self._norm_test(*params)
 
 
+class TestLayerNorm(CustomTestCase):
+
+    def _forward_native(
+        self,
+        x: torch.Tensor,
+        weight: torch.Tensor,
+        variance_epsilon: float,
+        residual: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        orig_dtype = x.dtype
+        x = x.to(torch.float32)
+        if residual is not None:
+            x = x + residual.to(torch.float32)
+            residual = x.to(orig_dtype)
+
+        (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
+        if residual is None:
+            return x
+        else:
+            return 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:
+        x_ln = torch.randn([m, n], dtype=dtype)
+        x_ln = make_non_contiguous(x_ln)
+        ref_x_ln = x_ln.clone()
+        hidden_size = x_ln.size(-1)
+        weight = torch.randn(hidden_size, dtype=dtype)
+        variance_epsilon = 1e-6
+
+        torch.ops.sgl_kernel.layernorm_cpu(x_ln, weight, variance_epsilon)
+        ref_ln_out = self._forward_native(ref_x_ln, weight, variance_epsilon)
+
+        atol = rtol = precision[ref_ln_out.dtype]
+        torch.testing.assert_close(x_ln, 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(
+            x_add_ln, residual, weight, variance_epsilon
+        )
+        ref_add_ln_out, ref_residual = self._forward_native(
+            ref_x_add_ln, weight, variance_epsilon, ref_residual
+        )
+
+        torch.testing.assert_close(x_add_ln, ref_add_ln_out, atol=atol, rtol=rtol)
+        torch.testing.assert_close(residual, ref_residual, atol=atol, rtol=rtol)
+
+
 if __name__ == "__main__":
     unittest.main()