diff --git a/.github/workflows/pr-test-jit-kernel.yml b/.github/workflows/pr-test-jit-kernel.yml index dfdb4cd72..f94ccb9d2 100644 --- a/.github/workflows/pr-test-jit-kernel.yml +++ b/.github/workflows/pr-test-jit-kernel.yml @@ -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 diff --git a/python/sglang/test/ci/ci_utils.py b/python/sglang/test/ci/ci_utils.py index 9e0ed0b2b..50a7da70c 100644 --- a/python/sglang/test/ci/ci_utils.py +++ b/python/sglang/test/ci/ci_utils.py @@ -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: diff --git a/python/sglang/test/ci/fork_test_worker.py b/python/sglang/test/ci/fork_test_worker.py new file mode 100644 index 000000000..cccb52e79 --- /dev/null +++ b/python/sglang/test/ci/fork_test_worker.py @@ -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()) diff --git a/test/registered/kernels/ops/attention/test_rope.py b/test/registered/kernels/ops/attention/test_rope.py index d28a022b3..a60d94349 100644 --- a/test/registered/kernels/ops/attention/test_rope.py +++ b/test/registered/kernels/ops/attention/test_rope.py @@ -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, diff --git a/test/registered/kernels/ops/diffusion/test_rope.py b/test/registered/kernels/ops/diffusion/test_rope.py index 4f9bb1427..ee1fa089f 100644 --- a/test/registered/kernels/ops/diffusion/test_rope.py +++ b/test/registered/kernels/ops/diffusion/test_rope.py @@ -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, diff --git a/test/registered/kernels/ops/mamba/test_sconv_extend_metadata.py b/test/registered/kernels/ops/mamba/test_sconv_extend_metadata.py index 73c34c442..20cc96984 100644 --- a/test/registered/kernels/ops/mamba/test_sconv_extend_metadata.py +++ b/test/registered/kernels/ops/mamba/test_sconv_extend_metadata.py @@ -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) diff --git a/test/registered/kernels/ops/quantization/test_hadamard_jit.py b/test/registered/kernels/ops/quantization/test_hadamard_jit.py index cefd5f68d..7860685da 100644 --- a/test/registered/kernels/ops/quantization/test_hadamard_jit.py +++ b/test/registered/kernels/ops/quantization/test_hadamard_jit.py @@ -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" diff --git a/test/registered/kernels/ops/quantization/test_per_token_group_quant.py b/test/registered/kernels/ops/quantization/test_per_token_group_quant.py index cdcc00086..7718a02ac 100644 --- a/test/registered/kernels/ops/quantization/test_per_token_group_quant.py +++ b/test/registered/kernels/ops/quantization/test_per_token_group_quant.py @@ -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 diff --git a/test/registered/unit/test_fork_test_worker.py b/test/registered/unit/test_fork_test_worker.py new file mode 100644 index 000000000..13ff61f6b --- /dev/null +++ b/test/registered/unit/test_fork_test_worker.py @@ -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() diff --git a/test/run_suite.py b/test/run_suite.py index 1963b01b9..c60b6e03c 100644 --- a/test/run_suite.py +++ b/test/run_suite.py @@ -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(