[Kernel] Strengthen kernel shape coverage (#29636)

Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>

sgl-kernel/tests/test_dsv3_fused_a_gemm.py

@@ -6,7 +6,7 @@ import torch.nn.functional as F
 from sgl_kernel import dsv3_fused_a_gemm
 
 
-@pytest.mark.parametrize("num_tokens", [i + 1 for i in range(16)])
+@pytest.mark.parametrize("num_tokens", [1, 8, 15, 16])
 def test_dsv3_fused_a_gemm(num_tokens):
     kHdIn = 7168
     kHdOut = 2112

sgl-kernel/tests/test_fp8_gemm.py

@@ -92,5 +92,20 @@ def test_accuracy_sm90_swap_ab(shape_mn, K, with_bias, out_dtype):
     _test_accuracy_once(M, N, K, with_bias, out_dtype, "cuda")
 
 
+PRODUCTION_LIKE_FP8_GEMM_CASES = [
+    (189, 4608, 8192, False, torch.bfloat16),
+    (3330, 256, 8192, False, torch.bfloat16),
+    (17, 9216, 2048, False, torch.bfloat16),
+]
+
+
+@pytest.mark.parametrize(
+    "M,N,K,with_bias,out_dtype",
+    PRODUCTION_LIKE_FP8_GEMM_CASES,
+)
+def test_accuracy_production_like_shapes(M, N, K, with_bias, out_dtype):
+    _test_accuracy_once(M, N, K, with_bias, out_dtype, "cuda")
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

sgl-kernel/tests/test_norm.py

@@ -52,6 +52,13 @@ def fused_add_rms_norm(x, residual, weight, eps):
     return x, residual
 
 
+def assert_close_norm(actual, expected, dtype):
+    if dtype is torch.bfloat16:
+        torch.testing.assert_close(actual, expected, rtol=1e-2, atol=2e-2)
+    else:
+        torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3)
+
+
 @pytest.mark.parametrize("batch_size", [1, 19, 99, 989])
 @pytest.mark.parametrize("hidden_size", [111, 500, 1024, 3072, 3584, 4096, 8192, 16384])
 @pytest.mark.parametrize("dtype", [torch.float16])
@@ -96,6 +103,57 @@ def test_fused_add_rmsnorm(batch_size, hidden_size, dtype):
     torch.testing.assert_close(residual_fused, residual_native, rtol=1e-3, atol=1e-3)
 
 
+PRODUCTION_LIKE_NORM_CASES = [
+    (38, 4096, torch.bfloat16),
+    (1240, 1536, torch.bfloat16),
+    (7807, 128, torch.bfloat16),
+]
+
+
+@pytest.mark.parametrize("batch_size,hidden_size,dtype", PRODUCTION_LIKE_NORM_CASES)
+def test_norm_production_like_shapes(batch_size, hidden_size, dtype):
+    x = torch.randn(batch_size, hidden_size, dtype=dtype, device="cuda")
+    w = torch.randn(hidden_size, dtype=dtype, device="cuda")
+
+    y_ref = llama_rms_norm(x, w)
+    enable_pdl = is_arch_support_pdl()
+    y = sgl_kernel.rmsnorm(x, w, enable_pdl=enable_pdl)
+
+    assert_close_norm(y_ref, y, dtype)
+
+
+PRODUCTION_LIKE_FUSED_ADD_RMSNORM_CASES = [
+    (39, 4096, torch.bfloat16),
+    (39, 8192, torch.bfloat16),
+    (89, 4096, torch.bfloat16),
+]
+
+
+@pytest.mark.parametrize(
+    "batch_size,hidden_size,dtype", PRODUCTION_LIKE_FUSED_ADD_RMSNORM_CASES
+)
+def test_fused_add_rmsnorm_production_like_shapes(batch_size, hidden_size, dtype):
+    eps = 1e-6
+
+    x = torch.randn(batch_size, hidden_size, dtype=dtype, device="cuda")
+    residual = torch.randn_like(x)
+    weight = torch.randn(hidden_size, dtype=dtype, device="cuda")
+
+    x_native, residual_native = fused_add_rms_norm(
+        x.clone(), residual.clone(), weight, eps
+    )
+
+    x_fused = x.clone()
+    residual_fused = residual.clone()
+    enable_pdl = is_arch_support_pdl()
+    sgl_kernel.fused_add_rmsnorm(
+        x_fused, residual_fused, weight, eps, enable_pdl=enable_pdl
+    )
+
+    assert_close_norm(x_fused, x_native, dtype)
+    torch.testing.assert_close(residual_fused, residual_native, rtol=1e-3, atol=1e-3)
+
+
 @pytest.mark.parametrize("batch_size", [1, 19, 99, 989])
 @pytest.mark.parametrize("hidden_size", [111, 500, 1024, 3072, 3584, 4096, 8192, 16384])
 @pytest.mark.parametrize("dtype", [torch.float16])

sgl-kernel/tests/test_per_token_quant_fp8.py

@@ -1,6 +1,6 @@
 import itertools
 import sys
-from typing import Optional, Tuple
+from typing import Tuple
 
 import pytest
 import torch
@@ -35,10 +35,16 @@ def sglang_per_token_quant_fp8(
     return output, scale
 
 
-@pytest.mark.parametrize(
-    "num_tokens,hidden_dim",
-    list(itertools.product([128, 256, 512], [512, 1076, 1368, 2048, 4096])),
-)
+PER_TOKEN_QUANT_CASES = list(
+    itertools.product([128, 256, 512], [512, 1076, 1368, 2048, 4096])
+) + [
+    (39, 1536),
+    (1392, 1536),
+    (7807, 1536),
+]
+
+
+@pytest.mark.parametrize("num_tokens,hidden_dim", PER_TOKEN_QUANT_CASES)
 def test_per_token_quant_compare_implementations(
     num_tokens: int,
     hidden_dim: int,