config: the derived parallel widths are computed from the leaves

python/sglang/srt/distributed/parallel_state.py

@@ -51,7 +51,9 @@
 )
 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

@@ -29,6 +29,7 @@
     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,21 +404,23 @@ 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!"
 
     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

@@ -131,6 +131,57 @@ def _parallel_config_leaves() -> 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 @@ def _v(self, name, getter):
         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 @@ def pp_rank(self) -> int:
 
     @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,15 +318,19 @@ def moe_dp_rank(self) -> int:
 
     @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:
         return self._v("moe_tp_rank", _ps().get_moe_tensor_parallel_rank)
 
     @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 @@ def getter():
                 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 @@ def attn_dcp_rank(self) -> int:
 
     @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,14 +1609,17 @@ 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
     _adaptive_draft_token_bound.cache_clear()
     _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()