diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 2e8c0bafd07c..d95235ceb80d 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/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() diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index 2d7efc68396b..c50242e7cd86 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/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,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 diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index 5c0f85e61e41..e496b9a43833 100644 --- a/python/sglang/srt/runtime_context.py +++ b/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,7 +318,9 @@ 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: @@ -233,7 +328,9 @@ def moe_tp_rank(self) -> int: @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,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() diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 9e2c10136f77..ab86286a6193 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -7,11 +7,13 @@ 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 @@ RuntimeContext, _FlagGroupBase, assert_published, + derive_parallel_widths, get_context, get_exec, get_flags, @@ -33,6 +36,7 @@ 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 @@ def test_an_unknown_name_is_still_an_attribute_error(self): 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()