diff --git a/sgl-kernel/tests/test_dsv3_fused_a_gemm.py b/sgl-kernel/tests/test_dsv3_fused_a_gemm.py index 1ec683f8a..661103837 100644 --- a/sgl-kernel/tests/test_dsv3_fused_a_gemm.py +++ b/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 diff --git a/sgl-kernel/tests/test_fp8_gemm.py b/sgl-kernel/tests/test_fp8_gemm.py index 079394e05..5f8fbdbf4 100644 --- a/sgl-kernel/tests/test_fp8_gemm.py +++ b/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__])) diff --git a/sgl-kernel/tests/test_norm.py b/sgl-kernel/tests/test_norm.py index 6ec0b1cf8..640878659 100644 --- a/sgl-kernel/tests/test_norm.py +++ b/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]) diff --git a/sgl-kernel/tests/test_per_token_quant_fp8.py b/sgl-kernel/tests/test_per_token_quant_fp8.py index e7b18555c..9073936ca 100644 --- a/sgl-kernel/tests/test_per_token_quant_fp8.py +++ b/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, diff --git a/test/registered/jit/benchmark/bench_sparse_mla_q8kv8_prefill_sm90.py b/test/registered/jit/benchmark/bench_sparse_mla_q8kv8_prefill_sm90.py index 2a89fb23f..6d9b6122a 100644 --- a/test/registered/jit/benchmark/bench_sparse_mla_q8kv8_prefill_sm90.py +++ b/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( - q, kv, indices, sm_scale, q_scale, kv_scale, D_V - ) + + 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}") diff --git a/test/registered/jit/test_activation.py b/test/registered/jit/test_activation.py index c4683f11b..4f7e92037 100644 --- a/test/registered/jit/test_activation.py +++ b/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)], ) diff --git a/test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py b/test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py index 398d965e4..96993fc01 100644 --- a/test/registered/jit/test_cutedsl_dsv3_fused_a_gemm.py +++ b/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_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") +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") diff --git a/test/registered/jit/test_dsv3_fused_a_gemm.py b/test/registered/jit/test_dsv3_fused_a_gemm.py index 01308d55e..c223b1d32 100644 --- a/test/registered/jit/test_dsv3_fused_a_gemm.py +++ b/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") diff --git a/test/registered/jit/test_dsv3_router_gemm.py b/test/registered/jit/test_dsv3_router_gemm.py index c8535a211..af4200209 100644 --- a/test/registered/jit/test_dsv3_router_gemm.py +++ b/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("hidden_dim", HIDDEN_DIMS) -@pytest.mark.parametrize("num_tokens", list(range(1, 17))) -@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize( + "num_experts,hidden_dim,num_tokens,out_dtype", ROUTER_GEMM_CASES +) 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") diff --git a/test/registered/jit/test_fused_add_rmsnorm.py b/test/registered/jit/test_fused_add_rmsnorm.py index 70f13c6f2..1408fefb1 100644 --- a/test/registered/jit/test_fused_add_rmsnorm.py +++ b/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 = get_ci_test_range( - [512, 1024, 1536, 2048, 3072, 4096, 5120, 6144, 7168, 8192], - [512, 2048, 8192], +HIDDEN_SIZE_LIST = [512, 1024, 1536, 2048, 3072, 4096, 5120, 6144, 7168, 8192] +FUSED_ADD_RMSNORM_CASES = get_ci_test_range( + list(itertools.product(BS_LIST, HIDDEN_SIZE_LIST)), + [ + (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 diff --git a/test/registered/jit/test_moe_align_block_size.py b/test/registered/jit/test_moe_align_block_size.py index 41ba65d02..95dbc963b 100644 --- a/test/registered/jit/test_moe_align_block_size.py +++ b/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( - "block_size,num_tokens,topk,num_experts,pad_sorted_token_ids", +MOE_ALIGN_CASES = get_ci_test_range( 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 diff --git a/test/registered/jit/test_per_tensor_quant_fp8.py b/test/registered/jit/test_per_tensor_quant_fp8.py index 8bae572ec..fdee171ce 100644 --- a/test/registered/jit/test_per_tensor_quant_fp8.py +++ b/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( - "num_tokens,hidden_dim", - list(itertools.product([128, 256, 512], [512, 2048, 4096])), +PER_TENSOR_QUANT_CASES = get_ci_test_range( + list( + 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, diff --git a/test/registered/jit/test_per_token_group_quant_8bit.py b/test/registered/jit/test_per_token_group_quant_8bit.py index d3c64bc15..25a139212 100644 --- a/test/registered/jit/test_per_token_group_quant_8bit.py +++ b/test/registered/jit/test_per_token_group_quant_8bit.py @@ -4,106 +4,124 @@ 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 ( - create_per_token_group_quant_fp8_output_scale, -) -from sglang.srt.layers.quantization.fp8_kernel import ( - per_token_group_quant_8bit as triton_per_token_group_quant_8bit, -) +from sglang.jit_kernel.utils import get_ci_test_range +from sglang.srt.utils import is_hip 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( - itertools.product( - [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 - [16, 32, 64, 128], # group_size - [None], # num_ranks - [fp8_type_], # dtype - [ - dict( - column_major_scales=False, - scale_tma_aligned=False, - scale_ue8m0=False, - fuse_silu_and_mul=False, - masked_layout_mode=None, - ), - dict( - column_major_scales=True, - scale_tma_aligned=False, - scale_ue8m0=False, - fuse_silu_and_mul=False, - masked_layout_mode=None, - ), - dict( - column_major_scales=True, - scale_tma_aligned=True, - scale_ue8m0=False, - fuse_silu_and_mul=False, - masked_layout_mode=None, - ), - dict( - column_major_scales=True, - scale_tma_aligned=True, - scale_ue8m0=True, - fuse_silu_and_mul=False, - masked_layout_mode=None, - ), - ], - ) -) + 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, - scale_ue8m0=True, - fuse_silu_and_mul=True, - masked_layout_mode=None, - ), - dict( - column_major_scales=True, - scale_tma_aligned=True, - scale_ue8m0=True, - fuse_silu_and_mul=True, - masked_layout_mode="balanced", - ), - dict( - column_major_scales=True, - scale_tma_aligned=True, - scale_ue8m0=True, - fuse_silu_and_mul=True, - masked_layout_mode="imbalanced", - ), - dict( - column_major_scales=True, - scale_tma_aligned=True, - scale_ue8m0=True, - fuse_silu_and_mul=True, - masked_layout_mode="extreme", - ), - ], +if not torch.cuda.is_available(): + pytest.skip("CUDA required", allow_module_level=True) + +from sgl_kernel.test_utils import ( # noqa: E402 + assert_all_close_or_tiny_diff, + create_per_token_group_quant_test_data, +) + +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, + scale_ue8m0=False, + fuse_silu_and_mul=False, + masked_layout_mode=None, + ), + dict( + column_major_scales=True, + scale_tma_aligned=False, + scale_ue8m0=False, + fuse_silu_and_mul=False, + masked_layout_mode=None, + ), + dict( + column_major_scales=True, + scale_tma_aligned=True, + scale_ue8m0=False, + fuse_silu_and_mul=False, + masked_layout_mode=None, + ), + dict( + column_major_scales=True, + scale_tma_aligned=True, + scale_ue8m0=True, + fuse_silu_and_mul=False, + masked_layout_mode=None, + ), +] +FUSED_FLAGS = [ + dict( + column_major_scales=True, + scale_tma_aligned=True, + scale_ue8m0=True, + fuse_silu_and_mul=True, + masked_layout_mode=None, + ), + dict( + column_major_scales=True, + scale_tma_aligned=True, + scale_ue8m0=True, + fuse_silu_and_mul=True, + masked_layout_mode="balanced", + ), + dict( + column_major_scales=True, + scale_tma_aligned=True, + scale_ue8m0=True, + fuse_silu_and_mul=True, + masked_layout_mode="imbalanced", + ), + dict( + column_major_scales=True, + scale_tma_aligned=True, + scale_ue8m0=True, + 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]), + ], ) diff --git a/test/registered/jit/test_per_token_group_quant_8bit_v2.py b/test/registered/jit/test_per_token_group_quant_8bit_v2.py index 12442cdce..fcb4e29f5 100644 --- a/test/registered/jit/test_per_token_group_quant_8bit_v2.py +++ b/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]) -@pytest.mark.parametrize("num_tokens", [1, 7, 64, 333]) -@pytest.mark.parametrize("hidden", [128, 2048, 4096]) -@pytest.mark.parametrize("fuse_silu_and_mul", [False, True]) -@pytest.mark.parametrize("scale_ue8m0", [False, True]) +V2_QUANT_CASES = get_ci_test_range( + list( + itertools.product( + [torch.bfloat16, torch.float16], + [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]) -@pytest.mark.parametrize("num_tokens", [1, 33, 128]) -@pytest.mark.parametrize("hidden", [128, 512, 4096, 7168]) +ROW_MAJOR_UE8M0_CASES = get_ci_test_range( + list( + 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]) -@pytest.mark.parametrize("hidden", [2048, 4096]) -@pytest.mark.parametrize("tokens_pad", [128, 384]) +MASKED_V2_CASES = get_ci_test_range( + # MaskedLayoutScheduler requires hidden / group_size to be divisible by 16. + 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( diff --git a/test/registered/jit/test_rmsnorm.py b/test/registered/jit/test_rmsnorm.py index 6baac1103..8ec878f35 100644 --- a/test/registered/jit/test_rmsnorm.py +++ b/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 = get_ci_test_range( - [64, 128, 256, 512, *range(1024, 8192 + 1, 1024), 2304, 2560, 12288, 16384], - [256, 1024, 16384], +SUPPORTED_HIDDEN_SIZE_LIST = [ + 64, + 128, + 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])