[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(
+PER_TOKEN_QUANT_CASES = list(
-    "num_tokens,hidden_dim",
+    itertools.product([128, 256, 512], [512, 1076, 1368, 2048, 4096])
-    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,

test/registered/jit/benchmark/bench_sparse_mla_q8kv8_prefill_sm90.py

@@ -126,16 +126,22 @@ def bench_sparse_mla_q8kv8_prefill_sm90(
                 "sgl_kernel.flash_mla.flash_mla_sparse_fwd is not available"
             )
         q, kv, indices, sm_scale = _make_q16_inputs(s_q, s_kv, h_q, d_qk, topk)
-        fn = lambda: flash_mla_sparse_fwd(q, kv, indices, sm_scale, D_V)
+
+        def fn():
+            return flash_mla_sparse_fwd(q, kv, indices, sm_scale, D_V)
+
     elif provider == "q8_fp8_jit":
         if not _sm90_available():
             raise RuntimeError("Q8KV8 sparse prefill benchmark requires SM90 CUDA")
         q, kv, indices, sm_scale, q_scale, kv_scale = _make_q8_inputs(
             s_q, s_kv, h_q, d_qk, topk
         )
-        fn = lambda: sparse_mla_q8kv8_prefill_fwd(
+
+        def fn():
+            return sparse_mla_q8kv8_prefill_fwd(
                 q, kv, indices, sm_scale, q_scale, kv_scale, D_V
             )
+
     else:
         raise ValueError(f"Unknown provider: {provider}")
 

test/registered/jit/test_activation.py

@@ -26,10 +26,13 @@ SHAPES = get_ci_test_range(
         (3, 5, 16),
         (2, 3, 512),
         (1, 17, 4096),
+        (48, 3072),
+        (38, 8192),
+        (39, 32768),
         *[(2**x, 2048) for x in range(0, 15, 2)],
         *[(2**x, 65536) for x in range(0, 5, 2)],
     ],
-    ci_range=[(7, 16), (2, 3, 512)],
+    ci_range=[(7, 16), (2, 3, 512), (48, 3072), (38, 8192)],
 )
 
 
@@ -167,9 +170,10 @@ UNARY_SHAPES = get_ci_test_range(
         (3, 5, 16),
         (2, 3, 512),
         (1, 17, 4096),
+        (38, 4096),
         *[(2**x, 2048) for x in range(0, 15, 2)],
     ],
-    ci_range=[(7, 16), (2, 3, 512)],
+    ci_range=[(7, 16), (2, 3, 512), (38, 4096)],
 )
 
 

test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py

@@ -5,22 +5,27 @@ import sys
 import pytest
 import torch
 
-from sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm import dsv3_fused_a_gemm
+from sglang.jit_kernel.utils import get_ci_test_range, get_jit_cuda_arch, is_hip_runtime
-from sglang.jit_kernel.utils import get_jit_cuda_arch, is_hip_runtime
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 
+if not torch.cuda.is_available():
+    pytest.skip("CUDA required", allow_module_level=True)
+
+from sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm import dsv3_fused_a_gemm  # noqa: E402
+
 # hd_in must be a multiple of 256; 6144/7168 cover the real fused-A shapes.
 HD_INS = [6144, 7168]
 # hd_out must be a multiple of 16; 2112 and 2624 cover real fused-A variants.
 HD_OUTS = [2112, 2624]
+NUM_TOKENS = get_ci_test_range(list(range(1, 17)), [1, 8, 16])
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 @pytest.mark.parametrize("hd_out", HD_OUTS)
 @pytest.mark.parametrize("hd_in", HD_INS)
-@pytest.mark.parametrize("num_tokens", list(range(1, 17)))
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
 def test_dsv3_fused_a_gemm(num_tokens, hd_in, hd_out):
     if is_hip_runtime() or get_jit_cuda_arch().major < 9:
         pytest.skip("SM90+ required")

test/registered/jit/test_dsv3_fused_a_gemm.py

@@ -7,7 +7,7 @@ import torch
 import torch.nn.functional as F
 
 from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm
-from sglang.jit_kernel.utils import get_jit_cuda_arch, is_hip_runtime
+from sglang.jit_kernel.utils import get_ci_test_range, get_jit_cuda_arch, is_hip_runtime
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@@ -16,12 +16,13 @@ register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-l
 HD_INS = [6144, 7168]
 # hd_out must be a multiple of 16; 2112 and 2624 cover real fused-A variants.
 HD_OUTS = [2112, 2624]
+NUM_TOKENS = get_ci_test_range(list(range(1, 17)), [1, 8, 16])
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 @pytest.mark.parametrize("hd_out", HD_OUTS)
 @pytest.mark.parametrize("hd_in", HD_INS)
-@pytest.mark.parametrize("num_tokens", list(range(1, 17)))
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
 def test_dsv3_fused_a_gemm(num_tokens, hd_in, hd_out):
     if is_hip_runtime() or get_jit_cuda_arch().major < 9:
         pytest.skip("SM90+ required")

test/registered/jit/test_dsv3_router_gemm.py

@@ -1,18 +1,37 @@
 """Tests for JIT dsv3_router_gemm kernel."""
 
+import itertools
 import sys
 
 import pytest
 import torch
 
 from sglang.jit_kernel.dsv3_router_gemm import dsv3_router_gemm
-from sglang.jit_kernel.utils import get_jit_cuda_arch, is_hip_runtime
+from sglang.jit_kernel.utils import get_ci_test_range, get_jit_cuda_arch, is_hip_runtime
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=37, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=148, suite="nightly-kernel-1-gpu", nightly=True)
 
 HIDDEN_DIMS = [1024, 4096, 5120, 6144, 7168]
+ROUTER_GEMM_CASES = get_ci_test_range(
+    list(
+        itertools.product(
+            [256, 384],
+            HIDDEN_DIMS,
+            list(range(1, 17)),
+            [torch.bfloat16, torch.float32],
+        )
+    ),
+    [
+        (256, 1024, 1, torch.bfloat16),
+        (256, 7168, 6, torch.bfloat16),
+        (256, 6144, 4, torch.float32),
+        (384, 7168, 8, torch.bfloat16),
+        (256, 7168, 16, torch.float32),
+        (384, 5120, 16, torch.float32),
+    ],
+)
 ATOL = 1e-2
 RTOL = 1e-2
 
@@ -22,10 +41,9 @@ def _ref(hidden_states, router_weights, out_dtype):
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@pytest.mark.parametrize("num_experts", [256, 384])
+@pytest.mark.parametrize(
-@pytest.mark.parametrize("hidden_dim", HIDDEN_DIMS)
+    "num_experts,hidden_dim,num_tokens,out_dtype", ROUTER_GEMM_CASES
-@pytest.mark.parametrize("num_tokens", list(range(1, 17)))
+)
-@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32])
 def test_dsv3_router_gemm(num_experts, hidden_dim, num_tokens, out_dtype):
     if is_hip_runtime() or get_jit_cuda_arch().major < 9:
         pytest.skip("SM90+ required")

test/registered/jit/test_fused_add_rmsnorm.py

@@ -46,10 +46,19 @@ def forward_native_hf_reference(
 
 BS_LIST = [2**n for n in range(0, 14)]
 BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
-BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109])
+HIDDEN_SIZE_LIST = [512, 1024, 1536, 2048, 3072, 4096, 5120, 6144, 7168, 8192]
-HIDDEN_SIZE_LIST = get_ci_test_range(
+FUSED_ADD_RMSNORM_CASES = get_ci_test_range(
-    [512, 1024, 1536, 2048, 3072, 4096, 5120, 6144, 7168, 8192],
+    list(itertools.product(BS_LIST, HIDDEN_SIZE_LIST)),
-    [512, 2048, 8192],
+    [
+        (1, 512),
+        (18, 4096),
+        (38, 4096),
+        (39, 4096),
+        (39, 5120),
+        (39, 8192),
+        (44, 8192),
+        (89, 4096),
+    ],
 )
 DEVICE = "cuda"
 DTYPE = torch.bfloat16
@@ -58,7 +67,7 @@ EPS = torch.finfo(torch.bfloat16).eps
 
 @pytest.mark.parametrize(
     "batch_size,hidden_size,cast_x_before_out_mul",
-    list(itertools.product(BS_LIST, HIDDEN_SIZE_LIST, [False, True])),
+    [(bs, hs, cast) for bs, hs in FUSED_ADD_RMSNORM_CASES for cast in [False, True]],
 )
 def test_fused_add_rmsnorm(
     batch_size: int, hidden_size: int, cast_x_before_out_mul: bool

test/registered/jit/test_moe_align_block_size.py

@@ -7,6 +7,7 @@ import triton
 import triton.language as tl
 
 from sglang.jit_kernel.moe_align import moe_align_block_size
+from sglang.jit_kernel.utils import get_ci_test_range
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=28, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@@ -143,17 +144,29 @@ def moe_align_block_size_triton(
     )
 
 
-@pytest.mark.parametrize(
+MOE_ALIGN_CASES = get_ci_test_range(
-    "block_size,num_tokens,topk,num_experts,pad_sorted_token_ids",
     list(
         itertools.product(
             [32, 64, 128, 256],  # block_size
             [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096],  # num_tokens
             [1, 2, 4, 8, 16, 32, 64],  # topk
-            [64, 160, 256, 257, 260, 264],  # num_experts
+            [64, 128, 160, 256, 257, 260, 264],  # num_experts
             [True, False],  # pad_sorted_token_ids
         )
     ),
+    [
+        (32, 1, 1, 64, True),
+        (128, 48, 1, 128, True),
+        (256, 103, 4, 256, False),
+        (128, 128, 1, 256, True),
+        (256, 4096, 8, 256, False),
+    ],
+)
+
+
+@pytest.mark.parametrize(
+    "block_size,num_tokens,topk,num_experts,pad_sorted_token_ids",
+    MOE_ALIGN_CASES,
 )
 def test_moe_align_block_size_compare_implementations(
     block_size, num_tokens, topk, num_experts, pad_sorted_token_ids

test/registered/jit/test_per_tensor_quant_fp8.py

@@ -6,6 +6,7 @@ import pytest
 import torch
 
 from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
+from sglang.jit_kernel.utils import get_ci_test_range
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=16, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@@ -44,10 +45,18 @@ def torch_scaled_fp8_quant(tensor, inv_scale):
     return qweight
 
 
-@pytest.mark.parametrize(
+PER_TENSOR_QUANT_CASES = get_ci_test_range(
-    "num_tokens,hidden_dim",
+    list(
-    list(itertools.product([128, 256, 512], [512, 2048, 4096])),
+        itertools.product(
+            [38, 39, 128, 256, 512, 1392],
+            [512, 1536, 2048, 4096, 7168],
         )
+    ),
+    [(38, 7168), (39, 1536), (128, 512), (512, 4096), (1392, 1536)],
+)
+
+
+@pytest.mark.parametrize("num_tokens,hidden_dim", PER_TENSOR_QUANT_CASES)
 def test_jit_per_tensor_quant_compare_implementations(
     num_tokens: int,
     hidden_dim: int,

test/registered/jit/test_per_token_group_quant_8bit.py

@@ -4,38 +4,35 @@ import sys
 import pytest
 import torch
 
-from sglang.srt.utils import is_hip
-
-_is_hip = is_hip()
-fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
-
-from sgl_kernel.test_utils import (
-    assert_all_close_or_tiny_diff,
-    create_per_token_group_quant_test_data,
-)
-
 from sglang.jit_kernel.per_token_group_quant_8bit import (
     per_token_group_quant_8bit as sglang_per_token_group_quant_8bit,
 )
-from sglang.srt.layers.quantization.fp8_kernel import (
+from sglang.jit_kernel.utils import get_ci_test_range
-    create_per_token_group_quant_fp8_output_scale,
+from sglang.srt.utils import is_hip
-)
-from sglang.srt.layers.quantization.fp8_kernel import (
-    per_token_group_quant_8bit as triton_per_token_group_quant_8bit,
-)
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=16, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
 
-configs = list(
+if not torch.cuda.is_available():
-    itertools.product(
+    pytest.skip("CUDA required", allow_module_level=True)
-        [1, 4, 16, 64, 127, 128, 512, 1024, 4096, 8192],  # num_tokens
+
-        [128, 256, 384, 512, 1024, 1536, 1664, 2048, 4096, 7168, 16384],  # hidden_dim
+from sgl_kernel.test_utils import (  # noqa: E402
-        [16, 32, 64, 128],  # group_size
+    assert_all_close_or_tiny_diff,
-        [None],  # num_ranks
+    create_per_token_group_quant_test_data,
-        [fp8_type_],  # dtype
+)
-        [
+
+from sglang.srt.layers.quantization.fp8_kernel import (  # noqa: E402
+    create_per_token_group_quant_fp8_output_scale,
+)
+from sglang.srt.layers.quantization.fp8_kernel import (  # noqa: E402
+    per_token_group_quant_8bit as triton_per_token_group_quant_8bit,
+)
+
+_is_hip = is_hip()
+fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
+
+BASE_FLAGS = [
     dict(
         column_major_scales=False,
         scale_tma_aligned=False,
@@ -64,16 +61,8 @@ configs = list(
         fuse_silu_and_mul=False,
         masked_layout_mode=None,
     ),
-        ],
+]
-    )
+FUSED_FLAGS = [
-) + list(
-    itertools.product(
-        [1, 4, 1 * 8, 4 * 8, 64 * 8, 256 * 8, 768 * 8],
-        [2048],
-        [128],
-        [8, 16, 32, 48],
-        [fp8_type_],
-        [
     dict(
         column_major_scales=True,
         scale_tma_aligned=True,
@@ -102,9 +91,38 @@ configs = list(
         fuse_silu_and_mul=True,
         masked_layout_mode="extreme",
     ),
-        ],
+]
+
+configs = get_ci_test_range(
+    list(
+        itertools.product(
+            [1, 4, 16, 17, 38, 51, 64, 127, 128, 512, 1024, 4096, 8192],
+            [128, 256, 384, 512, 768, 1024, 1536, 1664, 2048, 4096, 7168, 16384],
+            [16, 32, 64, 128],
+            [None],
+            [fp8_type_],
+            BASE_FLAGS,
         )
     )
+    + list(
+        itertools.product(
+            [1, 4, 1 * 8, 4 * 8, 64 * 8, 256 * 8, 768 * 8],
+            [2048],
+            [128],
+            [8, 16, 32, 48],
+            [fp8_type_],
+            FUSED_FLAGS,
+        )
+    ),
+    [
+        (1, 128, 128, None, fp8_type_, BASE_FLAGS[0]),
+        (17, 1536, 128, None, fp8_type_, BASE_FLAGS[2]),
+        (38, 4096, 128, None, fp8_type_, BASE_FLAGS[2]),
+        (51, 4096, 128, None, fp8_type_, BASE_FLAGS[2]),
+        (512, 2048, 128, 8, fp8_type_, FUSED_FLAGS[0]),
+        (2048, 2048, 128, 16, fp8_type_, FUSED_FLAGS[1]),
+    ],
+)
 
 
 @pytest.mark.parametrize(

test/registered/jit/test_per_token_group_quant_8bit_v2.py

@@ -1,20 +1,33 @@
+import itertools
+
 import pytest
 import torch
-from sgl_kernel import sgl_per_token_group_quant_8bit  # AOT v2 reference op
 
 from sglang.jit_kernel.per_token_group_quant_8bit_v2 import (
     per_token_group_quant_8bit_v2,
 )
-from sglang.srt.layers.quantization.fp8_kernel import (
+from sglang.jit_kernel.utils import get_ci_test_range
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+try:
+    from sgl_kernel import sgl_per_token_group_quant_8bit  # AOT v2 reference op
+except ImportError:
+    sgl_per_token_group_quant_8bit = None
+
+if sgl_per_token_group_quant_8bit is None and not torch.cuda.is_available():
+    pytest.skip("sgl_kernel AOT reference op is unavailable", allow_module_level=True)
+if sgl_per_token_group_quant_8bit is None:
+    raise ImportError("sgl_kernel AOT reference op is unavailable")
+
+from sglang.srt.layers.quantization.fp8_kernel import (  # noqa: E402
     create_per_token_group_quant_fp8_output_scale,
     fp8_dtype,
     fp8_max,
     fp8_min,
     sglang_per_token_group_quant_fp8,
 )
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 
 G = 128
 
@@ -35,11 +48,30 @@ def _alloc(x_shape, scale_ue8m0):
     return x_q, x_s
 
 
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+V2_QUANT_CASES = get_ci_test_range(
-@pytest.mark.parametrize("num_tokens", [1, 7, 64, 333])
+    list(
-@pytest.mark.parametrize("hidden", [128, 2048, 4096])
+        itertools.product(
-@pytest.mark.parametrize("fuse_silu_and_mul", [False, True])
+            [torch.bfloat16, torch.float16],
-@pytest.mark.parametrize("scale_ue8m0", [False, True])
+            [1, 7, 38, 64, 333],
+            [128, 2048, 4096, 7168],
+            [False, True],
+            [False, True],
+        )
+    ),
+    [
+        (torch.bfloat16, 1, 128, False, False),
+        (torch.bfloat16, 17, 1536, False, False),
+        (torch.bfloat16, 38, 7168, False, True),
+        (torch.bfloat16, 38, 4096, False, True),
+        (torch.bfloat16, 64, 4096, True, True),
+        (torch.float16, 333, 2048, True, False),
+    ],
+)
+
+
+@pytest.mark.parametrize(
+    "dtype,num_tokens,hidden,fuse_silu_and_mul,scale_ue8m0", V2_QUANT_CASES
+)
 def test_v2_jit_matches_aot(dtype, num_tokens, hidden, fuse_silu_and_mul, scale_ue8m0):
     """JIT v2 must be bit-exact with the AOT v2 across vanilla/silu and float/ue8m0
     scales (NaiveScheduler)."""
@@ -83,9 +115,23 @@ def test_v2_jit_matches_aot(dtype, num_tokens, hidden, fuse_silu_and_mul, scale_
     assert torch.equal(x_s, s_ref), "scales differ"
 
 
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+ROW_MAJOR_UE8M0_CASES = get_ci_test_range(
-@pytest.mark.parametrize("num_tokens", [1, 33, 128])
+    list(
-@pytest.mark.parametrize("hidden", [128, 512, 4096, 7168])
+        itertools.product(
+            [torch.bfloat16, torch.float16], [1, 33, 128], [128, 512, 4096, 7168]
+        )
+    ),
+    [
+        (torch.bfloat16, 1, 128),
+        (torch.bfloat16, 17, 1536),
+        (torch.bfloat16, 33, 7168),
+        (torch.bfloat16, 38, 4096),
+        (torch.float16, 128, 4096),
+    ],
+)
+
+
+@pytest.mark.parametrize("dtype,num_tokens,hidden", ROW_MAJOR_UE8M0_CASES)
 def test_sglang_per_token_group_quant_fp8_row_major_ue8m0(dtype, num_tokens, hidden):
     """Row-major scale_ue8m0=True quantizes WITH the rounded (power-of-2) scale.
     Verify: (1) scales are exact powers of 2, (2) dequant ≈ original within FP8 tolerance.
@@ -113,9 +159,14 @@ def test_sglang_per_token_group_quant_fp8_row_major_ue8m0(dtype, num_tokens, hid
 # column-major + ue8m0 + fused-silu+mul + masked combination. Input is 3D
 # [num_experts, tokens_padded, hidden*2]; only tokens < masked_m[e] are processed
 # (padding left untouched → zeros in both). Compare JIT vs AOT bit-exact.
-@pytest.mark.parametrize("num_experts", [2, 5])
+MASKED_V2_CASES = get_ci_test_range(
-@pytest.mark.parametrize("hidden", [2048, 4096])
+    # MaskedLayoutScheduler requires hidden / group_size to be divisible by 16.
-@pytest.mark.parametrize("tokens_pad", [128, 384])
+    list(itertools.product([2, 5], [2048, 4096, 8192], [128, 384])),
+    [(2, 2048, 128), (5, 4096, 384), (2, 8192, 128)],
+)
+
+
+@pytest.mark.parametrize("num_experts,hidden,tokens_pad", MASKED_V2_CASES)
 def test_v2_jit_masked_matches_aot(num_experts, hidden, tokens_pad):
     torch.manual_seed(num_experts * 1000 + hidden + tokens_pad)
     x = torch.randn(

test/registered/jit/test_rmsnorm.py

@@ -42,16 +42,36 @@ def flashinfer_rmsnorm(
 
 BS_LIST = [2**n for n in range(0, 14)]
 BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
-BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109])
+SUPPORTED_HIDDEN_SIZE_LIST = [
-SUPPORTED_HIDDEN_SIZE_LIST = get_ci_test_range(
+    64,
-    [64, 128, 256, 512, *range(1024, 8192 + 1, 1024), 2304, 2560, 12288, 16384],
+    128,
-    [256, 1024, 16384],
+    256,
+    512,
+    *range(1024, 8192 + 1, 1024),
+    1536,
+    2304,
+    2560,
+    8704,
+    12288,
+    16384,
+]
+RMSNORM_CASES = get_ci_test_range(
+    list(itertools.product(BS_LIST, SUPPORTED_HIDDEN_SIZE_LIST)),
+    [
+        (1, 256),
+        (18, 1024),
+        (38, 4096),
+        (1240, 1536),
+        (2500, 1024),
+        (4109, 1024),
+        (7807, 128),
+    ],
 )
 
 
 @pytest.mark.parametrize(
     "batch_size,hidden_size",
-    list(itertools.product(BS_LIST, SUPPORTED_HIDDEN_SIZE_LIST)),
+    RMSNORM_CASES,
 )
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("specify_out", [True, False])