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

config: route parallel config-leaf reads through get_parallel() (#33170)

Browse Source 
The parallel namespace joins the accessor migration: 106 config-leaf reads
(enable_dp_lm_head, enable_dp_attention, pp_async_batch_depth, dp_size,
ep_join_rank_offset, dwdp_size, ...) flip from get_server_args()/
self.server_args to get_parallel(), which serves config leaves from the
published parallel bag via __getattr__.

- ParallelContext.__getattr__ is restructured to stay dynamo-traceable
  (object.__getattribute__ graph-breaks): gate helpers such as
  enable_moe_dense_fully_dp() run inside compiled model forwards. A
  fullgraph regression test pins the pattern.
- The five live-shadowed topology sizes (tp/pp/dcp/attn_cp/moe_dp_size)
  keep their server_args reads: the live @property wins on the accessor,
  and conditionally-initialized groups would fail loud at unconditional
  call sites.
- Elastic-EP scale writers (ep_size/dp_size x4 in model_runner) reroute
  to get_context().override together with their remaining instance
  readers (expert_location gpus-per-node paths); the ServerArgs.override
  ratchet drops 39 -> 35.
- The expert placement helpers (compute_logical_to_rank_dispatch_
  physical_map, _compute_logical_to_all_physical_map,
  _prefer_same_node_experts) now read everything from the bags and drop
  their server_args parameter; their unit tests publish the config they
  need instead of stubbing it.
This commit is contained in:
Cheng Wan
2026-08-01 08:58:39 -07:00
committed by GitHub
parent df55e911d6
commit 47d8b5b749
81 changed files with 344 additions and 405 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/ops/test_aiter_allreduce_fusion_amd.py
+6 -3
View File
@@ -389,7 +389,6 @@ class TestAiterAllreduceFusionGate(CustomTestCase):
tp_size=8,
):
"""Run the gate with the aiter branch isolated (flashinfer forced off)."""
server_args = types.SimpleNamespace(enable_aiter_allreduce_fusion=aiter_enabled)
a2a_backend = types.SimpleNamespace(is_none=lambda: a2a_is_none)
with ExitStack() as stack:
@@ -417,10 +416,14 @@ class TestAiterAllreduceFusionGate(CustomTestCase):
lambda: types.SimpleNamespace(tp_size=tp_world_size),
)
)
# the gate reads get_exec().comm.enable_aiter_allreduce_fusion
from sglang.srt.runtime_context import get_context, get_flags
stack.enter_context(
mock.patch.object(comm, "get_server_args", lambda: server_args)
get_context().override_server_args(
enable_aiter_allreduce_fusion=aiter_enabled
)
)
from sglang.srt.runtime_context import get_flags
stack.enter_context(get_flags().dp.override(enabled=dp_attention))
stack.enter_context(
test/registered/unit/eplb/test_compute_logical_to_rank_dispatch_physical_map.py
+46 -46
View File
@@ -4,7 +4,6 @@ from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=7, suite="base-a-test-cpu")
import types
import unittest
import torch
@@ -14,20 +13,26 @@ from sglang.srt.eplb.expert_location import (
append_trivial_expert_slots,
compute_logical_to_rank_dispatch_physical_map,
)
from sglang.srt.runtime_context import get_context
from sglang.test.test_utils import CustomTestCase
def _make_server_args(ep_size: int, nnodes: int, moe_a2a_backend: str = "deepep"):
"""Minimal server_args stub for expert placement tests.
def _published(
ep_size: int,
nnodes: int,
moe_a2a_backend: str = "deepep",
ep_join_mode=None,
):
"""Scoped publish of the config the placement functions read.
`moe_a2a_backend` defaults to an a2a backend because these tests cover the
rank-local collapse, which is skipped when there is no a2a backend.
"""
return types.SimpleNamespace(
return get_context().override_server_args(
ep_size=ep_size,
nnodes=nnodes,
ep_join_mode=None,
moe_a2a_backend=moe_a2a_backend,
ep_join_mode=ep_join_mode,
)
@@ -74,7 +79,6 @@ class TestComputeLogicalToRankDispatchPhysicalMap(CustomTestCase):
NUM_LAYERS = 2
def setUp(self):
self.server_args = _make_server_args(self.EP_SIZE, self.NNODES)
self.logical_to_all_physical = _make_logical_to_all_physical_map(
num_layers=self.NUM_LAYERS,
num_logical_experts=self.NUM_LOGICAL,
@@ -83,14 +87,14 @@ class TestComputeLogicalToRankDispatchPhysicalMap(CustomTestCase):
)
def _call(self, ep_rank, seed=42):
return compute_logical_to_rank_dispatch_physical_map(
server_args=self.server_args,
logical_to_all_physical_map=self.logical_to_all_physical.clone(),
ep_size=self.EP_SIZE,
num_physical_experts=self.NUM_PHYSICAL,
ep_rank=ep_rank,
seed=seed,
)
with _published(self.EP_SIZE, self.NNODES):
return compute_logical_to_rank_dispatch_physical_map(
logical_to_all_physical_map=self.logical_to_all_physical.clone(),
ep_size=self.EP_SIZE,
num_physical_experts=self.NUM_PHYSICAL,
ep_rank=ep_rank,
seed=seed,
)
# ------------------------------------------------------------------ shape & range
@@ -164,26 +168,25 @@ class TestComputeLogicalToRankDispatchPhysicalMap(CustomTestCase):
num_physical_experts=self.NUM_PHYSICAL,
replicas_per_logical=2,
)
result = compute_logical_to_rank_dispatch_physical_map(
server_args=self.server_args,
logical_to_all_physical_map=logical_to_all_physical,
ep_size=self.EP_SIZE,
num_physical_experts=self.NUM_PHYSICAL,
ep_rank=0,
)
with _published(self.EP_SIZE, self.NNODES):
result = compute_logical_to_rank_dispatch_physical_map(
logical_to_all_physical_map=logical_to_all_physical,
ep_size=self.EP_SIZE,
num_physical_experts=self.NUM_PHYSICAL,
ep_rank=0,
)
self.assertEqual(result.shape, (1, self.NUM_LOGICAL))
self.assertTrue(torch.all(result >= 0))
def test_single_node(self):
"""With nnodes=1, all GPUs are on the same node."""
server_args = _make_server_args(ep_size=4, nnodes=1)
result = compute_logical_to_rank_dispatch_physical_map(
server_args=server_args,
logical_to_all_physical_map=self.logical_to_all_physical.clone(),
ep_size=self.EP_SIZE,
num_physical_experts=self.NUM_PHYSICAL,
ep_rank=0,
)
with _published(ep_size=4, nnodes=1):
result = compute_logical_to_rank_dispatch_physical_map(
logical_to_all_physical_map=self.logical_to_all_physical.clone(),
ep_size=self.EP_SIZE,
num_physical_experts=self.NUM_PHYSICAL,
ep_rank=0,
)
self.assertEqual(result.shape, (self.NUM_LAYERS, self.NUM_LOGICAL))
self.assertTrue(torch.all(result >= 0))
self.assertTrue(torch.all(result < self.NUM_PHYSICAL))
@@ -195,13 +198,13 @@ class TestComputeLogicalToRankDispatchPhysicalMap(CustomTestCase):
torch.arange(self.NUM_PHYSICAL, dtype=torch.int64).unsqueeze(0).unsqueeze(0)
)
mapping = mapping.expand(self.NUM_LAYERS, 1, self.NUM_PHYSICAL).clone()
result = compute_logical_to_rank_dispatch_physical_map(
server_args=self.server_args,
logical_to_all_physical_map=mapping,
ep_size=self.EP_SIZE,
num_physical_experts=self.NUM_PHYSICAL,
ep_rank=0,
)
with _published(self.EP_SIZE, self.NNODES):
result = compute_logical_to_rank_dispatch_physical_map(
logical_to_all_physical_map=mapping,
ep_size=self.EP_SIZE,
num_physical_experts=self.NUM_PHYSICAL,
ep_rank=0,
)
self.assertEqual(result.shape, (self.NUM_LAYERS, 1))
self.assertTrue(torch.all(result >= 0))
@@ -210,16 +213,13 @@ class TestComputeLogicalToRankDispatchPhysicalMap(CustomTestCase):
physical_to_logical = append_trivial_expert_slots(
physical_to_logical, count=16, num_logical_experts=64
)
server_args = _make_server_args(ep_size=5, nnodes=1)
server_args.ep_join_mode = "scale"
logical_to_physical = _compute_logical_to_all_physical_map(
server_args=server_args,
physical_to_logical_map=physical_to_logical,
num_logical_experts=64,
ep_size=5,
moe_ep_rank=4,
)
with _published(ep_size=5, nnodes=1, ep_join_mode="scale"):
logical_to_physical = _compute_logical_to_all_physical_map(
physical_to_logical_map=physical_to_logical,
num_logical_experts=64,
ep_size=5,
moe_ep_rank=4,
)
self.assertEqual(logical_to_physical[0, :16, 0].tolist(), list(range(64, 80)))
test/registered/unit/model_loader/test_presharded_loader.py
+8 -3
View File
@@ -787,9 +787,7 @@ class TestShardConfig(unittest.TestCase):
# moe_dense_tp_size / LM-head flags out of the cache key before.
loader = object.__new__(PreshardedModelLoader)
server_args = SimpleNamespace(
moe_dense_tp_size=1,
moe_dp_size=2,
enable_dp_lm_head=True,
enable_fp32_lm_head=True,
ep_num_redundant_experts=4,
enable_eplb=True,
@@ -812,7 +810,14 @@ class TestShardConfig(unittest.TestCase):
"init_expert_location",
"structural_signature",
}
parallel = SimpleNamespace(tp_size=8, moe_dp_size=2, moe_ep_size=4, pp_size=1)
parallel = SimpleNamespace(
tp_size=8,
moe_dp_size=2,
moe_ep_size=4,
pp_size=1,
moe_dense_tp_size=1,
enable_dp_lm_head=True,
)
with mock.patch(
"sglang.srt.model_loader.loader.get_server_args",
return_value=server_args,
test/registered/unit/test_runtime_context.py
+27
View File
@@ -760,6 +760,33 @@ class TestForwardFlags(_IsolatedServerArgs):
self.assertEqual(probe(torch.zeros(())).item(), 28)
self.assertEqual(probe(torch.zeros(())).item(), 0)
def test_parallel_config_leaves_trace_under_torch_compile(self):
# Regression: parallel config leaves resolve through
# ``ParallelContext.__getattr__`` (the bag fallback), and gate helpers
# such as ``enable_moe_dense_fully_dp()`` read them inside compiled
# model forwards — the fallback body must stay dynamo-traceable
# (``object.__getattribute__`` graph-breaks). fullgraph=True turns any
# graph break back into a failure.
import torch
from sglang.srt.runtime_context import get_parallel
reset_context()
with get_context().override_server_args(moe_dense_tp_size=1, dwdp_size=4):
@torch.compile(fullgraph=True, backend="eager", dynamic=False)
def probe(x):
par = get_parallel()
if par.enable_prefill_context_parallel:
x = x + 1
if par.moe_dense_tp_size == 1:
x = x + 2
if par.dwdp_size > 1:
x = x + 4
return x
self.assertEqual(probe(torch.zeros(())).item(), 6)
def test_graph_visible_flags_are_process_visible_across_threads(self):
# Documented divergence from the contextvar-backed flags: plain slots
# are process-global (the storage form these flags had before the
test/registered/unit/test_server_args_writer_ratchet.py
+1 -1
View File
@@ -49,7 +49,7 @@ _EXCLUDED = (
"multimodal_gen",
)
_BASELINE = 38
_BASELINE = 34
class TestServerArgsWriterRatchet(CustomTestCase):