[Refactor] Clean up JIT kernel utilites (#16884)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>

python/sglang/jit_kernel/tests/test_per_tensor_quant_fp8.py

@@ -58,7 +58,7 @@ def test_jit_per_tensor_quant_compare_implementations(
     )
 
 
-@pytest.mark.parametrize("shape", [(4, 8, 64), (2, 16, 128)])
+@pytest.mark.parametrize("shape", [(4, 8, 64), (2, 16, 128), (19260817, 1, 1)])
 def test_jit_per_tensor_quant_supports_3d(shape):
     device = torch.device("cuda")
     x = torch.rand(shape, dtype=torch.bfloat16, device=device)
@@ -74,7 +74,7 @@ def test_jit_per_tensor_quant_supports_3d(shape):
     torch.testing.assert_close(out.float(), out_ref.float(), rtol=1e-3, atol=1e-3)
 
     scale = torch.rand(1, dtype=torch.float32, device=device)
-    sglang_out, sglang_scale = sglang_scaled_fp8_quant(x, scale)
+    sglang_out, _ = sglang_scaled_fp8_quant(x, scale)
     torch_out = torch_scaled_fp8_quant(x, scale)
 
     torch.testing.assert_close(

python/sglang/jit_kernel/tests/test_qknorm.py

@@ -1,3 +1,6 @@
+import itertools
+
+import pytest
 import torch
 import triton
 
@@ -56,30 +59,35 @@ def torch_impl_qknorm(
     k.copy_(k.float() * k_norm * k_weight.float())
 
 
+BS_LIST = [2**n for n in range(0, 14)]
+BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
+N_K_LIST = [2, 4]
+N_Q_LIST = [8, 16]
+HEAD_DIM_LIST = [64, 128, 256]
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+
 # NOTE(dark): sgl_kernel use flashinfer template, which is bitwise identical to flashinfer impl.
 # However, sgl-jit-kernel, flashinfer, torch_impl, may have small numerical differences.
 # so we allow a small rel/abs tolerance in correctness test.
-def main():
-    N_K = 2
-    N_Q = 16
-    DEVICE = "cuda"
-    DTYPE = torch.bfloat16
-    BS_LIST = [2**n for n in range(0, 15)]
-    BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
-    for HEAD_DIM in [64, 128, 256]:
-        for BS in BS_LIST:
-            q = torch.randn(BS, N_Q, HEAD_DIM, device=DEVICE, dtype=DTYPE)
-            k = torch.randn(BS, N_K, HEAD_DIM, device=DEVICE, dtype=DTYPE)
-            q_weight = torch.randn(HEAD_DIM, device=DEVICE, dtype=DTYPE)
-            k_weight = torch.randn(HEAD_DIM, device=DEVICE, dtype=DTYPE)
-            q_k_aot = (q.clone(), k.clone())
-            q_k_jit = (q.clone(), k.clone())
-            sglang_aot_qknorm(q_k_aot[0], q_k_aot[1], q_weight, k_weight)
-            sglang_jit_qknorm(q_k_jit[0], q_k_jit[1], q_weight, k_weight)
-            triton.testing.assert_close(q_k_aot[0], q_k_jit[0], atol=1e-2, rtol=1e-2)
-            triton.testing.assert_close(q_k_aot[1], q_k_jit[1], atol=1e-2, rtol=1e-2)
-        print(f"HEAD_DIM={HEAD_DIM} correctness test passed.")
+
+
+@pytest.mark.parametrize(
+    "batch_size,n_k,n_q,head_dim",
+    list(itertools.product(BS_LIST, N_K_LIST, N_Q_LIST, HEAD_DIM_LIST)),
+)
+def test_qknorm(batch_size: int, n_k: int, n_q: int, head_dim: int) -> None:
+    q = torch.randn(batch_size, n_q, head_dim, device=DEVICE, dtype=DTYPE)
+    k = torch.randn(batch_size, n_k, head_dim, device=DEVICE, dtype=DTYPE)
+    q_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+    k_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+    q_k_aot = (q.clone(), k.clone())
+    q_k_jit = (q.clone(), k.clone())
+    sglang_aot_qknorm(q_k_aot[0], q_k_aot[1], q_weight, k_weight)
+    sglang_jit_qknorm(q_k_jit[0], q_k_jit[1], q_weight, k_weight)
+    triton.testing.assert_close(q_k_aot[0], q_k_jit[0], atol=1e-2, rtol=1e-2)
+    triton.testing.assert_close(q_k_aot[1], q_k_jit[1], atol=1e-2, rtol=1e-2)
 
 
 if __name__ == "__main__":
-    main()
+    pytest.main([__file__])

python/sglang/jit_kernel/tests/test_rmsnorm.py

@@ -0,0 +1,41 @@
+import itertools
+
+import pytest
+import torch
+import triton
+
+
+def sglang_jit_rmsnorm(input: torch.Tensor, weight: torch.Tensor) -> None:
+    from sglang.jit_kernel.norm import rmsnorm
+
+    rmsnorm(input, weight, output=input)
+
+
+def flashinfer_rmsnorm(input: torch.Tensor, weight: torch.Tensor) -> None:
+    from flashinfer.norm import rmsnorm
+
+    rmsnorm(input, weight, out=input)
+
+
+BS_LIST = [2**n for n in range(0, 14)]
+BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
+HIDDEN_SIZE_LIST = [512, 1024, 1536, 2048, 3072, 4096, 5120, 6144, 7168, 8192]
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+
+
+@pytest.mark.parametrize(
+    "batch_size,hidden_size", list(itertools.product(BS_LIST, HIDDEN_SIZE_LIST))
+)
+def test_rmsnorm(batch_size: int, hidden_size: int) -> None:
+    input = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=DTYPE)
+    weight = torch.randn(hidden_size, device=DEVICE, dtype=DTYPE)
+    input_sglang = input.clone()
+    input_flashinfer = input.clone()
+    sglang_jit_rmsnorm(input_sglang, weight)
+    flashinfer_rmsnorm(input_flashinfer, weight)
+    triton.testing.assert_close(input_sglang, input_flashinfer, atol=1e-2, rtol=1e-2)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])