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

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

python/sglang/srt/distributed/parallel_state.py

@@ -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

@@ -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

@@ -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()

test/registered/unit/test_runtime_context.py

@@ -7,11 +7,13 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
 import dataclasses
 import json
 import os
+import pathlib as _pathlib
 import shutil
 import tempfile
 import unittest
 from unittest.mock import patch
 
+import sglang as _sglang
 import sglang.srt.server_args as server_args_module
 from sglang.srt.arg_groups.arg_utils import NS, A, Arg
 from sglang.srt.runtime_context import (
@@ -20,6 +22,7 @@ from sglang.srt.runtime_context import (
     RuntimeContext,
     _FlagGroupBase,
     assert_published,
+    derive_parallel_widths,
     get_context,
     get_exec,
     get_flags,
@@ -33,6 +36,7 @@ from sglang.srt.runtime_context import (
 from sglang.srt.server_args import ServerArgs
 from sglang.test.test_utils import CustomTestCase
 
+_SRT = _pathlib.Path(next(iter(_sglang.__path__))).resolve() / "srt"
 _PS = "sglang.srt.distributed.parallel_state"
 _DP = "sglang.srt.layers.dp_attention"
 
@@ -1385,5 +1389,168 @@ class TestParallelLeafReads(_IsolatedServerArgs):
             getattr(ParallelContext(), "not_a_leaf")
 
 
+class TestDerivedWidths(_IsolatedOverrides):
+    """The widths no flag sets are computed from the leaves and stamped.
+
+    `attn_tp_size` and its siblings used to be read back off the group
+    coordinator that was built from them, which made the answer depend on
+    distributed init and, after an elastic scale, disagree with the leaves.
+    """
+
+    def setUp(self):
+        super().setUp()
+        parallel = get_parallel()
+        self._saved_derived = dict(parallel._derived)
+        parallel.clear_derived_widths()
+        self.addCleanup(
+            lambda: (
+                parallel.clear_derived_widths(),
+                parallel.stamp_derived_widths(**self._saved_derived),
+            )
+        )
+
+    def test_the_quotients_come_from_the_leaves(self):
+        widths = derive_parallel_widths(
+            tp_size=8,
+            attn_cp_size=1,
+            attn_dp_size=2,
+            moe_ep_size=4,
+            moe_dp_size=2,
+            dcp_size=1,
+            dcp_enabled=False,
+        )
+        self.assertEqual(widths["attn_tp_size"], 8 // 2 // 1)
+        self.assertEqual(widths["moe_tp_size"], 8 // 4 // 2)
+        self.assertEqual(widths["attn_dcp_size"], 1)
+
+    def test_the_world_size_is_not_stamped(self):
+        """It is not a quotient, and the live getter is right at every moment.
+        A stamp taken when the groups are built would answer with the launch
+        count after `try_admit_scale_ranks` expands WORLD, and with the joining
+        cohort's own width on a scale-joiner, which lays its groups out at
+        `tp * pp` while WORLD spans `ep_join_rank_offset + tp * pp`."""
+        widths = derive_parallel_widths(
+            tp_size=4,
+            attn_cp_size=1,
+            attn_dp_size=1,
+            moe_ep_size=1,
+            moe_dp_size=1,
+            dcp_size=1,
+            dcp_enabled=False,
+        )
+        self.assertNotIn("world_size", widths)
+        parallel = get_parallel()
+        parallel.stamp_derived_widths(attn_tp_size=4)
+        with patch(f"{_PS}.get_world_size", return_value=9):
+            self.assertEqual(parallel.world_size, 9)
+
+    def test_a_stamped_width_is_what_the_reader_answers_with(self):
+        parallel = get_parallel()
+        parallel.stamp_derived_widths(attn_tp_size=4, moe_tp_size=1)
+        with patch(
+            f"{_PS}.get_attn_tensor_model_parallel_world_size",
+            side_effect=AssertionError("the group must not be asked"),
+        ):
+            self.assertEqual(parallel.attn_tp_size, 4)
+
+    def test_an_override_still_wins_over_the_stamp(self):
+        parallel = get_parallel()
+        parallel.stamp_derived_widths(attn_tp_size=4)
+        with parallel.override(attn_tp_size=1):
+            self.assertEqual(parallel.attn_tp_size, 1)
+        self.assertEqual(parallel.attn_tp_size, 4)
+
+    def test_without_a_stamp_the_live_group_still_answers(self):
+        """A process that installed groups by hand keeps working."""
+        with patch(f"{_PS}.get_attn_tensor_model_parallel_world_size", return_value=2):
+            self.assertEqual(get_parallel().attn_tp_size, 2)
+
+    def test_with_neither_the_failure_names_the_cause(self):
+        with patch(
+            f"{_PS}.get_attn_tensor_model_parallel_world_size",
+            side_effect=AssertionError("attention tp group is not initialized"),
+        ):
+            with self.assertRaisesRegex(RuntimeError, r"derived parallel width"):
+                get_parallel().attn_tp_size
+
+    def test_a_temporary_disable_beats_the_stamp(self):
+        """`disable_dp_size()` runs a draft scope without DP attention. It moves
+        the module global the legacy getter reads, so it has to move the derived
+        width too -- the stamp wins over the live group, and a scope that left
+        it alone would answer with the target model's width for its duration."""
+        from sglang.srt.layers import dp_attention
+
+        parallel = get_parallel()
+        parallel.stamp_derived_widths(attn_dp_size=4)
+        with patch.object(dp_attention, "_ATTN_DP_SIZE", 4):
+            with dp_attention.disable_dp_size():
+                self.assertEqual(dp_attention.get_attention_dp_size(), 1)
+                self.assertEqual(parallel.attn_dp_size, 1)
+            self.assertEqual(parallel.attn_dp_size, 4)
+
+    def test_the_stamp_is_cleared_and_restamped(self):
+        parallel = get_parallel()
+        parallel.stamp_derived_widths(attn_dp_size=2)
+        self.assertEqual(parallel.attn_dp_size, 2)
+        # Elastic scaling restamps where it updates the live width.
+        parallel.stamp_derived_widths(attn_dp_size=4)
+        self.assertEqual(parallel.attn_dp_size, 4)
+        parallel.clear_derived_widths()
+        with patch(f"{_DP}.get_attention_dp_size", return_value=1):
+            self.assertEqual(parallel.attn_dp_size, 1)
+
+    def test_reset_context_drops_the_stamp(self):
+        """The stamp belongs to the lifecycle that made it.
+
+        `_derived_width` prefers the stamp over the live group, so a stamp that
+        outlived `reset_context()` would let the next test read the previous
+        topology.
+        """
+        from sglang.srt.runtime_context import reset_context
+
+        parallel = get_parallel()
+        parallel.stamp_derived_widths(attn_tp_size=4)
+        self.assertEqual(parallel.attn_tp_size, 4)
+        reset_context()
+        with patch(f"{_PS}.get_attn_tensor_model_parallel_world_size", return_value=1):
+            self.assertEqual(get_parallel().attn_tp_size, 1)
+
+    def test_the_arithmetic_has_one_home(self):
+        """`parallel_state` builds its groups from the same dict it stamps, and
+        `dp_attention` derives the pair it needs for the ranks, so a second copy
+        of a quotient would let two answers to one width drift apart."""
+        for rel, spelling in (
+            ("distributed/parallel_state.py", "derive_parallel_widths("),
+            ("layers/dp_attention.py", "derive_attention_widths("),
+        ):
+            source = (_SRT / rel).read_text(encoding="utf-8-sig")
+            self.assertNotIn("// attn_dp_size // attn_cp_size", source, rel)
+            self.assertNotIn("// attn_cp_size // attn_dp_size", source, rel)
+            self.assertNotIn("// moe_ep_size // moe_dp_size", source, rel)
+            self.assertNotIn("if enable_dp_attention else 1", source, rel)
+            self.assertIn(spelling, source, rel)
+
+    def test_the_rank_helper_agrees_with_the_stamp(self):
+        """`compute_dp_attention_world_info` keeps the ranks and takes the
+        widths from the same derivation the stamp uses."""
+        from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
+
+        for tp_size, dp_size, attn_cp_size in ((8, 2, 1), (8, 2, 2), (16, 4, 2)):
+            _, attn_tp_size, _, attn_dp_size = compute_dp_attention_world_info(
+                True, 0, tp_size, dp_size, attn_cp_size
+            )
+            widths = derive_parallel_widths(
+                tp_size=tp_size,
+                attn_cp_size=attn_cp_size,
+                attn_dp_size=attn_dp_size,
+                moe_ep_size=1,
+                moe_dp_size=1,
+                dcp_size=1,
+                dcp_enabled=False,
+            )
+            self.assertEqual(attn_tp_size, widths["attn_tp_size"])
+            self.assertEqual(attn_dp_size, widths["attn_dp_size"])
+
+
 if __name__ == "__main__":
     unittest.main()