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

This commit is contained in:
Cheng Wan
2026-09-21 12:25:32 -07:00
committed by GitHub
parent 65be3fa71a
commit 73f071db52
44 changed files with 933 additions and 501 deletions
@@ -97,9 +97,14 @@ def init_world_group(
def _sync_srt_world_group() -> None: def _sync_srt_world_group() -> None:
import sglang.srt.distributed.parallel_state as srt_parallel_state 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: if srt_parallel_state._WORLD is None:
srt_parallel_state._WORLD = _WORLD 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: def _clear_srt_world_group() -> None:
@@ -132,6 +137,27 @@ def _sync_srt_tp_group() -> None:
srt_parallel_state._ATTN_TP = _TP srt_parallel_state._ATTN_TP = _TP
if srt_parallel_state._ATTN_TP is _TP: if srt_parallel_state._ATTN_TP is _TP:
get_parallel().override_permanently( 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( **derive_parallel_widths(
tp_size=_TP.world_size, tp_size=_TP.world_size,
attn_cp_size=1, attn_cp_size=1,
@@ -151,6 +177,10 @@ def _clear_srt_tp_group() -> None:
if srt_parallel_state._ATTN_TP is _TP: if srt_parallel_state._ATTN_TP is _TP:
srt_parallel_state._ATTN_TP = None srt_parallel_state._ATTN_TP = None
get_parallel().clear_stamp() 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: if srt_parallel_state._TP is _TP:
srt_parallel_state._TP = None srt_parallel_state._TP = None
@@ -2,6 +2,7 @@ from contextlib import ExitStack
from types import SimpleNamespace from types import SimpleNamespace
from unittest.mock import call, patch from unittest.mock import call, patch
import pytest
import torch import torch
from torch import nn 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 # `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 # implies -- the shared `srt` vision layers ask for `attn_tp_size`, and this
# package publishes no `srt` config for that read to resolve against. # 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 ( with (
patch.object(parallel_state, "_TP", tp_group), 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._TP is tp_group
assert srt_parallel_state._ATTN_TP is tp_group assert srt_parallel_state._ATTN_TP is tp_group
assert get_parallel().attn_tp_size == 2 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() parallel_state._clear_srt_tp_group()
assert srt_parallel_state._TP is None assert srt_parallel_state._TP is None
assert srt_parallel_state._ATTN_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(): def test_srt_owned_groups_are_not_overwritten_or_cleared():
+6 -1
View File
@@ -107,7 +107,12 @@ class Derived(msgspec.Struct, frozen=True):
``publish`` and stored as an ordinary bag leaf -- a plain attribute load, ``publish`` and stored as an ordinary bag leaf -- a plain attribute load,
which is what a read inside compiled model code needs. 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 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 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`` computed but that a stamp can move one afterwards -- ``initialize_dp_attention``
@@ -320,3 +320,73 @@ class Parallel(msgspec.Struct):
doc="Whether decode context parallelism is in play: `dcp_size` is " doc="Whether decode context parallelism is in play: `dcp_size` is "
"wider than one rank, which is exactly when the group gets built.", "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.")
+140 -43
View File
@@ -24,10 +24,13 @@ If you only need to use the distributed environment without model/pipeline
""" """
import contextlib import contextlib
import functools
import gc import gc
import logging import logging
import os import os
import pickle import pickle
import sys
import warnings
import weakref import weakref
from collections import namedtuple from collections import namedtuple
from contextlib import contextmanager, nullcontext 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.platforms.device_mixin import _DEVICE_TO_DISTRIBUTED_BACKEND
from sglang.srt.runtime_context import ( from sglang.srt.runtime_context import (
_validate_parallel,
derive_parallel_widths, derive_parallel_widths,
get_global_dwdp_manager, get_global_dwdp_manager,
get_parallel, get_parallel,
@@ -2072,8 +2074,7 @@ _WORLD: Optional[GroupCoordinator] = None
def get_world_group() -> GroupCoordinator: def get_world_group() -> GroupCoordinator:
assert _WORLD is not None, "world group is not initialized" return get_parallel().world_group
return _WORLD
def init_world_group( def init_world_group(
@@ -2152,43 +2153,37 @@ _DCP: Optional[GroupCoordinator] = None
# duplicate GroupCoordinator for prefill in PD-Multiplexing # duplicate GroupCoordinator for prefill in PD-Multiplexing
_PDMUX_PREFILL_TP_GROUP: Optional[GroupCoordinator] = None _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): PD multiplexing builds a duplicate TP group -- the same ranks, a second
global _ENABLE_PDMUX_P_TP communicator -- so prefill and decode can occupy separate streams without
_ENABLE_PDMUX_P_TP = enable_prefill_multiplexing 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: def get_tp_group() -> GroupCoordinator:
if _ENABLE_PDMUX_P_TP: return get_parallel().tp_group
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
def get_attn_tp_group() -> GroupCoordinator: def get_attn_tp_group() -> GroupCoordinator:
assert _ATTN_TP is not None, ( return get_parallel().attn_tp_group
"attention tensor model parallel group is not initialized"
)
return _ATTN_TP
def get_shared_experts_tp_group() -> GroupCoordinator: def get_shared_experts_tp_group() -> GroupCoordinator:
assert _SHARED_EXPERTS_TP is not None, ( return get_parallel().shared_experts_tp_group
"shared-expert tensor model parallel group is not initialized"
)
return _SHARED_EXPERTS_TP
def get_attn_cp_group() -> GroupCoordinator: def get_attn_cp_group() -> GroupCoordinator:
assert _ATTN_CP is not None, ( return get_parallel().attn_cp_group
"attention context model parallel group is not initialized"
)
return _ATTN_CP
def get_dcp_group_no_assert() -> Optional[GroupCoordinator]: 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: def get_dcp_group() -> GroupCoordinator:
assert _DCP is not None, "decode context parallel group is not initialized" return get_parallel().dcp_group
return _DCP
_MOE_DP: Optional[GroupCoordinator] = None _MOE_DP: Optional[GroupCoordinator] = None
@@ -2206,18 +2200,15 @@ _MOE_TP: Optional[GroupCoordinator] = None
def get_moe_dp_group() -> GroupCoordinator: def get_moe_dp_group() -> GroupCoordinator:
assert _MOE_DP is not None, "moe data parallel group is not initialized" return get_parallel().moe_dp_group
return _MOE_DP
def get_moe_ep_group() -> GroupCoordinator: def get_moe_ep_group() -> GroupCoordinator:
assert _MOE_EP is not None, "expert model parallel group is not initialized" return get_parallel().moe_ep_group
return _MOE_EP
def get_moe_tp_group() -> GroupCoordinator: def get_moe_tp_group() -> GroupCoordinator:
assert _MOE_TP is not None, "expert model parallel group is not initialized" return get_parallel().moe_tp_group
return _MOE_TP
# kept for backward compatibility # kept for backward compatibility
@@ -2233,8 +2224,7 @@ def get_self_pp_group() -> GroupCoordinator:
def get_pp_group() -> GroupCoordinator: def get_pp_group() -> GroupCoordinator:
assert _PP is not None, "pipeline model parallel group is not initialized" return get_parallel().pp_group
return _PP
# kept for backward compatibility # kept for backward compatibility
@@ -2513,6 +2503,10 @@ def init_distributed_environment(
assert _WORLD.world_size == torch.distributed.get_world_size(), ( assert _WORLD.world_size == torch.distributed.get_world_size(), (
"world group already initialized with a different 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( def initialize_model_parallel(
@@ -2692,7 +2686,7 @@ def initialize_model_parallel(
rank_offset=rank_offset, rank_offset=rank_offset,
max_world_size=max_world_size, max_world_size=max_world_size,
) )
if get_tensor_model_parallel_rank() == 0: if _TP.rank_in_group == 0:
logger.info( logger.info(
f"DCP enabled, dcp_size={decode_context_parallel_size}, tp_size={tensor_model_parallel_size}" 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 # 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 # accounts of one layout, and this is where they meet: stating a group
# disagreement is still attributable, rather than letting a collective run # checks the identities, so a group built on the wrong peers is refused
# on the wrong peers. # here rather than hanging in a collective later.
_validate_parallel(get_parallel(), "group build") #
# 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( 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" assert not _PP_STATE_PATCHED, "Should not call when it's already patched"
_PP_STATE_PATCHED = True _PP_STATE_PATCHED = True
old_pp_group = get_pp_group()
global _PP global _PP
old_pp_group = _PP
_PP = pp_group _PP = pp_group
try: try:
# `pp_size` is a configured leaf: unlike the rank and the handle it # `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" assert not _TP_STATE_PATCHED, "Should not call when it's already patched"
_TP_STATE_PATCHED = True _TP_STATE_PATCHED = True
old_tp_group = get_tp_group()
global _TP global _TP
old_tp_group = _TP
_TP = tp_group _TP = tp_group
narrowed = dict( narrowed = dict(
tp_size=tp_group.world_size, 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_pp_group", get_pp_group)
setattr(vllm_parallel_state, "get_tp_group", get_tp_group) setattr(vllm_parallel_state, "get_tp_group", get_tp_group)
setattr(vllm_parallel_state, "get_world_group", get_world_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
@@ -58,7 +58,7 @@ def ranks_per_host() -> int:
return 1 return 1
try: try:
launch_world_size = get_parallel().launch_world_size launch_world_size = get_parallel().launch_world_size
except AssertionError: except (RuntimeError, ValueError):
return 1 return 1
if launch_world_size == 1: if launch_world_size == 1:
return 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 from sglang.srt.runtime_context import get_parallel
pp_group = get_parallel().pp_group pp_group = get_parallel().pp_group
except AssertionError: except RuntimeError:
return size return size
if pp_group.world_size <= 1: if pp_group.world_size <= 1:
@@ -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.compilation.torch_compile_decoration import set_torch_compile_config
from sglang.srt.distributed.parallel_state import ( from sglang.srt.distributed.parallel_state import (
graph_capture, graph_capture,
set_pdmux_status,
) )
from sglang.srt.dllm.config import DllmConfig from sglang.srt.dllm.config import DllmConfig
from sglang.srt.environ import envs from sglang.srt.environ import envs
@@ -1053,7 +1052,6 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
with self.backend.capture_session(self.stream): with self.backend.capture_session(self.stream):
self._capture_one_stream() self._capture_one_stream()
else: else:
set_pdmux_status(False)
for i, sg in enumerate(self.stream_groups): for i, sg in enumerate(self.stream_groups):
with ( with (
graph_capture(stream=sg[1]) as graph_capture_context, graph_capture(stream=sg[1]) as graph_capture_context,
+2 -2
View File
@@ -2073,7 +2073,7 @@ class PreshardedModelLoader(DefaultModelLoader):
try: try:
g = get_parallel().world_group g = get_parallel().world_group
return g.rank_in_group, g.world_size return g.rank_in_group, g.world_size
except (AssertionError, AttributeError): except (AssertionError, AttributeError, RuntimeError):
return 0, 1 return 0, 1
@staticmethod @staticmethod
@@ -2082,7 +2082,7 @@ class PreshardedModelLoader(DefaultModelLoader):
try: try:
get_parallel().world_group.barrier() get_parallel().world_group.barrier()
except (AssertionError, AttributeError): except (AssertionError, AttributeError, RuntimeError):
pass pass
@staticmethod @staticmethod
+2 -3
View File
@@ -23,7 +23,6 @@ from transformers import PretrainedConfig
import sglang.srt.models.deepseek_v2 as deepseek_v2 import sglang.srt.models.deepseek_v2 as deepseek_v2
from sglang.srt.configs.gigachat35 import GigaChat35Config 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.communicator import get_attn_tp_context
from sglang.srt.layers.layernorm import GemmaRMSNorm, RMSNorm from sglang.srt.layers.layernorm import GemmaRMSNorm, RMSNorm
from sglang.srt.layers.linear import ColumnParallelLinear from sglang.srt.layers.linear import ColumnParallelLinear
@@ -515,7 +514,7 @@ class GigaChat35Model(nn.Module):
self.config = config self.config = config
self.padding_idx = getattr(config, "pad_token_id", None) self.padding_idx = getattr(config, "pad_token_id", None)
self.vocab_size = config.vocab_size 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: if self.pp_group.is_first_rank:
self.embed_tokens = VocabParallelEmbedding( self.embed_tokens = VocabParallelEmbedding(
@@ -605,7 +604,7 @@ class GigaChat35ForCausalLM(DeepseekV2WeightLoaderMixin, nn.Module):
nn.Module.__init__(self) nn.Module.__init__(self)
self.config = config self.config = config
self.quant_config = quant_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.tp_size = get_parallel().tp_size
self.num_fused_shared_experts = 0 self.num_fused_shared_experts = 0
+1 -2
View File
@@ -18,7 +18,6 @@ from typing import Iterable, Optional
import torch import torch
from torch import nn from torch import nn
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.vocab_parallel_embedding import ( from sglang.srt.layers.vocab_parallel_embedding import (
@@ -140,7 +139,7 @@ class GigaChat35ForCausalLMNextN(DeepseekV2WeightLoaderMixin, nn.Module):
nn.Module.__init__(self) nn.Module.__init__(self)
self.config = config self.config = config
self.quant_config = quant_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.tp_size = get_parallel().tp_size
self.num_fused_shared_experts = 0 self.num_fused_shared_experts = 0
self.draft_model_idx = draft_model_idx or 0 self.draft_model_idx = draft_model_idx or 0
@@ -11,7 +11,7 @@ import torch
import torch.distributed as dist import torch.distributed as dist
from torch.cuda.streams import ExternalStream 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.model_executor.forward_batch_info import ForwardMode
from sglang.srt.multiplex.pdmux_context import ( from sglang.srt.multiplex.pdmux_context import (
get_current_stream_idx, get_current_stream_idx,
@@ -21,7 +21,7 @@ from sglang.srt.multiplex.pdmux_context import (
load_pdmux_config, load_pdmux_config,
set_current_stream_idx, 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: if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import ScheduleBatch 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 # 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) 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.stream_groups = get_stream_groups()
self.sm_counts = get_sm_counts() self.sm_counts = get_sm_counts()
self.real_sm_group_num = len(self.stream_groups) self.real_sm_group_num = len(self.stream_groups)
@@ -113,12 +113,10 @@ class SchedulerMultiplexMixin:
while True: while True:
with torch.cuda.stream(decode_stream): with torch.cuda.stream(decode_stream):
set_pdmux_status(False)
self.ingest_requests() self.ingest_requests()
running_batch = self.running_batch running_batch = self.running_batch
with torch.cuda.stream(prefill_stream): with torch.cuda.stream(prefill_stream), pdmux_prefill_tp_group():
set_pdmux_status(True)
sm_count = self.sm_counts[stream_idx][0] sm_count = self.sm_counts[stream_idx][0]
if not wait_prefill_kernel_done: if not wait_prefill_kernel_done:
created, running_batch = self.update_split_prefill_batch( created, running_batch = self.update_split_prefill_batch(
@@ -128,7 +126,6 @@ class SchedulerMultiplexMixin:
adjust_stream_group = created or adjust_stream_group adjust_stream_group = created or adjust_stream_group
with torch.cuda.stream(decode_stream): with torch.cuda.stream(decode_stream):
set_pdmux_status(False)
running_batch = self.update_running_batch(running_batch) running_batch = self.update_running_batch(running_batch)
self.running_batch = running_batch self.running_batch = running_batch
adjust_stream_group = adjust_stream_group or ( adjust_stream_group = adjust_stream_group or (
@@ -152,15 +149,13 @@ class SchedulerMultiplexMixin:
) )
with torch.cuda.stream(decode_stream): with torch.cuda.stream(decode_stream):
set_pdmux_status(False)
# process decode batch # process decode batch
if running_batch and not running_batch.is_empty(): if running_batch and not running_batch.is_empty():
decode_result = self.run_batch(running_batch) decode_result = self.run_batch(running_batch)
decode_done = True decode_done = True
else: else:
decode_done = False decode_done = False
with torch.cuda.stream(prefill_stream): with torch.cuda.stream(prefill_stream), pdmux_prefill_tp_group():
set_pdmux_status(True)
if ( if (
self.split_prefill_batch self.split_prefill_batch
and not self.split_prefill_batch.is_empty() and not self.split_prefill_batch.is_empty()
@@ -197,13 +192,11 @@ class SchedulerMultiplexMixin:
prefill_done = False prefill_done = False
with torch.cuda.stream(decode_stream): with torch.cuda.stream(decode_stream):
set_pdmux_status(False)
decode_stream.synchronize() decode_stream.synchronize()
if decode_done: if decode_done:
self.process_batch_result(running_batch, decode_result) self.process_batch_result(running_batch, decode_result)
with torch.cuda.stream(prefill_stream): with torch.cuda.stream(prefill_stream), pdmux_prefill_tp_group():
set_pdmux_status(True)
if prefill_done and self.split_prefill_batch.split_prefill_finished: if prefill_done and self.split_prefill_batch.split_prefill_finished:
wait_prefill_kernel_done = True wait_prefill_kernel_done = True
prefill_exe_done_flag = prefill_exe_done.query() prefill_exe_done_flag = prefill_exe_done.query()
@@ -214,7 +207,7 @@ class SchedulerMultiplexMixin:
) )
self.tp_cpu_group.allreduce(flags, dist.ReduceOp.SUM).wait() 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.process_batch_result(
self.split_prefill_batch, prefill_result self.split_prefill_batch, prefill_result
) )
+83 -170
View File
@@ -110,119 +110,16 @@ def _parallel_config_leaves() -> frozenset:
) )
# Ranks and group handles: the names no configuration carries, each with the # Ranks and group handles: the names no configuration carries. This table is
# canonical getter that answers it live. This table is their declaration, the # their declaration, the way `arg_groups/fields/parallel.py` is the leaves' and
# way `arg_groups/fields/parallel.py` is the leaves' and `Derived` is the # `Derived` is the widths'. A group handle names the getter that owns it,
# widths'. `None` marks a name only a stamp can answer: no coordinator knows # because the module that builds the groups is where it lives; a rank is a
# this process's attention-DP rank. # 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() _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) @functools.lru_cache(maxsize=1)
def _parallel_fields() -> frozenset: def _parallel_fields() -> frozenset:
"""Every name `ParallelContext` answers for, read from the declarations. """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; * configured leaves -- the `parallel` namespace of the record;
* derived widths -- the `Derived` declarations beside those leaves; * derived widths -- the `Derived` declarations beside those leaves;
* ranks and group handles -- `_LIVE_READS`, which is where they are * ranks and group handles -- declared beside the leaves with no `fn`,
declared because no configuration carries them. because nothing computes them; they are written at runtime.
The set is the union of those three, so `override()` cannot refuse a name The set is the union of those three, so `override()` cannot refuse a name
the class answers for. the class answers for.
@@ -244,7 +141,7 @@ def _parallel_fields() -> frozenset:
derived = { derived = {
name for name, decl in vars(Parallel).items() if isinstance(decl, 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( def derive_attention_widths(
@@ -331,6 +228,7 @@ def derive_spawn_ranks(
% (tp_size // moe_dp_size) % (tp_size // moe_dp_size)
// (tp_size // moe_dp_size // moe_ep_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): def attn_tp_size_of(cfg: Any):
"""`attn_tp_size`, computed at publish. See `parallel_widths_of`.""" """`attn_tp_size`, computed at publish. See `parallel_widths_of`."""
return parallel_widths_of(cfg)["attn_tp_size"] return parallel_widths_of(cfg)["attn_tp_size"]
@@ -586,24 +502,22 @@ def _validate_parallel(parallel, source: str) -> None:
class ParallelContext: class ParallelContext:
"""Parallel-topology namespace: one spelling per name. """Parallel-topology namespace: one spelling per name.
Ranks and group handles are read-through ``@property`` over the canonical Every name is answered by a lookup, never by asking a process group. A
getters, so they answer with the **live** process groups and raise before configured leaf and a width derived from one come off the published
distributed init. Every other name ``tp_size`` and its size siblings ``parallel`` bag; a rank is written by ``publish`` from the spawn bundle,
included, alongside config-only leaves such as ``nccl_port`` is answered and a group handle by ``initialize_model_parallel`` as it builds them. A
from the published ``parallel`` bag, in any process at any point after read before the write that answers it says which write is missing rather
publish. 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 That makes a scope a matter of stating names: ``patch_tensor_parallel_group``
exactly the configured widths. Two things do not follow that rule and are runs a draft worker under a different TP group by overriding the members it
asked of the group itself: ``initialize_model_parallel`` aliases ``_MOE_DP`` changes, and PD multiplexing points ``tp_group`` at the prefill
to ``_ATTN_CP`` when ``attn_cp_size > moe_dp_size``, so a reader that means communicator the same way. Elastic EP needs no rule at all: it scales
the MoE communicator's width calls ``get_moe_cp_size()``; and ``ep_size`` / ``dp_size`` on the published bag while ``launch_world_size``
``patch_tensor_parallel_group`` runs a scope under a different TP group, keeps the width the groups were built at, so the two are different names
which it declares by overriding ``tp_size``, ``tp_rank`` and ``tp_group`` rather than two answers to one name.
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.
""" """
__slots__ = ("_overrides", "_stamp", "_config") __slots__ = ("_overrides", "_stamp", "_config")
@@ -624,10 +538,9 @@ class ParallelContext:
def _read(self, name): def _read(self, name):
"""The one read path, for every kind of name in the namespace. """The one read path, for every kind of name in the namespace.
Scoped override, then the permanent stamp, then what the name is Scoped override, then the permanent stamp, then the published bag.
answered by when nobody has stated it: the published leaf for a A rank or a group handle is on no bag, so once those three are out the
configured value or a derived width, the canonical getter for a rank name has not been written yet and the read says so.
or a group handle.
The two override maps stay separate because they are taken down by The two override maps stay separate because they are taken down by
different things -- a `with` block and `clear_stamp()` -- and different things -- a `with` block and `clear_stamp()` -- and
@@ -643,21 +556,17 @@ class ParallelContext:
config = self._config config = self._config
if config is not None and name in config._fields: if config is not None and name in config._fields:
return getattr(config, name) 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(): if config is None and name in _parallel_config_leaves():
raise ValueError("config namespace 'parallel' not published") 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( raise RuntimeError(
f"derived parallel width {name!r} is not available: it is computed " f"derived parallel width {name!r} is not available: it is computed "
"from the configured leaves at publish, and permanently corrected " "from the configured leaves at publish, and permanently corrected "
@@ -731,33 +640,24 @@ def _derived_widths() -> dict:
def _install_parallel_properties() -> None: def _install_parallel_properties() -> None:
"""Give `ParallelContext` a property per name that is not a config leaf. """Give `ParallelContext` a property per name that is not a config leaf.
The quotients are declared in `arg_groups/fields/parallel.py`, beside the Every name the namespace answers that is not a plain leaf: the quotients,
leaves they are computed from; the ranks and group handles are declared in and the ranks and group handles declared beside them with no `fn`.
`_LIVE_READS`, because no configuration carries them. Properties rather
than names left to `__getattr__` because the class surface is what the Properties rather than names left to `__getattr__` because the class
guards introspect -- `hasattr(ParallelContext, "tp_group")` and surface is what the guards introspect -- `hasattr(ParallelContext,
`vars(ParallelContext)` are how the tests check the set from the class "tp_size")` and `vars(ParallelContext)` are how the tests read the set from
side -- and because each one carries its `Derived.doc`. the class side -- and because each one carries its `Derived.doc`.
Every one of them resolves through `_read`, so there is a single priority Every one of them resolves through `_read`, so there is a single priority
chain rather than one per kind of name. chain rather than one per kind of name.
""" """
docs = {name: decl.doc for name, decl in _derived_widths().items()} 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):
def getter(self, _name=name): def getter(self, _name=name):
return self._read(_name) return self._read(_name)
getter.__name__ = name getter.__name__ = name
getter.__doc__ = docs.get(name) getter.__doc__ = decl.doc
setattr(ParallelContext, name, property(getter)) setattr(ParallelContext, name, property(getter))
@@ -1912,6 +1812,12 @@ def publish(
), ),
) )
_CONTEXT._publish_role = role _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: if ranks is not None and ranks.gpu_id is not None:
_CONTEXT.override("spawn", gpu_id=ranks.gpu_id) _CONTEXT.override("spawn", gpu_id=ranks.gpu_id)
if ranks is not None: if ranks is not None:
@@ -1931,18 +1837,25 @@ def publish(
moe_dp_size=parallel.moe_dp_size, moe_dp_size=parallel.moe_dp_size,
moe_ep_size=parallel.moe_ep_size, moe_ep_size=parallel.moe_ep_size,
) )
# `moe_dp_rank` is a different quantity when the MoE-DP group is # `initialize_model_parallel` aliases the MoE-DP group to the
# aliased to the attention-CP one: the group answers the CP index, # attention-CP one when the CP dimension is the wider of the two, so
# while this computes the MoE-DP index. Leave it to the group there, so # this process's place in it is its CP index rather than the MoE-DP
# one name does not mean two things. # index the arithmetic above gives.
if parallel.moe_dp_size < parallel.attn_cp_size: 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 # `dp_rank` is recorded whatever it is, None included: replicas are
# separate WORLD groups, so no rank implies it and `None` is the answer # separate WORLD groups, so no rank implies it and `None` is the answer
# "no controller" rather than an absence. # "no controller" rather than an absence.
placement["dp_rank"] = ranks.dp_rank placement["dp_rank"] = ranks.dp_rank
placement["launch_world_rank"] = ranks.world_rank placement["launch_world_rank"] = ranks.world_rank
placement.update(_attention_ranks(parallel, placement["tp_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 # One stamp, not two: the identities are checked on every write, and a
# half-placed process satisfies none of them. # half-placed process satisfies none of them.
parallel.override_permanently(**placement) parallel.override_permanently(**placement)
@@ -24,8 +24,11 @@ from sglang.srt.runtime_context import get_context, get_parallel
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
# Unit tests run without distributed initialization. Backends that size buffers by # Unit tests run without distributed initialization. Backends that size buffers by
# attention tensor-parallel degree should see the single-rank default. # attention tensor-parallel degree should see the single-rank default, and a
_parallel_override = get_parallel().override(attn_tp_size=1) # 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__() _parallel_override.__enter__()
DEFAULT_HEAD_DIM = 16 DEFAULT_HEAD_DIM = 16
+14 -1
View File
@@ -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 build outlive any block, so the configuration describing them has to as
well. Callers that tear the groups down are already resetting the process. 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 from sglang.srt.server_args import ServerArgs
reset_context() reset_context()
@@ -2104,6 +2110,13 @@ def publish_build_topology(*, world_rank: int = 0, **server_args_fields):
role="test", role="test",
ranks=SpawnRanks(world_rank=world_rank), 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 _GPU_IDLE_TIMEOUT_SECS = 30.0
+3 -1
View File
@@ -134,7 +134,9 @@ class TestFilterDcpLocalChunkKvIndices(CustomTestCase):
def test_identity_without_dcp(self): def test_identity_without_dcp(self):
kv = torch.arange(37) 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( self.assertIs(
filter_dcp_local_chunk_kv_indices( filter_dcp_local_chunk_kv_indices(
kv, torch.tensor([0]), torch.tensor([37]) kv, torch.tensor([0]), torch.tensor([37])
@@ -25,6 +25,7 @@ import torch
import torch.distributed as dist import torch.distributed as dist
import sglang.srt.distributed.parallel_state as ps 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 from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci( register_cuda_ci(
@@ -130,6 +131,9 @@ def init_distributed():
local_rank=local_rank, local_rank=local_rank,
backend="nccl", 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 cpu_group = coord.cpu_group
nccl_group = coord.device_group nccl_group = coord.device_group
@@ -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.benchmark.utils import get_benchmark_range, multigpu_bench_main
from sglang.kernels.jit.utils import cache_once, is_arch_support_pdl from sglang.kernels.jit.utils import cache_once, is_arch_support_pdl
from sglang.kernels.ops.communication.mp import register_comm_cleanup 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 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci( register_cuda_ci(
@@ -61,6 +62,9 @@ def _init_cpu_group() -> dist.ProcessGroup:
local_rank=local_rank, local_rank=local_rank,
backend="nccl", 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) atexit.register(dist.destroy_process_group)
torch.cuda.set_stream(torch.cuda.Stream()) torch.cuda.set_stream(torch.cuda.Stream())
return coord.cpu_group return coord.cpu_group
@@ -33,6 +33,7 @@ from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import (
all_gather_inner, all_gather_inner,
create_state, create_state,
) )
from sglang.srt.runtime_context import get_parallel
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci( register_cuda_ci(
@@ -77,6 +78,9 @@ def _init_cpu_group() -> dist.ProcessGroup:
local_rank=local_rank, local_rank=local_rank,
backend="nccl", 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) atexit.register(dist.destroy_process_group)
logging.disable(logging.INFO) logging.disable(logging.INFO)
torch.cuda.set_stream(torch.cuda.Stream()) torch.cuda.set_stream(torch.cuda.Stream())
@@ -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 ( from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
CustomAllReduceV2, CustomAllReduceV2,
) )
from sglang.srt.runtime_context import get_parallel
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci( register_cuda_ci(
@@ -108,6 +109,9 @@ def _init_cpu_group() -> dist.ProcessGroup:
local_rank=local_rank, local_rank=local_rank,
backend="nccl", 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) atexit.register(dist.destroy_process_group)
logging.disable(logging.INFO) logging.disable(logging.INFO)
torch.cuda.set_stream(torch.cuda.Stream()) torch.cuda.set_stream(torch.cuda.Stream())
@@ -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 ( from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
CustomAllReduceV2, CustomAllReduceV2,
) )
from sglang.srt.runtime_context import get_parallel
from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kernels.utils import multigpu_pytest_main from sglang.test.kernels.utils import multigpu_pytest_main
@@ -129,6 +130,9 @@ def _init_cpu_group_once() -> dist.ProcessGroup:
local_rank=local_rank, local_rank=local_rank,
backend="nccl", 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) atexit.register(dist.destroy_process_group)
cpu_group = coord.cpu_group cpu_group = coord.cpu_group
assert isinstance(cpu_group, dist.ProcessGroup) assert isinstance(cpu_group, dist.ProcessGroup)
@@ -31,6 +31,7 @@ from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import (
all_gather_inner, all_gather_inner,
create_state, create_state,
) )
from sglang.srt.runtime_context import get_parallel
from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kernels.utils import multigpu_pytest_main from sglang.test.kernels.utils import multigpu_pytest_main
@@ -70,6 +71,9 @@ def _init_cpu_group_once() -> dist.ProcessGroup:
local_rank=local_rank, local_rank=local_rank,
backend="nccl", 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) atexit.register(dist.destroy_process_group)
logging.disable(logging.INFO) logging.disable(logging.INFO)
torch.cuda.set_stream(torch.cuda.Stream()) torch.cuda.set_stream(torch.cuda.Stream())
@@ -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 ( from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
CustomAllReduceV2, CustomAllReduceV2,
) )
from sglang.srt.runtime_context import get_parallel
from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kernels.utils import multigpu_pytest_main from sglang.test.kernels.utils import multigpu_pytest_main
@@ -92,6 +93,9 @@ def _init_cpu_group_once() -> dist.ProcessGroup:
local_rank=local_rank, local_rank=local_rank,
backend="nccl", 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) atexit.register(dist.destroy_process_group)
cpu_group = coord.cpu_group cpu_group = coord.cpu_group
assert isinstance(cpu_group, dist.ProcessGroup) assert isinstance(cpu_group, dist.ProcessGroup)
@@ -30,6 +30,7 @@ from sglang.kernels.ops.kimi_k3 import all_reduce
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
CustomAllReduceV2, CustomAllReduceV2,
) )
from sglang.srt.runtime_context import get_parallel
from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kernels.utils import multigpu_pytest_main from sglang.test.kernels.utils import multigpu_pytest_main
@@ -65,6 +66,9 @@ def _init_world():
local_rank=local_rank, local_rank=local_rank,
backend="nccl", 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) atexit.register(dist.destroy_process_group)
logging.disable(logging.INFO) logging.disable(logging.INFO)
torch.cuda.set_stream(torch.cuda.Stream()) torch.cuda.set_stream(torch.cuda.Stream())
@@ -20,6 +20,7 @@ from sglang.kernels.ops.kimi_k3 import (
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
CustomAllReduceV2, CustomAllReduceV2,
) )
from sglang.srt.runtime_context import get_parallel
from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kernels.utils import multigpu_pytest_main from sglang.test.kernels.utils import multigpu_pytest_main
@@ -53,6 +54,9 @@ def _init_world():
local_rank=local_rank, local_rank=local_rank,
backend="nccl", 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) atexit.register(dist.destroy_process_group)
cpu_group = coord.cpu_group cpu_group = coord.cpu_group
assert isinstance(cpu_group, dist.ProcessGroup) assert isinstance(cpu_group, dist.ProcessGroup)
@@ -48,6 +48,7 @@ from sglang.srt.mem_cache.multimodal_cache import (
EmbeddingResult, EmbeddingResult,
MultiModalStaticCache, MultiModalStaticCache,
) )
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils.common import safe_pickle_loads from sglang.srt.utils.common import safe_pickle_loads
from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase from sglang.test.test_utils import CustomTestCase
@@ -1011,10 +1012,7 @@ class TestEncoderDelivery(CustomTestCase):
statuses[1].copy_(torch.tensor([400, 1, 0, 0])) statuses[1].copy_(torch.tensor([400, 1, 0, 0]))
with ( with (
patch( get_parallel().override(tp_group=TPGroup()),
"sglang.srt.distributed.parallel_state.get_tp_group",
return_value=TPGroup(),
),
patch( patch(
"sglang.srt.disaggregation.encoder.server.torch.distributed.all_gather", "sglang.srt.disaggregation.encoder.server.torch.distributed.all_gather",
side_effect=all_gather, side_effect=all_gather,
@@ -1051,10 +1049,7 @@ class TestEncoderDelivery(CustomTestCase):
statuses[1][2] += 1 statuses[1][2] += 1
with ( with (
patch( get_parallel().override(tp_group=TPGroup()),
"sglang.srt.distributed.parallel_state.get_tp_group",
return_value=TPGroup(),
),
patch( patch(
"sglang.srt.disaggregation.encoder.server.torch.distributed.all_gather", "sglang.srt.disaggregation.encoder.server.torch.distributed.all_gather",
side_effect=all_gather, side_effect=all_gather,
@@ -1095,10 +1090,7 @@ class TestEncoderDelivery(CustomTestCase):
statuses[1].copy_(local_status) statuses[1].copy_(local_status)
with ( with (
patch( get_parallel().override(tp_group=TPGroup()),
"sglang.srt.distributed.parallel_state.get_tp_group",
return_value=TPGroup(),
),
patch( patch(
"sglang.srt.disaggregation.encoder.server.torch.distributed.all_gather", "sglang.srt.disaggregation.encoder.server.torch.distributed.all_gather",
side_effect=all_gather, side_effect=all_gather,
@@ -9,7 +9,7 @@ import unittest
from unittest.mock import MagicMock, call, patch from unittest.mock import MagicMock, call, patch
from sglang.srt.environ import envs 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 from sglang.test.test_utils import CustomTestCase
@@ -201,12 +201,12 @@ class TestRegisterToBootstrap(CustomTestCase):
self.assertIn("10.0.0.1", url_used) self.assertIn("10.0.0.1", url_used)
@patch("sglang.srt.disaggregation.common.conn.requests.put") @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( 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 = MagicMock()
success_resp.status_code = 200 success_resp.status_code = 200
mock_put.return_value = success_resp mock_put.return_value = success_resp
@@ -230,9 +230,12 @@ class TestRegisterToBootstrap(CustomTestCase):
for dp_rank, tp_rank, host, rank_port, _ in schedulers 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: for dp_rank, tp_rank, local_ip, _, rust_http_port in schedulers:
manager = self._make_manager() manager = self._make_manager()
manager.attn_dp_size = 2 manager.attn_dp_size = 2
@@ -277,7 +280,7 @@ class TestRegisterToBootstrap(CustomTestCase):
gather_call.args[0]["attn_dp_rank"], gather_call.args[0]["attn_dp_rank"],
gather_call.args[0]["attn_tp_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], [(dp, tp) for dp, tp, _, _, _ in schedulers],
) )
@@ -93,6 +93,9 @@ def _make_prefill_aware_swa_runner(
page_size=1, page_size=1,
attn_cp_size=1, attn_cp_size=1,
tp_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, is_draft_worker=False,
server_args=server_args, server_args=server_args,
attention_chunk_size=None, attention_chunk_size=None,
@@ -14,6 +14,7 @@ from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
maybe_stub_sgl_kernel() maybe_stub_sgl_kernel()
from sglang.srt.layers.moe import mega_moe 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") register_cpu_ci(est_time=8, suite="base-a-test-cpu")
@@ -271,11 +272,8 @@ class TestDeepGemmMegaMoeApi(CustomTestCase):
"init_new", "init_new",
return_value=object(), return_value=object(),
), ),
patch( get_parallel().override(
"sglang.srt.runtime_context.get_parallel", moe_ep_group=SimpleNamespace(device_group=object())
return_value=SimpleNamespace(
moe_ep_group=SimpleNamespace(device_group=object())
),
), ),
): ):
mega_moe._run_mega_routed( mega_moe._run_mega_routed(
@@ -11,6 +11,7 @@ from unittest.mock import Mock, patch
from parameterized import parameterized from parameterized import parameterized
from sglang.srt.distributed import parallel_state
from sglang.srt.managers import scheduler as scheduler_module from sglang.srt.managers import scheduler as scheduler_module
from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.scheduler import Scheduler from sglang.srt.managers.scheduler import Scheduler
@@ -168,12 +169,21 @@ class TestSchedulerIdleStepCounters(CustomTestCase):
) )
with ( with (
patch(f"{PDMUX_MODULE}.get_current_stream_idx", return_value=0), 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.empty_cache"),
patch( patch(
f"{PDMUX_MODULE}.torch.cuda.stream", f"{PDMUX_MODULE}.torch.cuda.stream",
side_effect=lambda stream: nullcontext(), 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( self.run_and_check(
scheduler, scheduler,
@@ -297,7 +297,9 @@ class TestPPMambaPoolSizing(unittest.TestCase):
server_args=SimpleNamespace(), server_args=SimpleNamespace(),
spec_algorithm=SimpleNamespace(is_none=lambda: True), spec_algorithm=SimpleNamespace(is_none=lambda: True),
layer_info=SimpleNamespace(start_layer=start, end_layer=end), 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, hybrid_gdn_config=None,
model_config=SimpleNamespace( model_config=SimpleNamespace(
hf_config=SimpleNamespace(), num_hidden_layers=cls.TOTAL_LAYERS hf_config=SimpleNamespace(), num_hidden_layers=cls.TOTAL_LAYERS
@@ -304,16 +304,13 @@ class TestHostMemoryBudget(CustomTestCase):
def test_ranks_per_host_divides_world_size_by_nodes(self): def test_ranks_per_host_divides_world_size_by_nodes(self):
# The launcher slices ranks uniformly across nodes, so the co-located # The launcher slices ranks uniformly across nodes, so the co-located
# rank count is world_size // nnodes — no hostname collective. # 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 ( with (
get_context().override_server_args(nnodes=2), get_context().override_server_args(nnodes=2, tp_size=16),
unittest.mock.patch.object( unittest.mock.patch.object(
torch.distributed, "is_initialized", return_value=True 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) self.assertEqual(base.ranks_per_host(), 8)
@@ -248,7 +248,7 @@ class TestUnifiedMLATokenToKVPool(unittest.TestCase):
] = float(layer + 1) ] = float(layer + 1)
with ( with (
get_parallel().override(dcp_enabled=False), get_parallel().override(dcp_enabled=False, attn_dcp_rank=0),
mock.patch( mock.patch(
"sglang.srt.mem_cache.memory_pool.current_platform.synchronize" "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("full")),
_FakeKVCache(pool.max_slots("mamba")), _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( allocator = UnifiedMambaTokenToKVPoolAllocator(
unified_buffer=pool, kvcache=kvcache, device=_DEV, page_size=ps unified_buffer=pool, kvcache=kvcache, device=_DEV, page_size=ps
) )
@@ -18,6 +18,7 @@ from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import (
FullCudaGraphBackend, FullCudaGraphBackend,
) )
from sglang.srt.model_executor.runner_utils import pool 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 import dflash_utils, dflash_worker_v2, eagle_utils
from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2 from sglang.srt.speculative.dflash_worker_v2 import DFlashWorkerV2
from sglang.test.ci.ci_register import register_cuda_ci 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", "sglang.srt.layers.dp_attention.is_dp_attention_enabled",
return_value=False, return_value=False,
), ),
# `parallel_state`, not the package re-export: a stub on the get_parallel().override(tp_group=tp_group),
# re-export is never consulted.
patch(
"sglang.srt.distributed.parallel_state.get_tp_group",
return_value=tp_group,
),
patch( patch(
"sglang.kernels.ops.speculative.sampling.tree_speculative_sampling_target_only", "sglang.kernels.ops.speculative.sampling.tree_speculative_sampling_target_only",
side_effect=fake_sampling, side_effect=fake_sampling,
@@ -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.model_runner_components import kv_pool_runtime
from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig 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.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase from sglang.test.test_utils import CustomTestCase
@@ -60,9 +60,8 @@ class TestCanaryHeadroom(CustomTestCase):
), ),
), ),
patch.object(kv_pool_runtime.torch.cuda, "synchronize"), patch.object(kv_pool_runtime.torch.cuda, "synchronize"),
patch( get_parallel().override(
"sglang.srt.distributed.parallel_state.get_world_group", world_group=SimpleNamespace(world_size=1, cpu_group=None)
return_value=SimpleNamespace(world_size=1, cpu_group=None),
), ),
patch.object(kv_pool_runtime, "get_available_gpu_memory", return_value=20), patch.object(kv_pool_runtime, "get_available_gpu_memory", return_value=20),
patch.object(kv_pool_runtime, "mambaish_config", return_value=None), patch.object(kv_pool_runtime, "mambaish_config", return_value=None),
@@ -36,7 +36,11 @@ def mock_cpu_env(kv_size=2, tp_size=1, swa_eviction_interval=4):
with ( with (
patch("torch._utils._element_size", return_value=kv_size), 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), envs.SGLANG_SWA_EVICTION_INTERVAL.override(swa_eviction_interval),
): ):
yield yield
@@ -53,9 +53,7 @@ class TestModelOptExport(unittest.TestCase):
self.mock_logger.start() self.mock_logger.start()
# Mock all distributed functions that might be called # Mock all distributed functions that might be called
self.mock_get_tp_group = patch( self.mock_get_tp_group = patch("sglang.srt.distributed.parallel_state._TP")
"sglang.srt.distributed.parallel_state.get_tp_group"
)
self.mock_get_tp_group.start() self.mock_get_tp_group.start()
# Mock model parallel initialization check # Mock model parallel initialization check
@@ -82,9 +82,7 @@ class TestModelOptModelLoader(CustomTestCase):
self.mock_logger.start() self.mock_logger.start()
# Mock all distributed functions that might be called # Mock all distributed functions that might be called
self.mock_get_tp_group = patch( self.mock_get_tp_group = patch("sglang.srt.distributed.parallel_state._TP")
"sglang.srt.distributed.parallel_state.get_tp_group"
)
self.mock_get_tp_group.start() self.mock_get_tp_group.start()
# Mock model parallel initialization check # Mock model parallel initialization check
@@ -25,6 +25,7 @@ from sglang.srt.model_loader.weight_utils import (
fastsafetensors_weights_iterator, fastsafetensors_weights_iterator,
safetensors_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.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase from sglang.test.test_utils import CustomTestCase
@@ -237,10 +238,7 @@ class TestPrefetchCheckpoints(CustomTestCase):
patch("threading.Thread", _InlineThread), patch("threading.Thread", _InlineThread),
patch("concurrent.futures.ThreadPoolExecutor", _InlineExecutor), patch("concurrent.futures.ThreadPoolExecutor", _InlineExecutor),
patch("concurrent.futures.wait", side_effect=_wait_all), patch("concurrent.futures.wait", side_effect=_wait_all),
patch( get_parallel().override(world_group=FakeWorldGroup()),
"sglang.srt.distributed.parallel_state.get_world_group",
return_value=FakeWorldGroup(),
),
patch( patch(
"sglang.srt.model_loader.weight_utils._prefetch_checkpoint_file", "sglang.srt.model_loader.weight_utils._prefetch_checkpoint_file",
side_effect=lambda path, cancel_event: loaded_paths.append(path), side_effect=lambda path, cancel_event: loaded_paths.append(path),
@@ -652,10 +652,7 @@ class TestStructuralSignature(unittest.TestCase):
# must see the same gathered list and thus the same aggregate. # must see the same gathered list and thus the same aggregate.
fake_group.all_gather_object.side_effect = lambda local: ["sig-pp0", "sig-pp1"] fake_group.all_gather_object.side_effect = lambda local: ["sig-pp0", "sig-pp1"]
with mock.patch( with get_parallel().override(world_group=fake_group):
"sglang.srt.distributed.parallel_state.get_world_group",
return_value=fake_group,
):
agg_from_rank0 = ( agg_from_rank0 = (
PreshardedModelLoader._make_rank_invariant_structural_signature( PreshardedModelLoader._make_rank_invariant_structural_signature(
"sig-pp0" "sig-pp0"
@@ -673,10 +670,7 @@ class TestStructuralSignature(unittest.TestCase):
"sig-pp0", "sig-pp0",
"sig-pp1-changed", "sig-pp1-changed",
] ]
with mock.patch( with get_parallel().override(world_group=fake_group):
"sglang.srt.distributed.parallel_state.get_world_group",
return_value=fake_group,
):
agg_changed = ( agg_changed = (
PreshardedModelLoader._make_rank_invariant_structural_signature( PreshardedModelLoader._make_rank_invariant_structural_signature(
"sig-pp0" "sig-pp0"
@@ -5,6 +5,7 @@ from types import SimpleNamespace
from unittest.mock import patch from unittest.mock import patch
from sglang.srt.models.transformers import TransformersBase 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.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase from sglang.test.test_utils import CustomTestCase
@@ -54,10 +55,9 @@ class TestTransformersFallbackSkipSubstrs(CustomTestCase):
pass pass
with ( with (
patch( # `__init__` only stashes the pipeline group, so an empty
"sglang.srt.distributed.parallel_state.get_pp_group", # stand-in carries it past the read.
return_value=SimpleNamespace(), get_parallel().override(pp_group=SimpleNamespace()),
),
patch( patch(
"sglang.srt.models.transformers.get_hf_text_config", "sglang.srt.models.transformers.get_hf_text_config",
return_value=SimpleNamespace(), return_value=SimpleNamespace(),
@@ -1,6 +1,5 @@
import sys import sys
from types import SimpleNamespace from types import SimpleNamespace
from unittest.mock import patch
import pytest import pytest
import torch import torch
@@ -13,6 +12,7 @@ from sglang.srt.multimodal.internvl_vit_cuda_graph_runner import (
InternViTCudaGraphRunner, InternViTCudaGraphRunner,
) )
from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner
from sglang.srt.runtime_context import get_parallel
class _Block: 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(): def test_dp_vit_graph_capture_does_not_enter_tp_communication_capture():
runner = _runner(use_data_parallel=True) runner = _runner(use_data_parallel=True)
with patch( # No tp_group is stated, so reading one would raise: the DP path must not
"sglang.srt.distributed.parallel_state.get_tp_group", # ask for the TP group at all.
side_effect=AssertionError("DP capture must be rank-local"), with runner._capture_context():
): pass
with runner._capture_context():
pass
def test_non_dp_vit_graph_capture_uses_tp_communication_capture(): 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())) group = SimpleNamespace(ca_comm=SimpleNamespace(capture=lambda: Capture()))
runner = _runner(use_data_parallel=False) runner = _runner(use_data_parallel=False)
with patch( with get_parallel().override(tp_group=group):
"sglang.srt.distributed.parallel_state.get_tp_group", return_value=group
):
with runner._capture_context(): with runner._capture_context():
pass pass
assert entered == [True] assert entered == [True]
@@ -132,18 +132,18 @@ def _stash_overlay(server_args):
def _live_topology_leaves(): 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. Read off the declarations that carry no `fn`, which is what those are.
Inferring them from "did the read raise" is wrong -- it only raises while Inferring them from "did the read raise" is wrong -- it raises only while
the process groups are missing, so in a process where an earlier test built nothing has written the name, so in a process where an earlier test stated
them the property answers the *live* size and a leaf check reads it as a one the property answers that value and a leaf check reads it as a config
config mismatch (`parallel.tp_size: bag=1 resolution=2`). Whether a name is mismatch (`parallel.tp_size: bag=1 resolution=2`). Whether a name is
shadowed is a property of the declaration, not of the process. 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): class TestResolutionDeclarations(CustomTestCase):
@@ -249,6 +249,8 @@ def test_worker_folds_a_gate_admitted_quantized_selector_head(monkeypatch):
block_size=8, block_size=8,
selector=object(), selector=object(),
model_runner=SimpleNamespace(tp_rank=0), 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), draft_model=SimpleNamespace(lm_head=None),
device="cpu", device="cpu",
_selector_sampling_enabled=True, _selector_sampling_enabled=True,
@@ -283,6 +285,8 @@ def test_worker_warns_once_when_selector_sampling_is_disabled(monkeypatch):
_selector_sampling_enabled=False, _selector_sampling_enabled=False,
_warned_sampling_fallback=False, _warned_sampling_fallback=False,
model_runner=SimpleNamespace(tp_rank=0), 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)) batch = SimpleNamespace(sampling_info=SimpleNamespace(is_all_greedy=False))
+439 -171
View File
@@ -10,8 +10,10 @@ import os
import pathlib as _pathlib import pathlib as _pathlib
import shutil import shutil
import tempfile import tempfile
import types
import unittest import unittest
import warnings import warnings
from types import SimpleNamespace
from unittest.mock import patch from unittest.mock import patch
import msgspec import msgspec
@@ -41,7 +43,6 @@ from sglang.srt.runtime_context import (
RuntimeContext, RuntimeContext,
SpawnRanks, SpawnRanks,
_FlagGroupBase, _FlagGroupBase,
_validate_parallel,
assert_published, assert_published,
derive_parallel_widths, derive_parallel_widths,
get_context, get_context,
@@ -111,39 +112,58 @@ def _scope_entries_that_say_nothing(paths):
_PS = "sglang.srt.distributed.parallel_state" _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" _DP = "sglang.srt.layers.dp_attention"
# Ranks and the launch width are asked of the group: they are not implied by #: The groups `initialize_model_parallel` states on the context, by the module
# anything, so there is nothing to derive them from. The quotients are not #: global it builds each one into. WORLD is not among them: it is built and
# here -- `attn_tp_size` and its siblings are functions of the configured #: stated by `init_distributed_environment`, one call earlier.
# leaves, and `TestDerivedWidths` pins them. `attn_dp_rank` is not here either: no group coordinator GROUP_STAMPS = {
# knows it, so it is stamped when the attention topology is initialized and "tp_group": "_TP",
# `TestStampedRanks` is what pins it. The other world width is not here "dcp_group": "_DCP",
# because the group does not know it; `TestTheTwoWorldWidths` pins it. "pp_group": "_PP",
SIZE_RANK_DELEGATIONS = [ "moe_ep_group": "_MOE_EP",
("launch_world_size", f"{_PS}.get_world_size"), "moe_dp_group": "_MOE_DP",
("launch_world_rank", f"{_PS}.get_world_rank"), "moe_tp_group": "_MOE_TP",
("tp_rank", f"{_PS}.get_tensor_model_parallel_rank"), "attn_tp_group": "_ATTN_TP",
("dcp_rank", f"{_PS}.get_dcp_rank"), "attn_cp_group": "_ATTN_CP",
("pp_rank", f"{_PS}.get_pipeline_model_parallel_rank"), "shared_experts_tp_group": "_SHARED_EXPERTS_TP",
("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"),
]
GROUP_DELEGATIONS = [
("world_group", f"{_PS}.get_world_group"), def _groups_the_build_states() -> dict:
("tp_group", f"{_PS}.get_tp_group"), """What `initialize_model_parallel` hands the context, read out of its
("dcp_group", f"{_PS}.get_dcp_group"), source: `{context name: the module global it passes}`.
("pp_group", f"{_PS}.get_pp_group"),
("moe_ep_group", f"{_PS}.get_moe_ep_group"), Out of the source because that is the only place the whole set appears at
("moe_dp_group", f"{_PS}.get_moe_dp_group"), once -- calling the function needs ten live process groups.
("moe_tp_group", f"{_PS}.get_moe_tp_group"), """
("attn_tp_group", f"{_PS}.get_attn_tp_group"), import ast
("attn_cp_group", f"{_PS}.get_attn_cp_group"), 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): class TestRuntimeContextSingletons(CustomTestCase):
@@ -171,29 +191,91 @@ class _IsolatedOverrides(CustomTestCase):
super().tearDown() 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): 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): def test_wrapper_holds_no_resolved_state(self):
# __slots__: no __dict__; the only instance state is the override hook. # __slots__: no __dict__; the only instance state is the override hook.
self.assertFalse(hasattr(get_parallel(), "__dict__")) self.assertFalse(hasattr(get_parallel(), "__dict__"))
@@ -204,50 +286,46 @@ class TestParallelDelegation(_IsolatedOverrides):
class TestTheTwoWorldWidths(_IsolatedOverrides): class TestTheTwoWorldWidths(_IsolatedOverrides):
"""Two questions about the WORLD group: what it was built at, and what it """Two questions about the WORLD group: what it was launched at, and what
has room for. it has room for.
Neither is stored here. How much of that room is serving after a scale-up Both are arithmetic over the configured leaves and are worked out at
is elastic-EP state, and is asked of the manager that owns it rather than publish, so they answer in a process that never builds the group -- which
mirrored onto this namespace. 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): def _published(self, **fields):
with patch(f"{_PS}.get_world_size", return_value=4): reset_context()
self.assertEqual(get_parallel().launch_world_size, 4) 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): def test_the_ceiling_is_the_configured_one_when_there_is_one(self):
parallel = get_parallel() self.assertEqual(self._published(tp_size=4, max_ep_size=32).max_world_size, 32)
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)
def test_without_a_configured_ceiling_the_room_is_the_launch_width(self): def test_without_a_configured_ceiling_the_room_is_the_launch_width(self):
parallel = get_parallel() parallel = self._published(tp_size=8)
with ( self.assertEqual(parallel.launch_world_size, 8)
parallel.override(max_ep_size=None), self.assertEqual(parallel.max_world_size, 8)
patch(f"{_PS}.get_world_size", return_value=8),
):
self.assertEqual(parallel.max_world_size, 8)
def test_each_width_can_be_stated_on_its_own(self): def test_each_width_can_be_stated_on_its_own(self):
"""Stating one must not answer for the other: they are two names.""" """Stating one must not answer for the other: they are two names."""
parallel = get_parallel() parallel = self._published(tp_size=8)
with ( with parallel.override(launch_world_size=2):
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"),
),
):
self.assertEqual(parallel.launch_world_size, 2) self.assertEqual(parallel.launch_world_size, 2)
self.assertEqual(parallel.max_world_size, 2) self.assertEqual(parallel.max_world_size, 8)
with parallel.override(max_ep_size=6): with parallel.override(max_world_size=6):
self.assertEqual(parallel.max_world_size, 6) self.assertEqual(parallel.max_world_size, 6)
self.assertEqual(parallel.launch_world_size, 2) self.assertEqual(parallel.launch_world_size, 2)
@@ -380,18 +458,20 @@ class TestAttentionRanksComeFromPublish(_IsolatedOverrides):
role="test", role="test",
ranks=SpawnRanks(world_rank=5), ranks=SpawnRanks(world_rank=5),
) )
with patch( with patch.object(_parallel_state(), "_ATTN_TP", None):
f"{_PS}.get_attn_tensor_model_parallel_rank",
side_effect=AssertionError("no group must be consulted"),
):
self.assertEqual(get_parallel().attn_tp_rank, 1) self.assertEqual(get_parallel().attn_tp_rank, 1)
self.assertEqual(get_parallel().attn_dp_rank, 1) self.assertEqual(get_parallel().attn_dp_rank, 1)
def test_without_a_bundle_it_still_asks_the_group(self): def test_without_a_bundle_a_rank_read_says_what_is_missing(self):
"""Unchanged for every process that publishes without a placement.""" """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") publish(ServerArgs(model_path="dummy", tp_size=8), role="test")
with patch(f"{_PS}.get_attn_tensor_model_parallel_rank", return_value=3): with self.assertRaises(RuntimeError) as caught:
self.assertEqual(get_parallel().attn_tp_rank, 3) 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): class TestStampedRanks(_IsolatedOverrides):
@@ -568,16 +648,14 @@ class TestEveryDeclaredParallelNameIsStatable(_IsolatedOverrides):
The bag carries the declared quotients as well as the operator's 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 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 the getter before publish and from the bag after."""
from sglang.srt.runtime_context import ( from sglang.srt.runtime_context import (
_LIVE_READS,
_derived_widths, _derived_widths,
_parallel_config_leaves, _parallel_config_leaves,
) )
self.assertEqual(set(_LIVE_READS) & _parallel_config_leaves(), set()) self.assertEqual(set(_derived_widths()) & _parallel_config_leaves(), set())
self.assertEqual(set(_LIVE_READS) & set(_derived_widths()), set())
def test_an_undeclared_name_is_refused(self): def test_an_undeclared_name_is_refused(self):
with self.assertRaises(ValueError): with self.assertRaises(ValueError):
@@ -741,19 +819,43 @@ class TestParallelDCP(_IsolatedOverrides):
self.assertTrue(parallel.dcp_enabled) self.assertTrue(parallel.dcp_enabled)
self.assertEqual(parallel.attn_dcp_size, 8) self.assertEqual(parallel.attn_dcp_size, 8)
def test_the_dcp_rank_still_reads_the_group(self): def _placed(self, world_rank, **fields):
"""A rank is not implied by the configuration, so it reads the group -- reset_context()
gated on a width that is.""" self.addCleanup(reset_context)
with ( publish(
get_parallel().override(tp_size=8, dcp_size=8, dcp_enabled=False), ServerArgs(model_path="dummy", **fields),
patch(f"{_PS}.get_dcp_rank", side_effect=AssertionError), role="test",
): ranks=SpawnRanks(world_rank=world_rank),
self.assertEqual(get_parallel().attn_dcp_rank, 0) )
with ( return get_parallel()
get_parallel().override(tp_size=8, dcp_size=8, dcp_enabled=True),
patch(f"{_PS}.get_dcp_rank", return_value=3), 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
self.assertEqual(get_parallel().attn_dcp_rank, 3) 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): def test_the_width_does_not_consult_the_platform(self):
with patch("sglang.srt.utils.is_cuda", return_value=False) as is_cuda: 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}, {name for name in widths if "world" in name},
set(), set(),
) )
parallel = get_parallel() # Its own arithmetic spans the offset, which is the leaf the quotients
parallel.override_permanently(attn_tp_size=4) # above are not given and could not account for.
with patch(f"{_PS}.get_world_size", return_value=9): from sglang.srt.runtime_context import launch_world_size_of
self.assertEqual(parallel.launch_world_size, 9)
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): def test_the_bare_name_is_gone(self):
"""It answered two questions, so every reader had to remember which. """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", f"{name} is declared but no property was installed",
) )
def test_the_declared_set_is_what_derive_parallel_widths_produces(self): def test_a_computed_name_names_the_function_that_computes_it(self):
"""The declaration is not a second list to keep in step: it names """The declaration is not a second list to keep in step: a name that
exactly the quotients the derivation returns.""" is a function of the leaves points at a function called after it, and
from sglang.srt.arg_groups.arg_utils import Derived that function existing is the whole of what publish needs."""
from sglang.srt.arg_groups.fields.parallel import Parallel 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( produced = set(
derive_parallel_widths( derive_parallel_widths(
tp_size=8, tp_size=8,
@@ -2335,7 +2455,7 @@ class TestTheDerivedHalfIsDeclared(CustomTestCase):
dcp_enabled=False, 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): def test_a_declared_quotient_is_not_a_record_field(self):
"""It has no operator input to preserve, and the record is what crosses """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. has its own parallel state.
""" """
#: Not topology. `get_self_pp_group` builds the single-rank group a draft #: May have callers. `get_self_pp_group` builds the single-rank group a
#: pipeline scope installs, so there is nothing for the context to answer #: draft pipeline scope installs, so there is nothing for the context to
#: with until the scope has installed it. #: answer with until the scope has installed it; the other two are not
#: topology at all.
ALLOWED = { ALLOWED = {
"get_self_pp_group", "get_self_pp_group",
"get_default_distributed_backend", "get_default_distributed_backend",
"get_mooncake_transfer_engine", "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): def _accessors(self):
"""Derived from the source, not listed here: a guard whose subject set """Derived from the source, not listed here: a guard whose subject set
is written by hand stops watching whatever gets added next.""" is written by hand stops watching whatever gets added next."""
@@ -2458,6 +2597,139 @@ class TestTheAccessorsHaveNoCallersOutsideTheirPackage(CustomTestCase):
"context cannot answer them", "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): def test_the_guard_would_notice_a_caller(self):
"""The subject set is derived, so this checks the search finds a real """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` 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): def test_a_group_built_at_another_width_is_refused(self):
"""The other end of the same identity: what the configuration says and """The other end of the same identity: what the configuration says and
what the coordinators were actually built at, checked where the what the coordinators were actually built at. Stating a group is how
disagreement is still attributable to the build.""" the build hands it over, so that write is where the disagreement
from sglang.srt.distributed import parallel_state surfaces -- still attributable to the build, and before a collective
runs on the wrong peers."""
from sglang.srt.distributed.parallel_state import GroupCoordinator from sglang.srt.distributed.parallel_state import GroupCoordinator
self._publish_square() self._publish_square()
wrong = GroupCoordinator.__new__(GroupCoordinator) wrong = GroupCoordinator.__new__(GroupCoordinator)
wrong.world_size = 8 wrong.world_size = 8
wrong.rank_in_group = 0 wrong.rank_in_group = 0
with patch.object(parallel_state, "_TP", wrong): with self.assertRaises(ValueError) as caught:
with self.assertRaises(ValueError) as caught: get_parallel().override_permanently(tp_group=wrong)
_validate_parallel(get_parallel(), "group build")
message = str(caught.exception) message = str(caught.exception)
self.assertIn("set by group build", message)
self.assertIn("tp_group.world_size == tp_size", message) self.assertIn("tp_group.world_size == tp_size", message)
self.assertIn("built 8, configured 4", 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): 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 from sglang.srt.distributed.parallel_state import GroupCoordinator
self._publish_square() self._publish_square()
right = GroupCoordinator.__new__(GroupCoordinator) right = GroupCoordinator.__new__(GroupCoordinator)
right.world_size = 4 right.world_size = 4
right.rank_in_group = 3 right.rank_in_group = 3
with patch.object(parallel_state, "_TP", right): get_parallel().override_permanently(tp_group=right)
_validate_parallel(get_parallel(), "group build") self.assertIs(get_parallel().tp_group, right)
def test_a_draft_scope_states_a_consistent_topology(self): def test_a_draft_scope_states_a_consistent_topology(self):
"""The scope narrows four names at once, so the identity applies to it """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 import parallel_state
from sglang.srt.distributed.parallel_state import GroupCoordinator from sglang.srt.distributed.parallel_state import GroupCoordinator
@@ -2609,9 +2883,8 @@ class TestTheTopologyIdentities(CustomTestCase):
group = GroupCoordinator.__new__(GroupCoordinator) group = GroupCoordinator.__new__(GroupCoordinator)
group.world_size = 2 group.world_size = 2
group.rank_in_group = 1 group.rank_in_group = 1
with patch.object(parallel_state, "_TP", group): with parallel_state.patch_tensor_parallel_group(group, owns_attention=True):
with parallel_state.patch_tensor_parallel_group(group, owns_attention=True): self.assertEqual(get_parallel().attn_tp_size, 2)
self.assertEqual(get_parallel().attn_tp_size, 2)
class TestWhoAnswersDuringADraftScope(CustomTestCase): class TestWhoAnswersDuringADraftScope(CustomTestCase):
@@ -2650,11 +2923,10 @@ class TestWhoAnswersDuringADraftScope(CustomTestCase):
group = self._single_member_group() group = self._single_member_group()
self._two_stage_pipeline() self._two_stage_pipeline()
self.assertEqual(get_parallel().pp_size, 2) self.assertEqual(get_parallel().pp_size, 2)
with patch.object(parallel_state, "_PP", group): with parallel_state.patch_pipeline_parallel_group(group):
with parallel_state.patch_pipeline_parallel_group(group): self.assertEqual(get_parallel().pp_size, 1)
self.assertEqual(get_parallel().pp_size, 1) self.assertEqual(get_parallel().pp_rank, 0)
self.assertEqual(get_parallel().pp_rank, 0) self.assertIs(get_parallel().pp_group, group)
self.assertIs(get_parallel().pp_group, group)
self.assertEqual(get_parallel().pp_size, 2) self.assertEqual(get_parallel().pp_size, 2)
self.assertEqual(get_parallel().pp_rank, 1) self.assertEqual(get_parallel().pp_rank, 1)
@@ -2685,28 +2957,25 @@ class TestWhoAnswersDuringADraftScope(CustomTestCase):
self.assertEqual(get_parallel().attn_tp_size, 2) self.assertEqual(get_parallel().attn_tp_size, 2)
group = self._group(world_size=2, rank=1) 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):
with parallel_state.patch_tensor_parallel_group(group, owns_attention=True): parallel = get_parallel()
parallel = get_parallel() self.assertEqual(parallel.tp_size, 2)
self.assertEqual(parallel.tp_size, 2) self.assertEqual(parallel.attn_tp_size, 2)
self.assertEqual(parallel.attn_tp_size, 2) self.assertEqual(parallel.attn_tp_rank, 1)
self.assertEqual(parallel.attn_tp_rank, 1) self.assertEqual(parallel.attn_dp_size, 1)
self.assertEqual(parallel.attn_dp_size, 1) self.assertEqual(parallel.attn_dp_rank, 0)
self.assertEqual(parallel.attn_dp_rank, 0) self.assertEqual(parallel.attn_cp_size, 1)
self.assertEqual(parallel.attn_cp_size, 1) self.assertEqual(parallel.attn_cp_rank, 0)
self.assertEqual(parallel.attn_cp_rank, 0) # `dp_size` is the deployment's replica count, not a property of
# `dp_size` is the deployment's replica count, not a property # the group being installed, so the scope leaves it alone --
# of the group being installed, so the scope leaves it alone -- # `require_mlp_tp_gather` asserts on it under dp attention.
# `require_mlp_tp_gather` asserts on it under dp attention. self.assertEqual(parallel.dp_size, 2)
self.assertEqual(parallel.dp_size, 2) # The whole point of stating the rest: the identity the override
# The whole point of stating the rest: the identity the # path and the group build both check holds in here.
# override path and the group build both check holds in here. self.assertEqual(
self.assertEqual( parallel.tp_size,
parallel.tp_size, parallel.attn_tp_size * parallel.attn_dp_size * parallel.attn_cp_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().attn_dp_size, 2)
self.assertEqual(get_parallel().dp_size, 2) self.assertEqual(get_parallel().dp_size, 2)
@@ -2790,9 +3059,8 @@ class TestWhoAnswersDuringADraftScope(CustomTestCase):
self._two_stage_pipeline() self._two_stage_pipeline()
self.assertEqual(get_parallel().pp_size, 2) self.assertEqual(get_parallel().pp_size, 2)
group = self._single_member_group() group = self._single_member_group()
with patch.object(parallel_state, "_PP", group): with parallel_state.patch_pipeline_parallel_group(group):
with parallel_state.patch_pipeline_parallel_group(group): checker = WeightChecker(get_model=lambda: None)
checker = WeightChecker(get_model=lambda: None)
# The scope has closed and the context answers the target's shape again. # The scope has closed and the context answers the target's shape again.
self.assertEqual(get_parallel().pp_size, 2) self.assertEqual(get_parallel().pp_size, 2)