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

.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

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 use_fork_worker:
-
+                if (
-            if capture_output:
+                    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

@@ -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

@@ -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)]
+_FULL_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
+_FULL_BS_LIST += [x + 1 for x in _FULL_BS_LIST]  # stress non-aligned paths
-BS_LIST = get_ci_test_range(BS_LIST, [1, 129, 2048, 2049])
+_FULL_NUM_KV_HEADS_LIST = [1, 2, 8]
-NUM_KV_HEADS_LIST = get_ci_test_range([1, 2, 8], [1, 8])
+_FULL_GQA_RATIO_LIST = [1, 4, 8]
-GQA_RATIO = get_ci_test_range([1, 4, 8], [1, 8])
+_FULL_ROPE_DIM_LIST = [64, 128, 256, 512]
-ROPE_DIM_LIST = get_ci_test_range([64, 128, 256, 512], [64, 256])
 IS_NEOX_LIST = [False, True]
-DTYPE_LIST = get_ci_test_range(
+_FULL_DTYPE_LIST = [torch.bfloat16, torch.float16]
-    [torch.bfloat16, torch.float16], [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(
-@pytest.mark.parametrize("gqa_ratio", GQA_RATIO)
+    "batch_size,gqa_ratio,num_kv_heads,rope_dim,is_neox,dtype", ROPE_CASES
-@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)
 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("batch_size,is_neox,rope_dim,head_dim", PARTIAL_ROPE_CASES)
-@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)
 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(
-@pytest.mark.parametrize("gqa_ratio", GQA_RATIO)
+    "batch_size,gqa_ratio,num_kv_heads,rope_dim,is_neox", FUSED_ROPE_STORE_CASES
-@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)
 def test_fused_rope_store(
     batch_size: int,
     gqa_ratio: int,

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)]
+_FULL_BS_LIST = [2**n for n in range(13)]
-BS_LIST += [x + 1 for x in BS_LIST]
+_FULL_BS_LIST += [x + 1 for x in _FULL_BS_LIST]
-BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 129, 257, 2049, 4097])
+_FULL_HEADS_LIST = [8, 16, 24, 32]
-HEADS_LIST = get_ci_test_range([8, 16, 24, 32], [8, 24])
+_FULL_HEAD_DIM_LIST = [64, 128, 256]
-HEAD_DIM_LIST = get_ci_test_range([64, 128, 256], [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],
 }
-
+QKNORM_ROPE_CASES = get_ci_test_range(
-
-@pytest.mark.parametrize(
-    "batch_size,num_heads,head_dim,is_neox,position_dtype",
     list(
         itertools.product(
-            BS_LIST,
+            _FULL_BS_LIST,
-            HEADS_LIST,
+            _FULL_HEADS_LIST,
-            HEAD_DIM_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

@@ -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("b,his_mode,lens_dtype", EXTEND_CASES)
-@pytest.mark.parametrize("his_mode", [HIS_ZEROS, HIS_PREFIX, HIS_SEQ_MINUS_EXT])
-@pytest.mark.parametrize("lens_dtype", [torch.int32, torch.int64])
 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("b,draft_token_num", VERIFY_CASES)
-@pytest.mark.parametrize("draft_token_num", [1, 9])
 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

@@ -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])
+_DTYPES = [torch.float32, torch.float16, torch.bfloat16]
-@pytest.mark.parametrize(
+_POWER_OF_TWO_DIMS = [
-    "dim",
+    1,
-    # Power-of-2 dims from python/sglang/kernels/aot/tests/test_hadamard.py (old AOT test)
+    2,
-    [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768],
+    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])
+_NON_POWER_OF_TWO_CASES = get_ci_test_range(
-@pytest.mark.parametrize(
+    list(itertools.product([137, 500, 1000], _DTYPES)),
-    "dim",
+    [
-    # Non-power-of-2 dims to test the padding path
+        (137, torch.float32),
-    # (137 from python/sglang/kernels/aot/tests/test_hadamard.py, 500/1000 added for coverage)
+        (137, torch.bfloat16),
-    [137, 500, 1000],
+        (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])
+def _mn_cases(dims):
-@pytest.mark.parametrize("dim", _12N_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,dtype", _mn_cases(_20N_DIMS))
-@pytest.mark.parametrize("dim", _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,dtype", _mn_cases(_28N_DIMS))
-@pytest.mark.parametrize("dim", _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,dtype", _mn_cases(_40N_DIMS))
-@pytest.mark.parametrize("dim", _40N_DIMS)
 def test_hadamard_transform_40n(dim, dtype):
     device = "cuda"
 

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])
+FP32_SCALE_CASES = get_ci_test_range(
-@pytest.mark.parametrize("column_major", [False, True])
+    list(
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+        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("shape,expected_m,masked_m_dtype", MASKED_TEST_CASES)
-@pytest.mark.parametrize("expected_m", [None, 4])
+def test_masked(shape, expected_m, masked_m_dtype):
-@pytest.mark.parametrize("num_experts,hidden,tokens_pad", MASKED_CASES)
-def test_masked(num_experts, hidden, tokens_pad, 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")])
+NON_FINITE_CASES = get_ci_test_range(
-@pytest.mark.parametrize("scale_ue8m0", [False, True])
+    list(
-@pytest.mark.parametrize("masked", [False, True])
+        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

@@ -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

@@ -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(