minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CI] Slim JIT kernel unit tests (#36887)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-08-29 07:26:49 +08:00
committed by GitHub
parent d12b313b93
commit 96a4dcdde8
10 changed files with 543 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/ops/quantization/test_hadamard_jit.py
+65 -20
View File
@@ -1,3 +1,4 @@
import itertools
import math
import sys
@@ -7,6 +8,7 @@ import torch
import torch.nn.functional as F
from scipy.linalg import hadamard
from sglang.kernels.jit.utils import get_ci_test_range
from sglang.kernels.ops.quantization.hadamard import (
hadamard_transform,
hadamard_transform_12n,
@@ -16,7 +18,8 @@ from sglang.kernels.ops.quantization.hadamard import (
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=128, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_cuda_ci(est_time=32, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_cuda_ci(est_time=128, stage="nightly", runner_config="1-gpu-large")
# Exact M×N Hadamard matrices (±1 entries) copied from
# python/sglang/kernels/jit/csrc/fast-hadamard-transform/code_gen.py.
@@ -218,12 +221,42 @@ def hadamard_transform_mn_ref(x, multiple, scale=1.0):
return x[..., : x_shape[-1]].reshape(*x_shape)
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize(
"dim",
# Power-of-2 dims from python/sglang/kernels/aot/tests/test_hadamard.py (old AOT test)
[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768],
_DTYPES = [torch.float32, torch.float16, torch.bfloat16]
_POWER_OF_TWO_DIMS = [
1,
2,
4,
8,
16,
32,
64,
128,
256,
512,
1024,
2048,
4096,
8192,
16384,
32768,
]
_POWER_OF_TWO_CASES = get_ci_test_range(
list(itertools.product(_POWER_OF_TWO_DIMS, _DTYPES)),
[
(1, torch.float32),
(2, torch.float16),
(4, torch.bfloat16),
(32, torch.bfloat16),
(256, torch.float32),
(2048, torch.float16),
(8192, torch.bfloat16),
(16384, torch.float32),
(32768, torch.float16),
],
)
@pytest.mark.parametrize("dim,dtype", _POWER_OF_TWO_CASES)
def test_hadamard_transform(dim, dtype):
device = "cuda"
@@ -248,13 +281,18 @@ def test_hadamard_transform(dim, dtype):
torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol)
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize(
"dim",
# Non-power-of-2 dims to test the padding path
# (137 from python/sglang/kernels/aot/tests/test_hadamard.py, 500/1000 added for coverage)
[137, 500, 1000],
_NON_POWER_OF_TWO_CASES = get_ci_test_range(
list(itertools.product([137, 500, 1000], _DTYPES)),
[
(137, torch.float32),
(137, torch.bfloat16),
(500, torch.float16),
(1000, torch.bfloat16),
],
)
@pytest.mark.parametrize("dim,dtype", _NON_POWER_OF_TWO_CASES)
def test_hadamard_transform_non_power_of_two(dim, dtype):
device = "cuda"
@@ -327,8 +365,18 @@ _28N_DIMS = [28 * (2**k) for k in range(2, 9)] # 112, 224, ... , 7168
_40N_DIMS = [40 * (2**k) for k in range(2, 9)] # 160, 320, ... , 10240
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize("dim", _12N_DIMS)
def _mn_cases(dims):
return get_ci_test_range(
list(itertools.product(dims, _DTYPES)),
[
(dims[0], torch.float32),
(dims[len(dims) // 2], torch.float16),
(dims[-1], torch.bfloat16),
],
)
@pytest.mark.parametrize("dim,dtype", _mn_cases(_12N_DIMS))
def test_hadamard_transform_12n(dim, dtype):
device = "cuda"
@@ -351,8 +399,7 @@ def test_hadamard_transform_12n(dim, dtype):
torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol)
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize("dim", _20N_DIMS)
@pytest.mark.parametrize("dim,dtype", _mn_cases(_20N_DIMS))
def test_hadamard_transform_20n(dim, dtype):
device = "cuda"
@@ -375,8 +422,7 @@ def test_hadamard_transform_20n(dim, dtype):
torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol)
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize("dim", _28N_DIMS)
@pytest.mark.parametrize("dim,dtype", _mn_cases(_28N_DIMS))
def test_hadamard_transform_28n(dim, dtype):
device = "cuda"
@@ -399,8 +445,7 @@ def test_hadamard_transform_28n(dim, dtype):
torch.testing.assert_close(out.float(), out_ref, rtol=rtol, atol=atol)
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize("dim", _40N_DIMS)
@pytest.mark.parametrize("dim,dtype", _mn_cases(_40N_DIMS))
def test_hadamard_transform_40n(dim, dtype):
device = "cuda"
test/registered/kernels/ops/quantization/test_per_token_group_quant.py
+48 -11
View File
@@ -32,8 +32,9 @@ from sglang.kernels.ops.quantization.per_token_group_quant import (
)
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")
register_cuda_ci(est_time=65, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
register_cuda_ci(est_time=120, stage="nightly", runner_config="1-gpu-large")
G = 128
FMAX = float(fp8_max) # 448 for e4m3
@@ -220,9 +221,20 @@ def test_ue8m0_row_packed_bitexact(hidden):
# fp32 / int8 scale paths: exact stored scale + dequant round-trip (the codes
# are not bit-reproducible under fast-math division).
# --------------------------------------------------------------------------- #
@pytest.mark.parametrize("hidden", [4096, 768])
@pytest.mark.parametrize("column_major", [False, True])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
FP32_SCALE_CASES = get_ci_test_range(
list(
itertools.product([torch.bfloat16, torch.float16], [False, True], [4096, 768])
),
[
(torch.bfloat16, False, 4096),
(torch.bfloat16, True, 768),
(torch.float16, False, 768),
(torch.float16, True, 4096),
],
)
@pytest.mark.parametrize("dtype,column_major,hidden", FP32_SCALE_CASES)
def test_fp32_scale(dtype, column_major, hidden):
"""fp32 scale (row-major contiguous / col-major TMA view): the stored scale
is amax/FMAX (a single multiply, bit-exact) and dequant round-trips within
@@ -342,12 +354,19 @@ MASKED_CASES = get_ci_test_range(
list(itertools.product([2, 5], [2048, 4096], [128, 384])),
[(2, 2048, 128), (5, 4096, 384)],
)
MASKED_TEST_CASES = get_ci_test_range(
list(itertools.product(MASKED_CASES, [None, 4], [torch.int32, torch.int64])),
[
((2, 2048, 128), None, torch.int32),
((2, 2048, 128), 4, torch.int64),
((5, 4096, 384), None, torch.int32),
((5, 4096, 384), 4, torch.int64),
],
)
@pytest.mark.parametrize("masked_m_dtype", [torch.int32, torch.int64])
@pytest.mark.parametrize("expected_m", [None, 4])
@pytest.mark.parametrize("num_experts,hidden,tokens_pad", MASKED_CASES)
def test_masked(num_experts, hidden, tokens_pad, expected_m, masked_m_dtype):
@pytest.mark.parametrize("shape,expected_m,masked_m_dtype", MASKED_TEST_CASES)
def test_masked(shape, expected_m, masked_m_dtype):
"""Masked EP-MoE schedule (col-packed ue8m0, plain quant -- no silu, so the
quant is bit-reproducible): rows < masked_m[e] are bit-exact vs the torch
reference; rows >= masked_m[e] stay zero (untouched). Fusion numerics are
@@ -359,6 +378,7 @@ def test_masked(num_experts, hidden, tokens_pad, expected_m, masked_m_dtype):
expected_m=4 shrinks the grid's token axis far below masked_m, so the
grid-stride token loop must still cover every valid token -- guards the
host-hint-only contract (a wrong hint can never drop tokens)."""
num_experts, hidden, tokens_pad = shape
torch.manual_seed(num_experts * 1000 + hidden + tokens_pad)
x = torch.randn(
num_experts, tokens_pad, hidden, device="cuda", dtype=torch.bfloat16
@@ -431,9 +451,26 @@ def test_masked_fused():
assert torch.all(x_q[e, m:].view(torch.int8) == 0), "padding touched"
@pytest.mark.parametrize("poison", [float("nan"), float("inf"), -float("inf")])
@pytest.mark.parametrize("scale_ue8m0", [False, True])
@pytest.mark.parametrize("masked", [False, True])
NON_FINITE_CASES = get_ci_test_range(
list(
itertools.product(
[float("nan"), float("inf"), -float("inf")],
[False, True],
[False, True],
)
),
[
(float("nan"), False, False),
(float("nan"), True, True),
(float("inf"), False, True),
(float("inf"), True, False),
(-float("inf"), False, False),
(-float("inf"), True, True),
],
)
@pytest.mark.parametrize("poison,scale_ue8m0,masked", NON_FINITE_CASES)
def test_non_finite_inputs_are_sanitized(poison, scale_ue8m0, masked):
"""CUDA-graph capture warmup runs the model on reused, uninitialized
buffers, so quant inputs can contain NaN/Inf bit patterns. The v1/v2/Triton