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

[refactor] Adopt get_parallel() everywhere and close out the parallel wrapper surface (#30492)

Browse Source
This commit is contained in:
Cheng Wan
2026-07-09 02:09:39 -07:00
committed by GitHub
parent 06eb1b1838
commit e703f9e566
71 changed files with 341 additions and 415 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/cp/test_cp_strategy_unit.py
+4 -6
View File
@@ -371,9 +371,8 @@ class TestCPZigzagStrategy(CustomTestCase):
fb = self._forward_batch(metas[rank], extend_seq_lens)
with (
patch(
"sglang.srt.layers.cp.zigzag.get_attention_cp_group",
return_value=_FakeCPGroup(padded_rank_tensors),
get_parallel().override(
attn_cp_group=_FakeCPGroup(padded_rank_tensors)
),
patch(
"sglang.srt.distributed.device_communicators.pynccl_allocator.use_symmetric_memory",
@@ -405,9 +404,8 @@ class TestCPZigzagStrategy(CustomTestCase):
fb = self._forward_batch(metas[rank], extend_seq_lens)
with (
patch(
"sglang.srt.layers.cp.zigzag.get_attention_cp_group",
return_value=_FakeCPGroup(padded_rank_tensors),
get_parallel().override(
attn_cp_group=_FakeCPGroup(padded_rank_tensors)
),
patch(
"sglang.srt.distributed.device_communicators.pynccl_allocator.use_symmetric_memory",
test/registered/debug_utils/comparator/test_entrypoint.py
-2
View File
@@ -4544,8 +4544,6 @@ class TestEntrypointDpAttentionMissingAlias:
"attn_tp_size": 1,
"attn_dp_rank": tp_rank,
"attn_dp_size": 2,
"local_attn_dp_rank": tp_rank,
"local_attn_dp_size": 2,
"attn_cp_rank": 0,
"attn_cp_size": 1,
}
test/registered/debug_utils/test_dumper.py
-2
View File
@@ -2344,8 +2344,6 @@ class TestDumperE2E:
"attn_tp_size",
"attn_dp_rank",
"attn_dp_size",
"local_attn_dp_rank",
"local_attn_dp_size",
"attn_cp_rank",
"attn_cp_size",
]
test/registered/unit/layers/test_dsv4_nonpaged_indexer.py
+2 -1
View File
@@ -9,6 +9,7 @@ from sglang.srt.environ import envs
from sglang.srt.layers.attention.dsv4.indexer import FP8_DTYPE, C4IndexerBackendMixin
from sglang.srt.layers.attention.dsv4.metadata import NonPagedIndexerPlan
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.runtime_context import get_parallel
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -39,7 +40,7 @@ class TestDSV4NonPagedIndexer(CustomTestCase):
envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.override(False),
patch(f"{_INDEXER}.is_cuda", return_value=True),
patch(f"{_INDEXER}.is_hip", return_value=False),
patch(f"{_INDEXER}.get_attention_cp_size", return_value=1),
get_parallel().override(attn_cp_size=1),
patch(
f"{_INDEXER}.is_in_tc_piecewise_cuda_graph",
return_value=overrides.get("piecewise_graph", False),
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+7 -24
View File
@@ -2832,33 +2832,16 @@ class UnifiedRadixCacheSuite:
storage_extra_config = None
if storage_backend == "file":
import sglang.srt.managers.cache_controller as cache_controller
from sglang.srt.runtime_context import get_parallel
# The file-backend storage config records TP/PP rank/size. These unit
# The file-backend storage config records TP/PP rank/size. These unit
# fixtures run without initializing distributed parallel state, so
# provide the local single-rank values that the fixture represents.
tp_rank_patcher = mock.patch.object(
cache_controller, "get_tensor_model_parallel_rank", return_value=0
# force the local single-rank topology the fixture represents.
parallel_override = get_parallel().override(
tp_rank=0, tp_size=1, pp_rank=0, pp_size=1
)
tp_size_patcher = mock.patch.object(
cache_controller, "get_tensor_model_parallel_world_size", return_value=1
)
pp_rank_patcher = mock.patch.object(
cache_controller, "get_pipeline_model_parallel_rank", return_value=0
)
pp_size_patcher = mock.patch.object(
cache_controller,
"get_pipeline_model_parallel_world_size",
return_value=1,
)
tp_rank_patcher.start()
tp_size_patcher.start()
pp_rank_patcher.start()
pp_size_patcher.start()
self.addCleanup(tp_rank_patcher.stop)
self.addCleanup(tp_size_patcher.stop)
self.addCleanup(pp_rank_patcher.stop)
self.addCleanup(pp_size_patcher.stop)
parallel_override.__enter__()
self.addCleanup(parallel_override.__exit__, None, None, None)
assert storage_dir is not None, "file backend needs a storage_dir"
# HiCacheFile reads the directory from this env var.
test/registered/unit/model_executor/test_pool_configurator.py
+2 -4
View File
@@ -10,6 +10,7 @@ import unittest
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from sglang.srt.runtime_context import get_parallel
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
@@ -27,10 +28,7 @@ def mock_cpu_env(kv_size=2, tp_size=1, swa_eviction_interval=4):
with (
patch("torch._utils._element_size", return_value=kv_size),
patch(
"sglang.srt.model_executor.pool_configurator.get_attention_tp_size",
return_value=tp_size,
),
get_parallel().override(attn_tp_size=tp_size),
envs.SGLANG_SWA_EVICTION_INTERVAL.override(swa_eviction_interval),
):
yield
test/registered/unit/test_module_state_ratchet.py
-2
View File
@@ -34,8 +34,6 @@ _PINNED_GLOBALS = {
# DP-attention topology (parallel vertical scope).
"_ATTN_DP_RANK",
"_ATTN_DP_SIZE",
"_LOCAL_ATTN_DP_SIZE",
"_LOCAL_ATTN_DP_RANK",
}
),
}
test/registered/unit/test_parallel_adoption_ratchet.py
+74
View File
@@ -0,0 +1,74 @@
"""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: ``layers/dp_attention.py`` is delegation
substrate (the wrapper's attn-DP dims delegate TO it), and ``layers/dcp/``
is the DCP subsystem's own plumbing, booked for a follow-up sweep. 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"\bget_(?:"
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"|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
"layers/dp_attention.py", # delegation substrate for the attn-DP dims
# 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()