minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Refactor] Clean up parallel runtime comments (#40632)

Browse Source
This commit is contained in:
Cheng Wan
2026-09-21 14:32:22 -07:00
committed by GitHub
parent f532ad1f9a
commit acac4dd9d9
63 changed files with 199 additions and 1122 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/benchmark/one_batch.py
-1
View File
@@ -327,7 +327,6 @@ def load_model(server_args, port_args, gpu_id, tp_rank):
server_args=server_args,
)
# Phase two: this entry has no scheduler to run it.
bootstrap.init_parallel_runtime(
server_args=server_args,
model_config=model_config,
python/sglang/multimodal_gen/runtime/distributed/parallel_state.py
+5 -27
View File
@@ -102,8 +102,6 @@ def _sync_srt_world_group() -> None:
if srt_parallel_state._WORLD is None:
srt_parallel_state._WORLD = _WORLD
if srt_parallel_state._WORLD is _WORLD:
# On the context too: that is where a handle is read from, and
# assigning the module global above does not reach it.
get_parallel().override_permanently(world_group=_WORLD)
@@ -115,18 +113,11 @@ def _clear_srt_world_group() -> None:
def _sync_srt_tp_group() -> None:
"""Lend this package's TP group to `srt`, and state the widths it implies.
"""Expose this package's TP group, widths, and ranks to shared SRT layers.
Shared `srt` layers run in this package and ask `get_parallel()` for how to
shard -- `srt/layers/attention/vision.py` reads `attn_tp_size`. The
published `srt` config cannot answer: `gpu_worker.py` publishes a dummy
carrying *this* package's `tp_size`, which a sequence-parallel launch sets
to 1 while the group lent here is as wide as the world. So the widths are
permanently overridden alongside the group -- this runs with no `srt`
config published at all, which is exactly why it cannot go through
`RuntimeContext.override` (it requires one).
Only tensor parallelism folds this way, so every other dimension is one.
Use the group's actual width: sequence parallelism can make it wider than
the TP size in the dummy SRT configuration. Other parallel dimensions are
one. Overrides also work before SRT configuration is published.
"""
import sglang.srt.distributed.parallel_state as srt_parallel_state
from sglang.srt.runtime_context import derive_parallel_widths, get_parallel
@@ -137,21 +128,9 @@ def _sync_srt_tp_group() -> None:
srt_parallel_state._ATTN_TP = _TP
if srt_parallel_state._ATTN_TP is _TP:
get_parallel().override_permanently(
# The group itself, because that is what the `srt` context answers
# a handle with -- assigning the module global above does not reach
# it. `tp_size` comes with them: the group is as wide as the world
# while the dummy carries this package's, and the widths below are
# quotients of one number, so stating a subset would describe a
# layout that does not exist.
tp_group=_TP,
attn_tp_group=_TP,
tp_size=_TP.world_size,
# The ranks too. The shared layers shard by them -- `vision.py`
# reads `attn_tp_rank`, every `srt` linear built without an
# explicit rank reads `tp_rank` -- and this package publishes no
# rank bundle, so nothing else writes one. The draft has no
# pipeline, context or expert dimension of its own, so those
# positions are zero.
tp_rank=_TP.rank_in_group,
attn_tp_rank=_TP.rank_in_group,
moe_tp_rank=_TP.rank_in_group,
@@ -178,8 +157,7 @@ def _clear_srt_tp_group() -> None:
srt_parallel_state._ATTN_TP = None
get_parallel().clear_stamp()
if srt_parallel_state._WORLD is not None:
# `clear_stamp` drops every stamped name; the WORLD group this
# package lent is still built, so hand it back.
# Restore the still-active WORLD handle after clearing TP overrides.
get_parallel().override_permanently(world_group=srt_parallel_state._WORLD)
if srt_parallel_state._TP is _TP:
srt_parallel_state._TP = None
python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py
-5
View File
@@ -148,13 +148,8 @@ def test_srt_attention_tp_group_tracks_diffusion_tp_group():
assert srt_parallel_state._TP is tp_group
assert srt_parallel_state._ATTN_TP is tp_group
assert get_parallel().attn_tp_size == 2
# The handle too: assigning the module global does not reach the `srt`
# context, which is what the shared layers ask for a group.
assert get_parallel().tp_group is tp_group
assert get_parallel().attn_tp_group is tp_group
# And the ranks the shared layers shard by. Nothing else writes one
# here: this package publishes no rank bundle, so a handle without a
# rank leaves every `srt` linear unable to say which shard it is.
assert get_parallel().tp_rank == 1
assert get_parallel().attn_tp_rank == 1
python/sglang/srt/arg_groups/arg_utils.py
+7 -29
View File
@@ -96,40 +96,18 @@ _NO_DEFAULT = object()
class Derived(msgspec.Struct, frozen=True):
"""Metadata for a field the configuration implies, not one anyone types.
"""Metadata for namespace fields that are not CLI inputs or record fields.
The other half of a namespace. An ``Arg`` field is the operator's input and
is collected into ``ServerArgs``; a ``Derived`` field carries no annotation,
so it is not a dataclass field and never reaches the record -- which is
right, because it has no input to preserve and the record is what crosses a
process boundary.
``fn`` names what computes it, as a dotted path resolved lazily so that a
declaration module stays free of runtime imports. Such a field is a pure
function of the published configuration, so it is computed once at
``publish`` and stored as an ordinary bag leaf -- a plain attribute load,
which is what a read inside compiled model code needs.
A declaration with no ``fn`` is one nothing can compute: a rank, or a
process group. Those are written into the namespace at runtime -- by
``publish`` from the spawn bundle, or by the build that creates the group --
and until then the name has no answer.
Most declarations carry ``fn``, the parallel quotients included:
they are a function of the configured leaves, so they are computed at
publish like the rest. What is special about them is not how they are
computed but that a stamp can move one afterwards -- ``initialize_dp_attention``
restamps ``attn_dp_size`` -- which ``ParallelContext`` answers above the
published leaf.
``fn`` is a lazily resolved dotted function path. It computes a value from
resolved configuration once at publication. Fields without ``fn``, such as
ranks and group handles, are set at runtime. Parallel overrides take
precedence over published values.
"""
doc: str = ""
fn: str = ""
# For a declaration with no ``fn`` whose absence is itself an answer:
# ``gpu_id`` is ``None`` in a process that runs on no device, and a reader
# wants that rather than an error. A rank has no such value -- the wrong
# one is a hang in a collective -- so it carries no default and a read
# before the write says so.
# Default for runtime-only fields, e.g. ``gpu_id=None`` without a device.
# Fields without a default raise if read before initialization.
default: Any = _NO_DEFAULT
python/sglang/srt/arg_groups/fields/parallel.py
+2 -19
View File
@@ -280,18 +280,7 @@ class Parallel(msgspec.Struct):
"Maximum EP size the server can scale to at runtime. Pre-allocates active-rank state and backend buffers to this size. Defaults to the launch-time world size.",
] = None
# ---- derived: the quotients of the leaves above -------------------------
#
# Declared here, beside what they are computed from, because a namespace is
# one file and one class. They are not annotated, so they are not dataclass
# fields and `collect_input_fields` does not put them on the record -- which
# is right: a quotient has no operator input to preserve, and the record is
# what crosses a process boundary, so a width put there would be a stale
# copy the moment an elastic scale-up restamps one. Every input is a leaf
# above, so all six are fixed once the configuration is: `publish` computes
# them through `parallel_widths_of` and stores them as ordinary bag leaves,
# and `ParallelContext` answers with the stamp when a scale-up has moved
# one.
# Derived fields are computed at publication and are not stored in ServerArgs.
attn_tp_size = Derived(
fn="sglang.srt.runtime_context.attn_tp_size_of",
doc="Attention tensor-parallel width: `tp_size` divided by the "
@@ -321,13 +310,7 @@ class Parallel(msgspec.Struct):
"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.
# Runtime fields: publish sets ranks; distributed initialization sets groups.
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.")
python/sglang/srt/disaggregation/encoder/server.py
+1 -6
View File
@@ -589,12 +589,7 @@ class MMEncoder:
distributed_init_method=dist_init_method,
local_rank=rank,
)
# The encoder serves the vision tower on a world of its own: `tp_size`
# ranks wide, with no pipeline, no expert or MoE-DP dimension and no
# decode context parallelism, whatever the generation side published.
# That has always been the layout it builds; stating it is what stops
# the context from answering with the other side's topology while these
# groups answer with this one.
# The encoder uses a separate WORLD with tensor and attention-CP parallelism.
parallel = get_parallel()
attn_cp_size = parallel.attn_cp_size
attn_tp_size = parallel.tp_size // attn_cp_size
python/sglang/srt/distributed/bootstrap.py
+10 -31
View File
@@ -59,23 +59,18 @@ _is_cpu_arm64 = is_host_cpu_arm64()
_TP_ALL_TO_ALL_WARMUP_BYTES_PER_PEER = 4 << 20
#: Set by `init_parallel`; `destroy_model_parallel` clears it, so a test that
#: tears the groups down can build them again.
#: Cleared by `destroy_model_parallel`.
_PARALLEL_INITIALISED = False
def reset_parallel_initialised() -> None:
"""Forget that the groups were built. Paired with tearing them down."""
"""Reset the initialization guard when model-parallel groups are destroyed."""
global _PARALLEL_INITIALISED
_PARALLEL_INITIALISED = False
def _bind_threads_if_cpu(*, device: str) -> "Optional[List[int]]":
"""Pin OpenMP threads to this process's NUMA node, on CPU.
A precondition of the CPU group build, which reads the binding, so it is
done here rather than left for a caller to remember.
"""
"""Bind OpenMP threads to the NUMA node before CPU group initialization."""
if device != "cpu":
return None
from sglang.srt.utils import numa_utils
@@ -99,23 +94,11 @@ def init_parallel_runtime(
device: str,
dist_port: int,
) -> None:
"""Phase two of startup: bring the parallel runtime up, once.
"""Initialize the parallel runtime once, after publishing configuration.
Publish says what the topology is; this makes it exist. Nothing returns,
because the groups are read through the runtime context -- a caller that
wants one asks `get_parallel()`, in this process or any later phase.
"Runtime" rather than "groups": two things have to be in place before the
groups can be built, and they are done here rather than left for every
entry to remember. The OpenMP/NUMA binding is what the CPU group build
reads, and the shared Mooncake transfer engine is what the Mooncake
process-group backend asks for -- create that one late and a second engine
appears. Both are preconditions of the build, not separate work.
Runs on the target worker only. A draft worker shares its target's groups,
which is why this is a phase the entry runs rather than something a runner
does on its way up: whether the groups exist must not depend on which
runner happened to be constructed first.
Set up CPU thread binding, the current device, and the shared Mooncake
engine before creating process groups. Draft workers reuse their target's
groups and must not call this function.
"""
global _PARALLEL_INITIALISED
if _PARALLEL_INITIALISED:
@@ -140,9 +123,7 @@ def init_parallel_runtime(
_set_all_reduce_flags()
local_omp_cpuid = _bind_threads_if_cpu(device=device)
# Everything below allocates on the current device -- the NCCL warm-up, the
# mooncake all-reduce buffer -- and without this every rank on a node would
# pick device 0, because the default is not to reindex the visible set.
# Select the local device before communicator allocation and warmup.
try:
torch.get_device_module(device).set_device(get_device().gpu_id)
except Exception:
@@ -199,11 +180,9 @@ def init_parallel_runtime(
def measure_pre_model_load_memory(*, device: str, is_draft_worker: bool) -> float:
"""Available memory after the groups exist and before the model loads.
"""Measure available memory for KV-cache sizing before this runner loads.
Sized into the KV cache later, so it has to be taken at exactly this point
-- which is why it stays with the runner rather than moving into the
parallel phase.
Call after parallel initialization and before allocating model weights.
"""
before_avail_memory = get_available_gpu_memory(device, get_device().gpu_id)
python/sglang/srt/distributed/parallel_state.py
+17 -85
View File
@@ -2156,12 +2156,10 @@ _PDMUX_PREFILL_TP_GROUP: Optional[GroupCoordinator] = None
@contextmanager
def pdmux_prefill_tp_group():
"""Run on the prefill stream's own tensor-parallel communicator.
"""Use the duplicate TP communicator for the prefill stream.
PD multiplexing builds a duplicate TP group -- the same ranks, a second
communicator -- so prefill and decode can occupy separate streams without
serialising on one. Nothing about the topology differs, so the scope states
the handle and nothing else.
PD multiplexing keeps prefill and decode on separate communicators with
the same ranks. Only the TP handle changes within this scope.
"""
assert _PDMUX_PREFILL_TP_GROUP is not None, (
"tensor model parallel group for PD-Multiplexing Prefill is not initialized"
@@ -2503,9 +2501,7 @@ def init_distributed_environment(
assert _WORLD.world_size == torch.distributed.get_world_size(), (
"world group already initialized with a different world size"
)
# Stated here rather than with the groups below it: WORLD is built in this
# function, and every group `initialize_model_parallel` builds is placed by
# reading it back.
# Publish WORLD before model-parallel initialization reads its local rank.
get_parallel().override_permanently(world_group=_WORLD)
@@ -2520,17 +2516,8 @@ def initialize_model_parallel(
"""
Initialize model parallel groups at the published widths.
Every width comes from the runtime context rather than from an argument:
the configuration already says how wide each dimension is, and a caller
that translates it again is a second place for the two to disagree. A
process that needs a narrower layout than the one it published -- the
media encoder is the case in the tree -- states that layout on the context
first, so what it builds and what it answers stay the same thing.
The remaining arguments are not topology. `backend` is decided by the
device, `duplicate_tp_group` and `enable_symm_mem` by other namespaces, and
`recovered_rank` / `rank_offset` / `max_world_size` describe this
particular join rather than the layout being joined.
Read topology widths from ``get_parallel()``. Callers needing a different
layout must override the context before building groups.
The widths this reads:
tp_size: GPUs used for tensor model parallelism.
@@ -2935,19 +2922,8 @@ def initialize_model_parallel(
max_world_size=max_world_size,
)
# The groups just built and the configuration they were built from are two
# accounts of one layout, and this is where they meet: stating a group
# checks the identities, so a group built on the wrong peers is refused
# here rather than hanging in a collective later.
#
# A dimension this configuration does not have is left unstated -- `_DCP`
# is None without decode context parallelism -- so reading it says the
# group was never built, which is what these getters have always said,
# rather than handing back a None to fail on at the collective.
#
# WORLD is not here: it is built and stated by
# `init_distributed_environment`, which is what lets every build above
# place its group by reading `get_world_group().local_rank`.
# Validate group widths against the context. Leave disabled groups unset
# so reading them raises. WORLD was published by distributed initialization.
built = {
"tp_group": _TP,
"pp_group": _PP,
@@ -3061,8 +3037,6 @@ def patch_pipeline_parallel_group(pp_group: GroupCoordinator):
old_pp_group = _PP
_PP = pp_group
try:
# `pp_size` is a configured leaf: unlike the rank and the handle it
# does not follow the group being swapped, so the scope has to name it.
with get_parallel().override(
pp_size=pp_group.world_size,
pp_rank=pp_group.rank_in_group,
@@ -3076,36 +3050,12 @@ def patch_pipeline_parallel_group(pp_group: GroupCoordinator):
@contextmanager
def patch_tensor_parallel_group(tp_group: GroupCoordinator, *, owns_attention: bool):
"""Run under a different tensor-parallel group until this scope ends.
"""Temporarily replace the TP group and its runtime-context values.
This is for draft workers of speculative decoding, which run the draft model
at the target's attention-TP width rather than its global TP width.
The scope replaces both the module global that ``get_tp_group()`` reads and
the members the runtime context answers with.
Which members depends on what the draft is, and only the worker knows: the
same call site hands over an attention-TP slice for one draft and the
target's whole TP group for another, so this cannot be read off the group.
``owns_attention`` says which. A draft that owns its attention topology
runs the whole model on the group being installed -- there is no
attention-DP replica inside it, so its attention identity is the group
itself, one replica, one context shard, and no expert dimension either.
Leaving those names on the target's answers is what lets a draft read
report a replica count the draft does not have.
A draft that does not own it was built outside any scope and keeps the
target's layout: the process is still one of several attention-DP replicas
and still gathers with them. Claiming one replica there is the same error
in the other direction, and the reader that acts on it is a collective -- a
DP gather takes its buffer size from the replica count and its communicator
from this group, so the two stop agreeing.
Args:
tp_group (GroupCoordinator): the tp group coordinator
owns_attention (bool): whether the draft's attention topology is this
group, decided by the worker where it builds its draft runner
For speculative drafts with ``owns_attention=True``, the installed group
is the draft's attention-TP group, with attention-DP, attention-CP, and
MoE-DP/EP widths set to one. Otherwise, retain the target's attention
topology. The worker must specify this based on how it constructed the draft.
"""
global _TP_STATE_PATCHED
@@ -3458,20 +3408,9 @@ def monkey_patch_vllm_parallel_state(reverse: bool = False):
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.
# Use `get_parallel()` outside this package. Warn once per deprecated getter.
_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",
@@ -3494,13 +3433,8 @@ _CONTEXT_NAME_OF = {
"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.
# Only deprecate width getters whose group widths are validated against
# configuration by `_WIDTH_AND_GROUP` in `runtime_context`.
_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"
@@ -3540,9 +3474,7 @@ del _name, _replacement, _fn
# What `from sglang.srt.distributed import *` re-exports: everything public
# except the deprecated getters. Business code reaches them through
# `get_parallel()`, and the package that defines them imports them from this
# module by name, so nothing needs the package path to reach one.
# except the deprecated getters.
__all__ = [
_public
for _public in list(globals())
python/sglang/srt/layers/dp_attention.py
+12 -30
View File
@@ -45,12 +45,10 @@ if TYPE_CHECKING:
def deployment_attn_dp_size() -> int:
"""Attention-DP replicas in the deployment, which no draft scope narrows.
"""Return the deployment's attention-DP replica count.
A draft runs on one attention-DP replica and its scope says so, but the
metadata a draft gathers is shaped by the replicas it gathers *with* --
the target's. Those come from the configuration, which the scope leaves
alone, so this answers the same number inside the scope and outside it.
Draft scopes retain this count because their metadata gathers include
the target's replicas.
"""
parallel = get_parallel()
attn_dp_size, _ = derive_attention_widths(
@@ -63,25 +61,20 @@ def deployment_attn_dp_size() -> int:
def dp_gather_width() -> int:
"""How many replicas the DP sync gathers over.
"""Return the DP gather width.
The attention-DP replicas, except after an elastic-EP scale-up, when the
gather spans the expanded WORLD -- whose width is the `dp_size` the
scale-up published. Read from the context either way: a scoped width has
to reach this, which is the whole reason the name has one home.
After elastic scale-up, the gather spans the expanded WORLD; otherwise
it spans the attention-DP replicas.
"""
parallel = get_parallel()
return parallel.dp_size if world_dp_gather_enabled() else parallel.attn_dp_size
def dp_gather_slot() -> int:
"""This process's index in the list the DP sync just gathered.
"""Return this process's index in the DP gather.
The gather spans the attention-DP replicas, except after an elastic-EP
scale-up, when it spans the expanded WORLD and the joining cohort is
numbered from its offset. Which list was gathered is what the flag below
says, so the index is read from there rather than kept as a second name on
the topology.
After elastic scale-up, use the TP rank plus the join offset; otherwise
use the attention-DP rank.
"""
parallel = get_parallel()
if world_dp_gather_enabled():
@@ -100,12 +93,10 @@ def enable_joiner_all_gather():
def update_dp_attention_post_scale(new_dp_size: int, new_dp_rank: int):
"""Point the DP gather at the expanded WORLD.
"""Switch DP gathers to the expanded WORLD.
The widths themselves are not written here: the caller scales `dp_size` on
the published bag, and the gather reads its width and this process's slot
from there. The arguments are the values the caller is about to publish,
kept so the log says which scale-up this was.
The caller updates the configured widths; these arguments identify the
scale-up in the log.
"""
get_flags().dp.use_world_group_for_gather = True
logger.debug(
@@ -433,9 +424,6 @@ def initialize_dp_attention(
if ep_scale_joiner_of(resolving_view(server_args)):
dp.joiner_skip_all_gather = True
# Stamped together, after the elastic adjustment: the width and the rank
# describe one topology, and a reader that caught them mid-update would
# see this process placed in a group it is not in.
get_parallel().override_permanently(
attn_dp_size=attn_dp_size, attn_dp_rank=attn_dp_rank
)
@@ -457,9 +445,6 @@ def is_allocation_symmetric() -> bool:
def get_dp_local_info(forward_batch: ForwardBatch) -> Tuple[torch.Tensor, torch.Tensor]:
# `get_dp_local_info` is only called in global DP gather and scatter. We use global DP rank here.
# The slot in the list that was gathered. A scale-up widens that list
# to WORLD, and this process's index in it is not its index among the
# launch replicas.
dp_rank = dp_gather_slot()
if forward_batch.dp_local_start_pos is None:
@@ -484,9 +469,6 @@ def get_dp_local_slice_cpu(
# CPU (start, length) slice for DP-local data in a rank-padded buffer.
# Returns Python ints (no D2H sync) and handles the cuda-graph-padded layout.
global_num_tokens = forward_batch.global_num_tokens_cpu
# The slot in the list that was gathered. A scale-up widens that list
# to WORLD, and this process's index in it is not its index among the
# launch replicas.
dp_rank = dp_gather_slot()
local_num_tokens = global_num_tokens[dp_rank]
if can_run_graph:
python/sglang/srt/layers/sampler.py
+1 -4
View File
@@ -112,10 +112,7 @@ class Sampler(nn.Module):
self.cp_sync_group = None
if is_dp_attention_enabled():
self.tp_sync_group = get_parallel().attn_tp_group.device_group
# Only when there is more than one context shard to reconcile. The
# sync below already short-circuits on that, and a model running on
# one shard -- a speculative draft, under the scope that says so --
# has no context-parallel communicator to name.
# Single-shard drafts may have no context-parallel group.
if get_parallel().attn_cp_size > 1:
self.cp_sync_group = get_parallel().attn_cp_group.device_group
python/sglang/srt/managers/scheduler.py
+1 -15
View File
@@ -517,8 +517,6 @@ class Scheduler(
# Init ZBAL, switch allocator should before any torch alloc action
self.init_zbal_on_npu()
# The groups are the first thing that allocates, so this comes after the
# allocator switch above and before anything that reads a group.
bootstrap.init_parallel_runtime(
server_args=server_args,
model_config=self.model_config,
@@ -5910,7 +5908,6 @@ def dispatch_event_loop(scheduler: Scheduler):
def _dispatch_event_loop_once(scheduler: Scheduler):
# The live PP property asserts before torch.distributed init (MLX stub).
disaggregation_mode: DisaggregationMode = scheduler.disaggregation_mode
if disaggregation_mode == DisaggregationMode.NULL:
if scheduler.enable_pdmux:
@@ -5940,15 +5937,8 @@ def _dispatch_event_loop_once(scheduler: Scheduler):
def resolve_spawn_dp_rank(dp_rank: Optional[int]) -> Optional[int]:
"""The `dp_rank` this process was spawned with, in either of its two forms.
A router does not pass it as an argument, it sets `SGLANG_DP_RANK`. Both
forms are the launcher naming this process's place, so both have to be in
hand before `publish` records the placement -- resolving one of them after
would leave the context answering `None` for a process that has a rank.
"""
"""Resolve the launcher DP rank, falling back to ``SGLANG_DP_RANK``."""
if dp_rank is None and "SGLANG_DP_RANK" in os.environ:
# [For Router] if env var "SGLANG_DP_RANK" exist, set dp_rank to the value of the env var
return int(os.environ["SGLANG_DP_RANK"])
return dp_rank
@@ -6034,10 +6024,6 @@ def run_scheduler_process(
# Load plugins so hooks can override Scheduler and its dependencies.
load_plugins()
dp_rank = resolve_spawn_dp_rank(dp_rank)
# Publish before anything in this process reads configuration, with the
# placement the launcher decided: from here on a rank read is answered
# without a process group, which is what every reader needs before
# `init_torch_distributed` has run.
publish(
server_args,
role="scheduler",
python/sglang/srt/managers/scheduler_components/dp_attn.py
+1 -3
View File
@@ -59,9 +59,7 @@ def _resolve_elastic_world_dp_size(
live_dp_size = dp_gather_width()
effective_ep_size = ElasticEPStateManager.get_effective_ep_size()
# The group's own membership, not the width it was built at: this is the
# one number an out-of-process join moves, and it is the upper bound the
# served width has to stay under.
# Query live membership because elastic joins can expand WORLD.
world_size = torch.distributed.get_world_size(group)
if live_dp_size != effective_ep_size:
python/sglang/srt/mem_cache/kv_cache_builder.py
-2
View File
@@ -267,8 +267,6 @@ def build_kv_cache(
enable_hierarchical_cache: bool,
hicache_draft_plan: Optional[HiCacheDraftPlan] = None,
) -> KVCacheBuildResult:
# Built from the scheduler loop, outside any draft scope, so the context
# answers for the process this cache belongs to.
parallel = get_parallel()
sliding_window_size: Optional[int] = None
full_tokens_per_layer: Optional[int] = None
python/sglang/srt/mem_cache/kv_cache_configurator.py
+1 -2
View File
@@ -259,8 +259,7 @@ class _PoolSizes(msgspec.Struct, frozen=True, kw_only=True):
class KVCacheConfigurator:
device: str
gpu_id: int
# Frozen at construction, not asked for later: this configurator is built
# inside the scope that describes a draft runner and used outside it.
# Capture draft placement at construction; the configurator outlives the scope.
attn_dp_size: int
pp_size: int
pp_group: Any
python/sglang/srt/mem_cache/pool_host/base.py
+2 -6
View File
@@ -46,13 +46,9 @@ def host_memory_budget_scope(budget_bytes: int):
def ranks_per_host() -> int:
"""Number of ranks of this job running on the same machine as this one.
"""Return the launch ranks per host, assuming uniform placement.
Derived as the launch width // nnodes: the launcher slices ranks
uniformly across nodes (resolution asserts divisibility), so no hostname
collective is needed — a collective here would have to be issued the same
number of times on every rank, and ranks build different numbers of host
pools.
Avoid a collective: ranks may construct different numbers of host pools.
"""
if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
return 1
python/sglang/srt/model_executor/model_runner.py
+1 -5
View File
@@ -1157,11 +1157,7 @@ class ModelRunner:
self.pre_model_load_memory = bootstrap.measure_pre_model_load_memory(
device=self.device, is_draft_worker=self.is_draft_worker
)
# Read once, here: a draft runner is constructed inside the scope that
# states its topology and used outside it, so what it holds has to be
# the placement it was built for rather than whatever the context
# answers later. Groups and widths alike -- a runner asked about its own
# shape after the scope has closed must still describe itself.
# Capture draft placement at construction; the runner outlives the scope.
parallel = get_parallel()
self.tp_group = parallel.tp_group
self.pp_group = parallel.pp_group
python/sglang/srt/models/inkling_common/moe.py
+2 -6
View File
@@ -671,12 +671,8 @@ class InklingSharedFusedMoE(FusedMoE):
quant_config: QuantizationConfig | None,
inference_moe_w13_interleaved: bool,
) -> None:
# FusedMoE.__init__ reads get_parallel() once and caches it on self, so
# scoping the override to just this call is sufficient for the module's lifetime.
# The shared experts are replicated rather than sharded, so there is no
# expert-parallel communication here and no group to name: a width of
# one with the wider group still installed would describe a layout that
# does not exist.
# FusedMoE caches this topology at construction. Shared experts are
# replicated, so they need no expert-parallel group.
with get_parallel().override(
moe_ep_size=1,
moe_ep_rank=0,
python/sglang/srt/runtime_context.py
+61 -257
View File
@@ -13,13 +13,10 @@
# ==============================================================================
"""A single structured accessor for process-static runtime state.
``get_parallel()`` returns a ``ParallelContext``. Ranks and process-group handles
read through **live** to the canonical getter in ``distributed.parallel_state`` /
``layers.dp_attention`` — exactly what those getters return, a read-through
wrapper and not a cache. Every other name, the sizes included, is a leaf of the
published ``parallel`` bag. It gives call-sites one import and one naming scheme
in place of a dozen free functions, plus an ``override()`` hook to force a
topology without monkeypatching the underlying getters.
``get_parallel()`` returns a ``ParallelContext`` for configuration, ranks, and
process-group handles. Reads use scoped overrides, then permanent overrides,
then the published configuration. ``publish`` records ranks from ``SpawnRanks``;
distributed initialization records group handles.
``get_server_args()`` returns the process-wide ``ServerArgs``. This is the
user's raw input, kept **read-only** for debug and reproduction; what
@@ -72,13 +69,7 @@ _PARALLEL_STATE = None
def _ps():
"""The module every rank and group read ends at.
Cached because the import statement dominated the read: a group read is
two attribute lookups plus this, and it runs per row-linear on an eager
forward. The getter is still resolved by name on the returned module, so
a test that patches `parallel_state.get_tp_group` is still seen.
"""
"""Lazily import and cache the parallel-state module."""
global _PARALLEL_STATE
if _PARALLEL_STATE is None:
from sglang.srt.distributed import parallel_state
@@ -95,11 +86,7 @@ def _dp():
@functools.lru_cache(maxsize=1)
def _parallel_config_leaves() -> frozenset:
"""Names under the ``parallel`` namespace, for the unpublished error path.
Read from the field metadata rather than the bag, which is what does not
exist yet when this is needed.
"""
"""Return configured parallel field names, including before publication."""
from sglang.srt.arg_groups.arg_utils import namespace_of
from sglang.srt.server_args import ServerArgs
@@ -110,31 +97,12 @@ def _parallel_config_leaves() -> frozenset:
)
# Ranks and group handles: the names no configuration carries. This table is
# their declaration, the way `arg_groups/fields/parallel.py` is the leaves' and
# `Derived` is the widths'. A group handle names the getter that owns it,
# because the module that builds the groups is where it lives; a rank is a
# position in one of those groups, so it is read off the handle. `None` marks a
# name only a stamp can answer: no coordinator knows this process's
# attention-DP rank.
_MISSING_READ = object()
@functools.lru_cache(maxsize=1)
def _parallel_fields() -> frozenset:
"""Every name `ParallelContext` answers for, read from the declarations.
Three sources, because the namespace has three kinds of name and each one
declares itself somewhere already:
* configured leaves -- the `parallel` namespace of the record;
* derived widths -- the `Derived` declarations beside those leaves;
* ranks and group handles -- declared beside the leaves with no `fn`,
because nothing computes them; they are written at runtime.
The set is the union of those three, so `override()` cannot refuse a name
the class answers for.
"""
"""Return configured and derived parallel names accepted by overrides."""
from sglang.srt.arg_groups.arg_utils import Derived
from sglang.srt.arg_groups.fields.parallel import Parallel
@@ -147,12 +115,7 @@ def _parallel_fields() -> frozenset:
def derive_attention_widths(
*, tp_size: int, attn_cp_size: int, dp_size: int, enable_dp_attention: bool
) -> tuple:
"""(attn_dp_size, attn_tp_size) from the leaves.
Split out because the rank computation in
`dp_attention.compute_dp_attention_world_info` needs the same two numbers
and must not carry a second copy of the arithmetic.
"""
"""Return (attn_dp_size, attn_tp_size) from the configured widths."""
attn_dp_size = dp_size if enable_dp_attention else 1
return attn_dp_size, tp_size // attn_dp_size // attn_cp_size
@@ -160,18 +123,12 @@ def derive_attention_widths(
def derive_attention_ranks(
*, tp_rank: int, attn_tp_size: int, attn_cp_size: int, enable_dp_attention: bool
) -> tuple:
"""(attn_tp_rank, attn_dp_rank) for a process at `tp_rank`.
"""Return (attn_tp_rank, attn_dp_rank) for a process at ``tp_rank``.
The rank layout is (dp, cp, tp) with tp the fastest-changing dimension::
The rank layout is (dp, cp, tp), with tp changing fastest::
tp_rank = (attn_dp_rank * attn_cp_size + attn_cp_rank) * attn_tp_size
+ attn_tp_rank
Split out beside `derive_attention_widths` because two places need it from
different inputs: `publish` has this process's `tp_rank` from the spawn and
the widths from the configuration, while `initialize_dp_attention` has them
from the groups it just built. They must not carry separate copies of the
arithmetic -- the point of computing it at publish is that the two agree.
"""
attn_tp_rank = tp_rank % attn_tp_size
if not enable_dp_attention:
@@ -180,13 +137,9 @@ def derive_attention_ranks(
def spawn_world_rank(server_args, *, tp_rank: int, pp_rank: int) -> int:
"""This process's place in WORLD, from the ranks its entry was given.
"""Compute the WORLD rank from launcher TP and PP ranks.
The inverse of `derive_spawn_ranks`, for the entries that have the pieces
but not the whole: the same expression `bootstrap` hands to
`init_distributed_environment` when it builds the group.
Reads a resolving view because it runs before `publish`.
Uses a resolving view because callers may run before ``publish``.
"""
from sglang.srt.arg_groups.model_override_base import resolving_view
@@ -204,17 +157,10 @@ def derive_spawn_ranks(
moe_dp_size: int,
moe_ep_size: int,
) -> dict:
"""Every rank a process group would answer, from its place in WORLD.
"""Derive TP, PP, attention, and MoE ranks without process groups.
WORLD is laid out `rank = ep_join_rank_offset + tp_size * pp_rank +
tp_rank`: `initialize_model_parallel` builds tensor-parallel groups as
contiguous blocks of `tp_size` and pipeline groups strided by it, so the
map is a bijection and this is its inverse. The attention and MoE ranks
are then positions inside the tensor-parallel block, which is what the
launcher computes when it decides what to spawn.
Pure arithmetic over the published widths: no group is consulted, which is
the point -- this runs at publish, before any of them exist.
WORLD uses ``rank = ep_join_rank_offset + tp_size * pp_rank + tp_rank``.
TP groups are contiguous blocks; PP groups are strided by ``tp_size``.
"""
local = world_rank - ep_join_rank_offset
tp_rank = local % tp_size
@@ -242,27 +188,10 @@ def derive_parallel_widths(
dcp_size: int,
dcp_enabled: bool,
) -> dict:
"""The parallel widths no flag sets, from the leaves that do.
`tp_size` and its siblings are configured; these are quotients of them, so
the arithmetic lives here rather than being read back off the group
coordinators.
The world widths are not among them: neither is a quotient, and they are
not one number. `launch_world_size` is what the WORLD group was built at
and is frozen there -- `GroupCoordinator.world_size` is `len(ranks)`, so it
does not move when mooncake admits ranks into an expandable WORLD;
`max_world_size` is what that group has room for. How much of that room is
serving right now is elastic-EP state and is asked of its owner. Deriving
either from the leaves would be wrong in a further way: on a scale joiner it
would answer with the joining cohort's own `tp * pp`, while that process's
WORLD spans `ep_join_rank_offset + tp * pp`.
"""
"""Derive attention and MoE widths and DCP settings from configuration."""
return {
"attn_dp_size": attn_dp_size,
# `attn_dp_size` is already the effective width (1 when DP attention is
# off), so the flag is spent here; a caller passing the raw `dp_size`
# leaf with the attention disabled would get tp/dp/cp instead of tp/1/cp.
# `attn_dp_size` already accounts for disabled DP attention.
"attn_tp_size": derive_attention_widths(
tp_size=tp_size,
attn_cp_size=attn_cp_size,
@@ -277,14 +206,7 @@ def derive_parallel_widths(
def parallel_widths_of(cfg: Any) -> dict:
"""The six quotients, from a resolved config.
Every input is a record field, so this is a function of the configuration
and nothing else -- which is why the six are declared `Derived(fn=...)` and
computed once at publish rather than on every read. `dcp_enabled` is
`dcp_size > 1` because that is exactly when `initialize_model_parallel`
builds the group.
"""
"""Return derived parallel settings from resolved configuration."""
attn_dp_size, _ = derive_attention_widths(
tp_size=cfg.tp_size,
attn_cp_size=cfg.attn_cp_size,
@@ -303,20 +225,15 @@ def parallel_widths_of(cfg: Any) -> dict:
def launch_world_size_of(cfg: Any):
"""`launch_world_size`, computed at publish.
"""Return the initial WORLD width, including existing ranks below a joiner.
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.
This value remains fixed after elastic scale-up.
"""
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 the WORLD capacity: ``max_ep_size`` or the launch width."""
return cfg.max_ep_size or launch_world_size_of(cfg)
@@ -351,25 +268,12 @@ def dcp_enabled_of(cfg: Any):
class SpawnRanks(msgspec.Struct, frozen=True):
"""Where the launcher put this process, in the two numbers only it knows.
"""Process placement supplied by the launcher.
`world_rank` is this process's place in the WORLD group, which fixes every
other rank: the groups are laid out from the published widths, so
`tp_rank`, `pp_rank` and the attention / MoE ranks are functions of it (see
`derive_spawn_ranks`). Passing them separately would be passing the same
fact five more times, with five more ways for an entry to contradict
itself.
`dp_rank` is the exception, because data-parallel replicas are separate
WORLD groups: with `--dp-size 2` each replica holds ranks `0 .. n-1`, so
the rank cannot say which replica this is. `None` means "no controller",
which is an answer rather than an absence, and it is recorded as one.
`gpu_id` is the device the parent picked for this process. It is not a
position in any group -- reindexing narrows the visible devices before the
spawn, and Ray allocates from its own pool -- but it is the same kind of
fact: something only the entry that spawned the process can state. `None`
for a process that runs on no device.
``world_rank`` determines TP, PP, attention, and MoE ranks from the
configured widths. ``dp_rank`` identifies the replica across separate
WORLD groups; ``None`` means no data-parallel controller. ``gpu_id`` is
the assigned device index, or ``None`` for a process without a device.
"""
world_rank: int
@@ -386,10 +290,8 @@ _RANK_AND_WIDTH = (
("moe_ep_rank", "moe_ep_size"),
)
# `moe_dp` is absent because `initialize_model_parallel` aliases the MoE-DP
# group to the attention-CP group when the latter is wider: there the group and
# the name are two facts, which is the same reason `moe_dp_rank` is left off the
# record at publish.
# MoE-DP may alias a wider attention-CP group, so its configured width
# need not match the group width.
_WIDTH_AND_GROUP = (
("tp_size", "tp_group"),
("pp_size", "pp_group"),
@@ -402,25 +304,14 @@ _UNREADABLE = object()
def _validate_parallel(parallel, source: str) -> None:
"""Fail on a topology that cannot describe a real process layout.
"""Check rank bounds, topology factorizations, and group widths.
Every identity holds unconditionally: a width and a rank are both
plausible small integers whichever way they are wrong, so an inconsistent
set is not caught by anything downstream -- it surfaces as a hang or a
wrong answer in a collective, far from the write. Stating one leaf without
the quotients that follow from it leaves the namespace describing no real
layout, and the caller that did so is the one that has to say what it meant.
Names that cannot be read are skipped rather than treated as zero: a
process that has published nothing can still stamp a rank, and a group that
has not been built answers nothing at all.
Skip unavailable or non-integer values so partially initialized contexts
can be validated.
"""
def read(name):
"""A width or a rank, or `_UNREADABLE` for anything these identities
cannot be stated about -- an absent name, `None`, or a stand-in a test
put in a group's place. Booleans are integers in Python and are not
widths, so they are out too."""
"""Return an integer topology value, or ``_UNREADABLE``; exclude booleans."""
try:
value = getattr(parallel, name)
except Exception:
@@ -500,24 +391,12 @@ def _validate_parallel(parallel, source: str) -> None:
class ParallelContext:
"""Parallel-topology namespace: one spelling per name.
"""Parallel configuration, process ranks, and group handles.
Every name is answered by a lookup, never by asking a process group. A
configured leaf and a width derived from one come off the published
``parallel`` bag; a rank is written by ``publish`` from the spawn bundle,
and a group handle by ``initialize_model_parallel`` as it builds them. A
read before the write that answers it says which write is missing rather
than deriving a number from whatever is installed -- the two answer
different questions, and a plausible wrong rank surfaces as a hang in a
collective far from here.
That makes a scope a matter of stating names: ``patch_tensor_parallel_group``
runs a draft worker under a different TP group by overriding the members it
changes, and PD multiplexing points ``tp_group`` at the prefill
communicator the same way. Elastic EP needs no rule at all: it scales
``ep_size`` / ``dp_size`` on the published bag while ``launch_world_size``
keeps the width the groups were built at, so the two are different names
rather than two answers to one name.
Configured and derived widths come from the published configuration.
``publish`` records ranks from the launcher; distributed initialization
records group handles. Scoped overrides take precedence over permanent
overrides and configuration. Uninitialized runtime fields raise on read.
"""
__slots__ = ("_overrides", "_stamp", "_config")
@@ -536,17 +415,7 @@ class ParallelContext:
return self._read(name)
def _read(self, name):
"""The one read path, for every kind of name in the namespace.
Scoped override, then the permanent stamp, then the published bag.
A rank or a group handle is on no bag, so once those three are out the
name has not been written yet and the read says so.
The two override maps stay separate because they are taken down by
different things -- a `with` block and `clear_stamp()` -- and
merging them would let a teardown of one drop the other, and would
turn "which wins" into whichever was written last.
"""
"""Read a scoped override, permanent override, or published value, in order."""
overrides = self._overrides
if name in overrides:
return overrides[name]
@@ -577,20 +446,10 @@ class ParallelContext:
raise AttributeError(f"ParallelContext has no {name!r}")
def override_permanently(self, **values) -> None:
"""Permanently record a width or rank the published bag can't answer
or no longer answers correctly -- not `RuntimeContext.override`,
because neither is a resolved config leaf and this must work with no
config published at all (`multimodal_gen` lends a TP group to `srt`
layers with no `srt` config to publish against).
"""Set parallel values until ``clear_stamp`` or ``reset_context``.
Widths are quotients of the configured leaves, so the bag can usually
answer and this only corrects it; a rank is a per-process fact the
configuration never carries, so for those this is the only source.
Lives beside, not inside, the `@contextmanager` `override` below -- a
name it cannot also have on this class -- because these are permanent
for the process, not scoped to a `with` block: none of the real
callers ever restore the value they set here.
Works without published configuration. Validate the combined topology
and restore the previous values if validation fails.
"""
unknown = set(values) - _parallel_fields()
if unknown:
@@ -628,7 +487,7 @@ class ParallelContext:
def _derived_widths() -> dict:
"""The declared quotients, by name -- `{name: Derived}`."""
"""Return parallel ``Derived`` declarations, including ranks and groups."""
from sglang.srt.arg_groups.arg_utils import Derived
from sglang.srt.arg_groups.fields.parallel import Parallel
@@ -638,18 +497,9 @@ def _derived_widths() -> dict:
def _install_parallel_properties() -> None:
"""Give `ParallelContext` a property per name that is not a config leaf.
"""Expose declared parallel fields as documented properties.
Every name the namespace answers that is not a plain leaf: the quotients,
and the ranks and group handles declared beside them with no `fn`.
Properties rather than names left to `__getattr__` because the class
surface is what the guards introspect -- `hasattr(ParallelContext,
"tp_size")` and `vars(ParallelContext)` are how the tests read the set from
the class side -- and because each one carries its `Derived.doc`.
Every one of them resolves through `_read`, so there is a single priority
chain rather than one per kind of name.
Properties support class-level introspection; all reads use ``_read``.
"""
for name, decl in _derived_widths().items():
@@ -755,10 +605,10 @@ class MoeFlags(_FlagGroupBase):
class DpFlags(_FlagGroupBase):
"""DP-attention runtime flags, materialized by ``initialize_dp_attention``
(after distributed setup; reads the model config). The topology values it
also computes -- the attention-DP width and rank -- are stamped on
``get_parallel()``, not kept here."""
"""DP-attention runtime flags set by ``initialize_dp_attention``.
Attention-DP width and rank are stored on ``get_parallel()``.
"""
enabled: bool = False
use_world_group_for_gather: bool = False
@@ -1134,9 +984,7 @@ def _install_derived_leaves(tops: dict, server_args: Any) -> None:
if not isinstance(decl, Derived):
continue
if not decl.fn:
# Nothing computes it. Seed the ones whose absence is itself an
# answer, so a process that never states one still reads it;
# the rest stay unwritten and say so when read.
# Initialize runtime-only fields that declare a default.
if decl.default is not _NO_DEFAULT:
bag = tops.get(path.split(".")[0])
for segment in path.split(".")[1:]:
@@ -1262,10 +1110,7 @@ class RuntimeContext:
# Snapshot resolved config into the namespace bags (the single source of
# truth for config reads). Placed by `namespace_of`; a mock/partial
# config that declares no namespace yields an empty tree (no bags).
# A name the configuration does not carry survives the re-projection.
# `gpu_id` is stated by the spawn, not derived, so rebuilding the bags
# from the record must not unwrite it -- the record has no field for it
# to be rebuilt from.
# Preserve the launcher-assigned device when rebuilding config bags.
stated = {}
if self._config_bags is not None:
device = self._config_bags.get("device")
@@ -1805,11 +1650,9 @@ def publish(
namespace-read enforcement (``record`` audits the reads instead).
``hf_config`` is accepted for forward-compat and currently unused.
``ranks`` is who this process is, from the entry that spawned it. It is
optional because most roles are not placed in the topology at all -- a
tokenizer has no ``tp_rank`` -- and those processes raise on a rank read
exactly as they do today, with a message naming the missing bundle rather
than an absent process group.
``ranks`` supplies launcher placement for roles that participate in the
parallel topology. Without it, rank reads require an explicit override,
except for ``attn_dcp_rank=0`` when DCP is disabled.
A process holds at most one live config: the bags always describe the
engine running now. Re-publish is allowed and is **last-publish-wins**
@@ -1836,26 +1679,14 @@ def publish(
),
)
_CONTEXT._publish_role = role
# Zero for every process when decode context parallelism is off, which is a
# fact about the configuration and not about the spawn -- so it answers
# without a rank bundle, the way it did when it stood for a group that was
# never built. With DCP on it is a position, and the bundle below states it.
# Disabled DCP has rank zero even in processes without a rank bundle.
if not _CONTEXT.parallel.dcp_enabled:
_CONTEXT.parallel.override_permanently(attn_dcp_rank=0)
# Stated on the bag directly: `gpu_id` is declared but not configured, so
# it is not a leaf `override` can route, and the spawn is the only thing
# that knows it. Written whatever it is, `None` included -- most roles run
# on no device, and that is an answer rather than a name nobody wrote.
# The device is assigned by the launcher; it is not a config field.
_CONTEXT.config_bag("device")._set(
"gpu_id", ranks.gpu_id if ranks is not None else None
)
if ranks is not None:
# The placement, worked out here rather than carried: the widths are on
# the bag a moment ago, and `world_rank` fixes the rest. A read of any
# of these then needs no process group, which is the point -- they are
# read long before one exists. The scoped overrides that swap a group
# for a draft worker sit above the record in the read chain, so a
# scope still wins.
parallel = _CONTEXT.parallel
placement = derive_spawn_ranks(
world_rank=ranks.world_rank,
@@ -1866,30 +1697,18 @@ def publish(
moe_dp_size=parallel.moe_dp_size,
moe_ep_size=parallel.moe_ep_size,
)
# `initialize_model_parallel` aliases the MoE-DP group to the
# attention-CP one when the CP dimension is the wider of the two, so
# this process's place in it is its CP index rather than the MoE-DP
# index the arithmetic above gives.
# MoE-DP aliases the attention-CP group when CP is wider.
if parallel.moe_dp_size < parallel.attn_cp_size:
placement["moe_dp_rank"] = placement["attn_cp_rank"]
# `dp_rank` is recorded whatever it is, None included: replicas are
# separate WORLD groups, so no rank implies it and `None` is the answer
# "no controller" rather than an absence.
# `None` means no data-parallel controller.
placement["dp_rank"] = ranks.dp_rank
placement["launch_world_rank"] = ranks.world_rank
placement.update(_attention_ranks(parallel, placement["tp_rank"]))
# A DCP group is a contiguous slice of a TP group, so this process's
# place in one is its TP rank folded by that width. `attn_dcp_rank` is
# the same number, and zero where decode context parallelism is off, so
# a reader does not have to ask whether it is on first.
# DCP groups are contiguous slices of a TP group.
if parallel.dcp_enabled:
placement["dcp_rank"] = placement["tp_rank"] % parallel.dcp_size
placement["attn_dcp_rank"] = placement.get("dcp_rank", 0)
# One stamp, not two: the identities are checked on every write, and a
# half-placed process satisfies none of them.
parallel.override_permanently(**placement)
# Publish established the whole layout, so every identity applies here,
# not just the ones the stamp happened to name.
_validate_parallel(parallel, "publish")
if _ROLE_NS_MODE == "record":
# The '-' marker distinguishes a zero-read role from a process where
@@ -1906,18 +1725,7 @@ def publish(
def _attention_ranks(parallel, tp_rank: int) -> dict:
"""Place this process in the attention topology, from the configuration.
The widths are already on the bag -- `publish` computed them a moment ago --
and the rank comes from the spawn, so the position is known here, before any
process group exists. That is the point: a rank read then works in a process
that never initialises distributed, which is what the per-runner record used to provide
by being a plain frozen record.
These are stamped rather than written as bag leaves because they are
per-process facts, and nothing about the configuration distinguishes one
rank from another.
"""
"""Derive attention ranks from the configured widths and the TP rank."""
attn_tp_rank, attn_dp_rank = derive_attention_ranks(
tp_rank=tp_rank,
attn_tp_size=parallel.attn_tp_size,
@@ -2065,13 +1873,9 @@ def restore_context(state: dict[str, Any]) -> None:
def reset_context() -> None:
"""Clear the context-owned store (unit-test teardown): drop the published
``server_args`` and install fresh ``Flags`` and ``Resources``.
"""Clear published configuration, parallel overrides, flags, and resources.
``parallel`` holds the permanently-overridden derived widths, which go
with the lifecycle that set them: `_read` prefers them over the
published leaves, so leaving one behind lets the next test read the
previous topology.
Used for test teardown and runtime lifecycle reset.
"""
_CONTEXT._server_args = None
_CONTEXT._config_bags = None
python/sglang/srt/speculative/dflash_worker_v2.py
+1 -5
View File
@@ -396,11 +396,7 @@ class DFlashWorkerV2(BaseSpecWorker):
self.draft_tp_context = (
draft_tp_context if get_parallel().enable_dp_attention else empty_context
)
# One decision, used twice: whether the draft runs on an attention-TP
# slice of its own. It picks how the runner is built, and then what the
# scope may say about attention every time it is entered -- a draft
# built outside the scope keeps the target's replica count and still
# gathers with it.
# Use the same attention topology during draft construction and execution.
self.draft_owns_attention = get_parallel().enable_dp_attention
if self.draft_owns_attention:
draft_init_ctx = draft_tp_context(
python/sglang/srt/speculative/eagle_worker_v2.py
+1 -6
View File
@@ -260,12 +260,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
self._rebuild_topk1_chain_buffers()
# Load draft model weights only.
# One decision, used twice: whether the draft runs on an attention-TP
# slice of its own. It picks how the runner is built, and then what the
# scope may say about attention every time it is entered -- a draft
# built outside the scope keeps the target's replica count and still
# gathers with it.
# Use the same attention topology during draft construction and execution.
self.draft_owns_attention = (
get_parallel().enable_dp_attention
and self.speculative_algorithm.is_eagle3()
python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py
+1 -3
View File
@@ -160,9 +160,7 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
self.kv_context: Optional[FrozenKVMTPContext] = None
# Built above under the pipeline scope only, so this runner carries the
# target's attention topology: entering the tensor scope later swaps the
# communicator without giving the draft a replica of its own.
# Retain the target's attention topology when swapping TP groups.
self.draft_owns_attention = False
self.draft_tp_context = (
draft_tp_context if get_parallel().enable_dp_attention else empty_context
python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+1 -4
View File
@@ -188,10 +188,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
"InklingForConditionalGenerationMTP",
"GigaChat35ForCausalLMNextN",
]
# The draft runner is built outside any tensor-parallel scope, so it
# carries the target's topology: entering the scope later swaps the
# communicator without making this process a draft with an attention
# replica of its own. It still gathers with the target's replicas.
# Retain the target's attention topology when swapping TP groups.
self.draft_owns_attention = False
self.draft_tp_context = (
draft_tp_context if get_parallel().enable_dp_attention else empty_context
python/sglang/srt/speculative/standalone_worker_v2.py
+1 -4
View File
@@ -88,10 +88,7 @@ class StandaloneDraftWorker(EagleDraftWorker):
# Alias for better readability
self.draft_runner = self.draft_worker.model_runner
# The draft runner is built outside any tensor-parallel scope, so it
# carries the target's topology: entering the scope later swaps the
# communicator without making this process a draft with an attention
# replica of its own. It still gathers with the target's replicas.
# Retain the target's attention topology when swapping TP groups.
self.draft_owns_attention = False
self.draft_tp_context = (
draft_tp_context if get_parallel().enable_dp_attention else empty_context
python/sglang/srt/utils/weight_checker.py
+2 -7
View File
@@ -70,10 +70,7 @@ def _is_non_persistent_buffer_name(name: str) -> bool:
class WeightChecker:
def __init__(self, *, get_model: Callable[[], Any]):
self._get_model = get_model
# A check is served on demand from the scheduler loop, which is outside
# the scope that describes a draft runner. The report has to name the
# runner it was built for, so the placement is read here, at
# construction, rather than asked for when the request arrives.
# Capture the runner placement before its draft scope exits.
parallel = get_parallel()
self._placement = ParallelismInfo(
tp_rank=parallel.tp_rank,
@@ -176,9 +173,7 @@ class WeightChecker:
return info.model_dump()
def _parallelism_info(self) -> ParallelismInfo:
# The WORLD position is asked for now rather than frozen: unlike the
# runner's placement it is a property of the process, and an elastic
# scale-up moves it.
# Read the current WORLD rank because elastic scale-up can change it.
return self._placement.model_copy(
update={
"rank": dist.get_rank() if dist.is_initialized() else 0,
python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py
+1 -5
View File
@@ -22,11 +22,7 @@ from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import get_context, get_parallel
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
# Unit tests run without distributed initialization. Backends that size buffers by
# attention tensor-parallel degree should see the single-rank default, and a
# backend that places itself in the decode context-parallel group needs a
# position: nothing publishes here, so there is no configuration to derive the
# zero these tests run at from.
# Use single-rank attention without distributed initialization.
_parallel_override = get_parallel().override(attn_tp_size=1, attn_dcp_rank=0)
_parallel_override.__enter__()
python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py
+1 -3
View File
@@ -5,9 +5,7 @@ import torch
import torch.nn.functional as F
from torch import nn
# State the topology before importing modules that read it at __init__. The
# group is stated too: `RowParallelLinear.forward` asks for it to manage
# symmetric memory, and `world_size=1` short-circuits that.
# Set the single-rank topology before importing the attention implementation.
from sglang.srt.runtime_context import get_context, get_parallel
_parallel_override = get_parallel().override(
python/sglang/test/test_utils.py
+7 -25
View File
@@ -2056,19 +2056,10 @@ def maybe_stub_sgl_kernel():
@contextlib.contextmanager
def published_topology(role: str = "test", *, ranks=None, **server_args_fields):
"""Publish a record describing the parallel topology a test wants.
"""Publish a test topology, defaulting to WORLD rank zero.
Replaces standing a per-process parallel record into the object under
test. The widths arrive the way production gets them -- from published
configuration -- and the per-process ranks the way a spawned process gets
them, so a rank read is answered without building a process group. Stating
the topology through the same door production uses also keeps the derived
widths honest: a hand-built double can claim an `attn_tp_size` the
configuration would never produce.
`ranks` overrides the spawn identities; by default this process is rank
zero of the world, which fixes every other rank. The context is reset on exit, including when the
test fails.
``ranks`` overrides the launcher placement. Reset the context before
publication and on exit, including when the test fails.
"""
from sglang.srt.runtime_context import SpawnRanks, publish, reset_context
from sglang.srt.server_args import ServerArgs
@@ -2085,15 +2076,10 @@ def published_topology(role: str = "test", *, ranks=None, **server_args_fields):
def publish_build_topology(*, world_rank: int = 0, **server_args_fields):
"""State the widths `initialize_model_parallel` is about to build at.
"""Publish the topology for a subsequent ``initialize_model_parallel`` call.
The build reads every width from the runtime context, so a test that wants
a particular topology publishes it here rather than passing it in -- the
same door production uses, which also keeps the derived widths honest.
Unlike `published_topology` this is not a scope: the groups it is about to
build outlive any block, so the configuration describing them has to as
well. Callers that tear the groups down are already resetting the process.
Preserve an existing WORLD group across the context reset. The caller is
responsible for tearing down groups and resetting the context afterward.
"""
from sglang.srt.distributed import parallel_state
from sglang.srt.runtime_context import (
@@ -2110,11 +2096,7 @@ def publish_build_topology(*, world_rank: int = 0, **server_args_fields):
role="test",
ranks=SpawnRanks(world_rank=world_rank),
)
# Callers that go on to build groups have already run
# `init_distributed_environment`, which states the WORLD group -- and the
# build below places every group it creates by reading that back. The reset
# above drops it, so hand it over again: publishing a configuration does not
# unbuild a process group.
# Restore the existing WORLD handle after resetting the context.
if parallel_state._WORLD is not None:
get_parallel().override_permanently(world_group=parallel_state._WORLD)