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
.github/workflows/pr-test-jit-kernel.yml
+2 -1
View File
@@ -92,7 +92,8 @@ jobs:
python3 run_suite.py --hw cuda \
--suite base-b-kernel-unit-test-1-gpu-large \
--auto-partition-id ${{ matrix.partition }} \
--auto-partition-size 2
--auto-partition-size 2 \
--fork-worker-batch-size 20
jit-kernel-multigpu-unit-test:
# Runs whenever call-jit-kernel-tests dispatches this workflow. That caller is the
python/sglang/test/ci/ci_utils.py
+97 -7
View File
@@ -23,6 +23,66 @@ class TestFile:
estimated_time: float = 60
class _ForkTestWorker:
"""Preloaded interpreter that forks an isolated child for each test file."""
def __init__(self):
result_read_fd, result_write_fd = os.pipe()
worker_path = os.path.join(os.path.dirname(__file__), "fork_test_worker.py")
self.process = subprocess.Popen(
["python3", worker_path, "--result-fd", str(result_write_fd)],
stdin=subprocess.PIPE,
stdout=None,
stderr=None,
text=True,
pass_fds=(result_write_fd,),
)
os.close(result_write_fd)
self.result_stream = os.fdopen(result_read_fd)
self.files_run = 0
def run(self, filename: str) -> tuple[int, float]:
tic = time.perf_counter()
if self.process.poll() is not None or self.process.stdin is None:
return 1, 0.0
try:
self.process.stdin.write(json.dumps({"filename": filename}) + "\n")
self.process.stdin.flush()
result_line = self.result_stream.readline()
except (BrokenPipeError, OSError):
return 1, time.perf_counter() - tic
if not result_line:
return 1, time.perf_counter() - tic
try:
result = json.loads(result_line)
except json.JSONDecodeError:
return 1, time.perf_counter() - tic
self.files_run += 1
return int(result["returncode"]), float(result["elapsed"])
def close(self, terminate: bool = False):
if self.process.poll() is None:
if terminate:
kill_process_tree(self.process.pid)
elif self.process.stdin is not None:
try:
self.process.stdin.write(json.dumps({"command": "stop"}) + "\n")
self.process.stdin.flush()
self.process.wait(timeout=10)
except (BrokenPipeError, subprocess.TimeoutExpired):
kill_process_tree(self.process.pid)
if self.process.poll() is None:
self.process.kill()
try:
self.process.wait(timeout=10)
except subprocess.TimeoutExpired:
self.process.kill()
self.process.wait()
if self.process.stdin is not None:
self.process.stdin.close()
self.result_stream.close()
# Patterns that indicate retriable accuracy/performance failures
RETRIABLE_PATTERNS = [
r"AssertionError:.*not greater than",
@@ -158,6 +218,7 @@ def run_unittest_files(
enable_retry: bool = False,
max_attempts: int = 2,
retry_wait_seconds: int = 60,
fork_worker_batch_size: int = 1,
):
"""
Run a list of test files.
@@ -172,6 +233,9 @@ def run_unittest_files(
assertion failures (not code errors).
max_attempts: Maximum number of attempts per file including initial run (default: 2).
retry_wait_seconds: Seconds to wait between retries (default: 60).
fork_worker_batch_size: Number of files served by one preloaded fork
worker. Each file still runs in a fresh child
process. One keeps the existing exec behavior.
"""
coredump_enabled = cuda_coredump.is_enabled()
if coredump_enabled:
@@ -185,6 +249,8 @@ def run_unittest_files(
# Per-file elapsed seconds, latest attempt wins. Consumed by the
# TIMINGS block emitted at the end of this function.
file_elapsed: Dict[str, float] = {}
fork_worker = None
use_fork_worker = fork_worker_batch_size > 1 and not enable_retry
for i, file in enumerate(files):
if isinstance(file, CIRegistry):
@@ -203,7 +269,7 @@ def run_unittest_files(
output_lines = []
def run_one_file(filename, capture_output=False):
nonlocal process, output_lines
nonlocal process, output_lines, fork_worker
full_path = os.path.join(os.getcwd(), filename)
logger.info(
@@ -211,10 +277,24 @@ def run_unittest_files(
)
file_tic = time.perf_counter()
cmd = ["python3", full_path, "-f"]
if capture_output:
if use_fork_worker:
if (
fork_worker is None
or fork_worker.files_run >= fork_worker_batch_size
):
if fork_worker is not None:
fork_worker.close()
fork_worker = _ForkTestWorker()
process = fork_worker.process
ret_code, _ = fork_worker.run(full_path)
if ret_code != 0 or fork_worker.files_run >= fork_worker_batch_size:
fork_worker.close()
fork_worker = None
process = None
elapsed = time.perf_counter() - file_tic
elif capture_output:
# Capture output for retry decision
cmd = ["python3", full_path, "-f"]
process = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
@@ -236,16 +316,19 @@ def run_unittest_files(
# Bounded wait for the reader to finish.
reader_thread.join(timeout=60)
else:
cmd = ["python3", full_path, "-f"]
process = subprocess.Popen(cmd, stdout=None, stderr=None)
process.wait()
elapsed = time.perf_counter() - file_tic
if not use_fork_worker:
elapsed = time.perf_counter() - file_tic
ret_code = process.returncode
file_elapsed[filename] = elapsed
logger.info(
f".\n.\nEnd ({i}/{len(files) - 1}):\n{filename=}, {elapsed=:.0f}, {estimated_time=}\n.\n.\n"
)
return process.returncode
return ret_code
# Retry loop for each file
attempt = 1
@@ -305,7 +388,11 @@ def run_unittest_files(
break
except TimeoutError:
kill_process_tree(process.pid)
if fork_worker is not None:
fork_worker.close(terminate=True)
fork_worker = None
elif process is not None:
kill_process_tree(process.pid)
time.sleep(5)
# TimeoutError aborts run_one_file before its elapsed write;
# record the timeout cap as an upper bound so the file still
@@ -333,6 +420,9 @@ def run_unittest_files(
if not continue_on_error:
break
if fork_worker is not None:
fork_worker.close()
elapsed_total = time.perf_counter() - tic
if coredump_enabled and not success:
python/sglang/test/ci/fork_test_worker.py
+105
View File
@@ -0,0 +1,105 @@
"""Execute test files in isolated fork children of a preloaded interpreter."""
import argparse
import json
import os
import runpy
import signal
import sys
import time
import traceback
def _preload_common_modules() -> None:
"""Load the expensive common stack without creating a CUDA context."""
import numpy # noqa: F401
import pytest # noqa: F401
import scipy # noqa: F401
import torch
import triton # noqa: F401
if torch.cuda.is_initialized():
raise RuntimeError("fork test worker initialized CUDA before forking")
def _normalize_exit_code(code) -> int:
if code is None:
return 0
if isinstance(code, int):
return code
print(code, file=sys.stderr, flush=True)
return 1
def _run_file(filename: str) -> int:
sys.argv = [filename, "-f"]
try:
runpy.run_path(filename, run_name="__main__")
return 0
except SystemExit as exc:
return _normalize_exit_code(exc.code)
except BaseException:
traceback.print_exc()
return 1
def _wait_status_to_returncode(status: int) -> int:
if os.WIFEXITED(status):
return os.WEXITSTATUS(status)
if os.WIFSIGNALED(status):
return 128 + os.WTERMSIG(status)
return 1
def run_file_in_fork(filename: str) -> tuple[int, float]:
sys.stdout.flush()
sys.stderr.flush()
tic = time.perf_counter()
child_pid = os.fork()
if child_pid == 0:
signal.signal(signal.SIGINT, signal.SIG_DFL)
returncode = _run_file(filename)
sys.stdout.flush()
sys.stderr.flush()
os._exit(max(0, min(returncode, 255)))
_, status = os.waitpid(child_pid, 0)
return _wait_status_to_returncode(status), time.perf_counter() - tic
def main() -> int:
if not hasattr(os, "fork"):
raise RuntimeError("fork test worker requires a POSIX platform")
parser = argparse.ArgumentParser()
parser.add_argument("--result-fd", type=int, required=True)
args = parser.parse_args()
_preload_common_modules()
with os.fdopen(args.result_fd, "w", buffering=1) as result_stream:
for line in sys.stdin:
command = json.loads(line)
if command.get("command") == "stop":
break
filename = command["filename"]
returncode, elapsed = run_file_in_fork(filename)
result_stream.write(
json.dumps(
{
"filename": filename,
"returncode": returncode,
"elapsed": elapsed,
}
)
+ "\n"
)
return 0
if __name__ == "__main__":
# Ignore Ctrl-C in the preloader. The active child keeps the default
# handler, while suite timeouts terminate the complete process tree.
signal.signal(signal.SIGINT, signal.SIG_IGN)
raise SystemExit(main())
test/registered/kernels/ops/attention/test_rope.py
+76 -26
View File
@@ -1,3 +1,4 @@
import itertools
import sys
import pytest
@@ -8,7 +9,7 @@ from sglang.kernels.jit.utils import get_ci_test_range
from sglang.srt.utils import is_hip
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci(est_time=64, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_cuda_ci(est_time=24, stage="base-b-kernel-unit", runner_config="1-gpu-large")
# Nightly is not redundant here: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1 to expand get_ci_test_range sweeps.
register_cuda_ci(est_time=390, stage="nightly", runner_config="1-gpu-large")
register_amd_ci(est_time=64, suite="jit-kernel-unit-test-amd")
@@ -142,26 +143,80 @@ def reference_rope(
# Test parameters
# ---------------------------------------------------------------------------
BS_LIST = [2**x for x in range(12)]
BS_LIST += [x + 1 for x in BS_LIST] # odd sizes to stress non-aligned paths
BS_LIST = get_ci_test_range(BS_LIST, [1, 129, 2048, 2049])
NUM_KV_HEADS_LIST = get_ci_test_range([1, 2, 8], [1, 8])
GQA_RATIO = get_ci_test_range([1, 4, 8], [1, 8])
ROPE_DIM_LIST = get_ci_test_range([64, 128, 256, 512], [64, 256])
_FULL_BS_LIST = [2**x for x in range(12)]
_FULL_BS_LIST += [x + 1 for x in _FULL_BS_LIST] # stress non-aligned paths
_FULL_NUM_KV_HEADS_LIST = [1, 2, 8]
_FULL_GQA_RATIO_LIST = [1, 4, 8]
_FULL_ROPE_DIM_LIST = [64, 128, 256, 512]
IS_NEOX_LIST = [False, True]
DTYPE_LIST = get_ci_test_range(
[torch.bfloat16, torch.float16], [torch.bfloat16, torch.float16]
_FULL_DTYPE_LIST = [torch.bfloat16, torch.float16]
_FULL_PARTIAL_ROPE_DIM_LIST = [64, 80, 96, 128]
_FULL_HEAD_DIM_LIST = [64, 128, 256]
ROPE_CASES = get_ci_test_range(
list(
itertools.product(
_FULL_BS_LIST,
_FULL_GQA_RATIO_LIST,
_FULL_NUM_KV_HEADS_LIST,
_FULL_ROPE_DIM_LIST,
IS_NEOX_LIST,
_FULL_DTYPE_LIST,
)
),
[
(1, 1, 1, 64, False, torch.bfloat16),
(129, 8, 1, 256, True, torch.float16),
(2048, 1, 8, 64, True, torch.float16),
(2049, 8, 8, 256, False, torch.bfloat16),
(1, 8, 8, 64, True, torch.bfloat16),
(129, 1, 1, 256, False, torch.float16),
(2048, 8, 1, 256, False, torch.bfloat16),
(2049, 1, 8, 64, True, torch.float16),
],
)
PARTIAL_ROPE_CASES = get_ci_test_range(
list(
itertools.product(
_FULL_BS_LIST,
IS_NEOX_LIST,
_FULL_PARTIAL_ROPE_DIM_LIST,
_FULL_HEAD_DIM_LIST,
)
),
[
(1, False, 64, 64),
(129, True, 96, 256),
(2048, False, 96, 128),
(2049, True, 64, 256),
],
)
FUSED_ROPE_STORE_CASES = get_ci_test_range(
list(
itertools.product(
_FULL_BS_LIST,
_FULL_GQA_RATIO_LIST,
_FULL_NUM_KV_HEADS_LIST,
_FULL_ROPE_DIM_LIST,
IS_NEOX_LIST,
)
),
[
(1, 1, 1, 64, False),
(129, 8, 1, 256, True),
(2048, 1, 8, 64, True),
(2049, 8, 8, 256, False),
(1, 8, 8, 64, True),
(129, 1, 1, 256, False),
(2048, 8, 1, 256, False),
(2049, 1, 8, 64, True),
],
)
PARTIAL_ROPE_DIM_LIST = get_ci_test_range([64, 80, 96, 128], [64, 96])
HEAD_DIM_LIST = get_ci_test_range([64, 128, 256], [64, 256])
@pytest.mark.parametrize("batch_size", BS_LIST)
@pytest.mark.parametrize("gqa_ratio", GQA_RATIO)
@pytest.mark.parametrize("num_kv_heads", NUM_KV_HEADS_LIST)
@pytest.mark.parametrize("rope_dim", ROPE_DIM_LIST)
@pytest.mark.parametrize("is_neox", IS_NEOX_LIST)
@pytest.mark.parametrize("dtype", DTYPE_LIST)
@pytest.mark.parametrize(
"batch_size,gqa_ratio,num_kv_heads,rope_dim,is_neox,dtype", ROPE_CASES
)
def test_rope(
batch_size: int,
gqa_ratio: int,
@@ -295,10 +350,7 @@ def test_rope_store_mixed_q_dtype(is_neox: bool) -> None:
assert torch.equal(vc_mixed, vc_ref)
@pytest.mark.parametrize("batch_size", BS_LIST)
@pytest.mark.parametrize("is_neox", IS_NEOX_LIST)
@pytest.mark.parametrize("rope_dim", PARTIAL_ROPE_DIM_LIST)
@pytest.mark.parametrize("head_dim", HEAD_DIM_LIST)
@pytest.mark.parametrize("batch_size,is_neox,rope_dim,head_dim", PARTIAL_ROPE_CASES)
def test_partial_rope(batch_size: int, is_neox: bool, rope_dim: int, head_dim: int):
if head_dim < rope_dim:
pytest.skip("Invalid config: head_dim must be >= rope_dim.")
@@ -321,11 +373,9 @@ def test_partial_rope(batch_size: int, is_neox: bool, rope_dim: int, head_dim: i
triton.testing.assert_close(k_fi, k_jit, atol=atol, rtol=rtol)
@pytest.mark.parametrize("batch_size", BS_LIST)
@pytest.mark.parametrize("gqa_ratio", GQA_RATIO)
@pytest.mark.parametrize("num_kv_heads", NUM_KV_HEADS_LIST)
@pytest.mark.parametrize("rope_dim", ROPE_DIM_LIST)
@pytest.mark.parametrize("is_neox", IS_NEOX_LIST)
@pytest.mark.parametrize(
"batch_size,gqa_ratio,num_kv_heads,rope_dim,is_neox", FUSED_ROPE_STORE_CASES
)
def test_fused_rope_store(
batch_size: int,
gqa_ratio: int,
test/registered/kernels/ops/diffusion/test_rope.py
+27 -13
View File
@@ -31,7 +31,7 @@ from sglang.kernels.ops.diffusion import (
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=44, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_cuda_ci(est_time=18, stage="base-b-kernel-unit", runner_config="1-gpu-large")
# Nightly is not redundant: it sets SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1, which
# expands the get_ci_test_range sweeps below.
register_cuda_ci(est_time=220, stage="nightly", runner_config="1-gpu-large")
@@ -115,11 +115,10 @@ def test_qknorm_rope_rejects_unsupported_dtypes() -> None:
)
BS_LIST = [2**n for n in range(13)]
BS_LIST += [x + 1 for x in BS_LIST]
BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 129, 257, 2049, 4097])
HEADS_LIST = get_ci_test_range([8, 16, 24, 32], [8, 24])
HEAD_DIM_LIST = get_ci_test_range([64, 128, 256], [64, 128, 256])
_FULL_BS_LIST = [2**n for n in range(13)]
_FULL_BS_LIST += [x + 1 for x in _FULL_BS_LIST]
_FULL_HEADS_LIST = [8, 16, 24, 32]
_FULL_HEAD_DIM_LIST = [64, 128, 256]
IS_NEOX_LIST = [False, True]
POSITION_DTYPES = [torch.int32, torch.int64]
ROPE_DIM_CHOICES = {
@@ -127,19 +126,34 @@ ROPE_DIM_CHOICES = {
128: [64, 128],
256: [64, 128, 256],
}
@pytest.mark.parametrize(
"batch_size,num_heads,head_dim,is_neox,position_dtype",
QKNORM_ROPE_CASES = get_ci_test_range(
list(
itertools.product(
BS_LIST,
HEADS_LIST,
HEAD_DIM_LIST,
_FULL_BS_LIST,
_FULL_HEADS_LIST,
_FULL_HEAD_DIM_LIST,
IS_NEOX_LIST,
POSITION_DTYPES,
)
),
[
(1, 8, 64, False, torch.int32),
(9, 24, 128, True, torch.int64),
(129, 8, 256, True, torch.int32),
(257, 24, 64, False, torch.int64),
(2049, 8, 128, True, torch.int32),
(4097, 24, 256, False, torch.int64),
(1, 24, 64, True, torch.int64),
(129, 8, 128, False, torch.int32),
(2049, 24, 256, True, torch.int64),
(4097, 8, 64, False, torch.int32),
],
)
@pytest.mark.parametrize(
"batch_size,num_heads,head_dim,is_neox,position_dtype",
QKNORM_ROPE_CASES,
)
def test_qknorm_rope(
batch_size: int,
test/registered/kernels/ops/mamba/test_sconv_extend_metadata.py
+33 -6
View File
@@ -9,6 +9,7 @@ clamp, long, to(int64), arange, searchsorted, clamp, to(int32).
import pytest
import torch
from sglang.kernels.jit.utils import get_ci_test_range
from sglang.srt.models.inkling_common.kernels.sconv import (
HIS_ONES,
HIS_PREFIX,
@@ -20,12 +21,41 @@ from sglang.srt.models.inkling_common.kernels.sconv import (
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_cuda_ci(est_time=12, stage="base-b-kernel-unit", runner_config="1-gpu-large")
# Nightly expands the representative PR cases below to the complete matrix.
register_cuda_ci(est_time=40, stage="nightly", runner_config="1-gpu-large")
requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only")
# cross si tiles (BLOCK_T=256) and the single-tile B bound
BATCH_SIZES = [1, 2, 7, 64, 257, 1023]
EXTEND_CASES = get_ci_test_range(
[
(b, his_mode, lens_dtype)
for lens_dtype in (torch.int32, torch.int64)
for his_mode in (HIS_ZEROS, HIS_PREFIX, HIS_SEQ_MINUS_EXT)
for b in BATCH_SIZES
],
[
(1, HIS_ZEROS, torch.int32),
(2, HIS_PREFIX, torch.int64),
(7, HIS_SEQ_MINUS_EXT, torch.int32),
(64, HIS_ZEROS, torch.int64),
(257, HIS_PREFIX, torch.int32),
(1023, HIS_SEQ_MINUS_EXT, torch.int64),
],
)
VERIFY_CASES = get_ci_test_range(
[(b, draft_token_num) for draft_token_num in (1, 9) for b in BATCH_SIZES],
[
(1, 1),
(2, 9),
(7, 1),
(64, 9),
(257, 1),
(1023, 9),
],
)
def _ref_extend(B, extend_seq_lens, his_mode, his_src, cache_indices, T):
@@ -88,9 +118,7 @@ def _cache_indices(b, idx_dtype):
@requires_cuda
@pytest.mark.parametrize("b", BATCH_SIZES)
@pytest.mark.parametrize("his_mode", [HIS_ZEROS, HIS_PREFIX, HIS_SEQ_MINUS_EXT])
@pytest.mark.parametrize("lens_dtype", [torch.int32, torch.int64])
@pytest.mark.parametrize("b,his_mode,lens_dtype", EXTEND_CASES)
def test_extend_matches_unfused(b, his_mode, lens_dtype):
torch.manual_seed(b * 10 + his_mode)
lens = torch.randint(0, 33, (b,), dtype=lens_dtype, device="cuda")
@@ -118,8 +146,7 @@ def test_extend_matches_unfused(b, his_mode, lens_dtype):
@requires_cuda
@pytest.mark.parametrize("b", BATCH_SIZES)
@pytest.mark.parametrize("draft_token_num", [1, 9])
@pytest.mark.parametrize("b,draft_token_num", VERIFY_CASES)
def test_verify_matches_unfused(b, draft_token_num):
torch.manual_seed(b)
cache_indices = _cache_indices(b, torch.int64)
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
test/registered/unit/test_fork_test_worker.py
+77
View File
@@ -0,0 +1,77 @@
import json
import os
import subprocess
import sys
import tempfile
import textwrap
import unittest
from pathlib import Path
from sglang.test.ci import fork_test_worker
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=15, suite="base-a-test-cpu")
@unittest.skipUnless(hasattr(os, "fork"), "fork requires a POSIX platform")
class TestForkTestWorker(CustomTestCase):
def test_files_run_in_isolated_children(self):
result_read_fd, result_write_fd = os.pipe()
process = subprocess.Popen(
[
sys.executable,
fork_test_worker.__file__,
"--result-fd",
str(result_write_fd),
],
stdin=subprocess.PIPE,
text=True,
pass_fds=(result_write_fd,),
)
os.close(result_write_fd)
try:
with (
tempfile.TemporaryDirectory() as tmpdir,
os.fdopen(result_read_fd) as result_stream,
):
first = Path(tmpdir) / "first.py"
first.write_text(textwrap.dedent("""
import builtins
import os
builtins._sglang_fork_worker_marker = 41
os.environ["SGLANG_FORK_WORKER_TEST"] = "leaked"
raise SystemExit(0)
"""))
second = Path(tmpdir) / "second.py"
second.write_text(textwrap.dedent("""
import builtins
import os
assert not hasattr(builtins, "_sglang_fork_worker_marker")
assert "SGLANG_FORK_WORKER_TEST" not in os.environ
raise SystemExit(3)
"""))
results = []
for filename in (first, second):
process.stdin.write(json.dumps({"filename": str(filename)}) + "\n")
process.stdin.flush()
results.append(json.loads(result_stream.readline()))
self.assertEqual([result["returncode"] for result in results], [0, 3])
self.assertTrue(all(result["elapsed"] >= 0 for result in results))
process.stdin.write(json.dumps({"command": "stop"}) + "\n")
process.stdin.flush()
self.assertEqual(process.wait(timeout=30), 0)
finally:
if process.poll() is None:
process.kill()
process.wait()
if __name__ == "__main__":
unittest.main()
test/run_suite.py
+13
View File
@@ -395,6 +395,7 @@ def run_a_suite(args):
enable_retry=args.enable_retry,
max_attempts=args.max_attempts,
retry_wait_seconds=args.retry_wait_seconds,
fork_worker_batch_size=args.fork_worker_batch_size,
)
@@ -487,8 +488,20 @@ def main():
default=None,
help="Path to sglang-ci-stats model.json for live LPT est; missing/malformed -> in-source est_time fallback.",
)
parser.add_argument(
"--fork-worker-batch-size",
type=int,
default=1,
help=(
"Preload common modules, then run this many files in isolated fork "
"children (default: 1, preserving one exec per file)."
),
)
args = parser.parse_args()
if args.fork_worker_batch_size <= 0:
parser.error("--fork-worker-batch-size must be positive")
# Validate auto-partition arguments
if (args.auto_partition_id is not None) != (args.auto_partition_size is not None):
parser.error(