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

[CI] Prune redundant CPU test overhead (#34309)

Browse Source 
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
This commit is contained in:
Xinyuan Tong
2026-08-13 19:51:17 -07:00
committed by GitHub
co-authored by hnyls2002
parent 6ad3f2d8fd
commit 85cdf1178d
78 changed files with 884 additions and 981 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/attention/unittests/conftest.py
+3 -4
View File
@@ -1,10 +1,9 @@
import sys
from pathlib import Path
# Add this directory to sys.path so that test files can do
# `sys.path.insert(0, str(Path(__file__).resolve().parents[1]))` equivalently,
# and so pytest can import subpackages (dense/, mla/, etc.) without
# confusing this directory with the Python stdlib `unittest` module.
# Put this directory on sys.path so pytest can import the subpackages
# (dense/, mla/, ...) without confusing this directory with the stdlib
# `unittest` module.
_here = str(Path(__file__).resolve().parent)
if _here not in sys.path:
sys.path.insert(0, _here)
test/registered/attention/unittests/dense/test_extend_init_contract.py
+1 -6
View File
@@ -23,23 +23,18 @@ specific regression #26735 introduced and then fixed
``breakable_cuda_graph_runner.py`` capture sites).
"""
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import get_device_sm
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
DenseAttentionCase,
build_dense_attention_fixture,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=10, stage="base-a", runner_config="1-gpu-small")
test/registered/attention/unittests/dense/test_fa3.py
+1 -6
View File
@@ -1,15 +1,9 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import get_device_sm
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
DenseAttentionCase,
@@ -34,6 +28,7 @@ from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_
from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
run_dense_split_op_extend_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=25, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/dense/test_fa4.py
+1 -6
View File
@@ -1,14 +1,8 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
DenseAttentionCase,
@@ -33,6 +27,7 @@ from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_
from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
run_dense_split_op_extend_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=45, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=45, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/dense/test_flashinfer.py
+1 -6
View File
@@ -1,15 +1,9 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import is_flashinfer_available
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
DenseAttentionCase,
@@ -30,6 +24,7 @@ from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_
from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
run_dense_split_op_extend_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=25, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=25, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/dense/test_flex_attention.py
+1 -6
View File
@@ -1,14 +1,8 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
DenseAttentionCase,
@@ -18,6 +12,7 @@ from sglang.test.kits.attention_unittest.attention_methods.dense_attention impor
from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
run_dense_split_op_extend_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=25, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=25, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/dense/test_hybrid_attn.py
+1 -6
View File
@@ -1,6 +1,4 @@
import sys
import unittest
from pathlib import Path
import torch
@@ -8,10 +6,6 @@ from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS
from sglang.srt.layers.attention.hybrid_attn_backend import HybridAttnBackend
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import is_flashinfer_available
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
DENSE_ATOL,
@@ -22,6 +16,7 @@ from sglang.test.kits.attention_unittest.attention_methods.dense_attention impor
replace_backend,
run_dense_fixture_eager,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/dense/test_tbo.py
+1 -6
View File
@@ -1,6 +1,4 @@
import sys
import unittest
from pathlib import Path
import torch
@@ -8,10 +6,6 @@ from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS
from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import get_device_sm
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
DENSE_ATOL,
@@ -25,6 +19,7 @@ from sglang.test.kits.attention_unittest.attention_methods.dense_attention impor
from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_runner import (
_prepare_spec_verify_batch,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/dense/test_torch_native.py
+1 -6
View File
@@ -1,20 +1,15 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
DenseAttentionCase,
make_dense_cases,
run_dense_attention_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/dense/test_triton.py
+1 -6
View File
@@ -1,15 +1,9 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import is_hip
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
DenseAttentionCase,
@@ -34,6 +28,7 @@ from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_
from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
run_dense_split_op_extend_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=25, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=25, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/dense/test_trtllm_mha.py
+1 -6
View File
@@ -1,6 +1,4 @@
import sys
import unittest
from pathlib import Path
import torch
@@ -11,10 +9,6 @@ from sglang.srt.utils.common import (
is_sm100_supported,
is_sm120_supported,
)
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
DenseAttentionCase,
@@ -31,6 +25,7 @@ from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_runner i
from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
run_dense_split_op_extend_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/dsa/test_dsa.py
+1 -6
View File
@@ -1,14 +1,8 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.dsa_attention import (
DSA_DECODE_IMPL_VARIANTS,
@@ -34,6 +28,7 @@ from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner i
from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_runner import (
run_dsa_eagle_draft_cuda_graph_runner_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=25, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=25, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/dsv4/test_deepseek_v4.py
-4
View File
@@ -13,9 +13,7 @@ gate+norm+rotate compression itself) is a deferred follow-up.
"""
import importlib.util
import sys
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest import mock
@@ -24,8 +22,6 @@ import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
_FLASH_MLA_AVAILABLE = (
importlib.util.find_spec("sgl_kernel") is not None
and importlib.util.find_spec("sgl_kernel.flash_mla") is not None
test/registered/attention/unittests/dual_chunk/test_dual_chunk_flash_attn.py
+1 -6
View File
@@ -1,14 +1,8 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.kits.attention_unittest.attention_methods.dual_chunk_attention import (
DualChunkAttentionCase,
make_dual_chunk_cases,
@@ -23,6 +17,7 @@ from sglang.test.kits.attention_unittest.attention_methods.dual_chunk_attention
from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner import (
run_dual_chunk_cuda_graph_decode_case,
)
from sglang.test.test_utils import CustomTestCase
# Container gate (KNOWN_FAILURES.md §1): `DualChunkFlashAttentionBackend` calls
test/registered/attention/unittests/gdn/test_flashinfer.py
+2 -7
View File
@@ -1,16 +1,10 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import is_flashinfer_available
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.gdn_attention import (
GDNAttentionCase,
@@ -29,6 +23,7 @@ from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_
from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
run_gdn_split_op_extend_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/gdn/test_linear_replayssm_decode.py
+1 -5
View File
@@ -33,21 +33,17 @@ L sweep:
Runnable as ``pytest`` and as ``__main__``.
"""
import sys
import unittest
from pathlib import Path
import torch
from sglang.test.test_utils import CustomTestCase
# Mirror sibling GDN unittests: register for CUDA/AMD CI. This is a kernel-math
# unit test; it lives with the other linear-attention kernel correctness tests.
# The registry calls MUST be module-level (the CI collector / check-registered-
# tests hook parses them statically via AST and only scans top-level statements
# -- a try/except wrapper hides them and fails the hook).
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
register_amd_ci(est_time=20, suite="stage-b-test-1-gpu-large-amd")
test/registered/attention/unittests/gdn/test_torch_native.py
+1 -6
View File
@@ -1,15 +1,9 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import is_hip
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.gdn_attention import (
GDNAttentionCase,
@@ -19,6 +13,7 @@ from sglang.test.kits.attention_unittest.attention_methods.gdn_attention import
from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
run_gdn_split_op_extend_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/gdn/test_triton.py
+1 -6
View File
@@ -1,6 +1,4 @@
import sys
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock
@@ -12,10 +10,6 @@ from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
)
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import is_hip
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.gdn_attention import (
GDNAttentionCase,
@@ -32,6 +26,7 @@ from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_
from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
run_gdn_split_op_extend_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/hybrid_linear/test_flashinfer_mla_chunk_metadata.py
-5
View File
@@ -13,7 +13,6 @@ so the chunked-MHA path never runs.
"""
import unittest
from pathlib import Path
from types import SimpleNamespace
import torch
@@ -143,8 +142,4 @@ class TestHybridLinearChunkMetadataDelegation(CustomTestCase):
if __name__ == "__main__":
sys_path_parent = str(Path(__file__).resolve().parents[1])
import sys
sys.path.insert(0, sys_path_parent)
unittest.main()
test/registered/attention/unittests/kda/test_triton.py
+1 -6
View File
@@ -1,15 +1,9 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import is_hip
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.kda_attention import (
KDAAttentionCase,
@@ -26,6 +20,7 @@ from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_
from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
run_kda_split_op_extend_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/lightning/test_triton.py
+1 -6
View File
@@ -1,15 +1,9 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.kernels.ops.attention.linear.seg_la import SegLaMeta, seg_la_fwd
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.lightning_attention import (
LightningAttentionCase,
@@ -23,6 +17,7 @@ from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_
run_lightning_eagle_verify_case,
run_lightning_eagle_verify_cuda_graph_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/mamba/test_mamba2.py
+1 -6
View File
@@ -1,6 +1,4 @@
import sys
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock
@@ -11,10 +9,6 @@ from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
MambaAttnBackendBase,
)
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.mamba2_attention import (
DEFAULT_CONV_KERNEL,
@@ -36,6 +30,7 @@ from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_
run_mamba2_eagle_verify_case,
run_mamba2_eagle_verify_cuda_graph_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/mla/test_cutlass_mla.py
+1 -6
View File
@@ -1,18 +1,13 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.kits.attention_unittest.attention_methods.mla_attention import (
MLAAttentionCase,
run_mla_attention_case,
)
from sglang.test.test_utils import CustomTestCase
# Cutlass MLA requires exactly Blackwell SM 10.0. The sgl-kernel
# `cutlass_mla_decode` checks `sm_version == 100` (major*10+minor), so
test/registered/attention/unittests/mla/test_flashinfer.py
+1 -6
View File
@@ -1,14 +1,8 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.kits.attention_unittest.attention_methods.mla_attention import (
MLAAttentionCase,
make_mla_cases,
@@ -27,6 +21,7 @@ from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_
from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
run_mla_split_op_extend_case,
)
from sglang.test.test_utils import CustomTestCase
MLA_SHAPE_KWARGS = dict(
kv_lora_rank=512,
test/registered/attention/unittests/mla/test_flashmla.py
+1 -6
View File
@@ -1,16 +1,10 @@
import sys
import unittest
from pathlib import Path
import torch
import triton
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.kits.attention_unittest.attention_methods.mla_attention import (
MLAAttentionCase,
build_mla_attention_fixture,
@@ -33,6 +27,7 @@ from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_
from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
run_mla_split_op_extend_case,
)
from sglang.test.test_utils import CustomTestCase
MLA_SHAPE_KWARGS = dict(
kv_lora_rank=512,
test/registered/attention/unittests/mla/test_tokenspeed_mla.py
+1 -6
View File
@@ -1,19 +1,14 @@
import importlib.util
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.kits.attention_unittest.attention_methods.mla_attention import (
MLAAttentionCase,
run_mla_attention_case,
)
from sglang.test.test_utils import CustomTestCase
# tokenspeed_mla is a CuTe DSL backend for Blackwell (SM100). It additionally
# enforces:
test/registered/attention/unittests/mla/test_triton.py
+1 -6
View File
@@ -1,14 +1,8 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.mla_attention import (
MLAAttentionCase,
@@ -32,6 +26,7 @@ from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_
from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
run_mla_split_op_extend_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=25, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=25, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/mla/test_trtllm_mla.py
+1 -6
View File
@@ -1,18 +1,13 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.kits.attention_unittest.attention_methods.mla_attention import (
MLAAttentionCase,
run_mla_attention_case,
)
from sglang.test.test_utils import CustomTestCase
# trtllm_mla goes through FlashInfer's XQA MLA path. Per PLAN.md and the
# project's is_sm120_supported helper (device_capability_majors=[12]), the
test/registered/attention/unittests/swa/test_flashinfer.py
+1 -6
View File
@@ -1,16 +1,10 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.environ import envs
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import is_flashinfer_available
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
DenseAttentionCase,
@@ -28,6 +22,7 @@ from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_
from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
run_dense_split_op_extend_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/swa/test_swa_out_cache_loc.py
+1 -6
View File
@@ -6,20 +6,15 @@ uses swa_loc directly for SWA layers and asserts it is provided. The per-backend
cuda-graph buffer plumbing is covered by the backend SWA integration tests.
"""
import sys
import unittest
from pathlib import Path
from types import SimpleNamespace
import torch
from sglang.srt.mem_cache.memory_pool import KVWriteLoc
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
test/registered/attention/unittests/swa/test_torch_native.py
+1 -6
View File
@@ -1,14 +1,8 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
DenseAttentionCase,
@@ -16,6 +10,7 @@ from sglang.test.kits.attention_unittest.attention_methods.dense_attention impor
make_swa_prefix_input_config_cases,
run_dense_attention_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
test/registered/attention/unittests/swa/test_triton.py
+1 -6
View File
@@ -1,15 +1,9 @@
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import is_hip
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
DenseAttentionCase,
@@ -27,6 +21,7 @@ from sglang.test.kits.attention_unittest.runner_modes.speculative_target_verify_
from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
run_dense_split_op_extend_case,
)
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200")
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
test/registered/cpu/arm64/test_moe.py
+1 -10
View File
@@ -11,17 +11,13 @@ register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64")
import itertools
import math
import os
import platform
import sys
import unittest
import torch
# Add parent dir (test/srt/cpu/) to path for utils import
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from sglang.srt.layers.amx_utils import CPUQuantMethod
from sglang.test.cpu_test_utils import precision, torch_w8a8_per_column_fused_moe
from sglang.test.test_utils import CustomTestCase
kernel = torch.ops.sgl_kernel
@@ -29,11 +25,6 @@ IS_ARM64 = platform.machine().lower() in ("aarch64", "arm64")
torch.manual_seed(128)
from utils import (
precision,
torch_w8a8_per_column_fused_moe,
)
class TestFusedExpertsInt8(CustomTestCase):
M = [1, 6, 32, 64]
test/registered/cpu/test_activation.py
+1 -1
View File
@@ -2,10 +2,10 @@ import sys
import pytest
import torch
from utils import GeluAndMul, SiluAndMul, precision
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.cpu_test_utils import GeluAndMul, SiluAndMul, precision
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64")
test/registered/cpu/test_bmm.py
+1 -1
View File
@@ -4,10 +4,10 @@ import unittest
# TODO: use interface in cpu.py
import torch
import torch.nn as nn
from utils import precision
from sglang.srt.layers.quantization.fp8_utils import input_to_float8
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.cpu_test_utils import precision
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
test/registered/cpu/test_causal_conv1d.py
+1 -1
View File
@@ -4,9 +4,9 @@ from typing import Optional
import sgl_kernel # noqa: F401
import torch
import torch.nn.functional as F
from utils import parametrize, precision
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.cpu_test_utils import parametrize, precision
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
test/registered/cpu/test_flash_attn.py
+1 -1
View File
@@ -3,9 +3,9 @@ import unittest
import sgl_kernel # noqa: F401
import torch
import torch.nn.functional as F
from utils import parametrize, precision
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.cpu_test_utils import parametrize, precision
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
test/registered/cpu/test_gemm.py
+3 -3
View File
@@ -3,7 +3,9 @@ import unittest
# TODO: use interface in cpu.py
import torch
import torch.nn as nn
from utils import (
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.cpu_test_utils import (
MXFP4QuantizeUtil,
convert_weight,
native_w8a8_per_token_matmul,
@@ -13,8 +15,6 @@ from utils import (
unpack_and_dequant_awq,
unpack_and_dequant_gptq,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
test/registered/cpu/test_mamba.py
+1 -1
View File
@@ -4,9 +4,9 @@ import pytest
import torch
import torch.nn.functional as F
from torch.nn.functional import softplus
from utils import precision
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.cpu_test_utils import precision
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
test/registered/cpu/test_mla.py
+1 -1
View File
@@ -2,9 +2,9 @@ import unittest
import torch
from torch.nn.functional import scaled_dot_product_attention
from utils import precision
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.cpu_test_utils import precision
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
test/registered/cpu/test_moe.py
+2 -3
View File
@@ -16,7 +16,8 @@ prepack = True
alpha = 1.702
limit = 7.0
from utils import (
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.cpu_test_utils import (
BLOCK_K,
BLOCK_N,
MXFP4QuantizeUtil,
@@ -32,8 +33,6 @@ from utils import (
unpack_and_dequant_awq,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
test/registered/cpu/test_norm.py
+1 -1
View File
@@ -3,9 +3,9 @@ from typing import Optional, Tuple, Union
import pytest
import torch
from utils import make_non_contiguous, precision
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.cpu_test_utils import make_non_contiguous, precision
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64")
test/registered/cpu/test_qkv_proj_with_rope.py
+5 -5
View File
@@ -1,16 +1,16 @@
import unittest
import torch
from utils import (
from sglang.srt.layers.quantization.fp8_utils import input_to_float8
from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.cpu_test_utils import (
convert_weight,
native_w8a8_per_token_matmul,
per_token_quant_int8,
precision,
)
from sglang.srt.layers.quantization.fp8_utils import input_to_float8
from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
test/registered/cpu/test_qwen3.py
+1 -1
View File
@@ -2,10 +2,10 @@ import sys
import pytest
import torch
from utils import precision
from sglang.srt.utils import is_host_cpu_arm64
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.cpu_test_utils import precision
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64")
test/registered/cpu/test_rope.py
+1 -1
View File
@@ -1,7 +1,6 @@
import unittest
import torch
from utils import precision
from sglang.srt.layers.rotary_embedding import (
MRotaryEmbedding,
@@ -14,6 +13,7 @@ from sglang.srt.layers.rotary_embedding.rope_variant import (
from sglang.srt.layers.rotary_embedding.utils import apply_rotary_pos_emb_native_eager
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.cpu_test_utils import precision
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
test/registered/cpu/test_shared_expert.py
+3 -3
View File
@@ -3,7 +3,9 @@ import math
import unittest
import torch
from utils import (
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.cpu_test_utils import (
BLOCK_K,
BLOCK_N,
factor_for_scale,
@@ -15,8 +17,6 @@ from utils import (
torch_naive_moe,
torch_w8a8_per_column_moe,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
test/registered/cpu/test_spec_kernels.py
+1 -1
View File
@@ -3,10 +3,10 @@ import unittest
import sgl_kernel # noqa: F401
import torch
import torch.nn.functional as F
from utils import precision
from sglang.srt.speculative.eagle_utils import TreeMaskMode, organize_draft_results
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.cpu_test_utils import precision
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=20, suite="base-b-test-cpu")
test/registered/cpu/utils.py
-659
View File
@@ -1,659 +0,0 @@
import itertools
import math
import torch
import torch.nn.functional as F
precision = {
torch.bfloat16: 1e-2,
torch.float16: 1e-3,
torch.float32: 1e-5,
}
BLOCK_N, BLOCK_K = 64, 128
factor_for_scale = 1e-3
fp8_max, fp8_min = 400, -400
def parametrize(**params):
def decorator(func):
def wrapper(self):
for combo in itertools.product(*params.values()):
kwargs = dict(zip(params.keys(), combo))
with self.subTest(**kwargs):
func(self, **kwargs)
return wrapper
return decorator
def SiluAndMul(x: torch.Tensor) -> torch.Tensor:
d = x.shape[-1] // 2
return F.silu(x[..., :d]) * x[..., d:]
def GeluAndMul(x: torch.Tensor, approximate="tanh") -> torch.Tensor:
d = x.shape[-1] // 2
return F.gelu(x[..., :d], approximate=approximate) * x[..., d:]
def per_token_quant_int8(x):
x = x.float()
absmax = x.abs().max(dim=-1).values
absmax = absmax.clamp_min(1e-10).unsqueeze(-1)
scale_x = absmax / 127
x_q = x.mul(127 / absmax)
x_q = torch.round(x_q).to(torch.int8)
return x_q, scale_x
def convert_weight(weight, scale_block_size, A_dtype):
N, K = weight.size()
fp8_max = 448.0
scale_block_size_N, scale_block_size_K = scale_block_size # (128, 128)
pad_N = (scale_block_size_N - (N % scale_block_size_N)) % scale_block_size_N
pad_K = (scale_block_size_K - (K % scale_block_size_K)) % scale_block_size_K
if pad_N > 0 or pad_K > 0:
weight = torch.nn.functional.pad(weight, (0, pad_K, 0, pad_N))
weight_blocks = weight.view(
math.ceil(N / scale_block_size_N),
scale_block_size_N,
math.ceil(K / scale_block_size_K),
scale_block_size_K,
) # (8, 128, 8, 128)
weight_blocks = weight_blocks.permute(0, 2, 1, 3).contiguous() # (8, 8, 128, 128)
# Step 2: compute per-block max abs values → scale
abs_max = weight_blocks.abs().amax(dim=(-2, -1), keepdim=True) # (8, 8, 1, 1)
scales = abs_max / fp8_max
scales = torch.where(
scales == 0, torch.ones_like(scales), scales
) # avoid division by zero
q_fp8 = (weight_blocks / scales).to(torch.float8_e4m3fn)
q_fp8_reshape = q_fp8.permute(0, 2, 1, 3).contiguous()
if pad_N > 0 or pad_K > 0:
q_fp8_reshape = q_fp8_reshape.view(N + pad_N, K + pad_K)
q_fp8_reshape = q_fp8_reshape[:N, :K].contiguous()
else:
q_fp8_reshape = q_fp8_reshape.view(N, K)
dq_weight = q_fp8.float() * scales
dq_weight = dq_weight.permute(0, 2, 1, 3).contiguous() # (8, 128, 8, 128)
if pad_N > 0 or pad_K > 0:
w_dq = dq_weight.view(N + pad_N, K + pad_K).to(A_dtype)
w_dq = w_dq[:N, :K].contiguous()
else:
w_dq = dq_weight.view(N, K).to(A_dtype)
scales = scales.view(
math.ceil(N / scale_block_size_N), math.ceil(K / scale_block_size_K)
)
return q_fp8_reshape, scales, w_dq
def native_w8a8_per_token_matmul(A, B, As, Bs, bias, output_dtype=torch.bfloat16):
"""Matrix multiplication function that supports per-token input quantization and per-column weight quantization"""
A = A.to(torch.float32)
B = B.to(torch.float32)
assert A.shape[-1] == B.shape[-1], "Dimension mismatch"
assert B.ndim == 2 and B.is_contiguous(), "B must be a 2D contiguous tensor"
# Reshape input
M = A.numel() // A.shape[-1]
B = B.t() # Transpose weight matrix
N, K = B.shape
origin_C_shape = A.shape[:-1] + (K,)
A = A.reshape(M, N)
# As is per-token [M, 1], Bs is per-column [1, K]
C = torch.matmul(A, B) # [M, K]
C = As * C * Bs.view(1, -1) # Broadcast per-column scale
if bias is not None:
C.add_(bias.view(1, -1))
return C.reshape(origin_C_shape).to(output_dtype)
def torch_naive_moe(a, w1, w2, b, routed_scaling_factor, output_dtype=torch.bfloat16):
a = a.to(torch.float32)
w1 = w1.to(torch.float32)
w2 = w2.to(torch.float32)
b = b.to(torch.float32) if b is not None else None
ic1 = torch.matmul(a, w1.transpose(0, 1))
ic2 = SiluAndMul(ic1)
ic3 = torch.matmul(ic2, w2.transpose(0, 1))
out = ic3 if b is None else ic3 + b * routed_scaling_factor
return out.to(output_dtype)
def torch_w8a8_per_column_moe(
a, w1_q, w2_q, w1_s, w2_s, b, routed_scaling_factor, output_dtype=torch.bfloat16
):
a = a.to(torch.float32)
b = b.to(torch.float32) if b is not None else None
# Perform per-token quantization
a_q, a_s = per_token_quant_int8(a)
ic1 = native_w8a8_per_token_matmul(
a_q, w1_q, a_s, w1_s, bias=None, output_dtype=torch.float32
)
ic2 = SiluAndMul(ic1)
a1_q, a1_s = per_token_quant_int8(ic2)
ic3 = native_w8a8_per_token_matmul(
a1_q, w2_q, a1_s, w2_s, bias=None, output_dtype=torch.float32
)
out = ic3 if b is None else ic3 + b * routed_scaling_factor
return out.to(output_dtype)
def scaled_weight(weight, scales):
E, N, K = weight.shape
pad_N = (BLOCK_N - (N % BLOCK_N)) % BLOCK_N
pad_K = (BLOCK_K - (K % BLOCK_K)) % BLOCK_K
if pad_N > 0 or pad_K > 0:
weight = torch.nn.functional.pad(weight, (0, pad_K, 0, pad_N))
weight_block = (
weight.view(E, math.ceil(N / BLOCK_N), BLOCK_N, math.ceil(K / BLOCK_K), BLOCK_K)
.permute(0, 1, 3, 2, 4)
.float()
.contiguous()
)
weight_scaled = (
(
weight_block
* scales.view(E, math.ceil(N / BLOCK_N), math.ceil(K / BLOCK_K), 1, 1)
)
.permute(0, 1, 3, 2, 4)
.contiguous()
)
if pad_N > 0 or pad_K > 0:
weight_scaled = weight_scaled.view(E, N + pad_N, K + pad_K)
weight_scaled = weight_scaled[..., :N, :K].contiguous()
else:
weight_scaled = weight_scaled.view(E, N, K)
return weight_scaled
def torch_naive_fused_moe(a, w1, w2, score, topk, renormalize):
B, D = a.shape
a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
if renormalize:
topk_weight = topk_weight / topk_weight.sum(dim=-1, keepdim=True)
topk_weight = topk_weight.view(-1)
topk_ids = topk_ids.view(-1)
for i in range(w1.shape[0]):
mask = topk_ids == i
if mask.sum():
out[mask] = SiluAndMul(a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(
0, 1
)
return (
out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
).sum(dim=1)
def moe_gptoss_act(x, alpha: float = 1.702, limit: float = 7.0):
x_glu, x_linear = x[..., ::2], x[..., 1::2]
# Clamp the input values
x_glu = x_glu.clamp(min=None, max=limit)
x_linear = x_linear.clamp(min=-limit, max=limit)
out_glu = x_glu * torch.sigmoid(alpha * x_glu)
# Note we add an extra bias of 1 to the linear layer
return out_glu * (x_linear + 1.0)
def torch_naive_gptoss_fused_moe(
x,
w1,
w2,
w1_bias,
w2_bias,
topk_weights,
topk_ids,
activation_alpha,
swiglu_limit,
len_experts,
) -> torch.Tensor:
# Ref code from https://huggingface.co/deepseek-ai/DeepSeek-V2/blob/e0828e3cc0a03408724b80c3cc92c8e072db8d01/modeling_deepseek.py#L589
cnts = topk_ids.new_zeros((topk_ids.shape[0], len_experts))
cnts.scatter_(1, topk_ids.to(torch.int64), 1)
tokens_per_expert = cnts.sum(dim=0)
idxs = topk_ids.view(-1).argsort()
sorted_tokens = x[idxs // topk_ids.shape[1]]
tokens_per_expert = tokens_per_expert.cpu().numpy()
outputs = []
start_idx = 0
for i, num_tokens in enumerate(tokens_per_expert):
end_idx = start_idx + num_tokens
if num_tokens == 0:
continue
tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
layer_w13_weight = w1[i]
layer_w13_weight_bias = w1_bias[i]
layer_w2_weight_bias = w2_bias[i]
layer_w2_weight = w2[i]
gate_up = F.linear(
tokens_for_this_expert,
layer_w13_weight,
bias=layer_w13_weight_bias.to(torch.bfloat16),
)
gate_up = moe_gptoss_act(gate_up, activation_alpha, swiglu_limit)
expert_out = F.linear(
gate_up, layer_w2_weight, bias=layer_w2_weight_bias.to(torch.bfloat16)
)
outputs.append(expert_out)
start_idx = end_idx
outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
new_x = torch.empty_like(outs)
new_x[idxs] = outs
final_out = (
new_x.view(*topk_ids.shape, -1)
.type(topk_weights.dtype)
.mul_(topk_weights.unsqueeze(dim=-1))
.sum(dim=1)
.type(new_x.dtype)
)
return final_out
def torch_naive_fused_moe_gptoss(
a,
w1,
w2,
w1_bias,
w2_bias,
topk_weight,
topk_ids,
renormalize,
activation_alpha,
swiglu_limit,
len_experts,
):
if renormalize:
topk_weight = topk_weight / topk_weight.sum(dim=-1, keepdim=True)
return torch_naive_gptoss_fused_moe(
a,
w1,
w2,
w1_bias,
w2_bias,
topk_weight,
topk_ids,
activation_alpha,
swiglu_limit,
len_experts,
)
def torch_w8a8_per_column_fused_moe(a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, topk):
"""This function performs fused moe with per-column int8 quantization using native torch."""
B, D = a.shape
# Perform per-token quantization
a_q, a_s = per_token_quant_int8(a)
# Repeat tokens to match topk
a_q = a_q.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
# Also repeat the scale
a_s = a_s.view(B, -1, 1).repeat(1, topk, 1).reshape(-1, 1) # [B*topk, 1]
out = torch.zeros(B * topk, w2.shape[1], dtype=torch.float32, device=a.device)
# Calculate routing
topk_weight = topk_weight.view(-1)
topk_ids = topk_ids.view(-1)
# Process each expert
for i in range(w1.shape[0]):
mask = topk_ids == i
if mask.sum():
# First MLP layer: note that a_s is now per-token
inter_out = native_w8a8_per_token_matmul(
a_q[mask],
w1[i],
a_s[mask],
w1_s[i],
bias=None,
output_dtype=torch.float32,
)
# Activation function
act_out = SiluAndMul(inter_out)
# Quantize activation output with per-token
act_out_q, act_out_s = per_token_quant_int8(act_out)
# Second MLP layer
out[mask] = native_w8a8_per_token_matmul(
act_out_q,
w2[i],
act_out_s,
w2_s[i],
bias=None,
output_dtype=torch.float32,
)
# Apply routing weights and sum
return (
(out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype))
.sum(dim=1)
.to(a.dtype)
)
def native_fp8_fused_moe(a, w1, w2, topk_weight, topk_ids, topk):
B, D = a.shape
a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D).float()
out = torch.zeros(B * topk, w2.shape[1], dtype=torch.float32, device=a.device)
# Calculate routing
topk_weight = topk_weight.view(-1)
topk_ids = topk_ids.view(-1)
for i in range(w1.shape[0]):
mask = topk_ids == i
if mask.sum():
ic0 = torch.matmul(a[mask], w1[i].transpose(0, 1))
ic1 = SiluAndMul(ic0)
out[mask] = torch.matmul(ic1, w2[i].transpose(0, 1))
return (
(out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype))
.sum(dim=1)
.to(a.dtype)
)
# https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/modelopt/torch/quantization/qtensor/mxfp4_tensor.py
class MXFP4QuantizeUtil:
E2M1_max = 6.0
E2M1_values = [0, 0.5, 1, 1.5, 2, 3, 4, 6]
E2M1_bounds = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5])
block_size = 32
@classmethod
def quantize(cls, input: torch.Tensor) -> tuple:
"""Converting a tensor to a quantized format based on MXFP4 quantization. Only E4M3 is supported.
Args:
input (torch.Tensor): The input tensor to be quantized.
"""
def cast_fp4(x):
sign = torch.sign(x)
sign_bit = (2 - sign) // 2
ord_ = torch.sum(
(x.abs().unsqueeze(-1) - cls.E2M1_bounds.to(x.device)) > 0, dim=-1
)
fp4_val = (sign_bit * 0b1000 + ord_).to(torch.uint8)
return fp4_val
def fuse_uint4_to_uint8(x):
# If the last dimension is odd, pad with zeros
# If this behavior is not desired, please modify the code accordingly
left_side = x[..., 0::2] # Even indices (0, 2, 4...)
right_side = x[..., 1::2] # Odd indices (1, 3, 5...)
new_data = (
right_side.clone() << 4
) # Put odd indices (higher addresses) in high bits
new_data[
..., : left_side.shape[-1]
] += left_side # Put even indices in low bits
return new_data
original_shape = input.shape
original_dtype = input.dtype
input = input.view(-1, cls.block_size)
# get scales
input_amax = input.abs().max(dim=-1, keepdim=True).values
descale = input_amax / cls.E2M1_max
min_value = torch.tensor(-127.0, device=descale.device)
e8m0_scale = torch.ceil(torch.maximum(torch.log2(descale), min_value))
input = (input / torch.exp2(e8m0_scale)).view(original_shape)
input_q = cast_fp4(input)
input_q = fuse_uint4_to_uint8(input_q)
e8m0_scale = (e8m0_scale + 127).to(torch.uint8)
return input_q, e8m0_scale
@classmethod
def dequantize(cls, quantized_data, dtype: torch.dtype, scale):
"""Dequantze MXFP4 packed tensor to a target dtype."""
def unfuse_uint8_to_uint4(x):
"""Unfuse uint8 values back to uint4 values.
This is the inverse operation of fuse_uint4_to_uint8.
"""
# Extract the lower 4 bits (even indices)
left_side = x & 0x0F
# Extract the upper 4 bits (odd indices)
right_side = (x >> 4) & 0x0F
# Create a new tensor with alternating values
shape = list(x.shape)
shape[-1] = shape[-1] * 2
result = torch.zeros(shape, dtype=torch.uint8, device=x.device)
# Fill in the values - even indices get low bits, odd indices get high bits
result[..., 0::2] = left_side # Even indices from low bits
result[..., 1::2] = right_side # Odd indices from high bits
return result
e8m0_scale = scale
# Unfuse the uint8 values back to uint4
x_unfused = unfuse_uint8_to_uint4(quantized_data)
# print("@@@ x_unfused: ", x_unfused)
# Extract sign and magnitude
sign = 1 - 2 * ((x_unfused & 0b1000) >> 3).to(
torch.float32
) # Extract sign bit and convert to +1/-1
magnitude = x_unfused & 0b0111 # Extract magnitude bits
magnitude = magnitude.to(torch.long)
# Create a tensor with the E2M1 values
values = torch.tensor(cls.E2M1_values, device=quantized_data.device)
# Use gather to index the values tensor properly
# We need to reshape magnitude to match the dimensions we want to gather along
original_shape = magnitude.shape
x_float = values[magnitude.reshape(-1)].reshape(original_shape)
# Apply sign and scale
x_float = sign.float() * x_float
# Reshape to apply block-wise scaling
x_float = x_float.reshape(-1, cls.block_size)
# Apply the E8M0 scale
scale_factor = torch.exp2(e8m0_scale.float() - 127)
scale_factor = scale_factor.reshape(-1, 1) # Reshape for proper broadcasting
# Apply scaling and reshape back to original shape
x_float = x_float * scale_factor
# Reshape back to the original shape
return x_float.reshape(original_shape).to(dtype)
def make_non_contiguous(x: torch.Tensor) -> torch.Tensor:
# Make a tensor non-contiguous without changing shape.
if not x.is_contiguous():
return x
last_dim = x.shape[-1]
expanded = torch.empty(*x.shape[:-1], last_dim + 32, dtype=x.dtype, device=x.device)
expanded[..., :last_dim].copy_(x)
return expanded.narrow(-1, 0, last_dim)
def awq_reverse_reorder_int_tensor(int_tensor, bits: int):
assert bits == 4
int_tensor = int_tensor.T.contiguous()
compress_ratio = 32 // bits
assert int_tensor.shape[-1] % compress_ratio == 0
order_map = [0, 2, 4, 6, 1, 3, 5, 7]
order_tensor = torch.tensor(
order_map, dtype=torch.int32, device=int_tensor.device
).reshape(1, -1)
order_tensor = order_tensor.repeat(int_tensor.shape[1] // compress_ratio, 1)
order_tensor = order_tensor + torch.arange(
0,
int_tensor.shape[1],
compress_ratio,
dtype=torch.int32,
device=int_tensor.device,
).reshape(-1, 1)
order_tensor = order_tensor.reshape(-1)
reverse_order_tensor = torch.arange(order_tensor.shape[0])[order_tensor]
reverse_order_tensor = reverse_order_tensor[order_tensor]
int_tensor = int_tensor[:, reverse_order_tensor]
return int_tensor
def unpack_and_dequant_awq(
awq_qweight: torch.Tensor,
awq_qzeros: torch.Tensor,
awq_scales: torch.Tensor,
bits: int,
group_size: int,
):
"""
Args:
awq_qweight (`torch.LongTensor`):
Expected shape: (in_features, out_features // (32 // bits))
awq_qzeros (`torch.LongTensor`):
Expected shape: (in_features // group_size, out_features // (32 // bits))
awq_scales (`torch.LongTensor`):
Expected shape: (in_features // group_size, out_features)
Returns:
fp16_weight (`torch.LongTensor`):
With shape (in_features, out_features).
zeros (`torch.LongTensor`):
With shape (in_features // group_size, out_features).
"""
assert bits == 4
qzeros = awq_qzeros
qweight = awq_qweight
qweight = qweight.T.contiguous()
scales = awq_scales
scales = scales.reshape(-1, 1, scales.shape[-1])
infeatures = awq_qweight.shape[0]
wf = torch.tensor(
list(range(0, 32, bits)), dtype=torch.int32, device=qzeros.device
).unsqueeze(0)
zeros = torch.bitwise_right_shift(torch.unsqueeze(qzeros, 2), wf.unsqueeze(0)).to(
torch.int16 if bits == 8 else torch.int8
)
torch.bitwise_and(zeros, (2**bits) - 1, out=zeros)
zeros = zeros.reshape(-1, 1, zeros.shape[1] * zeros.shape[2])
weight = torch.bitwise_right_shift(
torch.unsqueeze(qweight, 1), wf.unsqueeze(-1)
).to(torch.int16 if bits == 8 else torch.int8)
torch.bitwise_and(weight, (2**bits) - 1, out=weight)
weight = weight.reshape(-1, group_size, weight.shape[2])
weight = weight.view(-1, weight.shape[-1])
zeros = zeros.view(-1, zeros.shape[-1])
zeros = zeros.T.contiguous()
zeros = awq_reverse_reorder_int_tensor(zeros, bits)
weight = awq_reverse_reorder_int_tensor(weight, bits)
# Dequantize weights.
scales = awq_scales
zeros = zeros.contiguous()
scale_zeros = zeros * scales
g_idx = torch.tensor(
[i // group_size for i in range(infeatures)], dtype=torch.int32
)
scale_mat = scales[g_idx]
scale_zeros_mat = scale_zeros[g_idx].to(torch.bfloat16)
qdq_weight_T = weight * scale_mat - scale_zeros_mat.to(torch.bfloat16)
fp16_weight = qdq_weight_T.T
return fp16_weight, zeros
def unpack_4bit_to_32bit_signed(qweight, qzeros):
# Unpack 4-bit values and interpret them as signed integers
unpacked_weights = torch.zeros(
(qweight.shape[0] * 8, qweight.shape[1]),
dtype=torch.int8,
device=qweight.device,
requires_grad=False,
)
unpacked_zeros = torch.zeros(
(qzeros.shape[0], qzeros.shape[1] * 8),
dtype=torch.int8,
device=qzeros.device,
requires_grad=False,
)
for row in range(unpacked_weights.shape[0]):
i = row % 8
unpacked_weights[row, :] = (qweight[row // 8, :] >> (4 * i)) & 0xF
for col in range(unpacked_zeros.shape[1]):
i = col % 8
unpacked_zeros[:, col] = (qzeros[:, col // 8] >> (4 * i)) & 0xF
return unpacked_weights, unpacked_zeros + 1
def unpack_and_dequant_gptq(qweight, qzeros, scales):
unpacked_qweight, unpacked_qzeros = unpack_4bit_to_32bit_signed(qweight, qzeros)
group_size = unpacked_qweight.shape[0] // scales.shape[0]
scales = scales.repeat_interleave(group_size, dim=0)
unpacked_qzeros = unpacked_qzeros.repeat_interleave(group_size, dim=0)
unpacked_qweight = (unpacked_qweight - unpacked_qzeros) * scales
return unpacked_qweight.T
test/registered/rust/test_cargo_workspace.py
-76
View File
@@ -1,76 +0,0 @@
"""Run the `rust/` Cargo workspace's unit tests from the CPU CI suite.
The `rust/` workspace (sglang-grpc, sglang-mm, sglang-server) is compiled into
the wheel by setuptools-rust, but until now nothing ran `cargo test` in CI --
`.github/workflows/pr-test-rust.yml` and `pr-benchmark-rust.yml` are both
path-scoped to `sgl-model-gateway/**`, a different workspace. `lint.yml` covers
rustfmt/clippy via the pre-commit hooks, so this file only adds the test run.
The debug profile is deliberate: these are pure-logic tests (no timing or
codegen assertions), and the release profile costs a full LTO build for the
same coverage.
"""
import shutil
import subprocess
import unittest
from pathlib import Path
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
# base-c-test-cpu is where this was asked for, and it matches the repo's
# base-a + base-c dual-registration convention -- but base-c-test-cpu currently
# has no runner job in any workflow (it was carved out of base-b in #28623 to
# *reduce* CPU CI scope), so base-a-test-cpu is what actually executes.
register_cpu_ci(est_time=300, suite="base-a-test-cpu")
# repo root: test/registered/rust/<this file>
RUST_WORKSPACE = Path(__file__).resolve().parents[3] / "rust"
# Not `est_time`: that is a scheduling hint for partition balancing (a rough
# average), this is a hard ceiling for the worst case. The 136 tests run in ~1s;
# what varies is the build. Cache-warm the workspace crates recompile in ~15s,
# but a Swatinem/rust-cache miss rebuilds all ~370 dependencies -- measured at
# 48s on 4 fast cores, so several minutes on a hosted runner.
#
# Capped below the 600s `timeout-minutes` on the suite's "Run test" step so a
# hang fails here, with output, instead of being killed as an opaque job
# timeout. The harness `--timeout-per-file` (1200s) is looser still.
BUILD_AND_RUN_TIMEOUT_S = 300
class TestCargoWorkspace(CustomTestCase):
def test_cargo_test_workspace(self):
# Not skipUnless: cargo is a hard dependency of the editable install
# (setuptools-rust builds sglang-grpc), so a missing toolchain is a
# broken environment, and a silently-skipped CI test is worthless.
self.assertIsNotNone(
shutil.which("cargo"),
"cargo not found on PATH; install a Rust toolchain "
"(scripts/ci/utils/install_rust_protoc.sh)",
)
self.assertTrue(
(RUST_WORKSPACE / "Cargo.toml").is_file(),
f"rust workspace manifest not found at {RUST_WORKSPACE}",
)
proc = subprocess.run(
["cargo", "test", "--workspace"],
cwd=RUST_WORKSPACE,
capture_output=True,
text=True,
timeout=BUILD_AND_RUN_TIMEOUT_S,
)
# Print unconditionally so a green run still shows which tests ran.
print(proc.stdout)
self.assertEqual(
proc.returncode,
0,
f"`cargo test --workspace` failed in {RUST_WORKSPACE}\n"
f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}",
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/layers/test_logprob_chunk_stitching.py
+36 -33
View File
@@ -6,7 +6,6 @@ were skipped or double-emitted, drifting the per-request entry counts that
the scheduler asserts on.
"""
import itertools
import unittest
from types import SimpleNamespace
@@ -14,6 +13,7 @@ import torch
from sglang.srt.layers.logprob_processor import InputLogprobProcessor
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.logprob_test_utils import coverage_cases
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=30, suite="base-a-test-cpu")
@@ -23,6 +23,11 @@ VOCAB = 11
TOPK_CYCLE = [2, 0, 3]
# [] is a valid probe set distinct from None (opt-out).
TOKEN_IDS_CYCLE = [[0, 3], None, [1], []]
# start == extend_len is the zero-logprob-row shape. Order determines the cyclic
# width-3/4 heterogeneous coverage cases.
SEQ_SPEC_MENU = ((1, 1), (2, 2), (3, 0), (4, 1), (5, 5), (2, 0), (6, 2))
# 7 singletons + 7*7 ordered pairs + 4*7 wide cases each at width 3 and 4.
EXPECTED_CASES = 112
def _build_batch(seq_specs, with_token_ids):
@@ -92,40 +97,38 @@ class TestLogprobChunkStitching(CustomTestCase):
def _sweep(self, with_token_ids):
torch.manual_seed(0)
proc = InputLogprobProcessor()
# (extend_len, start); start == extend_len is the degenerate
# zero-logprob-row shape.
menu = [(1, 1), (2, 2), (3, 0), (4, 1), (5, 5), (2, 0), (6, 2)]
combos = list(coverage_cases(SEQ_SPEC_MENU, max_seqs=4))
self.assertEqual(len(combos), EXPECTED_CASES)
tried = 0
for n_seqs in (1, 2, 3, 4):
for combo in itertools.product(menu, repeat=n_seqs):
batch = _build_batch(list(combo), with_token_ids)
# Same unit as the production gate: grid rows, not logprob rows.
total_rows = batch[0].shape[0]
for chunk_size in (1, 2, 3, 5):
if total_rows <= chunk_size:
continue
tried += 1
ref, ref_sampled = _run(proc, batch, False, 10**9)
got, got_sampled = _run(proc, batch, True, chunk_size)
label = f"specs={list(combo)} chunk={chunk_size}"
self.assertEqual(ref.top_logprobs_val, got.top_logprobs_val, label)
self.assertEqual(ref.top_logprobs_idx, got.top_logprobs_idx, label)
if with_token_ids:
self.assertEqual(
ref.token_ids_logprobs_val,
got.token_ids_logprobs_val,
label,
)
self.assertEqual(
ref.token_ids_logprobs_idx,
got.token_ids_logprobs_idx,
label,
)
torch.testing.assert_close(
ref.token_logprobs, got.token_logprobs, msg=label
for combo in combos:
batch = _build_batch(list(combo), with_token_ids)
# Same unit as the production gate: grid rows, not logprob rows.
total_rows = batch[0].shape[0]
for chunk_size in (1, 2, 3, 5):
if total_rows <= chunk_size:
continue
tried += 1
ref, ref_sampled = _run(proc, batch, False, 10**9)
got, got_sampled = _run(proc, batch, True, chunk_size)
label = f"specs={list(combo)} chunk={chunk_size}"
self.assertEqual(ref.top_logprobs_val, got.top_logprobs_val, label)
self.assertEqual(ref.top_logprobs_idx, got.top_logprobs_idx, label)
if with_token_ids:
self.assertEqual(
ref.token_ids_logprobs_val,
got.token_ids_logprobs_val,
label,
)
torch.testing.assert_close(ref_sampled, got_sampled, msg=label)
self.assertGreater(tried, 1000)
self.assertEqual(
ref.token_ids_logprobs_idx,
got.token_ids_logprobs_idx,
label,
)
torch.testing.assert_close(
ref.token_logprobs, got.token_logprobs, msg=label
)
torch.testing.assert_close(ref_sampled, got_sampled, msg=label)
self.assertGreater(tried, 100)
def test_top_logprobs_stitching(self):
self._sweep(with_token_ids=False)
test/registered/unit/layers/test_logprob_fast_input.py
+55 -56
View File
@@ -7,7 +7,6 @@ agree with the reference path to floating-point tolerance, with identical
top-k indices, across chunk splits and heterogeneous per-sequence params.
"""
import itertools
import unittest
from types import SimpleNamespace
@@ -18,6 +17,7 @@ from sglang.srt.layers.logprob_processor import (
compute_row_log_normalizer,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.logprob_test_utils import coverage_cases
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=30, suite="base-a-test-cpu")
@@ -27,6 +27,11 @@ VOCAB = 11
TOPK_CYCLE = [2, 0, 3]
# [] is a valid probe set distinct from None (opt-out).
TOKEN_IDS_CYCLE = [[0, 3], None, [1], []]
# start == extend_len is the zero-logprob-row shape. Order determines the cyclic
# width-3 heterogeneous coverage cases.
SEQ_SPEC_MENU = ((1, 1), (3, 0), (4, 1), (5, 5), (6, 2))
# 5 singletons + 5*5 ordered pairs + 4*5 wide cases at width 3.
EXPECTED_CASES = 50
def _build_batch(seq_specs, dtype, vocab=VOCAB):
@@ -121,48 +126,46 @@ class TestFastInputLogprobs(CustomTestCase):
def _sweep(self, dtype, rtol, atol):
torch.manual_seed(0)
proc = InputLogprobProcessor()
# (extend_len, start); start == extend_len is the degenerate
# zero-logprob-row shape.
menu = [(1, 1), (3, 0), (4, 1), (5, 5), (6, 2)]
combos = list(coverage_cases(SEQ_SPEC_MENU, max_seqs=3))
self.assertEqual(len(combos), EXPECTED_CASES)
tried = 0
for n_seqs in (1, 2, 3):
for combo in itertools.product(menu, repeat=n_seqs):
batch = _build_batch(list(combo), dtype)
for chunk_size in (None, 1, 2, 3, 5):
tried += 1
ref, ref_sampled = _run(proc, batch, False, chunk_size)
got, got_sampled = _run(proc, batch, True, chunk_size)
label = f"specs={list(combo)} chunk={chunk_size} dtype={dtype}"
# Top-k order comes from the same values shifted by a
# per-row constant, so indices must match exactly.
self.assertEqual(ref.top_logprobs_idx, got.top_logprobs_idx, label)
self.assertEqual(
ref.token_ids_logprobs_idx, got.token_ids_logprobs_idx, label
)
_assert_nested_close(
self,
ref.top_logprobs_val,
got.top_logprobs_val,
label,
rtol,
atol,
)
_assert_nested_close(
self,
ref.token_ids_logprobs_val,
got.token_ids_logprobs_val,
label,
rtol,
atol,
)
torch.testing.assert_close(
ref.token_logprobs.float(),
got.token_logprobs.float(),
rtol=rtol,
atol=atol,
msg=label,
)
torch.testing.assert_close(ref_sampled, got_sampled, msg=label)
for combo in combos:
batch = _build_batch(list(combo), dtype)
for chunk_size in (None, 1, 2, 3, 5):
tried += 1
ref, ref_sampled = _run(proc, batch, False, chunk_size)
got, got_sampled = _run(proc, batch, True, chunk_size)
label = f"specs={list(combo)} chunk={chunk_size} dtype={dtype}"
# Top-k order comes from the same values shifted by a
# per-row constant, so indices must match exactly.
self.assertEqual(ref.top_logprobs_idx, got.top_logprobs_idx, label)
self.assertEqual(
ref.token_ids_logprobs_idx, got.token_ids_logprobs_idx, label
)
_assert_nested_close(
self,
ref.top_logprobs_val,
got.top_logprobs_val,
label,
rtol,
atol,
)
_assert_nested_close(
self,
ref.token_ids_logprobs_val,
got.token_ids_logprobs_val,
label,
rtol,
atol,
)
torch.testing.assert_close(
ref.token_logprobs.float(),
got.token_logprobs.float(),
rtol=rtol,
atol=atol,
msg=label,
)
torch.testing.assert_close(ref_sampled, got_sampled, msg=label)
self.assertGreater(tried, 100)
def test_fast_matches_reference_fp32(self):
@@ -176,20 +179,16 @@ class TestFastInputLogprobs(CustomTestCase):
# sits much closer to the truth than bf16 resolution.
torch.manual_seed(0)
proc = InputLogprobProcessor()
menu = [(1, 1), (3, 0), (4, 1), (5, 5), (6, 2)]
for n_seqs in (1, 2, 3):
for combo in itertools.product(menu, repeat=n_seqs):
batch = _build_batch(list(combo), torch.bfloat16)
pruned_states, _, input_logprob_indices, _, metadata = batch
truth = torch.log_softmax(pruned_states.double(), dim=-1)[
input_logprob_indices
]
for chunk_size in (None, 2, 5):
got, _ = _run(proc, batch, True, chunk_size)
label = f"specs={list(combo)} chunk={chunk_size}"
self._assert_rows_match_truth(
got, truth, metadata, label, atol=1e-4
)
for combo in coverage_cases(SEQ_SPEC_MENU, max_seqs=3):
batch = _build_batch(list(combo), torch.bfloat16)
pruned_states, _, input_logprob_indices, _, metadata = batch
truth = torch.log_softmax(pruned_states.double(), dim=-1)[
input_logprob_indices
]
for chunk_size in (None, 2, 5):
got, _ = _run(proc, batch, True, chunk_size)
label = f"specs={list(combo)} chunk={chunk_size}"
self._assert_rows_match_truth(got, truth, metadata, label, atol=1e-4)
def _assert_rows_match_truth(self, got, truth, metadata, label, atol):
pt = 0
test/registered/unit/spec/test_decode_bookkeeping_ownership.py
-239
View File
@@ -1,239 +0,0 @@
"""Ownership contract for per-request bookkeeping clocks.
Per-request accounting state (`decode_batch_idx` / `extend_batch_idx` iter
clocks, `kv_committed_len` / `kv_allocated_len` KV watermarks,
`spec_verify_ct`, and the `maybe_evict_swa()` call) must only be advanced by
the reviewed owner sites in _OWNER_SITES; spec-v2 draft workers must not
repeat any of them (the scheduler-driven free function / resolve path already
does).
A clock that runs fast fires SWA eviction in the overlap race window and
releases the SWA prefix lock early; neither shows up in e2e CI or the idle
leak checker, hence this AST-level guard.
"""
import ast
import unittest
import warnings
from collections import Counter
from pathlib import Path
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=8, suite="base-a-test-cpu")
_REPO_ROOT = Path(__file__).resolve().parents[4]
_SRT_DIR = _REPO_ROOT / "python" / "sglang" / "srt"
_SPECULATIVE_DIR = _SRT_DIR / "speculative"
assert _SRT_DIR.is_dir(), f"srt dir not found: {_SRT_DIR}"
_TRACKED_ATTRS = (
"decode_batch_idx",
"extend_batch_idx",
"kv_committed_len",
"kv_allocated_len",
"spec_verify_ct",
)
_EVICT_METHOD = "maybe_evict_swa"
# {(path relative to srt/, scope, kind): mutation count}. Kind is the mutated
# attribute (`= 0` resets exempt) or "evict" for a `maybe_evict_swa()` call.
# Any added/removed/recounted site fails until reviewed here.
_SB = "managers/schedule_batch.py"
_EAGLE_DECODE = ("speculative/eagle_utils.py", "eagle_prepare_for_decode")
_RESOLVE = (
"managers/scheduler_components/batch_result_processor.py",
"SchedulerBatchResultProcessor._resolve_spec_v2_tokens",
)
_SS = "session/streaming_session.py"
_OWNER_SITES = {
# non-spec scheduler
(_SB, "ScheduleBatch.prepare_for_decode", "decode_batch_idx"): 1,
(_SB, "ScheduleBatch.prepare_for_decode", "kv_committed_len"): 1,
(_SB, "ScheduleBatch.prepare_for_extend", "extend_batch_idx"): 1,
(_SB, "ScheduleBatch.prepare_for_extend", "kv_committed_len"): 1,
# kv_allocated_len is settled inside the owned-kv alloc functions (op28).
("mem_cache/allocation.py", "alloc_for_extend", "evict"): 1,
("mem_cache/allocation.py", "alloc_for_extend", "kv_allocated_len"): 1,
("mem_cache/allocation.py", "alloc_for_decode", "evict"): 1,
("mem_cache/allocation.py", "alloc_for_decode", "kv_allocated_len"): 1,
# spec v2: no pre-claim; resolve commits the full accepted run uniformly.
# kv_allocated_len for spec v2 draft decode (eagle + dflash) is settled
# inside the owned-kv alloc_for_spec_decode function (op42).
(*_EAGLE_DECODE, "decode_batch_idx"): 1,
(*_EAGLE_DECODE, "evict"): 1,
(
"mem_cache/allocation.py",
"alloc_for_spec_decode",
"kv_allocated_len",
): 1,
(*_RESOLVE, "kv_committed_len"): 1,
(*_RESOLVE, "spec_verify_ct"): 1,
# disaggregation decode prealloc: kv_allocated_len is settled inside the
# owned-kv alloc_for_decode_prealloc(_hisparse) functions (op42).
(
"disaggregation/decode.py",
"DecodePreallocQueue._pre_alloc",
"kv_committed_len",
): 1,
("disaggregation/decode.py", "alloc_for_decode_prealloc", "kv_allocated_len"): 1,
(
"disaggregation/decode.py",
"alloc_for_decode_prealloc_hisparse",
"kv_allocated_len",
): 1,
# streaming session slot save/restore and tail trimming
(_SS, "SessionSlot.save_from_req", "kv_committed_len"): 1,
(_SS, "SessionSlot.restore_to_req", "kv_committed_len"): 1,
(_SS, "StreamingSession._free_tail", "kv_committed_len"): 2,
(_SS, "StreamingSession._free_tail", "kv_allocated_len"): 2,
(_SS, "StreamingSession._trim_overshoot", "kv_committed_len"): 1,
(_SS, "StreamingSession._trim_overshoot", "kv_allocated_len"): 1,
(_SS, "StreamingSession.try_cache_finished_req", "kv_allocated_len"): 1,
# Inherit the authoritative finished length (not the lagging req clock).
(_SS, "StreamingSession.try_cache_finished_req", "kv_committed_len"): 1,
}
def _iter_scoped_nodes(tree):
"""Yield (node, dotted Class.method scope) for every node."""
scope_of = {}
def visit(node, scope):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
scope = f"{scope}.{node.name}" if scope else node.name
scope_of[node] = scope
for child in ast.iter_child_nodes(node):
visit(child, scope)
visit(tree, "")
return scope_of.items()
def _is_zero_reset(node):
return isinstance(node, ast.Assign) and (
isinstance(node.value, ast.Constant) and node.value.value == 0
)
def _scan_tree(tree):
"""Count bookkeeping sites in an AST as Counter[(scope, kind)]."""
sites = Counter()
for node, scope in _iter_scoped_nodes(tree):
if isinstance(node, (ast.AugAssign, ast.Assign)):
targets = node.targets if isinstance(node, ast.Assign) else [node.target]
for target in targets:
if (
isinstance(target, ast.Attribute)
and target.attr in _TRACKED_ATTRS
and not _is_zero_reset(node)
):
sites[(scope, target.attr)] += 1
if (
isinstance(node, ast.Call)
and isinstance(node.func, ast.Attribute)
and node.func.attr == _EVICT_METHOD
):
sites[(scope, "evict")] += 1
return sites
def _parse(path: Path):
# utf-8-sig: some srt files carry a BOM that breaks plain-utf-8 ast.parse.
with warnings.catch_warnings():
warnings.simplefilter("ignore", SyntaxWarning)
return ast.parse(path.read_text(encoding="utf-8-sig"))
def _scan_srt():
"""Count all bookkeeping sites in srt/ as Counter[(rel, scope, kind)]."""
found = Counter()
for path in sorted(_SRT_DIR.rglob("*.py")):
rel = path.relative_to(_SRT_DIR).as_posix()
for (scope, kind), count in _scan_tree(_parse(path)).items():
found[(rel, scope, kind)] += count
return found
def _draft_worker_classes():
"""All transitive EagleDraftWorkerBase subclasses under speculative/."""
by_name = {}
for path in sorted(_SPECULATIVE_DIR.glob("*.py")):
rel = path.relative_to(_SRT_DIR).as_posix()
for node in ast.walk(_parse(path)):
if isinstance(node, ast.ClassDef):
bases = {
b.id if isinstance(b, ast.Name) else getattr(b, "attr", None)
for b in node.bases
}
by_name[node.name] = (rel, node, bases)
workers = {"EagleDraftWorkerBase"}
changed = True
while changed:
changed = False
for name, (_, _, bases) in by_name.items():
if name not in workers and bases & workers:
workers.add(name)
changed = True
return [
(rel, node)
for name, (rel, node, _) in sorted(by_name.items())
if name in workers and name != "EagleDraftWorkerBase"
]
def _scan_class_subtree(class_node):
"""Scan one ClassDef subtree; returns (method_scope, kind) sites."""
module = ast.Module(body=[class_node], type_ignores=[])
sites = set()
for scope, kind in _scan_tree(module):
# Strip the leading class name; keep method-level scope.
sites.add((scope.split(".", 1)[1] if "." in scope else scope, kind))
return sites
class TestDecodeBookkeepingOwnership(CustomTestCase):
def test_bookkeeping_sites_match_owner_allowlist(self):
found = _scan_srt()
allow = Counter(_OWNER_SITES)
unexpected = found - allow
missing = allow - found
msg = []
if unexpected:
msg.append(
"New bookkeeping mutation(s) beyond the recorded counts:\n "
+ "\n ".join(f"{site} x{n}" for site, n in sorted(unexpected.items()))
+ "\nThese are owned by the sites in _OWNER_SITES -- do not "
"repeat them; a genuinely new owner must be recorded there."
)
if missing:
msg.append(
"Recorded site(s) no longer exist (update _OWNER_SITES):\n "
+ "\n ".join(f"{site} x{n}" for site, n in sorted(missing.items()))
)
self.assertFalse(msg, "\n\n".join(msg))
def test_spec_v2_draft_workers_do_no_scheduler_bookkeeping(self):
classes = _draft_worker_classes()
names = {node.name for _, node in classes}
# Discovery sanity: fail loudly instead of silently guarding nothing.
self.assertIn("EagleDraftWorker", names)
self.assertIn("FrozenKVMTPDraftWorker", names)
violations = []
for rel, node in classes:
for scope, kind in _scan_class_subtree(node):
violations.append((rel, f"{node.name}.{scope}", kind))
self.assertFalse(
violations,
"Spec-v2 draft worker(s) repeat scheduler-owned bookkeeping:\n "
+ "\n ".join(map(str, sorted(violations)))
+ "\nUnder spec v2 the iter-clock ticks, `maybe_evict_swa`, and "
"KV watermark settlement are owned by the scheduler-driven "
"free function / resolve path. Remove these from the worker.",
)
if __name__ == "__main__":
unittest.main(verbosity=3)
test/registered/unit/test_global_config_read_ratchet.py
-520
View File
@@ -1,520 +0,0 @@
"""Ratchet guard: process-global config reads may only decrease.
``get_server_args()`` returns the published ``ServerArgs`` — one process's
startup record. Config decisions read the namespace accessors instead
(``get_exec()`` / ``get_memory()`` / …), which carry the resolved value
including post-publish overrides, and per-runner values come from the runner
that owns them.
Business code no longer reads the published record for a config value at all:
the baselines are zero for both shapes, over the whole package minus the two
modules that own the slot.
Where the remaining reads live (``runtime_context.py``, exempt by module):
- **Derived members.** ``@property`` / method members of ``ServerArgs``
(``mamba_cache_chunk_size``, ``max_speculative_num_draft_tokens``,
``use_mla_backend()``, ``get_attention_backends()``, ``get_model_config()``,
``cutedsl_moe_max_num_tokens()``) are computed from several fields plus the HF
config, so they are not namespace leaves and ``ServerArgs`` is their only
home. ``runtime_context`` exposes each one as a named accessor
(``mamba_cache_chunk_size()`` …) and is the only module that reads the slot
for them.
- **Config-intent reads of live-shadowed sizes.** ``get_parallel()`` shadows
``tp/pp/dcp/attn_cp/moe_dp_size`` with the live topology, and a few call sites
need what was *configured*: the ``configured_*_size()`` accessors. Their
reasons, per call site:
- ``dsa_indexer.pp_size`` gates ``pp_size > 1 and not get_pp_group()...``, and
the short circuit is the point: with PP off the group is never touched, which
is what lets the ``Indexer`` be constructed before distributed init. The live
property would demand the group either way.
- ``dp_attention.attn_cp_size`` / ``moe_dp_size``: the configuration the
predicate detects (``attn_cp_size > moe_dp_size``) is the one where
``initialize_model_parallel`` aliases ``_MOE_DP`` to ``_ATTN_CP``, so the live
sizes are equal there and a live comparison is always false.
- ``model_loader/loader.py`` reports both: the same dict carries the live
``moe_dp_size`` under ``"dp"``, so this entry is the configured intent.
What the ratchet sees, syntactically: ``get_server_args().field``,
``sa = get_server_args()`` followed by ``sa.field`` (function-local, module-level,
or parked on an instance attribute -- ``self._sa = get_server_args()`` read from
another method of the same class), function-local copies of an alias to a
fixpoint (``cfg = sa`` then ``cfg.field``), and the dynamic form of each --
``getattr(<either>, "field")`` -- since a string-named read reaches the same
slot. What it cannot see is a name computed at runtime (``getattr(sa, name)``)
or indirection deeper than a local name copy (through a container, an
attribute of another object, a cross-scope copy); the census tool in the
context repo is what audits those.
A whole-object pass (``def f(server_args)``) is not a global read and is not
counted -- there the caller decided which instance to hand over. An optional
parameter that falls back to the global (``f(server_args=None)``) hides one,
so those fallbacks were removed; the ratchet cannot see them and the census
tool in the context repo is what audits that shape.
"""
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
import ast
import unittest
from pathlib import Path
import sglang
from sglang.test.test_utils import CustomTestCase
# srt is the migrated surface; the rest of the package has no reads today and is
# scanned so a new one cannot appear there unnoticed.
_PACKAGE_ROOT = Path(next(iter(sglang.__path__)))
# The modules that own the slot: runtime_context publishes it and exposes the
# named accessors for the derived members, server_args/arg_groups ARE the
# resolution pipeline.
_SLOT_OWNERS = ("srt/runtime_context.py", "srt/server_args.py", "srt/arg_groups/")
# Every call site of a ``configured_*_size()`` accessor, with the reason the
# live topology cannot answer there. The test below asserts this map is exactly
# the set of call sites, so the reasons cannot drift away from the code.
_CONFIGURED_SIZE_CALL_SITES = {
("srt/layers/attention/dsa/dsa_indexer.py", "configured_pp_size"): (
"gates `pp_size > 1 and not get_pp_group()...`; the short circuit is the "
"point, since with PP off the group is never touched, which is what lets "
"the Indexer be constructed before distributed init"
),
("srt/mem_cache/kv_cache_configurator.py", "configured_pp_size"): (
"decides whether the token capacity needs a cross-PP all-reduce at all; "
"asking the configured size keeps that decision independent of whether a "
"PP group is installed in this process"
),
("srt/layers/dp_attention.py", "configured_attn_cp_size"): (
"compared against the configured moe_dp_size below"
),
("srt/layers/dp_attention.py", "configured_moe_dp_size"): (
"the configuration this predicate detects (attn_cp_size > moe_dp_size) is "
"the one where initialize_model_parallel aliases _MOE_DP to _ATTN_CP, so "
"the live sizes are equal there and a live comparison is always false"
),
("srt/model_loader/loader.py", "configured_moe_dp_size"): (
"the same dict already carries the live moe_dp_size under 'dp'; this entry "
"is the configured intent"
),
("srt/models/kimi_k25.py", "configured_tp_size"): (
"the IPC refcount must match the configured TP consumer count captured "
"when the tokenizer creates MmItemMemoryPool; a live attention subgroup "
"size could strand leases in the bounded pool"
),
("srt/models/kimi_k3.py", "configured_tp_size"): (
"same as kimi_k25: the IPC refcount must agree with the recycler's waiter"
),
}
# A dynamic read whose name is set nowhere in the tree, so the predicate it
# feeds is inert (the ``getattr`` default decides it). Converting it would mean
# choosing what it should have named, which is the CP path's call, not this
# sweep's -- so it is listed here rather than silently counted or "fixed".
_INERT_DYNAMIC_READS = frozenset({("srt/layers/cp/base.py", "_is_dsa_model_arch")})
_DIRECT_BASELINE = 0
_ALIAS_BASELINE = 0
def _is_global_call(node) -> bool:
"""``get_server_args()`` however it is spelled: bare, or module-qualified
(``ctx.get_server_args()``), which an ast.Name check alone would miss."""
if not isinstance(node, ast.Call):
return False
func = node.func
if isinstance(func, ast.Name):
return func.id == "get_server_args"
return isinstance(func, ast.Attribute) and func.attr == "get_server_args"
def _collect(rel: str, tree: ast.AST, inert: frozenset = frozenset()):
"""The (direct, alias) field reads in one module.
``inert`` names the fields listed in ``_INERT_DYNAMIC_READS`` for this file;
they are dropped here, at the point the read is recognized, so the filter
matches on the field name rather than on the rendered message.
"""
direct, alias = [], []
def counted(attr: str) -> bool:
return attr not in inert
def _getattr_name(node):
"""``getattr(<record>, "field")`` names a field just as ``.field`` does;
matching only ast.Attribute would let a dynamic read walk past."""
if not (
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "getattr"
and len(node.args) >= 2
and isinstance(node.args[1], ast.Constant)
and isinstance(node.args[1].value, str)
):
return None
return node.args[1].value
for node in ast.walk(tree):
if (
isinstance(node, ast.Attribute)
and _is_global_call(node.value)
and counted(node.attr)
):
direct.append(f"{rel}:{node.lineno}: get_server_args().{node.attr}")
name = _getattr_name(node)
if name is not None and _is_global_call(node.args[0]) and counted(name):
direct.append(f"{rel}:{node.lineno}: getattr(get_server_args(), {name!r})")
if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
continue
params = {a.arg for a in list(node.args.args) + list(node.args.kwonlyargs)}
bound = {}
for inner in ast.walk(node):
# ``sa = get_server_args()`` and its annotated form
# ``sa: ServerArgs = get_server_args()``.
if isinstance(inner, (ast.Assign, ast.AnnAssign)) and _is_global_call(
getattr(inner, "value", None)
):
targets = (
inner.targets if isinstance(inner, ast.Assign) else [inner.target]
)
for target in targets:
if not isinstance(target, ast.Name):
continue
# A parameter reassigned from the global is the
# optional-injection shape (``f(server_args=None)`` then
# ``server_args = get_server_args()``): the reads that
# follow are global reads wearing a parameter's name, so
# they count from the bind on.
bound.setdefault(target.id, inner.lineno)
if not bound:
continue
# A copy of an alias reaches the same record (``cfg = sa`` after
# ``sa = get_server_args()``), so follow Name-to-Name assignments to a
# fixpoint. Deeper indirection (through containers, attributes of
# other objects, cross-scope copies) stays census-tool territory.
changed = True
while changed:
changed = False
for inner in ast.walk(node):
if not isinstance(inner, (ast.Assign, ast.AnnAssign)):
continue
value = getattr(inner, "value", None)
if not (isinstance(value, ast.Name) and value.id in bound):
continue
targets = (
inner.targets if isinstance(inner, ast.Assign) else [inner.target]
)
for target in targets:
if isinstance(target, ast.Name) and target.id not in bound:
bound[target.id] = inner.lineno
changed = True
for inner in ast.walk(node):
if (
isinstance(inner, ast.Attribute)
and isinstance(inner.value, ast.Name)
and inner.value.id in bound
and inner.lineno >= bound[inner.value.id]
and counted(inner.attr)
):
alias.append(
f"{rel}:{inner.lineno}: {inner.value.id}.{inner.attr} "
f"(bound from get_server_args() at line {bound[inner.value.id]})"
)
name = _getattr_name(inner)
if (
name is not None
and isinstance(inner.args[0], ast.Name)
and inner.args[0].id in bound
and inner.lineno >= bound[inner.args[0].id]
and counted(name)
):
alias.append(
f"{rel}:{inner.lineno}: getattr({inner.args[0].id}, {name!r}) "
f"(bound from get_server_args() at line {bound[inner.args[0].id]})"
)
# A module-level alias is visible to every function in the file, so it needs
# its own pass -- the per-function scan above deliberately does not reach
# across scopes.
module_bound = {}
module_stack = list(tree.body)
while module_stack:
stmt = module_stack.pop()
# A module-level bind can sit inside an `if` / `try` / `with`, so the
# walk descends into those bodies -- but not into a nested function or
# class, whose binds are that scope's own.
if isinstance(
stmt, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)
):
continue
module_stack.extend(ast.iter_child_nodes(stmt))
if isinstance(stmt, (ast.Assign, ast.AnnAssign)) and _is_global_call(
getattr(stmt, "value", None)
):
targets = stmt.targets if isinstance(stmt, ast.Assign) else [stmt.target]
for target in targets:
if isinstance(target, ast.Name):
module_bound.setdefault(target.id, stmt.lineno)
if module_bound:
# Shadowing is per lexical scope: a function with its own `sa` hides the
# module alias *inside that function only*. Aggregating the names
# file-wide would suppress every read in the module, including the
# top-level ones and the ones in functions that do resolve to the alias.
parents = {}
scope_binds = {}
stack = [tree]
while stack:
node = stack.pop()
enclosing = parents.get(id(node))
for child in ast.iter_child_nodes(node):
parents[id(child)] = (
node
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
else enclosing
)
stack.append(child)
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
names = {
a.arg for a in list(node.args.args) + list(node.args.kwonlyargs)
}
# Only this scope's own stores: a nested function's local `sa`
# shadows the alias inside *that* function, not in its parent.
pending = list(node.body)
while pending:
inner = pending.pop()
if isinstance(
inner,
(
ast.FunctionDef,
ast.AsyncFunctionDef,
ast.Lambda,
ast.ClassDef,
),
):
continue
if isinstance(inner, ast.Name) and isinstance(inner.ctx, ast.Store):
names.add(inner.id)
pending.extend(ast.iter_child_nodes(inner))
scope_binds[id(node)] = names
def _shadowed(node, name):
scope = parents.get(id(node))
while scope is not None:
if name in scope_binds.get(id(scope), ()):
return True
scope = parents.get(id(scope))
return False
for node in ast.walk(tree):
base = attr = None
if (
isinstance(node, ast.Attribute)
and isinstance(node.value, ast.Name)
and node.value.id in module_bound
):
base, attr = node.value.id, node.attr
shown = f"{base}.{attr}"
else:
attr_name = _getattr_name(node)
if (
attr_name is not None
and isinstance(node.args[0], ast.Name)
and node.args[0].id in module_bound
):
base, attr = node.args[0].id, attr_name
shown = f"getattr({base}, {attr!r})"
if base and not _shadowed(node, base) and counted(attr):
alias.append(
f"{rel}:{node.lineno}: {shown} "
f"(module-level bind from get_server_args() at line "
f"{module_bound[base]})"
)
# An alias parked on an instance attribute (``self._sa = get_server_args()``
# in one method, ``self._sa.field`` in another) reaches the same slot and
# crosses function scopes, so it is collected per class rather than per
# function.
for node in ast.walk(tree):
if not isinstance(node, ast.ClassDef):
continue
attr_bound = {}
for inner in ast.walk(node):
if isinstance(inner, (ast.Assign, ast.AnnAssign)) and _is_global_call(
getattr(inner, "value", None)
):
targets = (
inner.targets if isinstance(inner, ast.Assign) else [inner.target]
)
for target in targets:
if (
isinstance(target, ast.Attribute)
and isinstance(target.value, ast.Name)
and target.value.id in ("self", "cls")
):
attr_bound.setdefault(
(target.value.id, target.attr), inner.lineno
)
if not attr_bound:
continue
def _bound_attr(value):
"""``self._sa`` when that attribute was bound from the global."""
if (
isinstance(value, ast.Attribute)
and isinstance(value.value, ast.Name)
and (value.value.id, value.attr) in attr_bound
):
return (value.value.id, value.attr)
return None
for inner in ast.walk(node):
key = shown = None
if isinstance(inner, ast.Attribute):
key = _bound_attr(inner.value)
if key is not None and counted(inner.attr):
shown = f"{key[0]}.{key[1]}.{inner.attr}"
else:
name = _getattr_name(inner)
if name is not None:
key = _bound_attr(inner.args[0])
if key is not None and counted(name):
shown = f"getattr({key[0]}.{key[1]}, {name!r})"
if shown is not None:
alias.append(
f"{rel}:{inner.lineno}: {shown} "
f"(attribute bind from get_server_args() at line "
f"{attr_bound[key]})"
)
return direct, alias
def _field_reads():
direct, alias = [], []
for path in sorted(_PACKAGE_ROOT.rglob("*.py")):
rel = path.relative_to(_PACKAGE_ROOT).as_posix()
if rel.startswith(_SLOT_OWNERS):
continue
try:
tree = ast.parse(path.read_text())
except SyntaxError:
continue
inert = frozenset(name for path_, name in _INERT_DYNAMIC_READS if path_ == rel)
module_direct, module_alias = _collect(rel, tree, inert)
direct += module_direct
alias += module_alias
return direct, alias
class TestGlobalConfigReadRatchet(CustomTestCase):
def _check(self, kind, reads, baseline):
if len(reads) > baseline:
self.fail(
f"{kind} process-global config field reads grew: {len(reads)} > "
f"baseline {baseline}. Read the namespace accessor for the "
"field's namespace, or the owning runner for a per-runner "
"field:\n" + "\n".join(reads)
)
if len(reads) < baseline:
self.fail(
f"{kind} process-global config field reads shrank: {len(reads)} < "
f"baseline {baseline}. Lower the baseline in this file to lock "
"in the progress."
)
def test_global_field_reads_match_the_baseline(self):
direct, alias = _field_reads()
self._check("direct", direct, _DIRECT_BASELINE)
self._check("alias-form", alias, _ALIAS_BASELINE)
class TestConfiguredSizeCallSites(CustomTestCase):
"""The configured-vs-live exceptions are enumerated, with reasons.
``configured_*_size()`` answers what the user asked for where
``get_parallel()`` would answer what the process ended up with. Each such
exception is listed above with why the live property cannot serve it, and
this case fails if the code and that list disagree.
The unit is **(file, accessor)**, not the individual call: a second
`configured_pp_size()` in a file already registered for it collapses into
the same entry, so the reason has to cover the file's use of that accessor
rather than one line. A new file, or a new accessor in a listed file, is
what this catches -- in either call form (bare or module-qualified).
"""
def test_the_call_sites_match_the_documented_set(self):
found = set()
for path in sorted(_PACKAGE_ROOT.rglob("*.py")):
rel = path.relative_to(_PACKAGE_ROOT).as_posix()
if rel.startswith(_SLOT_OWNERS):
continue
try:
tree = ast.parse(path.read_text())
except SyntaxError:
continue
for node in ast.walk(tree):
if not isinstance(node, ast.Call):
continue
func = node.func
name = (
func.id
if isinstance(func, ast.Name)
else (func.attr if isinstance(func, ast.Attribute) else None)
)
if name and name.startswith("configured_") and name.endswith("_size"):
found.add((rel, name))
documented = set(_CONFIGURED_SIZE_CALL_SITES)
self.assertEqual(
documented,
found,
"configured-size call sites drifted from their documented reasons.\n"
f" undocumented: {sorted(found - documented)}\n"
f" stale entries: {sorted(documented - found)}",
)
class TestNoRenamedAccessorImports(CustomTestCase):
"""The scanners above match ``get_server_args`` and ``configured_*_size``
by their literal names, so an ``import ... as`` rename would walk a read
straight past both the zero baseline and the call-site registry. Renaming
these accessors buys nothing (the names are already short and unambiguous),
so it is banned outright — which is exactly what makes literal-name
matching sound."""
def test_the_scanned_accessors_are_never_import_renamed(self):
offenders = []
for path in sorted(_PACKAGE_ROOT.rglob("*.py")):
rel = path.relative_to(_PACKAGE_ROOT).as_posix()
try:
tree = ast.parse(path.read_text())
except SyntaxError:
continue
for node in ast.walk(tree):
if not isinstance(node, (ast.ImportFrom, ast.Import)):
continue
for imported in node.names:
if imported.asname is None or imported.asname == imported.name:
continue
base = imported.name.rsplit(".", 1)[-1]
if base == "get_server_args" or (
base.startswith("configured_") and base.endswith("_size")
):
offenders.append(
f"{rel}:{node.lineno}: {imported.name} as "
f"{imported.asname}"
)
self.assertFalse(
offenders,
"get_server_args / configured_*_size imported under another name; "
"the read ratchet and the configured-size registry match these "
"accessors by their literal names, so a rename silently escapes "
"both:\n" + "\n".join(offenders),
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/test_legacy_global_ratchet.py
-65
View File
@@ -1,65 +0,0 @@
"""Ratchet guard: legacy global-accessor call-sites may only decrease.
The process-wide ``ServerArgs`` is owned by the runtime context; the legacy
``get_global_server_args`` / ``set_global_server_args_for_*`` names survive as
thin shims for the existing call-sites. New code should use the
``sglang.srt.runtime_context`` accessors (``get_server_args()`` /
``get_context().set_server_args()``), so the shim call-site counts below must
never grow. When your change removes call-sites, lower the matching baseline
to the new count.
"""
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
import re
import unittest
from pathlib import Path
import sglang.srt
from sglang.test.test_utils import CustomTestCase
_SRT_ROOT = Path(next(iter(sglang.srt.__path__)))
# Baselines counted over python/sglang/srt/**/*.py, including each function's
# own def line. Ratchet: decrease-only.
_RATCHETS = [
# Down to the shim definition itself; every call-site now goes through
# runtime_context.get_server_args().
("get_global_server_args", r"\bget_global_server_args\s*\(", 1),
(
"set_global_server_args_for_*",
r"\bset_global_server_args_for_(?:scheduler|tokenizer)\s*\(",
4,
),
]
class TestLegacyGlobalRatchet(CustomTestCase):
def test_legacy_accessor_call_sites_match_the_baselines(self):
# Exact pin, failing in BOTH directions: a grown count means new code
# bypassed the runtime_context accessors; a shrunk count means a
# removal forgot to lower the baseline, which would let later changes
# silently re-add call-sites up to the stale ceiling.
sources = [
path.read_text(encoding="utf-8", errors="replace")
for path in sorted(_SRT_ROOT.rglob("*.py"))
]
for name, pattern, baseline in _RATCHETS:
count = sum(len(re.findall(pattern, source)) for source in sources)
if count > baseline:
self.fail(
f"{name} call-sites grew: {count} > baseline {baseline}. "
"New code must use the sglang.srt.runtime_context accessors "
"(get_server_args() / get_context().set_server_args())."
)
if count < baseline:
self.fail(
f"{name} call-sites shrank: {count} < baseline {baseline}. "
"Lower the baseline in this file to lock in the progress."
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/test_module_state_ratchet.py
-68
View File
@@ -1,68 +0,0 @@
"""Ratchet guard: module-level runtime state in the flag-owning layers may
only shrink.
Runtime flags live on ``get_flags()`` groups (``moe`` / ``dp`` / ``capture``),
where they get lifecycle reset, typo-safe writes, and the transactional
test-override primitive. A new module-level global written through a
``global`` statement in these modules recreates the pattern this replaced:
state with ad-hoc lifecycle that leaks across unit-test teardowns and cannot
be overridden scoped.
The pin lists the survivors by name: the DP-attention topology values (owned
by the parallel vertical) and the TBO comm stream (a resource, owned by the
resources vertical). Migrating one of them must shrink its pin; adding a name
fails the ratchet.
"""
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
import ast
import unittest
from pathlib import Path
import sglang.srt
from sglang.test.test_utils import CustomTestCase
_SRT_ROOT = Path(next(iter(sglang.srt.__path__)))
_PINNED_GLOBALS = {
"layers/moe/utils.py": frozenset(),
"layers/dp_attention.py": frozenset(
{
# DP-attention topology (parallel vertical scope).
"_ATTN_DP_RANK",
"_ATTN_DP_SIZE",
}
),
}
class TestModuleStateRatchet(CustomTestCase):
def test_global_statements_match_the_pins(self):
for rel, pinned in _PINNED_GLOBALS.items():
tree = ast.parse((_SRT_ROOT / rel).read_text())
declared = {
name
for node in ast.walk(tree)
if isinstance(node, ast.Global)
for name in node.names
}
grown = declared - pinned
self.assertFalse(
grown,
f"{rel} declares new module-level runtime state {sorted(grown)}; "
"put runtime flags on a get_flags() group instead "
"(see runtime_context.MoeFlags / DpFlags).",
)
shrunk = pinned - declared
self.assertFalse(
shrunk,
f"{rel} no longer declares {sorted(shrunk)}; "
"shrink the pin in this file to lock in the progress.",
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/test_no_bare_pytest_main.py
-90
View File
@@ -1,90 +0,0 @@
import ast
import pathlib
import unittest
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
_REPO_ROOT = pathlib.Path(__file__).resolve().parents[3]
_SCAN_ROOTS = [_REPO_ROOT / "python", _REPO_ROOT / "test"]
class TestNoBarePytestMain(CustomTestCase):
def test_no_bare_pytest_main_in_repo(self):
offenders = []
for root in _SCAN_ROOTS:
if not root.exists():
continue
for path in root.rglob("*.py"):
violation = _find_bare_pytest_main(path)
if violation is not None:
offenders.append(violation)
self.assertFalse(
offenders,
msg=(
"Found bare `pytest.main(...)` in __main__ blocks (must be "
"wrapped in sys.exit(...) so failing tests propagate the exit "
"code to the CI runner):\n " + "\n ".join(offenders)
),
)
def _find_bare_pytest_main(path: pathlib.Path):
"""Return `<rel_path>:<lineno>` if `path` has a bare pytest.main(...) call
inside `if __name__ == "__main__":`, else None."""
try:
source = path.read_text(encoding="utf-8")
except (OSError, UnicodeDecodeError):
return None
try:
tree = ast.parse(source, filename=str(path))
except SyntaxError:
return None
for node in ast.walk(tree):
if not isinstance(node, ast.If):
continue
if not _is_main_guard(node.test):
continue
for stmt in node.body:
if _is_bare_pytest_main_call(stmt):
rel = path.relative_to(_REPO_ROOT)
return f"{rel}:{stmt.lineno}"
return None
def _is_main_guard(test: ast.expr) -> bool:
"""Match `__name__ == "__main__"` (either side)."""
if not isinstance(test, ast.Compare) or len(test.ops) != 1:
return False
if not isinstance(test.ops[0], ast.Eq):
return False
sides = [test.left, *test.comparators]
has_name = any(isinstance(s, ast.Name) and s.id == "__name__" for s in sides)
has_main = any(isinstance(s, ast.Constant) and s.value == "__main__" for s in sides)
return has_name and has_main
def _is_bare_pytest_main_call(stmt: ast.stmt) -> bool:
"""Match `pytest.main(...)` whose return value is discarded.
`sys.exit(pytest.main(...))` and `code = pytest.main(...)` are fine."""
if not isinstance(stmt, ast.Expr):
return False
call = stmt.value
if not isinstance(call, ast.Call):
return False
func = call.func
return (
isinstance(func, ast.Attribute)
and func.attr == "main"
and isinstance(func.value, ast.Name)
and func.value.id == "pytest"
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/test_parallel_adoption_ratchet.py
-78
View File
@@ -1,78 +0,0 @@
"""Ratchet guard: legacy parallel-getter calls in swept directories may only
shrink.
``models/`` and ``layers/`` read parallel topology through
``get_parallel().<dim>`` (the read-through wrapper in ``runtime_context``),
which gives one import, one naming scheme, and the scoped ``override()``
test primitive. Direct calls to the ``parallel_state`` size/rank getters in
these directories are regressions against that sweep.
Exemptions, pinned by path: ``runtime_context.py`` and
``layers/dp_attention.py`` are delegation substrate, while
``layers/dcp/comm.py`` retains deprecated DCP compatibility shims for
out-of-tree callers. Sweeping an exempt path must remove it from the pin.
"""
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
import re
import unittest
from pathlib import Path
import sglang.srt
from sglang.test.test_utils import CustomTestCase
_SRT_ROOT = Path(next(iter(sglang.srt.__path__)))
_BANNED_CALLS = re.compile(
r"\b(?:dcp_enabled|get_(?:"
r"tensor_model_parallel_(?:world_size|rank)"
r"|pipeline_model_parallel_(?:world_size|rank)"
r"|moe_expert_parallel_(?:world_size|rank)"
r"|moe_tensor_parallel_(?:world_size|rank)"
r"|moe_data_parallel_(?:world_size|rank)"
r"|attn_tensor_model_parallel_(?:world_size|rank)"
r"|attn_context_model_parallel_(?:world_size|rank)"
r"|dcp_(?:world_size|rank)"
r"|dcp_group(?:_no_assert)?"
r"|attention_dcp_(?:world_size|rank)"
r"|attention_(?:tp|cp)_(?:group|rank|size)"
r"))\(\)"
)
# The whole package is swept; the exemptions are the substrate itself.
_SWEPT_DIRS = ("",)
_EXEMPT = (
"distributed/", # parallel_state: defines the canonical getters
"runtime_context.py", # delegates DCP reads to canonical getters
"layers/dp_attention.py", # delegation substrate for the attn-DP dims
"layers/dcp/comm.py", # deprecated out-of-tree DCP compatibility shims
# The dumper's megatron plugin calls third-party getters that share the
# parallel_state names (self._mpu.get_tensor_model_parallel_rank()).
"debug_utils/dumper.py",
)
class TestParallelAdoptionRatchet(CustomTestCase):
def test_no_legacy_parallel_getters_in_swept_dirs(self):
offenders = []
for top in _SWEPT_DIRS:
for path in sorted((_SRT_ROOT / top).rglob("*.py")):
rel = path.relative_to(_SRT_ROOT).as_posix()
if rel.startswith(_EXEMPT):
continue
for i, line in enumerate(path.read_text().split("\n"), 1):
if _BANNED_CALLS.search(line):
offenders.append(f"{rel}:{i}")
self.assertFalse(
offenders,
"legacy parallel-getter calls in swept directories (use "
f"get_parallel().<dim> instead): {offenders}",
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/test_server_args_mutation_ratchet.py
-85
View File
@@ -1,85 +0,0 @@
"""Ratchet guard: server_args mutations outside the resolution pipeline may
only decrease.
After ``ServerArgs.__post_init__`` returns, the instance carries the resolved
configuration; the resolution pipeline (``server_args.py`` and
``arg_groups/``) is the only place that computes it. Every assignment to a
``server_args`` field elsewhere weakens that contract, so the count below is
an exact pin: new mutations must not appear, and removals must lower the
baseline to lock in the progress.
There is no post-resolution mutation entry point on the instance any more:
resolved config changes go to the context bags via
``get_context().override(source, **fields)``, and a value that differs for one
runner or worker travels as a constructor argument to it. The baseline is
therefore zero. ``ServerArgs.__setattr__`` raises
on a bare assignment after resolution; this ratchet catches the sites the tests
never execute.
"""
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
import re
import unittest
from pathlib import Path
import sglang
from sglang.test.test_utils import CustomTestCase
_SGLANG_ROOT = Path(next(iter(sglang.__path__)))
# Assignments to a server_args attribute (``server_args.x = ...``,
# ``self.server_args.x = ...``, and the ``sa`` alias used by a few helpers).
# ``==`` comparisons are excluded by the negative lookahead.
_MUTATION_PATTERNS = [
# (?![=}]) skips ``==`` comparisons and f-string ``{x=}`` debug specs.
re.compile(r"\bserver_args\.[a-z0-9_]+\s*=(?![=}])"),
re.compile(r"\bsa\.[a-z0-9_]+\s*=(?![=}])"),
re.compile(r"get_(?:global_)?server_args\(\)\.[a-z0-9_]+\s*=(?![=}])"),
# setattr is the same write with the attribute name behind a variable.
re.compile(
r"setattr\(\s*(?:[\w.]+\.)?(?:server_args|sa|get_(?:global_)?server_args\(\))\s*,"
),
]
# The resolution pipeline itself (mutation is its job) and multimodal_gen,
# whose ServerArgs is a different class outside this contract.
_EXCLUDED = (
"srt/server_args.py",
"srt/arg_groups",
"multimodal_gen",
)
_BASELINE = 0
class TestServerArgsMutationRatchet(CustomTestCase):
def test_out_of_pipeline_mutations_match_the_baseline(self):
count = 0
for path in sorted(_SGLANG_ROOT.rglob("*.py")):
rel = path.relative_to(_SGLANG_ROOT).as_posix()
if rel.startswith(_EXCLUDED):
continue
source = path.read_text()
count += sum(len(p.findall(source)) for p in _MUTATION_PATTERNS)
if count > _BASELINE:
self.fail(
f"server_args mutations outside the resolution pipeline grew: "
f"{count} > baseline {_BASELINE}. Configuration is resolved in "
"ServerArgs.__post_init__; declare through the pipeline "
"(passes / declare_late_resolution), change resolved config "
"with get_context().override(source, ...), or hand the value "
"to its runner as a constructor argument — do not assign fields."
)
if count < _BASELINE:
self.fail(
f"server_args mutations outside the resolution pipeline "
f"shrank: {count} < baseline {_BASELINE}. Lower the baseline "
"in this file to lock in the progress."
)
if __name__ == "__main__":
unittest.main()
test/run_suite.py
+3 -4
View File
@@ -319,10 +319,9 @@ def run_a_suite(args):
for f in glob.glob(
os.path.join(script_dir, "registered", "**", "*.py"), recursive=True
)
if not f.endswith("/conftest.py")
and not f.endswith("/__init__.py")
and not f.endswith("/cpu/utils.py")
and not f.endswith("/run_tests.py")
# conftest.py / __init__.py are pytest+package structure, never
# registered tests, and must not be executed as one.
if os.path.basename(f) not in ("conftest.py", "__init__.py")
]
# Strict: all discovered files must have proper registration