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

config: the derived parallel widths are computed from the leaves (#36790)

Browse Source 
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-08-28 10:18:46 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent c2928e86d7
commit 43c63a22ff
4 changed files with 316 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/parallel_state.py
+18 -2
View File
@@ -51,7 +51,9 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
)
from sglang.srt.platforms.device_mixin import _DEVICE_TO_DISTRIBUTED_BACKEND
from sglang.srt.runtime_context import (
derive_parallel_widths,
get_global_dwdp_manager,
get_parallel,
set_global_dwdp_manager,
)
from sglang.srt.utils import (
@@ -2504,7 +2506,18 @@ def initialize_model_parallel(
attn_dp_size = attention_data_parallel_size
attn_cp_size = attention_context_model_parallel_size
attn_tp_size = tensor_model_parallel_size // attn_cp_size // attn_dp_size
# The groups below are built at these numbers, and the same dict is stamped
# once they exist.
derived_widths = derive_parallel_widths(
tp_size=tensor_model_parallel_size,
attn_cp_size=attn_cp_size,
attn_dp_size=attn_dp_size,
moe_ep_size=expert_model_parallel_size,
moe_dp_size=moe_data_model_parallel_size,
dcp_size=decode_context_parallel_size,
dcp_enabled=_DCP is not None,
)
attn_tp_size = derived_widths["attn_tp_size"]
global _ATTN_CP
assert (
@@ -2590,7 +2603,7 @@ def initialize_model_parallel(
moe_ep_size = expert_model_parallel_size
moe_dp_size = moe_data_model_parallel_size
moe_tp_size = tensor_model_parallel_size // moe_ep_size // moe_dp_size
moe_tp_size = derived_widths["moe_tp_size"]
global _MOE_DP
assert _MOE_DP is None, "moe data parallel group is already initialized"
@@ -2703,6 +2716,8 @@ def initialize_model_parallel(
max_world_size=max_world_size,
)
get_parallel().stamp_derived_widths(**derived_widths)
def create_custom_parallel_group(
group_ranks: List[int], backend: str = "gloo"
@@ -2930,6 +2945,7 @@ def get_moe_tensor_parallel_rank():
def destroy_model_parallel():
"""Set the groups to none and destroy them."""
get_parallel().clear_derived_widths()
dwdp_mgr = get_global_dwdp_manager()
if dwdp_mgr is not None:
dwdp_mgr.cleanup()
python/sglang/srt/layers/dp_attention.py
+23 -10
View File
@@ -29,6 +29,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.runtime_context import (
derive_attention_widths,
get_device,
get_exec,
get_flags,
@@ -63,6 +64,7 @@ def update_dp_attention_post_scale(new_dp_size: int, new_dp_rank: int):
global _ATTN_DP_SIZE, _ATTN_DP_RANK
_ATTN_DP_SIZE = new_dp_size
_ATTN_DP_RANK = new_dp_rank
get_parallel().stamp_derived_widths(attn_dp_size=new_dp_size)
get_flags().dp.use_world_group_for_gather = True
logger.debug(
"[Elastic EP] dp_attention switched to WORLD: dp_size=%d dp_rank=%d",
@@ -324,8 +326,17 @@ def is_dp_max_padding() -> bool:
def compute_dp_attention_world_info(
enable_dp_attention, tp_rank, tp_size, dp_size, attn_cp_size: int = 1
):
attn_dp_size = dp_size if enable_dp_attention else 1
attn_tp_size = tp_size // attn_dp_size // attn_cp_size
"""This rank's place in the attention topology, plus the widths it sits in.
The widths come from `derive_attention_widths`; what this adds is the two
ranks, which are per-process and so are not part of the stamped set.
"""
attn_dp_size, attn_tp_size = derive_attention_widths(
tp_size=tp_size,
attn_cp_size=attn_cp_size,
dp_size=dp_size,
enable_dp_attention=enable_dp_attention,
)
attn_tp_rank = tp_rank % attn_tp_size
if not enable_dp_attention:
@@ -356,10 +367,10 @@ def initialize_dp_attention(
tp_rank = get_tensor_model_parallel_rank()
tp_size = get_tensor_model_parallel_world_size()
_, _, _ATTN_DP_RANK, _ = compute_dp_attention_world_info(
_, _, _ATTN_DP_RANK, _ATTN_DP_SIZE = compute_dp_attention_world_info(
enable_dp_attention, tp_rank, tp_size, dp_size, attn_cp_size
)
_ATTN_DP_SIZE = dp_size if enable_dp_attention else 1
get_parallel().stamp_derived_widths(attn_dp_size=_ATTN_DP_SIZE)
if get_exec().moe.elastic_ep_backend is not None and get_parallel().max_ep_size:
_ATTN_DP_RANK = tp_rank + get_parallel().ep_join_rank_offset
@@ -393,13 +404,14 @@ def get_attention_dp_size() -> int:
@contextmanager
def disable_dp_size():
"""Patch the tp group temporarily until this function ends.
"""Run without DP attention until this scope ends.
This method is for draft workers of speculative decoding to run draft model
with different tp degree from that of target model workers.
This is for draft workers of speculative decoding, which run the draft model
at a different width from the target model's workers.
Args:
tp_group (GroupCoordinator): the tp group coordinator
The scope replaces both the module global that ``get_attention_dp_size()``
reads and the derived width the runtime context answers with, so the two
spellings of the name cannot disagree inside it.
"""
global _ATTN_DP_SIZE
assert _ATTN_DP_SIZE is not None, "dp attention not initialized!"
@@ -407,7 +419,8 @@ def disable_dp_size():
old_dp_size = _ATTN_DP_SIZE
_ATTN_DP_SIZE = 1
try:
yield
with get_parallel().override(attn_dp_size=1):
yield
finally:
_ATTN_DP_SIZE = old_dp_size
python/sglang/srt/runtime_context.py
+108 -8
View File
@@ -131,6 +131,57 @@ _PARALLEL_FIELDS = frozenset(
)
def derive_attention_widths(
*, tp_size: int, attn_cp_size: int, dp_size: int, enable_dp_attention: bool
) -> tuple:
"""(attn_dp_size, attn_tp_size) from the leaves.
Split out because the rank computation in
`dp_attention.compute_dp_attention_world_info` needs the same two numbers
and must not carry a second copy of the arithmetic.
"""
attn_dp_size = dp_size if enable_dp_attention else 1
return attn_dp_size, tp_size // attn_dp_size // attn_cp_size
def derive_parallel_widths(
*,
tp_size: int,
attn_cp_size: int,
attn_dp_size: int,
moe_ep_size: int,
moe_dp_size: int,
dcp_size: int,
dcp_enabled: bool,
) -> dict:
"""The parallel widths no flag sets, from the leaves that do.
`tp_size` and its siblings are configured; these are quotients of them, so
the arithmetic lives here rather than being read back off the group
coordinators.
`world_size` is not among them: it is not a quotient, and `get_world_size()`
answers with the live WORLD group, which stays right through an elastic
scale-up that a stamp taken at group build would not survive.
"""
return {
"attn_dp_size": attn_dp_size,
# `attn_dp_size` is already the effective width (1 when DP attention is
# off), so the flag is spent here; a caller passing the raw `dp_size`
# leaf with the attention disabled would get tp/dp/cp instead of tp/1/cp.
"attn_tp_size": derive_attention_widths(
tp_size=tp_size,
attn_cp_size=attn_cp_size,
dp_size=attn_dp_size,
enable_dp_attention=True,
)[1],
"moe_ep_size": moe_ep_size,
"moe_tp_size": tp_size // moe_ep_size // moe_dp_size,
"dcp_enabled": dcp_enabled,
"attn_dcp_size": dcp_size if dcp_enabled else 1,
}
class ParallelContext:
"""Parallel-topology namespace: one spelling per name.
@@ -154,11 +205,12 @@ class ParallelContext:
different names rather than two answers to one name.
"""
__slots__ = ("_overrides", "_config")
__slots__ = ("_overrides", "_config", "_derived")
def __init__(self):
self._overrides = {}
self._config = None # parallel config bag, wired at publish
self._derived = {} # widths stamped when the groups are built
def __getattr__(self, name):
if name.startswith("_"):
@@ -181,6 +233,45 @@ class ParallelContext:
overrides = self._overrides
return overrides[name] if name in overrides else getter()
def stamp_derived_widths(self, **widths) -> None:
"""Record the widths derived from the leaves, as the groups are built.
`initialize_model_parallel` computes the set through
`derive_parallel_widths` and hands it here; `initialize_dp_attention`
stamps `attn_dp_size` again once it knows the effective width, and
elastic EP restamps it where it already updates the live one. A stamped
width is what the readers answer with.
"""
self._derived.update(widths)
def clear_derived_widths(self) -> None:
self._derived.clear()
def _derived_width(self, name, getter):
"""A width the leaves imply: the stamp, else the live group.
The fallback keeps a process that installed groups without going
through `initialize_model_parallel` working. When neither is there,
the failure says which of the two is missing rather than surfacing a
group getter's bare assertion.
"""
overrides = self._overrides
if name in overrides:
return overrides[name]
derived = self._derived
if name in derived:
return derived[name]
try:
return getter()
except (AssertionError, AttributeError, RuntimeError) as exc:
raise RuntimeError(
f"derived parallel width {name!r} is not available: it is "
"computed from the configured leaves when the process groups "
"are built (initialize_model_parallel / "
"initialize_dp_attention), and neither a stamp nor a live "
"group is present"
) from exc
@contextmanager
def override(self, **kwargs):
"""Temporarily force parallel values, restoring on exit. Validates keys and
@@ -213,7 +304,9 @@ class ParallelContext:
@property
def moe_ep_size(self) -> int:
return self._v("moe_ep_size", _ps().get_moe_expert_parallel_world_size)
return self._derived_width(
"moe_ep_size", _ps().get_moe_expert_parallel_world_size
)
@property
def moe_ep_rank(self) -> int:
@@ -225,7 +318,9 @@ class ParallelContext:
@property
def moe_tp_size(self) -> int:
return self._v("moe_tp_size", _ps().get_moe_tensor_parallel_world_size)
return self._derived_width(
"moe_tp_size", _ps().get_moe_tensor_parallel_world_size
)
@property
def moe_tp_rank(self) -> int:
@@ -233,7 +328,9 @@ class ParallelContext:
@property
def attn_tp_size(self) -> int:
return self._v("attn_tp_size", _ps().get_attn_tensor_model_parallel_world_size)
return self._derived_width(
"attn_tp_size", _ps().get_attn_tensor_model_parallel_world_size
)
@property
def attn_tp_rank(self) -> int:
@@ -254,11 +351,11 @@ class ParallelContext:
return False
return _ps().get_dcp_world_size() > 1
return self._v("dcp_enabled", getter)
return self._derived_width("dcp_enabled", getter)
@property
def attn_dcp_size(self) -> int:
return self._v(
return self._derived_width(
"attn_dcp_size",
lambda: _ps().get_dcp_world_size() if self.dcp_enabled else 1,
)
@@ -271,7 +368,7 @@ class ParallelContext:
@property
def attn_dp_size(self) -> int:
return self._v("attn_dp_size", _dp().get_attention_dp_size)
return self._derived_width("attn_dp_size", _dp().get_attention_dp_size)
@property
def attn_dp_rank(self) -> int:
@@ -1512,7 +1609,9 @@ def reset_context() -> None:
"""Clear the context-owned store (unit-test teardown): drop the published
``server_args`` and install fresh ``Flags`` and ``Resources``.
Wrapper subsystems (``parallel``) hold no state and are unaffected.
``parallel`` holds the stamped derived widths, which go with the lifecycle
that stamped them: `_derived_width` prefers the stamp over the live group,
so leaving one behind lets the next test read the previous topology.
"""
_CONTEXT._server_args = None
_CONTEXT._config_bags = None
@@ -1520,6 +1619,7 @@ def reset_context() -> None:
_CONTEXT._overrides_log = []
_CONTEXT._publish_role = None
_CONTEXT.parallel._config = None
_CONTEXT.parallel.clear_derived_widths()
_CONTEXT.flags = Flags()
_CONTEXT.resources = Resources()
_CONTEXT.forward = ForwardFlags()