Deprecate the parallel getters the context answers, and ratchet them shut (#40342)

python/sglang/multimodal_gen/runtime/distributed/parallel_state.py

@@ -97,9 +97,14 @@ def init_world_group(
 
 def _sync_srt_world_group() -> None:
     import sglang.srt.distributed.parallel_state as srt_parallel_state
+    from sglang.srt.runtime_context import get_parallel
 
     if srt_parallel_state._WORLD is None:
         srt_parallel_state._WORLD = _WORLD
+    if srt_parallel_state._WORLD is _WORLD:
+        # On the context too: that is where a handle is read from, and
+        # assigning the module global above does not reach it.
+        get_parallel().override_permanently(world_group=_WORLD)
 
 
 def _clear_srt_world_group() -> None:
@@ -132,6 +137,27 @@ def _sync_srt_tp_group() -> None:
         srt_parallel_state._ATTN_TP = _TP
     if srt_parallel_state._ATTN_TP is _TP:
         get_parallel().override_permanently(
+            # The group itself, because that is what the `srt` context answers
+            # a handle with -- assigning the module global above does not reach
+            # it. `tp_size` comes with them: the group is as wide as the world
+            # while the dummy carries this package's, and the widths below are
+            # quotients of one number, so stating a subset would describe a
+            # layout that does not exist.
+            tp_group=_TP,
+            attn_tp_group=_TP,
+            tp_size=_TP.world_size,
+            # The ranks too. The shared layers shard by them -- `vision.py`
+            # reads `attn_tp_rank`, every `srt` linear built without an
+            # explicit rank reads `tp_rank` -- and this package publishes no
+            # rank bundle, so nothing else writes one. The draft has no
+            # pipeline, context or expert dimension of its own, so those
+            # positions are zero.
+            tp_rank=_TP.rank_in_group,
+            attn_tp_rank=_TP.rank_in_group,
+            moe_tp_rank=_TP.rank_in_group,
+            attn_cp_rank=0,
+            pp_rank=0,
+            moe_ep_rank=0,
             **derive_parallel_widths(
                 tp_size=_TP.world_size,
                 attn_cp_size=1,
@@ -151,6 +177,10 @@ def _clear_srt_tp_group() -> None:
     if srt_parallel_state._ATTN_TP is _TP:
         srt_parallel_state._ATTN_TP = None
         get_parallel().clear_stamp()
+        if srt_parallel_state._WORLD is not None:
+            # `clear_stamp` drops every stamped name; the WORLD group this
+            # package lent is still built, so hand it back.
+            get_parallel().override_permanently(world_group=srt_parallel_state._WORLD)
     if srt_parallel_state._TP is _TP:
         srt_parallel_state._TP = None
 

python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py

@@ -2,6 +2,7 @@ from contextlib import ExitStack
 from types import SimpleNamespace
 from unittest.mock import call, patch
 
+import pytest
 import torch
 from torch import nn
 
@@ -135,7 +136,7 @@ def test_srt_attention_tp_group_tracks_diffusion_tp_group():
     # `world_size`, because lending the group also states the parallel widths it
     # implies -- the shared `srt` vision layers ask for `attn_tp_size`, and this
     # package publishes no `srt` config for that read to resolve against.
-    tp_group = SimpleNamespace(world_size=2)
+    tp_group = SimpleNamespace(world_size=2, rank_in_group=1)
 
     with (
         patch.object(parallel_state, "_TP", tp_group),
@@ -147,11 +148,22 @@ def test_srt_attention_tp_group_tracks_diffusion_tp_group():
         assert srt_parallel_state._TP is tp_group
         assert srt_parallel_state._ATTN_TP is tp_group
         assert get_parallel().attn_tp_size == 2
+        # The handle too: assigning the module global does not reach the `srt`
+        # context, which is what the shared layers ask for a group.
+        assert get_parallel().tp_group is tp_group
+        assert get_parallel().attn_tp_group is tp_group
+        # And the ranks the shared layers shard by. Nothing else writes one
+        # here: this package publishes no rank bundle, so a handle without a
+        # rank leaves every `srt` linear unable to say which shard it is.
+        assert get_parallel().tp_rank == 1
+        assert get_parallel().attn_tp_rank == 1
 
         parallel_state._clear_srt_tp_group()
 
         assert srt_parallel_state._TP is None
         assert srt_parallel_state._ATTN_TP is None
+        with pytest.raises(RuntimeError):
+            get_parallel().tp_group
 
 
 def test_srt_owned_groups_are_not_overwritten_or_cleared():

python/sglang/srt/arg_groups/arg_utils.py

@@ -107,7 +107,12 @@ class Derived(msgspec.Struct, frozen=True):
     ``publish`` and stored as an ordinary bag leaf -- a plain attribute load,
     which is what a read inside compiled model code needs.
 
-    Every declaration carries ``fn`` today, the parallel quotients included:
+    A declaration with no ``fn`` is one nothing can compute: a rank, or a
+    process group. Those are written into the namespace at runtime -- by
+    ``publish`` from the spawn bundle, or by the build that creates the group --
+    and until then the name has no answer.
+
+    Most declarations carry ``fn``, the parallel quotients included:
     they are a function of the configured leaves, so they are computed at
     publish like the rest. What is special about them is not how they are
     computed but that a stamp can move one afterwards -- ``initialize_dp_attention``

python/sglang/srt/arg_groups/fields/parallel.py

@@ -320,3 +320,73 @@ class Parallel(msgspec.Struct):
         doc="Whether decode context parallelism is in play: `dcp_size` is "
         "wider than one rank, which is exactly when the group gets built.",
     )
+
+    # -- written at runtime, not carried by any configuration --------------
+    #
+    # No `fn`: nothing here is a function of the leaves above. A rank is
+    # written by `publish` from the spawn bundle; a group by the build that
+    # creates it. Until one of them has run there is no answer, and a read
+    # says so rather than deriving something that would answer a different
+    # question.
+    tp_rank = Derived(doc="This process's place in the tensor-parallel group.")
+    pp_rank = Derived(doc="This process's place in the pipeline group.")
+    moe_ep_rank = Derived(doc="This process's place in the expert-parallel group.")
+    moe_dp_rank = Derived(doc=("This process's place in the MoE data-parallel group."))
+    moe_tp_rank = Derived(
+        doc=("This process's place in the MoE tensor-parallel group.")
+    )
+    attn_tp_rank = Derived(
+        doc=("This process's place in the attention tensor-parallel group.")
+    )
+    attn_cp_rank = Derived(
+        doc=("This process's place in the attention context-parallel group.")
+    )
+    dcp_rank = Derived(
+        doc=("This process's place in the decode context-parallel group.")
+    )
+    attn_dcp_rank = Derived(
+        doc=(
+            "Decode context-parallel rank inside the attention TP group, "
+            "zero where decode context parallelism is off."
+        )
+    )
+    attn_dp_rank = Derived(
+        doc=(
+            "This process's index in the attention-DP group, computed from "
+            "`tp_rank` when `initialize_dp_attention` runs."
+        )
+    )
+    dp_rank = Derived(
+        doc=(
+            "Which data-parallel replica this process serves, as the data "
+            "parallel controller numbered them at spawn. `None` when there "
+            "is no controller: unlike the other ranks it is a position in "
+            "no group, which is why the spawn states it."
+        )
+    )
+    launch_world_rank = Derived(
+        doc=(
+            "This process's rank in the WORLD group as built. A scale-up "
+            "does not renumber it."
+        )
+    )
+    launch_world_size = Derived(
+        fn="sglang.srt.runtime_context.launch_world_size_of",
+        doc="Width the WORLD group was built at -- what a scale-up leaves "
+        "behind rather than updates.",
+    )
+    max_world_size = Derived(
+        fn="sglang.srt.runtime_context.max_world_size_of",
+        doc="Ranks the WORLD group has room for: `--max-ep-size` when set, "
+        "otherwise the launch width.",
+    )
+    world_group = Derived(doc="The WORLD group.")
+    tp_group = Derived(doc="The tensor-parallel group.")
+    pp_group = Derived(doc="The pipeline group.")
+    moe_ep_group = Derived(doc="The expert-parallel group.")
+    moe_dp_group = Derived(doc="The MoE data-parallel group.")
+    moe_tp_group = Derived(doc="The MoE tensor-parallel group.")
+    attn_tp_group = Derived(doc="The attention tensor-parallel group.")
+    attn_cp_group = Derived(doc="The attention context-parallel group.")
+    shared_experts_tp_group = Derived(doc=("The shared-expert tensor-parallel group."))
+    dcp_group = Derived(doc="The decode context-parallel group.")

python/sglang/srt/distributed/parallel_state.py

@@ -24,10 +24,13 @@ If you only need to use the distributed environment without model/pipeline
 """
 
 import contextlib
+import functools
 import gc
 import logging
 import os
 import pickle
+import sys
+import warnings
 import weakref
 from collections import namedtuple
 from contextlib import contextmanager, nullcontext
@@ -50,7 +53,6 @@ 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 (
-    _validate_parallel,
     derive_parallel_widths,
     get_global_dwdp_manager,
     get_parallel,
@@ -2072,8 +2074,7 @@ _WORLD: Optional[GroupCoordinator] = None
 
 
 def get_world_group() -> GroupCoordinator:
-    assert _WORLD is not None, "world group is not initialized"
-    return _WORLD
+    return get_parallel().world_group
 
 
 def init_world_group(
@@ -2152,43 +2153,37 @@ _DCP: Optional[GroupCoordinator] = None
 # duplicate GroupCoordinator for prefill in PD-Multiplexing
 _PDMUX_PREFILL_TP_GROUP: Optional[GroupCoordinator] = None
 
-_ENABLE_PDMUX_P_TP: bool = False
 
+@contextmanager
+def pdmux_prefill_tp_group():
+    """Run on the prefill stream's own tensor-parallel communicator.
 
-def set_pdmux_status(enable_prefill_multiplexing: bool):
-    global _ENABLE_PDMUX_P_TP
-    _ENABLE_PDMUX_P_TP = enable_prefill_multiplexing
+    PD multiplexing builds a duplicate TP group -- the same ranks, a second
+    communicator -- so prefill and decode can occupy separate streams without
+    serialising on one. Nothing about the topology differs, so the scope states
+    the handle and nothing else.
+    """
+    assert _PDMUX_PREFILL_TP_GROUP is not None, (
+        "tensor model parallel group for PD-Multiplexing Prefill is not initialized"
+    )
+    with get_parallel().override(tp_group=_PDMUX_PREFILL_TP_GROUP):
+        yield
 
 
 def get_tp_group() -> GroupCoordinator:
-    if _ENABLE_PDMUX_P_TP:
-        assert _PDMUX_PREFILL_TP_GROUP is not None, (
-            "tensor model parallel group for PD-Multiplexing Prefill is not initialized"
-        )
-        return _PDMUX_PREFILL_TP_GROUP
-    assert _TP is not None, "tensor model parallel group is not initialized"
-    return _TP
+    return get_parallel().tp_group
 
 
 def get_attn_tp_group() -> GroupCoordinator:
-    assert _ATTN_TP is not None, (
-        "attention tensor model parallel group is not initialized"
-    )
-    return _ATTN_TP
+    return get_parallel().attn_tp_group
 
 
 def get_shared_experts_tp_group() -> GroupCoordinator:
-    assert _SHARED_EXPERTS_TP is not None, (
-        "shared-expert tensor model parallel group is not initialized"
-    )
-    return _SHARED_EXPERTS_TP
+    return get_parallel().shared_experts_tp_group
 
 
 def get_attn_cp_group() -> GroupCoordinator:
-    assert _ATTN_CP is not None, (
-        "attention context model parallel group is not initialized"
-    )
-    return _ATTN_CP
+    return get_parallel().attn_cp_group
 
 
 def get_dcp_group_no_assert() -> Optional[GroupCoordinator]:
@@ -2196,8 +2191,7 @@ def get_dcp_group_no_assert() -> Optional[GroupCoordinator]:
 
 
 def get_dcp_group() -> GroupCoordinator:
-    assert _DCP is not None, "decode context parallel group is not initialized"
-    return _DCP
+    return get_parallel().dcp_group
 
 
 _MOE_DP: Optional[GroupCoordinator] = None
@@ -2206,18 +2200,15 @@ _MOE_TP: Optional[GroupCoordinator] = None
 
 
 def get_moe_dp_group() -> GroupCoordinator:
-    assert _MOE_DP is not None, "moe data parallel group is not initialized"
-    return _MOE_DP
+    return get_parallel().moe_dp_group
 
 
 def get_moe_ep_group() -> GroupCoordinator:
-    assert _MOE_EP is not None, "expert model parallel group is not initialized"
-    return _MOE_EP
+    return get_parallel().moe_ep_group
 
 
 def get_moe_tp_group() -> GroupCoordinator:
-    assert _MOE_TP is not None, "expert model parallel group is not initialized"
-    return _MOE_TP
+    return get_parallel().moe_tp_group
 
 
 # kept for backward compatibility
@@ -2233,8 +2224,7 @@ def get_self_pp_group() -> GroupCoordinator:
 
 
 def get_pp_group() -> GroupCoordinator:
-    assert _PP is not None, "pipeline model parallel group is not initialized"
-    return _PP
+    return get_parallel().pp_group
 
 
 # kept for backward compatibility
@@ -2513,6 +2503,10 @@ def init_distributed_environment(
         assert _WORLD.world_size == torch.distributed.get_world_size(), (
             "world group already initialized with a different world size"
         )
+    # Stated here rather than with the groups below it: WORLD is built in this
+    # function, and every group `initialize_model_parallel` builds is placed by
+    # reading it back.
+    get_parallel().override_permanently(world_group=_WORLD)
 
 
 def initialize_model_parallel(
@@ -2692,7 +2686,7 @@ def initialize_model_parallel(
             rank_offset=rank_offset,
             max_world_size=max_world_size,
         )
-        if get_tensor_model_parallel_rank() == 0:
+        if _TP.rank_in_group == 0:
             logger.info(
                 f"DCP enabled, dcp_size={decode_context_parallel_size}, tp_size={tensor_model_parallel_size}"
             )
@@ -2942,10 +2936,32 @@ def initialize_model_parallel(
         )
 
     # The groups just built and the configuration they were built from are two
-    # accounts of one layout. Check them against each other here, where the
-    # disagreement is still attributable, rather than letting a collective run
-    # on the wrong peers.
-    _validate_parallel(get_parallel(), "group build")
+    # accounts of one layout, and this is where they meet: stating a group
+    # checks the identities, so a group built on the wrong peers is refused
+    # here rather than hanging in a collective later.
+    #
+    # A dimension this configuration does not have is left unstated -- `_DCP`
+    # is None without decode context parallelism -- so reading it says the
+    # group was never built, which is what these getters have always said,
+    # rather than handing back a None to fail on at the collective.
+    #
+    # WORLD is not here: it is built and stated by
+    # `init_distributed_environment`, which is what lets every build above
+    # place its group by reading `get_world_group().local_rank`.
+    built = {
+        "tp_group": _TP,
+        "pp_group": _PP,
+        "moe_ep_group": _MOE_EP,
+        "moe_dp_group": _MOE_DP,
+        "moe_tp_group": _MOE_TP,
+        "attn_tp_group": _ATTN_TP,
+        "attn_cp_group": _ATTN_CP,
+        "shared_experts_tp_group": _SHARED_EXPERTS_TP,
+        "dcp_group": _DCP,
+    }
+    get_parallel().override_permanently(
+        **{name: group for name, group in built.items() if group is not None}
+    )
 
 
 def create_custom_parallel_group(
@@ -3041,8 +3057,8 @@ def patch_pipeline_parallel_group(pp_group: GroupCoordinator):
     assert not _PP_STATE_PATCHED, "Should not call when it's already patched"
 
     _PP_STATE_PATCHED = True
-    old_pp_group = get_pp_group()
     global _PP
+    old_pp_group = _PP
     _PP = pp_group
     try:
         # `pp_size` is a configured leaf: unlike the rank and the handle it
@@ -3096,8 +3112,8 @@ def patch_tensor_parallel_group(tp_group: GroupCoordinator, *, owns_attention: b
     assert not _TP_STATE_PATCHED, "Should not call when it's already patched"
 
     _TP_STATE_PATCHED = True
-    old_tp_group = get_tp_group()
     global _TP
+    old_tp_group = _TP
     _TP = tp_group
     narrowed = dict(
         tp_size=tp_group.world_size,
@@ -3437,3 +3453,84 @@ def monkey_patch_vllm_parallel_state(reverse: bool = False):
         setattr(vllm_parallel_state, "get_pp_group", get_pp_group)
         setattr(vllm_parallel_state, "get_tp_group", get_tp_group)
         setattr(vllm_parallel_state, "get_world_group", get_world_group)
+
+
+# --- deprecation ---------------------------------------------------------
+#
+# These getters are the definition of a name, not a second spelling of it.
+# Business code asks `get_parallel()`, which answers by calling them and which
+# a scope can redirect; a call that arrives here directly cannot be redirected,
+# so a draft worker's scope does not reach it. The package that defines them
+# keeps calling them -- a read there would go through the context back into
+# itself -- so the warning fires only for callers outside it, and once per
+# name, because the point is to name the replacement rather than to fill a log.
+_EXEMPT_CALLERS = ("sglang.srt.distributed.",)
+
+# Which context name each getter here answers. The shim's own bookkeeping --
+# what a getter was replaced by is of no interest to whoever declares the field
+# -- so it is written next to the warning that uses it.
+_CONTEXT_NAME_OF = {
+    "get_world_group": "world_group",
+    "get_tp_group": "tp_group",
+    "get_pp_group": "pp_group",
+    "get_moe_ep_group": "moe_ep_group",
+    "get_moe_dp_group": "moe_dp_group",
+    "get_moe_tp_group": "moe_tp_group",
+    "get_attn_tp_group": "attn_tp_group",
+    "get_attn_cp_group": "attn_cp_group",
+    "get_shared_experts_tp_group": "shared_experts_tp_group",
+    "get_dcp_group": "dcp_group",
+    "get_world_size": "launch_world_size",
+    "get_world_rank": "launch_world_rank",
+    "get_tensor_model_parallel_rank": "tp_rank",
+    "get_pipeline_model_parallel_rank": "pp_rank",
+    "get_moe_expert_parallel_rank": "moe_ep_rank",
+    "get_moe_data_parallel_rank": "moe_dp_rank",
+    "get_moe_tensor_parallel_rank": "moe_tp_rank",
+    "get_attn_tensor_model_parallel_rank": "attn_tp_rank",
+    "get_attn_context_model_parallel_rank": "attn_cp_rank",
+    "get_dcp_rank": "dcp_rank",
+}
+# The width getters read a built group; the context answers the same names from
+# the configuration. Those are one answer rather than two only for the groups
+# the build checks against the configuration -- `_WIDTH_AND_GROUP` in
+# `runtime_context` -- so only those are listed here. `moe_dp`, `moe_tp` and
+# `dcp` are not on that list and are deliberately absent: the MoE-DP group is
+# the attention-CP group when the latter is wider, and the other two are simply
+# not pinned yet.
+_CONTEXT_NAME_OF["get_tensor_model_parallel_world_size"] = "tp_size"
+_CONTEXT_NAME_OF["get_attn_tensor_model_parallel_world_size"] = "attn_tp_size"
+_CONTEXT_NAME_OF["get_attn_context_model_parallel_world_size"] = "attn_cp_size"
+_CONTEXT_NAME_OF["get_pipeline_model_parallel_world_size"] = "pp_size"
+_CONTEXT_NAME_OF["get_moe_expert_parallel_world_size"] = "moe_ep_size"
+
+_ALREADY_WARNED: set = set()
+
+
+def _warn_if_called_from_outside(name: str, replacement: str):
+    def decorate(fn):
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            if name not in _ALREADY_WARNED:
+                caller = sys._getframe(1).f_globals.get("__name__", "")
+                if not caller.startswith(_EXEMPT_CALLERS):
+                    _ALREADY_WARNED.add(name)
+                    warnings.warn(
+                        f"{name}() is deprecated; read "
+                        f"get_parallel().{replacement} instead, which answers the "
+                        "same thing and can be redirected by a scope",
+                        DeprecationWarning,
+                        stacklevel=2,
+                    )
+            return fn(*args, **kwargs)
+
+        return wrapper
+
+    return decorate
+
+
+for _name, _replacement in _CONTEXT_NAME_OF.items():
+    _fn = globals().get(_name)
+    if _fn is not None:
+        globals()[_name] = _warn_if_called_from_outside(_name, _replacement)(_fn)
+del _name, _replacement, _fn

python/sglang/srt/mem_cache/pool_host/base.py

@@ -58,7 +58,7 @@ def ranks_per_host() -> int:
         return 1
     try:
         launch_world_size = get_parallel().launch_world_size
-    except AssertionError:
+    except (RuntimeError, ValueError):
         return 1
     if launch_world_size == 1:
         return 1
@@ -104,7 +104,7 @@ def sync_fixed_hicache_size(size: int, host_size: int) -> int:
         from sglang.srt.runtime_context import get_parallel
 
         pp_group = get_parallel().pp_group
-    except AssertionError:
+    except RuntimeError:
         return size
 
     if pp_group.world_size <= 1:

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py

@@ -41,7 +41,6 @@ from sglang.srt.compilation import torch_compile_decoration
 from sglang.srt.compilation.torch_compile_decoration import set_torch_compile_config
 from sglang.srt.distributed.parallel_state import (
     graph_capture,
-    set_pdmux_status,
 )
 from sglang.srt.dllm.config import DllmConfig
 from sglang.srt.environ import envs
@@ -1053,7 +1052,6 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                     with self.backend.capture_session(self.stream):
                         self._capture_one_stream()
             else:
-                set_pdmux_status(False)
                 for i, sg in enumerate(self.stream_groups):
                     with (
                         graph_capture(stream=sg[1]) as graph_capture_context,

python/sglang/srt/model_loader/loader.py

@@ -2073,7 +2073,7 @@ class PreshardedModelLoader(DefaultModelLoader):
         try:
             g = get_parallel().world_group
             return g.rank_in_group, g.world_size
-        except (AssertionError, AttributeError):
+        except (AssertionError, AttributeError, RuntimeError):
             return 0, 1
 
     @staticmethod
@@ -2082,7 +2082,7 @@ class PreshardedModelLoader(DefaultModelLoader):
 
         try:
             get_parallel().world_group.barrier()
-        except (AssertionError, AttributeError):
+        except (AssertionError, AttributeError, RuntimeError):
             pass
 
     @staticmethod

python/sglang/srt/models/gigachat35.py

@@ -23,7 +23,6 @@ from transformers import PretrainedConfig
 
 import sglang.srt.models.deepseek_v2 as deepseek_v2
 from sglang.srt.configs.gigachat35 import GigaChat35Config
-from sglang.srt.distributed import get_pp_group
 from sglang.srt.layers.communicator import get_attn_tp_context
 from sglang.srt.layers.layernorm import GemmaRMSNorm, RMSNorm
 from sglang.srt.layers.linear import ColumnParallelLinear
@@ -515,7 +514,7 @@ class GigaChat35Model(nn.Module):
         self.config = config
         self.padding_idx = getattr(config, "pad_token_id", None)
         self.vocab_size = config.vocab_size
-        self.pp_group = get_pp_group()
+        self.pp_group = get_parallel().pp_group
 
         if self.pp_group.is_first_rank:
             self.embed_tokens = VocabParallelEmbedding(
@@ -605,7 +604,7 @@ class GigaChat35ForCausalLM(DeepseekV2WeightLoaderMixin, nn.Module):
         nn.Module.__init__(self)
         self.config = config
         self.quant_config = quant_config
-        self.pp_group = get_pp_group()
+        self.pp_group = get_parallel().pp_group
         self.tp_size = get_parallel().tp_size
         self.num_fused_shared_experts = 0
 

python/sglang/srt/models/gigachat35_mtp.py

@@ -18,7 +18,6 @@ from typing import Iterable, Optional
 import torch
 from torch import nn
 
-from sglang.srt.distributed import get_pp_group
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.vocab_parallel_embedding import (
@@ -140,7 +139,7 @@ class GigaChat35ForCausalLMNextN(DeepseekV2WeightLoaderMixin, nn.Module):
         nn.Module.__init__(self)
         self.config = config
         self.quant_config = quant_config
-        self.pp_group = get_pp_group()
+        self.pp_group = get_parallel().pp_group
         self.tp_size = get_parallel().tp_size
         self.num_fused_shared_experts = 0
         self.draft_model_idx = draft_model_idx or 0

python/sglang/srt/multiplex/multiplexing_mixin.py

@@ -11,7 +11,7 @@ import torch
 import torch.distributed as dist
 from torch.cuda.streams import ExternalStream
 
-from sglang.srt.distributed.parallel_state import set_pdmux_status
+from sglang.srt.distributed.parallel_state import pdmux_prefill_tp_group
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.multiplex.pdmux_context import (
     get_current_stream_idx,
@@ -21,7 +21,7 @@ from sglang.srt.multiplex.pdmux_context import (
     load_pdmux_config,
     set_current_stream_idx,
 )
-from sglang.srt.runtime_context import get_disagg
+from sglang.srt.runtime_context import get_device, get_disagg, get_parallel
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import ScheduleBatch
@@ -37,7 +37,7 @@ class SchedulerMultiplexMixin:
 
         # for pd_multiplexing, Init stream_groups, exclude normal stream for prefill only and decode only
         self.pdmux_config = load_pdmux_config(get_disagg().pdmux_config_path)
-        initialize_stream_groups(self.gpu_id, self.pdmux_config)
+        initialize_stream_groups(get_device().gpu_id, self.pdmux_config)
         self.stream_groups = get_stream_groups()
         self.sm_counts = get_sm_counts()
         self.real_sm_group_num = len(self.stream_groups)
@@ -113,12 +113,10 @@ class SchedulerMultiplexMixin:
 
         while True:
             with torch.cuda.stream(decode_stream):
-                set_pdmux_status(False)
                 self.ingest_requests()
                 running_batch = self.running_batch
 
-            with torch.cuda.stream(prefill_stream):
-                set_pdmux_status(True)
+            with torch.cuda.stream(prefill_stream), pdmux_prefill_tp_group():
                 sm_count = self.sm_counts[stream_idx][0]
                 if not wait_prefill_kernel_done:
                     created, running_batch = self.update_split_prefill_batch(
@@ -128,7 +126,6 @@ class SchedulerMultiplexMixin:
                     adjust_stream_group = created or adjust_stream_group
 
             with torch.cuda.stream(decode_stream):
-                set_pdmux_status(False)
                 running_batch = self.update_running_batch(running_batch)
                 self.running_batch = running_batch
                 adjust_stream_group = adjust_stream_group or (
@@ -152,15 +149,13 @@ class SchedulerMultiplexMixin:
                 )
 
             with torch.cuda.stream(decode_stream):
-                set_pdmux_status(False)
                 # process decode batch
                 if running_batch and not running_batch.is_empty():
                     decode_result = self.run_batch(running_batch)
                     decode_done = True
                 else:
                     decode_done = False
-            with torch.cuda.stream(prefill_stream):
-                set_pdmux_status(True)
+            with torch.cuda.stream(prefill_stream), pdmux_prefill_tp_group():
                 if (
                     self.split_prefill_batch
                     and not self.split_prefill_batch.is_empty()
@@ -197,13 +192,11 @@ class SchedulerMultiplexMixin:
                     prefill_done = False
 
             with torch.cuda.stream(decode_stream):
-                set_pdmux_status(False)
                 decode_stream.synchronize()
                 if decode_done:
                     self.process_batch_result(running_batch, decode_result)
 
-            with torch.cuda.stream(prefill_stream):
-                set_pdmux_status(True)
+            with torch.cuda.stream(prefill_stream), pdmux_prefill_tp_group():
                 if prefill_done and self.split_prefill_batch.split_prefill_finished:
                     wait_prefill_kernel_done = True
                     prefill_exe_done_flag = prefill_exe_done.query()
@@ -214,7 +207,7 @@ class SchedulerMultiplexMixin:
                     )
 
                     self.tp_cpu_group.allreduce(flags, dist.ReduceOp.SUM).wait()
-                    if flags.item() == self.tp_size:
+                    if flags.item() == get_parallel().tp_size:
                         self.process_batch_result(
                             self.split_prefill_batch, prefill_result
                         )

python/sglang/srt/runtime_context.py

@@ -110,119 +110,16 @@ def _parallel_config_leaves() -> frozenset:
     )
 
 
-# Ranks and group handles: the names no configuration carries, each with the
-# canonical getter that answers it live. This table is their declaration, the
-# way `arg_groups/fields/parallel.py` is the leaves' and `Derived` is the
-# widths'. `None` marks a name only a stamp can answer: no coordinator knows
-# this process's attention-DP rank.
+# Ranks and group handles: the names no configuration carries. This table is
+# their declaration, the way `arg_groups/fields/parallel.py` is the leaves' and
+# `Derived` is the widths'. A group handle names the getter that owns it,
+# because the module that builds the groups is where it lives; a rank is a
+# position in one of those groups, so it is read off the handle. `None` marks a
+# name only a stamp can answer: no coordinator knows this process's
+# attention-DP rank.
 _MISSING_READ = object()
 
 
-class Live(msgspec.Struct, frozen=True):
-    """How a rank / group / world width is answered, and what it means.
-
-    `source` is the canonical getter's name in `parallel_state`, a callable
-    taking the context, or `None` for a name only a stamp can answer.
-
-    Most entries in the table below are a bare getter name: a rank or a group
-    handle is its own explanation. This shape is for a name whose meaning is
-    not in its getter, and it carries the prose with the declaration rather
-    than in a second table keyed by the same names.
-    """
-
-    source: Any = None
-    doc: str = ""
-    # For a stamp-only name (`source=None`): what a reader should be told when
-    # nothing has stamped it. These names have no fallback by construction, so
-    # the message is the only thing pointing at what did not happen.
-    unstamped: str = ""
-
-
-_LIVE_READS: dict = {
-    # Two widths of the WORLD group: what it was built at, and what it has
-    # room for. Both are properties of the group itself. How much of that room
-    # is currently serving is elastic-EP state, owned by `ElasticEPStateManager`
-    # and asked of it directly -- a width that lives somewhere else does not
-    # become a WORLD fact by being readable from here.
-    "launch_world_size": Live(
-        source="get_world_size",
-        doc=(
-            "Width the WORLD group was built at: `len(ranks)`, frozen when the "
-            "coordinator was constructed. What every startup reader wants -- "
-            "memory accounting, KV cache sizing, graph capture, weight loading "
-            "-- and what a scale-up leaves behind rather than updates."
-        ),
-    ),
-    "max_world_size": Live(
-        source=lambda self: self.max_ep_size or self.launch_world_size,
-        doc=(
-            "Ranks the WORLD group has room for: `--max-ep-size` when it is "
-            "set, otherwise the launch width. This is the ceiling the process "
-            "group was pre-allocated to -- mooncake sizes its active-rank mask "
-            "to it -- which is why `init_distributed_environment` takes it "
-            "under this name. Whether the group can grow at all is a separate "
-            "question, answered by the leaf being set rather than by this width."
-        ),
-    ),
-    "launch_world_rank": Live(
-        source="get_world_rank",
-        doc=(
-            "This process's rank in the WORLD group as built. Frozen with the "
-            "coordinator, exactly like `launch_world_size`, and named for the "
-            "same reason: a scale-up does not renumber it."
-        ),
-    ),
-    "tp_rank": "get_tensor_model_parallel_rank",
-    "pp_rank": "get_pipeline_model_parallel_rank",
-    "moe_ep_rank": "get_moe_expert_parallel_rank",
-    "moe_dp_rank": "get_moe_data_parallel_rank",
-    "moe_tp_rank": "get_moe_tensor_parallel_rank",
-    "attn_tp_rank": "get_attn_tensor_model_parallel_rank",
-    "attn_cp_rank": "get_attn_context_model_parallel_rank",
-    "dcp_rank": "get_dcp_rank",
-    "attn_dcp_rank": lambda self: self.dcp_rank if self.dcp_enabled else 0,
-    "attn_dp_rank": Live(
-        source=None,
-        doc=(
-            "This process's index in the attention-DP group. Computed from "
-            "`tp_rank` when the attention topology is initialized, and moved "
-            "by an elastic scale-up, so no coordinator can answer it."
-        ),
-        unstamped=(
-            "it is computed from this process's `tp_rank` when the attention "
-            "topology is initialized, so a process that never ran "
-            "`initialize_dp_attention` has no answer to give"
-        ),
-    ),
-    "dp_rank": Live(
-        source=None,
-        doc=(
-            "Which data-parallel replica this process serves, as the data "
-            "parallel controller numbered them at spawn. `None` when there is "
-            "no controller. Unlike `attn_dp_rank` and `moe_dp_rank` it is not "
-            "a position in any process group -- no group has one member per "
-            "replica -- which is why nothing can derive it and the spawn "
-            "states it instead."
-        ),
-        unstamped=(
-            "it is a spawn identity, handed to `publish(..., ranks=...)` by "
-            "the process entry; a process that published without a rank "
-            "bundle has no replica index to report"
-        ),
-    ),
-    "world_group": "get_world_group",
-    "tp_group": "get_tp_group",
-    "pp_group": "get_pp_group",
-    "moe_ep_group": "get_moe_ep_group",
-    "moe_dp_group": "get_moe_dp_group",
-    "moe_tp_group": "get_moe_tp_group",
-    "attn_tp_group": "get_attn_tp_group",
-    "attn_cp_group": "get_attn_cp_group",
-    "shared_experts_tp_group": "get_shared_experts_tp_group",
-    "dcp_group": "get_dcp_group",
-}
-
-
 @functools.lru_cache(maxsize=1)
 def _parallel_fields() -> frozenset:
     """Every name `ParallelContext` answers for, read from the declarations.
@@ -232,8 +129,8 @@ def _parallel_fields() -> frozenset:
 
     * configured leaves -- the `parallel` namespace of the record;
     * derived widths -- the `Derived` declarations beside those leaves;
-    * ranks and group handles -- `_LIVE_READS`, which is where they are
-      declared because no configuration carries them.
+    * ranks and group handles -- declared beside the leaves with no `fn`,
+      because nothing computes them; they are written at runtime.
 
     The set is the union of those three, so `override()` cannot refuse a name
     the class answers for.
@@ -244,7 +141,7 @@ def _parallel_fields() -> frozenset:
     derived = {
         name for name, decl in vars(Parallel).items() if isinstance(decl, Derived)
     }
-    return frozenset(_parallel_config_leaves() | derived | set(_LIVE_READS))
+    return frozenset(_parallel_config_leaves() | derived)
 
 
 def derive_attention_widths(
@@ -331,6 +228,7 @@ def derive_spawn_ranks(
             % (tp_size // moe_dp_size)
             // (tp_size // moe_dp_size // moe_ep_size)
         ),
+        "moe_tp_rank": tp_rank % (tp_size // moe_dp_size // moe_ep_size),
     }
 
 
@@ -404,6 +302,24 @@ def parallel_widths_of(cfg: Any) -> dict:
     )
 
 
+def launch_world_size_of(cfg: Any):
+    """`launch_world_size`, computed at publish.
+
+    The width `bootstrap` builds the WORLD at: one rank per pipeline stage of
+    each tensor-parallel group, above the offset a scale joiner comes in at --
+    zero for everyone else, which is the convention `spawn_world_rank` uses for
+    the same arithmetic. A scale-up does not move it, which is the point of the
+    name.
+    """
+    return cfg.ep_join_rank_offset + cfg.tp_size * cfg.pp_size
+
+
+def max_world_size_of(cfg: Any):
+    """`max_world_size`, computed at publish. The ceiling the group is
+    pre-allocated to: `--max-ep-size` when set, otherwise the launch width."""
+    return cfg.max_ep_size or launch_world_size_of(cfg)
+
+
 def attn_tp_size_of(cfg: Any):
     """`attn_tp_size`, computed at publish. See `parallel_widths_of`."""
     return parallel_widths_of(cfg)["attn_tp_size"]
@@ -586,24 +502,22 @@ def _validate_parallel(parallel, source: str) -> None:
 class ParallelContext:
     """Parallel-topology namespace: one spelling per name.
 
-    Ranks and group handles are read-through ``@property`` over the canonical
-    getters, so they answer with the **live** process groups and raise before
-    distributed init. Every other name — ``tp_size`` and its size siblings
-    included, alongside config-only leaves such as ``nccl_port`` — is answered
-    from the published ``parallel`` bag, in any process at any point after
-    publish.
+    Every name is answered by a lookup, never by asking a process group. A
+    configured leaf and a width derived from one come off the published
+    ``parallel`` bag; a rank is written by ``publish`` from the spawn bundle,
+    and a group handle by ``initialize_model_parallel`` as it builds them. A
+    read before the write that answers it says which write is missing rather
+    than deriving a number from whatever is installed -- the two answer
+    different questions, and a plausible wrong rank surfaces as a hang in a
+    collective far from here.
 
-    A size is read from the configuration because the groups are built at
-    exactly the configured widths. Two things do not follow that rule and are
-    asked of the group itself: ``initialize_model_parallel`` aliases ``_MOE_DP``
-    to ``_ATTN_CP`` when ``attn_cp_size > moe_dp_size``, so a reader that means
-    the MoE communicator's width calls ``get_moe_cp_size()``; and
-    ``patch_tensor_parallel_group`` runs a scope under a different TP group,
-    which it declares by overriding ``tp_size``, ``tp_rank`` and ``tp_group``
-    for its duration. Elastic EP is a third case, and it needs no rule here: it
-    scales ``ep_size`` / ``dp_size`` on the published bag while the group
-    coordinators keep the width they were constructed with, so the two are
-    different names rather than two answers to one name.
+    That makes a scope a matter of stating names: ``patch_tensor_parallel_group``
+    runs a draft worker under a different TP group by overriding the members it
+    changes, and PD multiplexing points ``tp_group`` at the prefill
+    communicator the same way. Elastic EP needs no rule at all: it scales
+    ``ep_size`` / ``dp_size`` on the published bag while ``launch_world_size``
+    keeps the width the groups were built at, so the two are different names
+    rather than two answers to one name.
     """
 
     __slots__ = ("_overrides", "_stamp", "_config")
@@ -624,10 +538,9 @@ class ParallelContext:
     def _read(self, name):
         """The one read path, for every kind of name in the namespace.
 
-        Scoped override, then the permanent stamp, then what the name is
-        answered by when nobody has stated it: the published leaf for a
-        configured value or a derived width, the canonical getter for a rank
-        or a group handle.
+        Scoped override, then the permanent stamp, then the published bag.
+        A rank or a group handle is on no bag, so once those three are out the
+        name has not been written yet and the read says so.
 
         The two override maps stay separate because they are taken down by
         different things -- a `with` block and `clear_stamp()` -- and
@@ -643,21 +556,17 @@ class ParallelContext:
         config = self._config
         if config is not None and name in config._fields:
             return getattr(config, name)
-        live = _LIVE_READS.get(name, _MISSING_READ)
-        if live is not _MISSING_READ:
-            source = live.source if isinstance(live, Live) else live
-            if isinstance(source, str):
-                return getattr(_ps(), source)()
-            if source is not None:
-                return source(self)
-            why = live.unstamped if isinstance(live, Live) else ""
-            raise RuntimeError(
-                f"parallel name {name!r} is not available: "
-                + (why or "nothing has stamped it in this process")
-            )
         if config is None and name in _parallel_config_leaves():
             raise ValueError("config namespace 'parallel' not published")
-        if name in _derived_widths():
+        declared = _derived_widths().get(name)
+        if declared is not None and not declared.fn:
+            raise RuntimeError(
+                f"parallel name {name!r} has not been written in this process. "
+                + declared.doc
+                + f" Write it by publishing a rank bundle or building the groups, "
+                f"or state it with get_parallel().override({name}=...)"
+            )
+        if declared is not None:
             raise RuntimeError(
                 f"derived parallel width {name!r} is not available: it is computed "
                 "from the configured leaves at publish, and permanently corrected "
@@ -731,33 +640,24 @@ def _derived_widths() -> dict:
 def _install_parallel_properties() -> None:
     """Give `ParallelContext` a property per name that is not a config leaf.
 
-    The quotients are declared in `arg_groups/fields/parallel.py`, beside the
-    leaves they are computed from; the ranks and group handles are declared in
-    `_LIVE_READS`, because no configuration carries them. Properties rather
-    than names left to `__getattr__` because the class surface is what the
-    guards introspect -- `hasattr(ParallelContext, "tp_group")` and
-    `vars(ParallelContext)` are how the tests check the set from the class
-    side -- and because each one carries its `Derived.doc`.
+    Every name the namespace answers that is not a plain leaf: the quotients,
+    and the ranks and group handles declared beside them with no `fn`.
+
+    Properties rather than names left to `__getattr__` because the class
+    surface is what the guards introspect -- `hasattr(ParallelContext,
+    "tp_size")` and `vars(ParallelContext)` are how the tests read the set from
+    the class side -- and because each one carries its `Derived.doc`.
 
     Every one of them resolves through `_read`, so there is a single priority
     chain rather than one per kind of name.
     """
-    docs = {name: decl.doc for name, decl in _derived_widths().items()}
-    docs.update(
-        {
-            name: live.doc
-            for name, live in _LIVE_READS.items()
-            if isinstance(live, Live) and live.doc
-        }
-    )
-
-    for name in list(_derived_widths()) + list(_LIVE_READS):
+    for name, decl in _derived_widths().items():
 
         def getter(self, _name=name):
             return self._read(_name)
 
         getter.__name__ = name
-        getter.__doc__ = docs.get(name)
+        getter.__doc__ = decl.doc
         setattr(ParallelContext, name, property(getter))
 
 
@@ -1912,6 +1812,12 @@ def publish(
             ),
         )
     _CONTEXT._publish_role = role
+    # Zero for every process when decode context parallelism is off, which is a
+    # fact about the configuration and not about the spawn -- so it answers
+    # without a rank bundle, the way it did when it stood for a group that was
+    # never built. With DCP on it is a position, and the bundle below states it.
+    if not _CONTEXT.parallel.dcp_enabled:
+        _CONTEXT.parallel.override_permanently(attn_dcp_rank=0)
     if ranks is not None and ranks.gpu_id is not None:
         _CONTEXT.override("spawn", gpu_id=ranks.gpu_id)
     if ranks is not None:
@@ -1931,18 +1837,25 @@ def publish(
             moe_dp_size=parallel.moe_dp_size,
             moe_ep_size=parallel.moe_ep_size,
         )
-        # `moe_dp_rank` is a different quantity when the MoE-DP group is
-        # aliased to the attention-CP one: the group answers the CP index,
-        # while this computes the MoE-DP index. Leave it to the group there, so
-        # one name does not mean two things.
+        # `initialize_model_parallel` aliases the MoE-DP group to the
+        # attention-CP one when the CP dimension is the wider of the two, so
+        # this process's place in it is its CP index rather than the MoE-DP
+        # index the arithmetic above gives.
         if parallel.moe_dp_size < parallel.attn_cp_size:
-            placement.pop("moe_dp_rank")
+            placement["moe_dp_rank"] = placement["attn_cp_rank"]
         # `dp_rank` is recorded whatever it is, None included: replicas are
         # separate WORLD groups, so no rank implies it and `None` is the answer
         # "no controller" rather than an absence.
         placement["dp_rank"] = ranks.dp_rank
         placement["launch_world_rank"] = ranks.world_rank
         placement.update(_attention_ranks(parallel, placement["tp_rank"]))
+        # A DCP group is a contiguous slice of a TP group, so this process's
+        # place in one is its TP rank folded by that width. `attn_dcp_rank` is
+        # the same number, and zero where decode context parallelism is off, so
+        # a reader does not have to ask whether it is on first.
+        if parallel.dcp_enabled:
+            placement["dcp_rank"] = placement["tp_rank"] % parallel.dcp_size
+        placement["attn_dcp_rank"] = placement.get("dcp_rank", 0)
         # One stamp, not two: the identities are checked on every write, and a
         # half-placed process satisfies none of them.
         parallel.override_permanently(**placement)

python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py

@@ -24,8 +24,11 @@ from sglang.srt.runtime_context import get_context, get_parallel
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 
 # Unit tests run without distributed initialization. Backends that size buffers by
-# attention tensor-parallel degree should see the single-rank default.
-_parallel_override = get_parallel().override(attn_tp_size=1)
+# attention tensor-parallel degree should see the single-rank default, and a
+# backend that places itself in the decode context-parallel group needs a
+# position: nothing publishes here, so there is no configuration to derive the
+# zero these tests run at from.
+_parallel_override = get_parallel().override(attn_tp_size=1, attn_dcp_rank=0)
 _parallel_override.__enter__()
 
 DEFAULT_HEAD_DIM = 16

python/sglang/test/test_utils.py

@@ -2095,7 +2095,13 @@ def publish_build_topology(*, world_rank: int = 0, **server_args_fields):
     build outlive any block, so the configuration describing them has to as
     well. Callers that tear the groups down are already resetting the process.
     """
-    from sglang.srt.runtime_context import SpawnRanks, publish, reset_context
+    from sglang.srt.distributed import parallel_state
+    from sglang.srt.runtime_context import (
+        SpawnRanks,
+        get_parallel,
+        publish,
+        reset_context,
+    )
     from sglang.srt.server_args import ServerArgs
 
     reset_context()
@@ -2104,6 +2110,13 @@ def publish_build_topology(*, world_rank: int = 0, **server_args_fields):
         role="test",
         ranks=SpawnRanks(world_rank=world_rank),
     )
+    # Callers that go on to build groups have already run
+    # `init_distributed_environment`, which states the WORLD group -- and the
+    # build below places every group it creates by reading that back. The reset
+    # above drops it, so hand it over again: publishing a configuration does not
+    # unbuild a process group.
+    if parallel_state._WORLD is not None:
+        get_parallel().override_permanently(world_group=parallel_state._WORLD)
 
 
 _GPU_IDLE_TIMEOUT_SECS = 30.0