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

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

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-08-29 07:26:49 +08:00
committed by GitHub
parent d12b313b93
commit 96a4dcdde8
10 changed files with 543 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/ops/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(