Enable mixed type LayerNorm kernel for NSA indexer (#12044)

python/sglang/test/test_layernorm.py

@@ -3,7 +3,7 @@ import unittest
 
 import torch
 
-from sglang.srt.layers.layernorm import GemmaRMSNorm, RMSNorm
+from sglang.srt.layers.layernorm import GemmaRMSNorm, LayerNorm, RMSNorm
 from sglang.test.test_utils import CustomTestCase
 
 
@@ -109,5 +109,77 @@ class TestGemmaRMSNorm(CustomTestCase):
                 self._run_gemma_rms_norm_test(*params)
 
 
+class TestLayerNorm(CustomTestCase):
+    DTYPES = [torch.half, torch.bfloat16]
+    PARAM_DTYPES = [torch.bfloat16, torch.float32]
+    NUM_TOKENS = [7, 83, 1024]
+    HIDDEN_SIZES = [128, 512, 1536, 5120, 5124, 5125, 5126, 7168]
+    USE_AFFINE = [False, True]
+    USE_BIAS = [False, True]
+    SEEDS = [0]
+
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA is not available")
+        torch.set_default_device("cuda")
+
+    def _run_layer_norm_test(
+        self, num_tokens, hidden_size, use_affine, use_bias, dtype, seed, param_dtype
+    ):
+        torch.manual_seed(seed)
+
+        layer = LayerNorm(
+            hidden_size, elementwise_affine=use_affine, bias=use_bias, dtype=param_dtype
+        )
+        if use_affine:
+            layer.weight.data.normal_(mean=1.0, std=0.1)
+            if use_bias:
+                layer.bias.data.normal_(mean=0.0, std=0.1)
+
+        scale = 1 / (2 * hidden_size)
+        x = torch.randn(num_tokens, hidden_size, dtype=dtype) * scale
+
+        with torch.inference_mode():
+            ref_out = layer.forward_native(x)
+            out = layer(x)
+
+        self.assertTrue(torch.allclose(out, ref_out, atol=1e-2, rtol=1e-3))
+
+        if (
+            use_affine
+            and use_bias
+            and not (dtype == torch.bfloat16 and param_dtype == torch.float32)
+        ):
+            layer.dtype = torch.float32
+            layer.weight.data = layer.weight.data.to(torch.float32)
+            layer.bias.data = layer.bias.data.to(torch.float32)
+            with torch.inference_mode():
+                cuda_out = layer(x.to(torch.bfloat16)).to(x.dtype)
+
+            self.assertTrue(torch.allclose(cuda_out, ref_out, atol=2e-2, rtol=1e-3))
+
+    def test_layer_norm(self):
+        for params in itertools.product(
+            self.NUM_TOKENS,
+            self.HIDDEN_SIZES,
+            self.USE_AFFINE,
+            self.USE_BIAS,
+            self.DTYPES,
+            self.SEEDS,
+            self.PARAM_DTYPES,
+        ):
+            with self.subTest(
+                num_tokens=params[0],
+                hidden_size=params[1],
+                use_affine=params[2],
+                use_bias=params[3],
+                dtype=params[4],
+                seed=params[5],
+                param_dtype=params[6],
+            ):
+                self._run_layer_norm_test(*params)
+
+
 if __name__ == "__main__":
     unittest.main(verbosity=2)