From 73f071db52e87f56f9d44546fc81b7dbc2aaa72e Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:25:32 -0700 Subject: [PATCH] Deprecate the parallel getters the context answers, and ratchet them shut (#40342) --- .../runtime/distributed/parallel_state.py | 30 + ...est_component_accuracy_parallel_runtime.py | 14 +- python/sglang/srt/arg_groups/arg_utils.py | 7 +- .../sglang/srt/arg_groups/fields/parallel.py | 70 ++ .../sglang/srt/distributed/parallel_state.py | 183 ++++-- python/sglang/srt/mem_cache/pool_host/base.py | 4 +- .../runner/decode_cuda_graph_runner.py | 2 - python/sglang/srt/model_loader/loader.py | 4 +- python/sglang/srt/models/gigachat35.py | 5 +- python/sglang/srt/models/gigachat35_mtp.py | 3 +- .../srt/multiplex/multiplexing_mixin.py | 21 +- python/sglang/srt/runtime_context.py | 253 +++----- .../attention_methods/dense_attention.py | 7 +- python/sglang/test/test_utils.py | 15 +- test/registered/dcp/test_dcp_layout_unit.py | 4 +- .../dcp/test_reduce_scatter_along_dim.py | 4 + .../communication/bench_custom_all_reduce.py | 4 + .../bench_symm_mem_all_gather.py | 4 + .../communication/bench_tp_qknorm.py | 4 + .../communication/test_custom_all_reduce.py | 4 + .../communication/test_symm_mem_all_gather.py | 4 + .../ops/communication/test_tp_qknorm.py | 4 + .../kernels/ops/kimi_k3/test_ar_fusion.py | 4 + .../kernels/ops/kimi_k3/test_collectives.py | 4 + .../unit/disaggregation/test_encode_server.py | 16 +- .../test_register_to_bootstrap.py | 19 +- ...flashattention_pa_swa_prefill_lens_size.py | 3 + .../layers/moe/test_mega_moe_deepgemm_api.py | 8 +- .../test_disagg_idle_step_counters.py | 12 +- .../test_mamba_donated_alloc_ratio.py | 4 +- .../unit/mem_cache/test_mem_pool_host.py | 9 +- .../unit/mem_cache/test_unified_mla_views.py | 6 +- .../runner_utils/test_graph_pool_borrow.py | 8 +- .../model_executor/test_kv_canary_headroom.py | 7 +- .../model_executor/test_pool_configurator.py | 6 +- .../unit/model_loader/test_modelopt_export.py | 4 +- .../unit/model_loader/test_modelopt_loader.py | 4 +- .../model_loader/test_prefetch_checkpoints.py | 6 +- .../model_loader/test_presharded_loader.py | 10 +- .../test_transformers_fallback.py | 8 +- .../multimodal/test_vit_cuda_graph_runner.py | 16 +- .../test_resolution_declarations.py | 16 +- .../unit/spec/test_dflash_logits.py | 4 + test/registered/unit/test_runtime_context.py | 610 +++++++++++++----- 44 files changed, 933 insertions(+), 501 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index 5123d0dbb..f3be2f255 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/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 diff --git a/python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py b/python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py index a7e6d0ac5..9539f9e87 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py +++ b/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(): diff --git a/python/sglang/srt/arg_groups/arg_utils.py b/python/sglang/srt/arg_groups/arg_utils.py index 1fbf53b30..4e25707d6 100644 --- a/python/sglang/srt/arg_groups/arg_utils.py +++ b/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`` diff --git a/python/sglang/srt/arg_groups/fields/parallel.py b/python/sglang/srt/arg_groups/fields/parallel.py index 1b4df5bdf..378a52220 100644 --- a/python/sglang/srt/arg_groups/fields/parallel.py +++ b/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.") diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 4c02ebb21..f127f2fa5 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/pool_host/base.py b/python/sglang/srt/mem_cache/pool_host/base.py index 6d3418193..924429484 100644 --- a/python/sglang/srt/mem_cache/pool_host/base.py +++ b/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: diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index b5dda425f..076ba0c2e 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/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, diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 71ecd9ac5..31f26a33b 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/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 diff --git a/python/sglang/srt/models/gigachat35.py b/python/sglang/srt/models/gigachat35.py index 724791800..0fa4c460d 100644 --- a/python/sglang/srt/models/gigachat35.py +++ b/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 diff --git a/python/sglang/srt/models/gigachat35_mtp.py b/python/sglang/srt/models/gigachat35_mtp.py index 77924d7cf..aa3703414 100644 --- a/python/sglang/srt/models/gigachat35_mtp.py +++ b/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 diff --git a/python/sglang/srt/multiplex/multiplexing_mixin.py b/python/sglang/srt/multiplex/multiplexing_mixin.py index d2f4024fa..b3c0f4b87 100644 --- a/python/sglang/srt/multiplex/multiplexing_mixin.py +++ b/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 ) diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index 81736cabb..da4320175 100644 --- a/python/sglang/srt/runtime_context.py +++ b/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) diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py index 9cbbfad9b..d26824efd 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py +++ b/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 diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index aae6098b0..dac26d3e1 100644 --- a/python/sglang/test/test_utils.py +++ b/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 diff --git a/test/registered/dcp/test_dcp_layout_unit.py b/test/registered/dcp/test_dcp_layout_unit.py index a98ad8325..7c63381f2 100644 --- a/test/registered/dcp/test_dcp_layout_unit.py +++ b/test/registered/dcp/test_dcp_layout_unit.py @@ -134,7 +134,9 @@ class TestFilterDcpLocalChunkKvIndices(CustomTestCase): def test_identity_without_dcp(self): kv = torch.arange(37) - with rc.get_parallel().override(dcp_enabled=False, dcp_size=1, dcp_rank=0): + with rc.get_parallel().override( + dcp_enabled=False, dcp_size=1, dcp_rank=0, attn_dcp_rank=0 + ): self.assertIs( filter_dcp_local_chunk_kv_indices( kv, torch.tensor([0]), torch.tensor([37]) diff --git a/test/registered/dcp/test_reduce_scatter_along_dim.py b/test/registered/dcp/test_reduce_scatter_along_dim.py index 481b46dc4..9bc6fa74b 100644 --- a/test/registered/dcp/test_reduce_scatter_along_dim.py +++ b/test/registered/dcp/test_reduce_scatter_along_dim.py @@ -25,6 +25,7 @@ import torch import torch.distributed as dist import sglang.srt.distributed.parallel_state as ps +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci( @@ -130,6 +131,9 @@ def init_distributed(): local_rank=local_rank, backend="nccl", ) + # The context answers a handle from what was stated on it, not from + # this module global, so a rank stood up by hand says so itself. + get_parallel().override_permanently(world_group=coord) cpu_group = coord.cpu_group nccl_group = coord.device_group diff --git a/test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py b/test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py index d42846f2c..71de22850 100644 --- a/test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py +++ b/test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py @@ -14,6 +14,7 @@ from sglang.kernels.jit.benchmark import marker from sglang.kernels.jit.benchmark.utils import get_benchmark_range, multigpu_bench_main from sglang.kernels.jit.utils import cache_once, is_arch_support_pdl from sglang.kernels.ops.communication.mp import register_comm_cleanup +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci( @@ -61,6 +62,9 @@ def _init_cpu_group() -> dist.ProcessGroup: local_rank=local_rank, backend="nccl", ) + # The context answers a handle from what was stated on it, not from + # this module global, so a rank stood up by hand says so itself. + get_parallel().override_permanently(world_group=coord) atexit.register(dist.destroy_process_group) torch.cuda.set_stream(torch.cuda.Stream()) return coord.cpu_group diff --git a/test/registered/kernels/benchmark/communication/bench_symm_mem_all_gather.py b/test/registered/kernels/benchmark/communication/bench_symm_mem_all_gather.py index e2fb7d7aa..4d3e23892 100644 --- a/test/registered/kernels/benchmark/communication/bench_symm_mem_all_gather.py +++ b/test/registered/kernels/benchmark/communication/bench_symm_mem_all_gather.py @@ -33,6 +33,7 @@ from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import ( all_gather_inner, create_state, ) +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci( @@ -77,6 +78,9 @@ def _init_cpu_group() -> dist.ProcessGroup: local_rank=local_rank, backend="nccl", ) + # The context answers a handle from what was stated on it, not from + # this module global, so a rank stood up by hand says so itself. + get_parallel().override_permanently(world_group=ps._WORLD) atexit.register(dist.destroy_process_group) logging.disable(logging.INFO) torch.cuda.set_stream(torch.cuda.Stream()) diff --git a/test/registered/kernels/benchmark/communication/bench_tp_qknorm.py b/test/registered/kernels/benchmark/communication/bench_tp_qknorm.py index e5f4be66f..0d0452749 100644 --- a/test/registered/kernels/benchmark/communication/bench_tp_qknorm.py +++ b/test/registered/kernels/benchmark/communication/bench_tp_qknorm.py @@ -36,6 +36,7 @@ from sglang.kernels.ops.communication.mp import register_comm_cleanup from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( CustomAllReduceV2, ) +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci( @@ -108,6 +109,9 @@ def _init_cpu_group() -> dist.ProcessGroup: local_rank=local_rank, backend="nccl", ) + # The context answers a handle from what was stated on it, not from + # this module global, so a rank stood up by hand says so itself. + get_parallel().override_permanently(world_group=coord) atexit.register(dist.destroy_process_group) logging.disable(logging.INFO) torch.cuda.set_stream(torch.cuda.Stream()) diff --git a/test/registered/kernels/ops/communication/test_custom_all_reduce.py b/test/registered/kernels/ops/communication/test_custom_all_reduce.py index 7ae2c2366..fd7984128 100644 --- a/test/registered/kernels/ops/communication/test_custom_all_reduce.py +++ b/test/registered/kernels/ops/communication/test_custom_all_reduce.py @@ -37,6 +37,7 @@ from sglang.kernels.ops.communication.mp import register_comm_cleanup from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( CustomAllReduceV2, ) +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.kernels.utils import multigpu_pytest_main @@ -129,6 +130,9 @@ def _init_cpu_group_once() -> dist.ProcessGroup: local_rank=local_rank, backend="nccl", ) + # The context answers a handle from what was stated on it, not from + # this module global, so a rank stood up by hand says so itself. + get_parallel().override_permanently(world_group=coord) atexit.register(dist.destroy_process_group) cpu_group = coord.cpu_group assert isinstance(cpu_group, dist.ProcessGroup) diff --git a/test/registered/kernels/ops/communication/test_symm_mem_all_gather.py b/test/registered/kernels/ops/communication/test_symm_mem_all_gather.py index 5cdd16470..2a2f48d5e 100644 --- a/test/registered/kernels/ops/communication/test_symm_mem_all_gather.py +++ b/test/registered/kernels/ops/communication/test_symm_mem_all_gather.py @@ -31,6 +31,7 @@ from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import ( all_gather_inner, create_state, ) +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.kernels.utils import multigpu_pytest_main @@ -70,6 +71,9 @@ def _init_cpu_group_once() -> dist.ProcessGroup: local_rank=local_rank, backend="nccl", ) + # The context answers a handle from what was stated on it, not from + # this module global, so a rank stood up by hand says so itself. + get_parallel().override_permanently(world_group=ps._WORLD) atexit.register(dist.destroy_process_group) logging.disable(logging.INFO) torch.cuda.set_stream(torch.cuda.Stream()) diff --git a/test/registered/kernels/ops/communication/test_tp_qknorm.py b/test/registered/kernels/ops/communication/test_tp_qknorm.py index 712c24ecd..a70c6b211 100644 --- a/test/registered/kernels/ops/communication/test_tp_qknorm.py +++ b/test/registered/kernels/ops/communication/test_tp_qknorm.py @@ -24,6 +24,7 @@ from sglang.kernels.ops.communication.mp import register_comm_cleanup from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( CustomAllReduceV2, ) +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.kernels.utils import multigpu_pytest_main @@ -92,6 +93,9 @@ def _init_cpu_group_once() -> dist.ProcessGroup: local_rank=local_rank, backend="nccl", ) + # The context answers a handle from what was stated on it, not from + # this module global, so a rank stood up by hand says so itself. + get_parallel().override_permanently(world_group=coord) atexit.register(dist.destroy_process_group) cpu_group = coord.cpu_group assert isinstance(cpu_group, dist.ProcessGroup) diff --git a/test/registered/kernels/ops/kimi_k3/test_ar_fusion.py b/test/registered/kernels/ops/kimi_k3/test_ar_fusion.py index 0bbcadd5a..229bad61d 100644 --- a/test/registered/kernels/ops/kimi_k3/test_ar_fusion.py +++ b/test/registered/kernels/ops/kimi_k3/test_ar_fusion.py @@ -30,6 +30,7 @@ from sglang.kernels.ops.kimi_k3 import all_reduce from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( CustomAllReduceV2, ) +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.kernels.utils import multigpu_pytest_main @@ -65,6 +66,9 @@ def _init_world(): local_rank=local_rank, backend="nccl", ) + # The context answers a handle from what was stated on it, not from + # this module global, so a rank stood up by hand says so itself. + get_parallel().override_permanently(world_group=coord) atexit.register(dist.destroy_process_group) logging.disable(logging.INFO) torch.cuda.set_stream(torch.cuda.Stream()) diff --git a/test/registered/kernels/ops/kimi_k3/test_collectives.py b/test/registered/kernels/ops/kimi_k3/test_collectives.py index dc286d3b7..1ccb32fb9 100644 --- a/test/registered/kernels/ops/kimi_k3/test_collectives.py +++ b/test/registered/kernels/ops/kimi_k3/test_collectives.py @@ -20,6 +20,7 @@ from sglang.kernels.ops.kimi_k3 import ( from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( CustomAllReduceV2, ) +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.kernels.utils import multigpu_pytest_main @@ -53,6 +54,9 @@ def _init_world(): local_rank=local_rank, backend="nccl", ) + # The context answers a handle from what was stated on it, not from + # this module global, so a rank stood up by hand says so itself. + get_parallel().override_permanently(world_group=coord) atexit.register(dist.destroy_process_group) cpu_group = coord.cpu_group assert isinstance(cpu_group, dist.ProcessGroup) diff --git a/test/registered/unit/disaggregation/test_encode_server.py b/test/registered/unit/disaggregation/test_encode_server.py index 757dbfb39..6f0c093db 100644 --- a/test/registered/unit/disaggregation/test_encode_server.py +++ b/test/registered/unit/disaggregation/test_encode_server.py @@ -48,6 +48,7 @@ from sglang.srt.mem_cache.multimodal_cache import ( EmbeddingResult, MultiModalStaticCache, ) +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils.common import safe_pickle_loads from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -1011,10 +1012,7 @@ class TestEncoderDelivery(CustomTestCase): statuses[1].copy_(torch.tensor([400, 1, 0, 0])) with ( - patch( - "sglang.srt.distributed.parallel_state.get_tp_group", - return_value=TPGroup(), - ), + get_parallel().override(tp_group=TPGroup()), patch( "sglang.srt.disaggregation.encoder.server.torch.distributed.all_gather", side_effect=all_gather, @@ -1051,10 +1049,7 @@ class TestEncoderDelivery(CustomTestCase): statuses[1][2] += 1 with ( - patch( - "sglang.srt.distributed.parallel_state.get_tp_group", - return_value=TPGroup(), - ), + get_parallel().override(tp_group=TPGroup()), patch( "sglang.srt.disaggregation.encoder.server.torch.distributed.all_gather", side_effect=all_gather, @@ -1095,10 +1090,7 @@ class TestEncoderDelivery(CustomTestCase): statuses[1].copy_(local_status) with ( - patch( - "sglang.srt.distributed.parallel_state.get_tp_group", - return_value=TPGroup(), - ), + get_parallel().override(tp_group=TPGroup()), patch( "sglang.srt.disaggregation.encoder.server.torch.distributed.all_gather", side_effect=all_gather, diff --git a/test/registered/unit/disaggregation/test_register_to_bootstrap.py b/test/registered/unit/disaggregation/test_register_to_bootstrap.py index 196a5cafe..c7692cbac 100644 --- a/test/registered/unit/disaggregation/test_register_to_bootstrap.py +++ b/test/registered/unit/disaggregation/test_register_to_bootstrap.py @@ -9,7 +9,7 @@ import unittest from unittest.mock import MagicMock, call, patch from sglang.srt.environ import envs -from sglang.srt.runtime_context import get_context +from sglang.srt.runtime_context import get_context, get_parallel from sglang.test.test_utils import CustomTestCase @@ -201,12 +201,12 @@ class TestRegisterToBootstrap(CustomTestCase): self.assertIn("10.0.0.1", url_used) @patch("sglang.srt.disaggregation.common.conn.requests.put") - # The consumer reads the group through `get_parallel()`, which reads - # through to the canonical getter, so that is where the stub belongs. - @patch("sglang.srt.distributed.parallel_state.get_world_group") def test_rust_attention_dp_replicates_complete_topology_across_hosts( - self, mock_world_group, mock_put + self, mock_put ): + # The consumer reads the group through `get_parallel()`, so the + # stub is stated there rather than in the module the build writes. + mock_world_group = MagicMock() success_resp = MagicMock() success_resp.status_code = 200 mock_put.return_value = success_resp @@ -230,9 +230,12 @@ class TestRegisterToBootstrap(CustomTestCase): for dp_rank, tp_rank, host, rank_port, _ in schedulers ] - mock_world_group.return_value.all_gather_object.side_effect = gather_topology + mock_world_group.all_gather_object.side_effect = gather_topology - with envs.SGLANG_RUST_SERVER.override(True): + with ( + get_parallel().override(world_group=mock_world_group), + envs.SGLANG_RUST_SERVER.override(True), + ): for dp_rank, tp_rank, local_ip, _, rust_http_port in schedulers: manager = self._make_manager() manager.attn_dp_size = 2 @@ -277,7 +280,7 @@ class TestRegisterToBootstrap(CustomTestCase): gather_call.args[0]["attn_dp_rank"], gather_call.args[0]["attn_tp_rank"], ) - for gather_call in mock_world_group.return_value.all_gather_object.call_args_list + for gather_call in mock_world_group.all_gather_object.call_args_list ], [(dp, tp) for dp, tp, _, _, _ in schedulers], ) diff --git a/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py b/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py index 2094e8c0c..2ecf5011d 100644 --- a/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py +++ b/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py @@ -93,6 +93,9 @@ def _make_prefill_aware_swa_runner( page_size=1, attn_cp_size=1, tp_size=1, + # The backend still reads the runner's frozen record for these two; + # same single-rank placement, stated where it looks for it. + ps=SimpleNamespace(attn_cp_size=1, tp_size=1), is_draft_worker=False, server_args=server_args, attention_chunk_size=None, diff --git a/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py b/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py index e7ca95663..445d567ac 100644 --- a/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py +++ b/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py @@ -14,6 +14,7 @@ from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel maybe_stub_sgl_kernel() from sglang.srt.layers.moe import mega_moe +from sglang.srt.runtime_context import get_parallel register_cpu_ci(est_time=8, suite="base-a-test-cpu") @@ -271,11 +272,8 @@ class TestDeepGemmMegaMoeApi(CustomTestCase): "init_new", return_value=object(), ), - patch( - "sglang.srt.runtime_context.get_parallel", - return_value=SimpleNamespace( - moe_ep_group=SimpleNamespace(device_group=object()) - ), + get_parallel().override( + moe_ep_group=SimpleNamespace(device_group=object()) ), ): mega_moe._run_mega_routed( diff --git a/test/registered/unit/managers/test_disagg_idle_step_counters.py b/test/registered/unit/managers/test_disagg_idle_step_counters.py index d6aa44aa4..8764e4f97 100644 --- a/test/registered/unit/managers/test_disagg_idle_step_counters.py +++ b/test/registered/unit/managers/test_disagg_idle_step_counters.py @@ -11,6 +11,7 @@ from unittest.mock import Mock, patch from parameterized import parameterized +from sglang.srt.distributed import parallel_state from sglang.srt.managers import scheduler as scheduler_module from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import Scheduler @@ -168,12 +169,21 @@ class TestSchedulerIdleStepCounters(CustomTestCase): ) with ( patch(f"{PDMUX_MODULE}.get_current_stream_idx", return_value=0), - patch(f"{PDMUX_MODULE}.set_pdmux_status"), patch(f"{PDMUX_MODULE}.torch.cuda.empty_cache"), patch( f"{PDMUX_MODULE}.torch.cuda.stream", side_effect=lambda stream: nullcontext(), ), + # The prefill section runs under the duplicate communicator + # `--enable-pdmux` builds, in place of the module flag this + # replaces. The loop has no process groups at all, so stand + # one in: the scope refuses to open without it rather than + # letting prefill quietly share the decode communicator. + patch.object( + parallel_state, + "_PDMUX_PREFILL_TP_GROUP", + SimpleNamespace(world_size=1, rank_in_group=0), + ), ): self.run_and_check( scheduler, diff --git a/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py b/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py index b684d867f..ce4205096 100644 --- a/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py +++ b/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py @@ -297,7 +297,9 @@ class TestPPMambaPoolSizing(unittest.TestCase): server_args=SimpleNamespace(), spec_algorithm=SimpleNamespace(is_none=lambda: True), layer_info=SimpleNamespace(start_layer=start, end_layer=end), - ps=SimpleNamespace(attn_dp_size=1, pp_size=pp_size), + # The runner carries its placement as plain attributes. + attn_dp_size=1, + pp_size=pp_size, hybrid_gdn_config=None, model_config=SimpleNamespace( hf_config=SimpleNamespace(), num_hidden_layers=cls.TOTAL_LAYERS diff --git a/test/registered/unit/mem_cache/test_mem_pool_host.py b/test/registered/unit/mem_cache/test_mem_pool_host.py index 909b8719a..99ec163af 100644 --- a/test/registered/unit/mem_cache/test_mem_pool_host.py +++ b/test/registered/unit/mem_cache/test_mem_pool_host.py @@ -304,16 +304,13 @@ class TestHostMemoryBudget(CustomTestCase): def test_ranks_per_host_divides_world_size_by_nodes(self): # The launcher slices ranks uniformly across nodes, so the co-located # rank count is world_size // nnodes — no hostname collective. - fake_group = unittest.mock.Mock(world_size=16) + # tp_size=16 states the launch width the count divides -- the + # published configuration is where ranks_per_host reads it from. with ( - get_context().override_server_args(nnodes=2), + get_context().override_server_args(nnodes=2, tp_size=16), unittest.mock.patch.object( torch.distributed, "is_initialized", return_value=True ), - unittest.mock.patch( - "sglang.srt.distributed.parallel_state.get_world_group", - return_value=fake_group, - ), ): self.assertEqual(base.ranks_per_host(), 8) diff --git a/test/registered/unit/mem_cache/test_unified_mla_views.py b/test/registered/unit/mem_cache/test_unified_mla_views.py index e295fcedb..0c0347327 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_views.py +++ b/test/registered/unit/mem_cache/test_unified_mla_views.py @@ -248,7 +248,7 @@ class TestUnifiedMLATokenToKVPool(unittest.TestCase): ] = float(layer + 1) with ( - get_parallel().override(dcp_enabled=False), + get_parallel().override(dcp_enabled=False, attn_dcp_rank=0), mock.patch( "sglang.srt.mem_cache.memory_pool.current_platform.synchronize" ), @@ -411,7 +411,9 @@ class TestMambaAllocatorCpuCopyIsPhysical(unittest.TestCase): _FakeKVCache(pool.max_slots("full")), _FakeKVCache(pool.max_slots("mamba")), ) - with get_parallel().override(dcp_enabled=False, attn_dcp_size=1): + with get_parallel().override( + dcp_enabled=False, attn_dcp_size=1, attn_dcp_rank=0 + ): allocator = UnifiedMambaTokenToKVPoolAllocator( unified_buffer=pool, kvcache=kvcache, device=_DEV, page_size=ps ) diff --git a/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py b/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py index 2d49cb46f..ff120edec 100644 --- a/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py +++ b/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py @@ -18,6 +18,7 @@ from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import ( FullCudaGraphBackend, ) from sglang.srt.model_executor.runner_utils import pool +from sglang.srt.runtime_context import get_parallel from sglang.srt.speculative import dflash_utils, dflash_worker_v2, eagle_utils from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2 from sglang.test.ci.ci_register import register_cuda_ci @@ -222,12 +223,7 @@ class TestGraphPoolBorrow(CustomTestCase): "sglang.srt.layers.dp_attention.is_dp_attention_enabled", return_value=False, ), - # `parallel_state`, not the package re-export: a stub on the - # re-export is never consulted. - patch( - "sglang.srt.distributed.parallel_state.get_tp_group", - return_value=tp_group, - ), + get_parallel().override(tp_group=tp_group), patch( "sglang.kernels.ops.speculative.sampling.tree_speculative_sampling_target_only", side_effect=fake_sampling, diff --git a/test/registered/unit/model_executor/test_kv_canary_headroom.py b/test/registered/unit/model_executor/test_kv_canary_headroom.py index 901e6057b..132092ad4 100644 --- a/test/registered/unit/model_executor/test_kv_canary_headroom.py +++ b/test/registered/unit/model_executor/test_kv_canary_headroom.py @@ -11,7 +11,7 @@ from sglang.srt.model_executor.cuda_graph_config import ( ) from sglang.srt.model_executor.model_runner_components import kv_pool_runtime from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig -from sglang.srt.runtime_context import get_context +from sglang.srt.runtime_context import get_context, get_parallel from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -60,9 +60,8 @@ class TestCanaryHeadroom(CustomTestCase): ), ), patch.object(kv_pool_runtime.torch.cuda, "synchronize"), - patch( - "sglang.srt.distributed.parallel_state.get_world_group", - return_value=SimpleNamespace(world_size=1, cpu_group=None), + get_parallel().override( + world_group=SimpleNamespace(world_size=1, cpu_group=None) ), patch.object(kv_pool_runtime, "get_available_gpu_memory", return_value=20), patch.object(kv_pool_runtime, "mambaish_config", return_value=None), diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 7fbf006f8..5f6f8fd67 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -36,7 +36,11 @@ def mock_cpu_env(kv_size=2, tp_size=1, swa_eviction_interval=4): with ( patch("torch._utils._element_size", return_value=kv_size), - get_parallel().override(attn_tp_size=tp_size), + # A width is a whole topology: state the TP siblings the identities + # relate it to, not the attention share alone. + get_parallel().override( + tp_size=tp_size, attn_tp_size=tp_size, moe_tp_size=tp_size + ), envs.SGLANG_SWA_EVICTION_INTERVAL.override(swa_eviction_interval), ): yield diff --git a/test/registered/unit/model_loader/test_modelopt_export.py b/test/registered/unit/model_loader/test_modelopt_export.py index 9f9711ece..25313f282 100644 --- a/test/registered/unit/model_loader/test_modelopt_export.py +++ b/test/registered/unit/model_loader/test_modelopt_export.py @@ -53,9 +53,7 @@ class TestModelOptExport(unittest.TestCase): self.mock_logger.start() # Mock all distributed functions that might be called - self.mock_get_tp_group = patch( - "sglang.srt.distributed.parallel_state.get_tp_group" - ) + self.mock_get_tp_group = patch("sglang.srt.distributed.parallel_state._TP") self.mock_get_tp_group.start() # Mock model parallel initialization check diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index ca2740f26..5bd16c741 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/test/registered/unit/model_loader/test_modelopt_loader.py @@ -82,9 +82,7 @@ class TestModelOptModelLoader(CustomTestCase): self.mock_logger.start() # Mock all distributed functions that might be called - self.mock_get_tp_group = patch( - "sglang.srt.distributed.parallel_state.get_tp_group" - ) + self.mock_get_tp_group = patch("sglang.srt.distributed.parallel_state._TP") self.mock_get_tp_group.start() # Mock model parallel initialization check diff --git a/test/registered/unit/model_loader/test_prefetch_checkpoints.py b/test/registered/unit/model_loader/test_prefetch_checkpoints.py index 13e992be0..c2f4439d2 100644 --- a/test/registered/unit/model_loader/test_prefetch_checkpoints.py +++ b/test/registered/unit/model_loader/test_prefetch_checkpoints.py @@ -25,6 +25,7 @@ from sglang.srt.model_loader.weight_utils import ( fastsafetensors_weights_iterator, safetensors_weights_iterator, ) +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -237,10 +238,7 @@ class TestPrefetchCheckpoints(CustomTestCase): patch("threading.Thread", _InlineThread), patch("concurrent.futures.ThreadPoolExecutor", _InlineExecutor), patch("concurrent.futures.wait", side_effect=_wait_all), - patch( - "sglang.srt.distributed.parallel_state.get_world_group", - return_value=FakeWorldGroup(), - ), + get_parallel().override(world_group=FakeWorldGroup()), patch( "sglang.srt.model_loader.weight_utils._prefetch_checkpoint_file", side_effect=lambda path, cancel_event: loaded_paths.append(path), diff --git a/test/registered/unit/model_loader/test_presharded_loader.py b/test/registered/unit/model_loader/test_presharded_loader.py index 527259ee7..f10a44d0d 100644 --- a/test/registered/unit/model_loader/test_presharded_loader.py +++ b/test/registered/unit/model_loader/test_presharded_loader.py @@ -652,10 +652,7 @@ class TestStructuralSignature(unittest.TestCase): # must see the same gathered list and thus the same aggregate. fake_group.all_gather_object.side_effect = lambda local: ["sig-pp0", "sig-pp1"] - with mock.patch( - "sglang.srt.distributed.parallel_state.get_world_group", - return_value=fake_group, - ): + with get_parallel().override(world_group=fake_group): agg_from_rank0 = ( PreshardedModelLoader._make_rank_invariant_structural_signature( "sig-pp0" @@ -673,10 +670,7 @@ class TestStructuralSignature(unittest.TestCase): "sig-pp0", "sig-pp1-changed", ] - with mock.patch( - "sglang.srt.distributed.parallel_state.get_world_group", - return_value=fake_group, - ): + with get_parallel().override(world_group=fake_group): agg_changed = ( PreshardedModelLoader._make_rank_invariant_structural_signature( "sig-pp0" diff --git a/test/registered/unit/model_loader/test_transformers_fallback.py b/test/registered/unit/model_loader/test_transformers_fallback.py index 9989c1569..36de833ee 100644 --- a/test/registered/unit/model_loader/test_transformers_fallback.py +++ b/test/registered/unit/model_loader/test_transformers_fallback.py @@ -5,6 +5,7 @@ from types import SimpleNamespace from unittest.mock import patch from sglang.srt.models.transformers import TransformersBase +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -54,10 +55,9 @@ class TestTransformersFallbackSkipSubstrs(CustomTestCase): pass with ( - patch( - "sglang.srt.distributed.parallel_state.get_pp_group", - return_value=SimpleNamespace(), - ), + # `__init__` only stashes the pipeline group, so an empty + # stand-in carries it past the read. + get_parallel().override(pp_group=SimpleNamespace()), patch( "sglang.srt.models.transformers.get_hf_text_config", return_value=SimpleNamespace(), diff --git a/test/registered/unit/multimodal/test_vit_cuda_graph_runner.py b/test/registered/unit/multimodal/test_vit_cuda_graph_runner.py index 09a72e4fd..c1c59de38 100644 --- a/test/registered/unit/multimodal/test_vit_cuda_graph_runner.py +++ b/test/registered/unit/multimodal/test_vit_cuda_graph_runner.py @@ -1,6 +1,5 @@ import sys from types import SimpleNamespace -from unittest.mock import patch import pytest import torch @@ -13,6 +12,7 @@ from sglang.srt.multimodal.internvl_vit_cuda_graph_runner import ( InternViTCudaGraphRunner, ) from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner +from sglang.srt.runtime_context import get_parallel class _Block: @@ -32,12 +32,10 @@ def _runner(*, use_data_parallel: bool) -> ViTCudaGraphRunner: def test_dp_vit_graph_capture_does_not_enter_tp_communication_capture(): runner = _runner(use_data_parallel=True) - with patch( - "sglang.srt.distributed.parallel_state.get_tp_group", - side_effect=AssertionError("DP capture must be rank-local"), - ): - with runner._capture_context(): - pass + # No tp_group is stated, so reading one would raise: the DP path must not + # ask for the TP group at all. + with runner._capture_context(): + pass def test_non_dp_vit_graph_capture_uses_tp_communication_capture(): @@ -52,9 +50,7 @@ def test_non_dp_vit_graph_capture_uses_tp_communication_capture(): group = SimpleNamespace(ca_comm=SimpleNamespace(capture=lambda: Capture())) runner = _runner(use_data_parallel=False) - with patch( - "sglang.srt.distributed.parallel_state.get_tp_group", return_value=group - ): + with get_parallel().override(tp_group=group): with runner._capture_context(): pass assert entered == [True] diff --git a/test/registered/unit/server_args/test_resolution_declarations.py b/test/registered/unit/server_args/test_resolution_declarations.py index 132792c7e..8a95bb41d 100644 --- a/test/registered/unit/server_args/test_resolution_declarations.py +++ b/test/registered/unit/server_args/test_resolution_declarations.py @@ -132,18 +132,18 @@ def _stash_overlay(server_args): def _live_topology_leaves(): - """Names `ParallelContext` serves from the live topology, not the config. + """Names `ParallelContext` answers from a runtime write, not the config. - Read out of `_LIVE_READS`, which is where those names are declared. - Inferring them from "did the read raise" is wrong -- it only raises while - the process groups are missing, so in a process where an earlier test built - them the property answers the *live* size and a leaf check reads it as a - config mismatch (`parallel.tp_size: bag=1 resolution=2`). Whether a name is + Read off the declarations that carry no `fn`, which is what those are. + Inferring them from "did the read raise" is wrong -- it raises only while + nothing has written the name, so in a process where an earlier test stated + one the property answers that value and a leaf check reads it as a config + mismatch (`parallel.tp_size: bag=1 resolution=2`). Whether a name is shadowed is a property of the declaration, not of the process. """ - from sglang.srt.runtime_context import _LIVE_READS + from sglang.srt.runtime_context import _derived_widths - return frozenset(_LIVE_READS) + return frozenset(n for n, d in _derived_widths().items() if not d.fn) class TestResolutionDeclarations(CustomTestCase): diff --git a/test/registered/unit/spec/test_dflash_logits.py b/test/registered/unit/spec/test_dflash_logits.py index 9ebbc9c2e..b01ce6572 100644 --- a/test/registered/unit/spec/test_dflash_logits.py +++ b/test/registered/unit/spec/test_dflash_logits.py @@ -249,6 +249,8 @@ def test_worker_folds_a_gate_admitted_quantized_selector_head(monkeypatch): block_size=8, selector=object(), model_runner=SimpleNamespace(tp_rank=0), + # The worker rank-gates its logging on its own frozen record. + ps=SimpleNamespace(tp_rank=0), draft_model=SimpleNamespace(lm_head=None), device="cpu", _selector_sampling_enabled=True, @@ -283,6 +285,8 @@ def test_worker_warns_once_when_selector_sampling_is_disabled(monkeypatch): _selector_sampling_enabled=False, _warned_sampling_fallback=False, model_runner=SimpleNamespace(tp_rank=0), + # The worker rank-gates its logging on its own frozen record. + ps=SimpleNamespace(tp_rank=0), ) batch = SimpleNamespace(sampling_info=SimpleNamespace(is_all_greedy=False)) diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 00c43ba70..fd022d5a5 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -10,8 +10,10 @@ import os import pathlib as _pathlib import shutil import tempfile +import types import unittest import warnings +from types import SimpleNamespace from unittest.mock import patch import msgspec @@ -41,7 +43,6 @@ from sglang.srt.runtime_context import ( RuntimeContext, SpawnRanks, _FlagGroupBase, - _validate_parallel, assert_published, derive_parallel_widths, get_context, @@ -111,39 +112,58 @@ def _scope_entries_that_say_nothing(paths): _PS = "sglang.srt.distributed.parallel_state" + + +def _parallel_state(): + from sglang.srt.distributed import parallel_state + + return parallel_state + + _DP = "sglang.srt.layers.dp_attention" -# Ranks and the launch width are asked of the group: they are not implied by -# anything, so there is nothing to derive them from. The quotients are not -# here -- `attn_tp_size` and its siblings are functions of the configured -# leaves, and `TestDerivedWidths` pins them. `attn_dp_rank` is not here either: no group coordinator -# knows it, so it is stamped when the attention topology is initialized and -# `TestStampedRanks` is what pins it. The other world width is not here -# because the group does not know it; `TestTheTwoWorldWidths` pins it. -SIZE_RANK_DELEGATIONS = [ - ("launch_world_size", f"{_PS}.get_world_size"), - ("launch_world_rank", f"{_PS}.get_world_rank"), - ("tp_rank", f"{_PS}.get_tensor_model_parallel_rank"), - ("dcp_rank", f"{_PS}.get_dcp_rank"), - ("pp_rank", f"{_PS}.get_pipeline_model_parallel_rank"), - ("moe_ep_rank", f"{_PS}.get_moe_expert_parallel_rank"), - ("moe_dp_rank", f"{_PS}.get_moe_data_parallel_rank"), - ("moe_tp_rank", f"{_PS}.get_moe_tensor_parallel_rank"), - ("attn_tp_rank", f"{_PS}.get_attn_tensor_model_parallel_rank"), - ("attn_cp_rank", f"{_PS}.get_attn_context_model_parallel_rank"), -] +#: The groups `initialize_model_parallel` states on the context, by the module +#: global it builds each one into. WORLD is not among them: it is built and +#: stated by `init_distributed_environment`, one call earlier. +GROUP_STAMPS = { + "tp_group": "_TP", + "dcp_group": "_DCP", + "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", +} -GROUP_DELEGATIONS = [ - ("world_group", f"{_PS}.get_world_group"), - ("tp_group", f"{_PS}.get_tp_group"), - ("dcp_group", f"{_PS}.get_dcp_group"), - ("pp_group", f"{_PS}.get_pp_group"), - ("moe_ep_group", f"{_PS}.get_moe_ep_group"), - ("moe_dp_group", f"{_PS}.get_moe_dp_group"), - ("moe_tp_group", f"{_PS}.get_moe_tp_group"), - ("attn_tp_group", f"{_PS}.get_attn_tp_group"), - ("attn_cp_group", f"{_PS}.get_attn_cp_group"), -] + +def _groups_the_build_states() -> dict: + """What `initialize_model_parallel` hands the context, read out of its + source: `{context name: the module global it passes}`. + + Out of the source because that is the only place the whole set appears at + once -- calling the function needs ten live process groups. + """ + import ast + import inspect + import textwrap + + from sglang.srt.distributed import parallel_state + + body = textwrap.dedent(inspect.getsource(parallel_state.initialize_model_parallel)) + for node in ast.walk(ast.parse(body)): + keys = getattr(node, "keys", None) + if ( + isinstance(node, ast.Dict) + and keys + and all( + isinstance(k, ast.Constant) and str(k.value).endswith("_group") + for k in keys + ) + ): + return {k.value: v.id for k, v in zip(node.keys, node.values)} + raise AssertionError("initialize_model_parallel states no group at all") class TestRuntimeContextSingletons(CustomTestCase): @@ -171,29 +191,91 @@ class _IsolatedOverrides(CustomTestCase): super().tearDown() +class TestTheBuildStatesEveryGroup(_IsolatedOverrides): + """Nothing derives a group, so one the build forgets to state is a name + that answers "not written" for the rest of the process -- and the reader + that finds out is a model layer, a long way from here.""" + + def test_the_build_states_every_group_the_context_declares(self): + from sglang.srt.runtime_context import _parallel_fields + + declared = {name for name in _parallel_fields() if name.endswith("_group")} + self.assertEqual(declared, set(GROUP_STAMPS) | {"world_group"}) + + def test_the_world_group_is_stated_where_it_is_built(self): + """`initialize_model_parallel` places every group it builds by reading + `get_world_group().local_rank`, so WORLD has to be answerable before it + runs -- one function earlier, where it is constructed.""" + import ast + import inspect + import textwrap + + from sglang.srt.distributed import parallel_state + + body = textwrap.dedent( + inspect.getsource(parallel_state.init_distributed_environment) + ) + stated = { + kw.arg + for node in ast.walk(ast.parse(body)) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "override_permanently" + for kw in node.keywords + } + self.assertIn("world_group", stated) + + def test_nothing_reads_a_name_this_build_has_not_stated_yet(self): + """The stamp is at the end, so a getter called before it answers a name + nothing has written -- a crash at startup, in a process no unit test + runs. What the function may read is what a *previous* call stated, and + that is WORLD alone: a rank belongs to the spawn, and this build does + not get to assume the spawn ran first.""" + import ast + import inspect + import textwrap + + from sglang.srt.distributed import parallel_state + + body = textwrap.dedent( + inspect.getsource(parallel_state.initialize_model_parallel) + ) + read = set() + for node in ast.walk(ast.parse(body)): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name): + name = parallel_state._CONTEXT_NAME_OF.get(node.func.id) + if name is not None: + read.add(name) + self.assertTrue(read, "no getter is called here; this proves nothing") + self.assertEqual(read - {"world_group"}, set()) + + def test_each_group_is_stated_from_the_global_it_was_built_into(self): + self.assertEqual(_groups_the_build_states(), GROUP_STAMPS) + + def test_a_dimension_the_configuration_has_not_got_is_left_unstated(self): + """`_DCP` is None without decode context parallelism, and every one of + these getters has always refused to answer for a group that was never + built rather than handing back a None to fail on at the collective.""" + import ast + import inspect + import textwrap + + from sglang.srt.distributed import parallel_state + + body = textwrap.dedent( + inspect.getsource(parallel_state.initialize_model_parallel) + ) + stamp = next( + node + for node in ast.walk(ast.parse(body)) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "override_permanently" + ) + self.assertIn("is not None", ast.unparse(stamp)) + + class TestParallelDelegation(_IsolatedOverrides): - def test_size_rank_delegate_to_canonical_getters(self): - # Patch each getter to a distinct sentinel: a miswired attribute would read - # a different (unpatched) getter and fail. - for i, (attr, target) in enumerate(SIZE_RANK_DELEGATIONS): - sentinel = 1000 + i - with patch(target, return_value=sentinel): - self.assertEqual( - getattr(get_parallel(), attr), - sentinel, - msg=f"{attr} must delegate to {target}", - ) - - def test_groups_delegate_to_canonical_getters(self): - for attr, target in GROUP_DELEGATIONS: - sentinel = object() - with patch(target, return_value=sentinel): - self.assertIs( - getattr(get_parallel(), attr), - sentinel, - msg=f"{attr} must delegate to {target}", - ) - def test_wrapper_holds_no_resolved_state(self): # __slots__: no __dict__; the only instance state is the override hook. self.assertFalse(hasattr(get_parallel(), "__dict__")) @@ -204,50 +286,46 @@ class TestParallelDelegation(_IsolatedOverrides): class TestTheTwoWorldWidths(_IsolatedOverrides): - """Two questions about the WORLD group: what it was built at, and what it - has room for. + """Two questions about the WORLD group: what it was launched at, and what + it has room for. - Neither is stored here. How much of that room is serving after a scale-up - is elastic-EP state, and is asked of the manager that owns it rather than - mirrored onto this namespace. + Both are arithmetic over the configured leaves and are worked out at + publish, so they answer in a process that never builds the group -- which + is where a good half of the readers are. How much of that room is serving + after a scale-up is neither of them: that is elastic-EP state, asked of the + manager that owns it rather than mirrored onto this namespace. """ - def test_the_launch_width_is_what_the_group_was_built_at(self): - with patch(f"{_PS}.get_world_size", return_value=4): - self.assertEqual(get_parallel().launch_world_size, 4) + def _published(self, **fields): + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy", **fields), role="test") + return get_parallel() + + def test_the_launch_width_is_a_rank_per_stage_of_each_group(self): + self.assertEqual(self._published(tp_size=4, pp_size=2).launch_world_size, 8) + + def test_the_launch_width_spans_the_ranks_a_joiner_came_in_above(self): + """A scale joiner lays its own groups out at `tp * pp`, while its WORLD + spans the cohort already running underneath it as well.""" + parallel = self._published(tp_size=4, pp_size=1, ep_join_rank_offset=8) + self.assertEqual(parallel.launch_world_size, 12) def test_the_ceiling_is_the_configured_one_when_there_is_one(self): - parallel = get_parallel() - with ( - parallel.override(max_ep_size=32), - patch( - f"{_PS}.get_world_size", - side_effect=AssertionError("the built group must not be asked"), - ), - ): - self.assertEqual(parallel.max_world_size, 32) + self.assertEqual(self._published(tp_size=4, max_ep_size=32).max_world_size, 32) def test_without_a_configured_ceiling_the_room_is_the_launch_width(self): - parallel = get_parallel() - with ( - parallel.override(max_ep_size=None), - patch(f"{_PS}.get_world_size", return_value=8), - ): - self.assertEqual(parallel.max_world_size, 8) + parallel = self._published(tp_size=8) + self.assertEqual(parallel.launch_world_size, 8) + self.assertEqual(parallel.max_world_size, 8) def test_each_width_can_be_stated_on_its_own(self): """Stating one must not answer for the other: they are two names.""" - parallel = get_parallel() - with ( - parallel.override(launch_world_size=2, max_ep_size=None), - patch( - f"{_PS}.get_world_size", - side_effect=AssertionError("the built group must not be asked"), - ), - ): + parallel = self._published(tp_size=8) + with parallel.override(launch_world_size=2): self.assertEqual(parallel.launch_world_size, 2) - self.assertEqual(parallel.max_world_size, 2) - with parallel.override(max_ep_size=6): + self.assertEqual(parallel.max_world_size, 8) + with parallel.override(max_world_size=6): self.assertEqual(parallel.max_world_size, 6) self.assertEqual(parallel.launch_world_size, 2) @@ -380,18 +458,20 @@ class TestAttentionRanksComeFromPublish(_IsolatedOverrides): role="test", ranks=SpawnRanks(world_rank=5), ) - with patch( - f"{_PS}.get_attn_tensor_model_parallel_rank", - side_effect=AssertionError("no group must be consulted"), - ): + with patch.object(_parallel_state(), "_ATTN_TP", None): self.assertEqual(get_parallel().attn_tp_rank, 1) self.assertEqual(get_parallel().attn_dp_rank, 1) - def test_without_a_bundle_it_still_asks_the_group(self): - """Unchanged for every process that publishes without a placement.""" + def test_without_a_bundle_a_rank_read_says_what_is_missing(self): + """There is nothing to fall back to. Deriving one from whatever group + happens to be installed would answer a different question -- where this + process sits in that group, not where the launcher put it.""" publish(ServerArgs(model_path="dummy", tp_size=8), role="test") - with patch(f"{_PS}.get_attn_tensor_model_parallel_rank", return_value=3): - self.assertEqual(get_parallel().attn_tp_rank, 3) + with self.assertRaises(RuntimeError) as caught: + get_parallel().attn_tp_rank + message = str(caught.exception) + self.assertIn("has not been written in this process", message) + self.assertIn("override(attn_tp_rank=...)", message) class TestStampedRanks(_IsolatedOverrides): @@ -568,16 +648,14 @@ class TestEveryDeclaredParallelNameIsStatable(_IsolatedOverrides): The bag carries the declared quotients as well as the operator's leaves, and both are ahead of the live getter once a configuration is - published: a name in `_LIVE_READS` and in either of them would answer + published: a declared name that is also a leaf would answer from the getter before publish and from the bag after.""" from sglang.srt.runtime_context import ( - _LIVE_READS, _derived_widths, _parallel_config_leaves, ) - self.assertEqual(set(_LIVE_READS) & _parallel_config_leaves(), set()) - self.assertEqual(set(_LIVE_READS) & set(_derived_widths()), set()) + self.assertEqual(set(_derived_widths()) & _parallel_config_leaves(), set()) def test_an_undeclared_name_is_refused(self): with self.assertRaises(ValueError): @@ -741,19 +819,43 @@ class TestParallelDCP(_IsolatedOverrides): self.assertTrue(parallel.dcp_enabled) self.assertEqual(parallel.attn_dcp_size, 8) - def test_the_dcp_rank_still_reads_the_group(self): - """A rank is not implied by the configuration, so it reads the group -- - gated on a width that is.""" - with ( - get_parallel().override(tp_size=8, dcp_size=8, dcp_enabled=False), - patch(f"{_PS}.get_dcp_rank", side_effect=AssertionError), - ): - self.assertEqual(get_parallel().attn_dcp_rank, 0) - with ( - get_parallel().override(tp_size=8, dcp_size=8, dcp_enabled=True), - patch(f"{_PS}.get_dcp_rank", return_value=3), - ): - self.assertEqual(get_parallel().attn_dcp_rank, 3) + def _placed(self, world_rank, **fields): + reset_context() + self.addCleanup(reset_context) + publish( + ServerArgs(model_path="dummy", **fields), + role="test", + ranks=SpawnRanks(world_rank=world_rank), + ) + return get_parallel() + + def test_the_dcp_rank_is_where_the_tp_rank_falls_in_its_slice(self): + """A DCP group is a contiguous slice of the TP group, so the place in + one is the TP rank folded by the width.""" + parallel = self._placed(5, tp_size=8, dcp_size=4) + self.assertEqual(parallel.dcp_rank, 1) + self.assertEqual(parallel.attn_dcp_rank, 1) + + def test_the_gated_off_rank_answers_without_a_spawn_bundle(self): + """Zero for every process when decode context parallelism is off, so a + reader on a path that never publishes a bundle -- a memory pool, an + attention backend built in a unit test -- still gets an answer. It + stood for a group that was never built before, and it has to keep + answering the same way.""" + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy", tp_size=8), role="test") + self.assertEqual(get_parallel().attn_dcp_rank, 0) + + def test_the_dcp_rank_is_gated_on_a_width_the_configuration_carries(self): + """Zero where decode context parallelism is off, so a reader does not + have to ask whether it is on before asking where it sits -- and the + gated-off name is not answered at all, because no group holds it.""" + parallel = self._placed(5, tp_size=8, dcp_size=1) + self.assertFalse(parallel.dcp_enabled) + self.assertEqual(parallel.attn_dcp_rank, 0) + with self.assertRaises(RuntimeError): + parallel.dcp_rank def test_the_width_does_not_consult_the_platform(self): with patch("sglang.srt.utils.is_cuda", return_value=False) as is_cuda: @@ -2074,10 +2176,16 @@ class TestDerivedWidths(_IsolatedOverrides): {name for name in widths if "world" in name}, set(), ) - parallel = get_parallel() - parallel.override_permanently(attn_tp_size=4) - with patch(f"{_PS}.get_world_size", return_value=9): - self.assertEqual(parallel.launch_world_size, 9) + # Its own arithmetic spans the offset, which is the leaf the quotients + # above are not given and could not account for. + from sglang.srt.runtime_context import launch_world_size_of + + self.assertEqual( + launch_world_size_of( + SimpleNamespace(ep_join_rank_offset=8, tp_size=4, pp_size=1) + ), + 12, + ) def test_the_bare_name_is_gone(self): """It answered two questions, so every reader had to remember which. @@ -2315,15 +2423,27 @@ class TestTheDerivedHalfIsDeclared(CustomTestCase): f"{name} is declared but no property was installed", ) - def test_the_declared_set_is_what_derive_parallel_widths_produces(self): - """The declaration is not a second list to keep in step: it names - exactly the quotients the derivation returns.""" - from sglang.srt.arg_groups.arg_utils import Derived - from sglang.srt.arg_groups.fields.parallel import Parallel + def test_a_computed_name_names_the_function_that_computes_it(self): + """The declaration is not a second list to keep in step: a name that + is a function of the leaves points at a function called after it, and + that function existing is the whole of what publish needs.""" + import importlib + + from sglang.srt.runtime_context import _derived_widths + + computed = {n: d.fn for n, d in _derived_widths().items() if d.fn} + self.assertTrue(computed, "nothing is computed from the leaves") + for name, fn in computed.items(): + module, _, attr = fn.rpartition(".") + self.assertEqual(attr, f"{name}_of", f"{name} is computed by {attr}") + self.assertTrue(callable(getattr(importlib.import_module(module), attr))) + + def test_the_arithmetic_produces_nothing_that_is_not_declared(self): + """The other side of it: a key the derivation returns and no + declaration names is a width the namespace never answers with, and the + group build re-states it into a name nobody can read.""" + from sglang.srt.runtime_context import _derived_widths - declared = { - name for name, value in vars(Parallel).items() if isinstance(value, Derived) - } produced = set( derive_parallel_widths( tp_size=8, @@ -2335,7 +2455,7 @@ class TestTheDerivedHalfIsDeclared(CustomTestCase): dcp_enabled=False, ) ) - self.assertEqual(declared, produced) + self.assertEqual(produced - set(_derived_widths()), set()) def test_a_declared_quotient_is_not_a_record_field(self): """It has no operator input to preserve, and the record is what crosses @@ -2405,15 +2525,34 @@ class TestTheAccessorsHaveNoCallersOutsideTheirPackage(CustomTestCase): has its own parallel state. """ - #: Not topology. `get_self_pp_group` builds the single-rank group a draft - #: pipeline scope installs, so there is nothing for the context to answer - #: with until the scope has installed it. + #: May have callers. `get_self_pp_group` builds the single-rank group a + #: draft pipeline scope installs, so there is nothing for the context to + #: answer with until the scope has installed it; the other two are not + #: topology at all. ALLOWED = { "get_self_pp_group", "get_default_distributed_backend", "get_mooncake_transfer_engine", } + #: Zero callers required, but not deprecated either: the context has no + #: name that answers the same question. + #: + #: The three widths read a group the build does not check against the + #: configuration, so "the group's width" and "the configured width" are two + #: facts -- the MoE-DP group is the attention-CP group when the latter is + #: wider, and the other two are simply not pinned yet. Pinning them in + #: `_WIDTH_AND_GROUP` is what would let them move. + NOT_ANSWERED_BY_THE_CONTEXT = { + "get_moe_data_parallel_world_size", + "get_moe_tensor_parallel_world_size", + "get_dcp_world_size", + # Answers `None` where the context asserts, which is the whole point of + # the caller that wants it. + "get_dcp_group_no_assert", + "get_torch_distributed_pg_options", + } + def _accessors(self): """Derived from the source, not listed here: a guard whose subject set is written by hand stops watching whatever gets added next.""" @@ -2458,6 +2597,139 @@ class TestTheAccessorsHaveNoCallersOutsideTheirPackage(CustomTestCase): "context cannot answer them", ) + #: How many callers each exempt accessor has outside the defining package. + #: A ratchet, not a description: these may go down and never up, and a name + #: that reaches zero comes off the list. Anything not here must have none. + ALLOWED_CALLERS = { + "get_self_pp_group": 1, + "get_default_distributed_backend": 1, + "get_mooncake_transfer_engine": 6, + } + + def test_the_exempt_accessors_do_not_grow_new_callers(self): + """The zero-caller rule above cannot cover the three that are not + topology, so they get a count instead. Ratchets only turn one way: a + number that has to go up means a new business-code reader of a name the + context should be answering.""" + for name, allowed in sorted(self.ALLOWED_CALLERS.items()): + callers = self._callers(name) + self.assertLessEqual( + len(callers), + allowed, + f"{name} grew a caller: {callers}. Read it through " + f"get_parallel() if the context can answer it; if it truly " + f"cannot, lower this number only when one goes away.", + ) + + def test_every_getter_the_context_answers_is_deprecated(self): + """The other half of the ratchet: the deprecation set is derived from + the table that maps a context name to the getter behind it, so dropping + a getter out of that table would quietly take it off the list. This + fails if one of them stops being marked.""" + from sglang.srt.distributed import parallel_state + + marked = set(parallel_state._CONTEXT_NAME_OF) + unclassified = ( + self._accessors() - self.ALLOWED - self.NOT_ANSWERED_BY_THE_CONTEXT + ) + for name in sorted(unclassified): + if name in marked: + continue + # Not answered by the context and not exempt: a getter that is + # neither is a name with no home, which is what this module exists + # to prevent. + self.assertIn( + name, + marked, + f"{name} is neither deprecated nor listed as exempt -- give it " + "a context name or say here why it has none", + ) + + def test_calling_one_from_outside_the_package_is_deprecated(self): + """The getters stay -- they are the definition -- but a call that comes + from outside the package that defines them cannot be redirected by a + scope, so it says what to read instead.""" + import warnings + + from sglang.srt.distributed import parallel_state + + parallel_state._ALREADY_WARNED.discard("get_tensor_model_parallel_rank") + self.addCleanup( + parallel_state._ALREADY_WARNED.discard, "get_tensor_model_parallel_rank" + ) + with warnings.catch_warnings(record=True) as seen: + warnings.simplefilter("always") + try: + parallel_state.get_tensor_model_parallel_rank() + except Exception: + pass + messages = [str(w.message) for w in seen] + self.assertTrue( + any("get_parallel().tp_rank" in m for m in messages), + f"expected the replacement to be named, got {messages}", + ) + + def test_nothing_the_context_answers_with_calls_back_into_the_package(self): + """The read path reaches the stored group, not the getter that used to + wrap it -- which is what lets the getters be deprecated without the + replacement tripping the warning meant for people who bypass it.""" + + import inspect + + from sglang.srt.distributed import parallel_state + from sglang.srt.runtime_context import ParallelContext, _derived_widths + + written = {n for n, d in _derived_widths().items() if not d.fn} + self.assertTrue(written, "no written-at-runtime names; this proves nothing") + self.assertIn("tp_group", written) + + # The read path is lookups only -- override, stamp, bag. Nothing in it + # can reach a getter, which is what lets them be deprecated without the + # replacement tripping the warning meant for people who bypass it. + body = inspect.getsource(ParallelContext._read) + self.assertNotIn("_ps()", body) + self.assertNotIn("parallel_state", body) + self.assertNotIn("sglang.srt.runtime_context", parallel_state._EXEMPT_CALLERS) + + def test_a_scope_reaches_callers_that_went_straight_to_the_getter(self): + """The getters read the context, so redirecting a group redirects them + too. PD multiplexing needs exactly this: the in-package readers have to + follow the prefill communicator, not just the ones asking the context.""" + from sglang.srt.distributed import parallel_state + + stand_in = SimpleNamespace(world_size=1, rank_in_group=0) + with get_parallel().override(tp_group=stand_in): + self.assertIs(get_parallel().tp_group, stand_in) + self.assertIs(parallel_state.get_tp_group(), stand_in) + + def test_the_package_that_defines_them_is_not_warned_at(self): + """`srt/distributed/` keeps calling them: a read there would go through + the context back into itself.""" + import warnings + + from sglang.srt.distributed import parallel_state + + parallel_state._ALREADY_WARNED.discard("get_tensor_model_parallel_rank") + self.addCleanup( + parallel_state._ALREADY_WARNED.discard, "get_tensor_model_parallel_rank" + ) + caller = types.ModuleType("sglang.srt.distributed.pretend_internal") + caller.__dict__["call"] = lambda: ( + parallel_state.get_tensor_model_parallel_rank() + ) + exec( + "def call():\n from sglang.srt.distributed import parallel_state\n" + " return parallel_state.get_tp_group()", + caller.__dict__, + ) + with warnings.catch_warnings(record=True) as seen: + warnings.simplefilter("always") + try: + caller.call() + except Exception: + pass + self.assertEqual([str(w.message) for w in seen], []) + def test_the_guard_would_notice_a_caller(self): """The subject set is derived, so this checks the search finds a real call rather than that the list happens to be empty: `get_self_pp_group` @@ -2571,37 +2843,39 @@ class TestTheTopologyIdentities(CustomTestCase): def test_a_group_built_at_another_width_is_refused(self): """The other end of the same identity: what the configuration says and - what the coordinators were actually built at, checked where the - disagreement is still attributable to the build.""" - from sglang.srt.distributed import parallel_state + what the coordinators were actually built at. Stating a group is how + the build hands it over, so that write is where the disagreement + surfaces -- still attributable to the build, and before a collective + runs on the wrong peers.""" from sglang.srt.distributed.parallel_state import GroupCoordinator self._publish_square() wrong = GroupCoordinator.__new__(GroupCoordinator) wrong.world_size = 8 wrong.rank_in_group = 0 - with patch.object(parallel_state, "_TP", wrong): - with self.assertRaises(ValueError) as caught: - _validate_parallel(get_parallel(), "group build") + with self.assertRaises(ValueError) as caught: + get_parallel().override_permanently(tp_group=wrong) message = str(caught.exception) - self.assertIn("set by group build", message) self.assertIn("tp_group.world_size == tp_size", message) self.assertIn("built 8, configured 4", message) + # The refused write left nothing behind: the name is unwritten, not + # holding a group no identity accepts. + with self.assertRaises(RuntimeError): + get_parallel().tp_group def test_a_group_built_at_the_configured_width_is_quiet(self): - from sglang.srt.distributed import parallel_state from sglang.srt.distributed.parallel_state import GroupCoordinator self._publish_square() right = GroupCoordinator.__new__(GroupCoordinator) right.world_size = 4 right.rank_in_group = 3 - with patch.object(parallel_state, "_TP", right): - _validate_parallel(get_parallel(), "group build") + get_parallel().override_permanently(tp_group=right) + self.assertIs(get_parallel().tp_group, right) def test_a_draft_scope_states_a_consistent_topology(self): """The scope narrows four names at once, so the identity applies to it - -- and holds, which is what step lets the guard stay on.""" + -- and holds, which is what lets the guard stay on.""" from sglang.srt.distributed import parallel_state from sglang.srt.distributed.parallel_state import GroupCoordinator @@ -2609,9 +2883,8 @@ class TestTheTopologyIdentities(CustomTestCase): group = GroupCoordinator.__new__(GroupCoordinator) group.world_size = 2 group.rank_in_group = 1 - with patch.object(parallel_state, "_TP", group): - with parallel_state.patch_tensor_parallel_group(group, owns_attention=True): - self.assertEqual(get_parallel().attn_tp_size, 2) + with parallel_state.patch_tensor_parallel_group(group, owns_attention=True): + self.assertEqual(get_parallel().attn_tp_size, 2) class TestWhoAnswersDuringADraftScope(CustomTestCase): @@ -2650,11 +2923,10 @@ class TestWhoAnswersDuringADraftScope(CustomTestCase): group = self._single_member_group() self._two_stage_pipeline() self.assertEqual(get_parallel().pp_size, 2) - with patch.object(parallel_state, "_PP", group): - with parallel_state.patch_pipeline_parallel_group(group): - self.assertEqual(get_parallel().pp_size, 1) - self.assertEqual(get_parallel().pp_rank, 0) - self.assertIs(get_parallel().pp_group, group) + with parallel_state.patch_pipeline_parallel_group(group): + self.assertEqual(get_parallel().pp_size, 1) + self.assertEqual(get_parallel().pp_rank, 0) + self.assertIs(get_parallel().pp_group, group) self.assertEqual(get_parallel().pp_size, 2) self.assertEqual(get_parallel().pp_rank, 1) @@ -2685,28 +2957,25 @@ class TestWhoAnswersDuringADraftScope(CustomTestCase): self.assertEqual(get_parallel().attn_tp_size, 2) group = self._group(world_size=2, rank=1) - with patch.object(parallel_state, "_TP", group): - with parallel_state.patch_tensor_parallel_group(group, owns_attention=True): - parallel = get_parallel() - self.assertEqual(parallel.tp_size, 2) - self.assertEqual(parallel.attn_tp_size, 2) - self.assertEqual(parallel.attn_tp_rank, 1) - self.assertEqual(parallel.attn_dp_size, 1) - self.assertEqual(parallel.attn_dp_rank, 0) - self.assertEqual(parallel.attn_cp_size, 1) - self.assertEqual(parallel.attn_cp_rank, 0) - # `dp_size` is the deployment's replica count, not a property - # of the group being installed, so the scope leaves it alone -- - # `require_mlp_tp_gather` asserts on it under dp attention. - self.assertEqual(parallel.dp_size, 2) - # The whole point of stating the rest: the identity the - # override path and the group build both check holds in here. - self.assertEqual( - parallel.tp_size, - parallel.attn_tp_size - * parallel.attn_dp_size - * parallel.attn_cp_size, - ) + with parallel_state.patch_tensor_parallel_group(group, owns_attention=True): + parallel = get_parallel() + self.assertEqual(parallel.tp_size, 2) + self.assertEqual(parallel.attn_tp_size, 2) + self.assertEqual(parallel.attn_tp_rank, 1) + self.assertEqual(parallel.attn_dp_size, 1) + self.assertEqual(parallel.attn_dp_rank, 0) + self.assertEqual(parallel.attn_cp_size, 1) + self.assertEqual(parallel.attn_cp_rank, 0) + # `dp_size` is the deployment's replica count, not a property of + # the group being installed, so the scope leaves it alone -- + # `require_mlp_tp_gather` asserts on it under dp attention. + self.assertEqual(parallel.dp_size, 2) + # The whole point of stating the rest: the identity the override + # path and the group build both check holds in here. + self.assertEqual( + parallel.tp_size, + parallel.attn_tp_size * parallel.attn_dp_size * parallel.attn_cp_size, + ) self.assertEqual(get_parallel().attn_dp_size, 2) self.assertEqual(get_parallel().dp_size, 2) @@ -2790,9 +3059,8 @@ class TestWhoAnswersDuringADraftScope(CustomTestCase): self._two_stage_pipeline() self.assertEqual(get_parallel().pp_size, 2) group = self._single_member_group() - with patch.object(parallel_state, "_PP", group): - with parallel_state.patch_pipeline_parallel_group(group): - checker = WeightChecker(get_model=lambda: None) + with parallel_state.patch_pipeline_parallel_group(group): + checker = WeightChecker(get_model=lambda: None) # The scope has closed and the context answers the target's shape again. self.assertEqual(get_parallel().pp_size, 2)