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

[Config] Round 6.5: a namespace declares what it derives, next to what it derives it from (#38113)

Browse Source 
Fifth of five; stacked on #38049. The split gave every namespace a file, but
only for the half an operator types. This is the other half.

## The parallel quotients are declared, not written out

`attn_tp_size` and its five siblings were sixty lines of near-identical
properties in the runtime context, a file away from the leaves they are
quotients of, so reading `parallel.py` told you what you could set and nothing
about what that decides.

They are declared in `Parallel` now, in the same class as those leaves. They
carry no annotation, so they are not dataclass fields and
`collect_input_fields` never puts them on the record -- the same mechanism that
already keeps `_NS_PATH` off it. That is the right exclusion: a quotient has no
operator input to preserve, and the record is what crosses a process boundary,
where a stamped width is one an elastic scale-up will not refresh.

## A quotient is a value in the bag, like every other derived one

`_derived_width` answered from a stamp or, failing that, a live process group.
The group read could never disagree with the stamp:

- `initialize_model_parallel` stamps all six as its last statement,
  unconditionally;
- an elastic scale-up restamps `attn_dp_size` through
  `update_dp_attention_post_scale` -- the comment claiming it does *not* was
  wrong;
- no hardware backend builds groups of its own;
- `multimodal_gen`, which has its own `initialize_model_parallel` and does not
  stamp, never reads a quotient.

So a built group was always already stamped, and the group read goes -- and with
it the last reason for a quotient to be resolved on every read.

Every input to `derive_parallel_widths` is a record field. `dcp_enabled` is
`decode_context_parallel_size > 1`, not a fact about a built group; it was
spelled `_DCP is not None`, which is a longer way to say the same thing. So the
six are fixed once the configuration is fixed -- the same test every other
`Derived(fn=...)` in this PR passes. They are declared the same way and computed
the same way: once, at publish, into ordinary bag leaves.

What remains is override -> stamp -> published leaf. The stamp stays above the
leaf because an elastic scale-up restamps `attn_dp_size`; the override stays on
top because that is how a test names a width.

## One answer for the config-derived predicates

`enable_mamba_extra_buffer` and its lazy variant, `is_ep_joiner`,
`is_ep_scale_joiner`, `is_startup_weight_load_overlap`: each existed as a
`ServerArgs` member for the resolution pipeline and, for most of them, again as
a `runtime_context` function for readers after publish. Three places to keep
saying the same thing.

A `Derived(fn=...)` is a pure function of the published configuration, so
`publish` computes it once and stores it as an ordinary bag leaf -- a plain
attribute load, which is what a read inside compiled model code needs. The
function is handed the whole resolved config rather than the bag it lands in,
because a derivation is free to span namespaces and the mamba one does: it
reads `memory.disable_radix_cache` alongside its own `exec.mamba` strategy,
which is why it could never have been a method on either bag.

The pre-publish helpers stay -- resolution needs the predicate before there is
a bag to read -- and three readers keep them, because they run before their own
process publishes: `initialize_dp_attention`, which the weight-cache daemon
calls while building its groups thirty lines before its `publish`, and
`PortArgs.init_new`, a factory handed the record that already reads eighteen
other fields off it.

## Notes for a reviewer

**Overriding a leaf does not move its quotient.** `override(tp_size=2)` leaves
`attn_tp_size` where the published config put it, because nothing is recomputed
on read. A test states a topology by publishing a config -- which is what a
real process does -- or by naming the width it wants, `override(attn_tp_size=2)`.
Six tests say it that way now. This is the price of having one answer computed
once, and it is the same price every other derived value in the config already
carries.

A caller that reads a quotient without publishing or overriding now gets an
explicit error naming the field, instead of a default that an uninitialised
group happened to supply. One fixture was in that state --
`TestMlaWriteDoorsUnderDcp` built a bare pool and asked whether DCP was on --
and it publishes a config now, which is what the process it stands in for
does.

Eighteen sites read these predicates without calling them. That is correct --
they are properties -- but it is worth saying they were checked, because a
census that assumes otherwise reports eighteen always-true conditions.

## The skill that documents this subsystem is updated with it

`.claude/rules/modify-component-must-read.md` points at
`.claude/skills/sglang-runtime-context/SKILL.md` before anyone touches these
files, so a stale sentence there is a wrong instruction rather than a stale
note. Four of its load-bearing statements stopped being true across this series
and are corrected here: `NS(...)` is no longer how a field states its namespace
(the declaring class is); the DCP degrade rule is gone, because the quotients
are not live reads; `mamba_extra_buffer_enabled()` and the other predicate
functions it named as the shape to copy no longer exist; and the
namespace-coverage ratchet is described in terms of the marker. The docstring of
`test_server_args_namespaces.py` said the same thing and is fixed too.

The consequence a test author actually trips over is stated there as well:
overriding a leaf no longer moves its quotient, so a topology is stated by
publishing a config or by naming the width.

## Verification

A full registered-unit sweep (648 files) against this stack's merge-base:
19 failures on both sides, the same 19 -- AMD `gfx950`, `modelopt`,
`cuda_vmm`, `weight_checker` and friends, none of them config. The narrower 139-file config sweep used earlier in this series
does not contain the files this change reaches -- `test_kv_index_translator`
never names `get_parallel()`, it constructs an object that does -- which is why
the baseline differential over everything is what is quoted here.
This commit is contained in:
Cheng Wan
2026-09-06 21:44:24 -07:00
committed by GitHub
parent b99175dc7d
commit aaf9a95763
56 changed files with 1118 additions and 379 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/benchmark/one_batch.py
+8 -3
View File
@@ -87,7 +87,12 @@ from sglang.srt.model_executor.cuda_graph_config import (
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import get_parallel, get_schedule, publish
from sglang.srt.runtime_context import (
get_model,
get_parallel,
get_schedule,
publish,
)
from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.srt.server_args import PortArgs, ServerArgs
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
@@ -356,12 +361,12 @@ def load_model(server_args, port_args, gpu_id, tp_rank):
model_runner = MlxModelRunnerStub(**runner_kwargs)
else:
model_runner = ModelRunner(**runner_kwargs)
if cfg.is_startup_weight_load_overlap:
if get_model().is_startup_weight_load_overlap:
model_runner.start_startup_weight_load()
model_runner.alloc_memory_pool()
model_runner.init_attention_backends()
model_runner.init_cuda_graphs()
if cfg.is_startup_weight_load_overlap:
if get_model().is_startup_weight_load_overlap:
model_runner.finalize_startup_weight_load()
rank_print(f"max_total_num_tokens={model_runner.max_total_num_tokens}")
tokenizer = get_tokenizer(
python/sglang/multimodal_gen/runtime/distributed/parallel_state.py
+47 -4
View File
@@ -164,19 +164,45 @@ 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.
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
stamped alongside the group, as `srt.initialize_model_parallel` does.
Only tensor parallelism folds this way, so every other dimension is one.
"""
import sglang.srt.distributed.parallel_state as srt_parallel_state
from sglang.srt.runtime_context import derive_parallel_widths, get_parallel
if srt_parallel_state._TP is None:
srt_parallel_state._TP = _TP
if srt_parallel_state._ATTN_TP is None:
srt_parallel_state._ATTN_TP = _TP
if srt_parallel_state._ATTN_TP is _TP:
get_parallel().stamp_derived_widths(
**derive_parallel_widths(
tp_size=_TP.world_size,
attn_cp_size=1,
attn_dp_size=1,
moe_ep_size=1,
moe_dp_size=1,
dcp_size=1,
dcp_enabled=False,
)
)
def _clear_srt_tp_group() -> None:
import sglang.srt.distributed.parallel_state as srt_parallel_state
from sglang.srt.runtime_context import get_parallel
if srt_parallel_state._ATTN_TP is _TP:
srt_parallel_state._ATTN_TP = None
get_parallel().clear_derived_widths()
if srt_parallel_state._TP is _TP:
srt_parallel_state._TP = None
@@ -704,11 +730,17 @@ def use_tensor_parallel_group(tp_group: GroupCoordinator):
The scope replaces the module globals that ``get_tp_group()`` and srt's
``get_tp_group()`` / ``get_attention_tp_group()`` read, and — like srt's
``patch_tensor_parallel_group`` — the three members the runtime context
answers with, so that a size read from the published bag cannot disagree
with a rank read from the swapped group.
``patch_tensor_parallel_group`` — the members the runtime context answers
with, so that a size read from the published bag cannot disagree with a rank
read from the swapped group.
The parallel quotients are part of that set: the published config describes
the launch (`tp=1` for a sequence-parallel run), not the group folded in
here. Left out, an encoder built in this scope keeps its heads whole on
`attn_tp_size == 1` while the `QKVParallelLinear` beside it shards on
`tp_size == 2`, and the weight loader narrows past the end of the tensor.
"""
from sglang.srt.runtime_context import get_parallel
from sglang.srt.runtime_context import derive_parallel_widths, get_parallel
old_tp_group = get_tp_group()
import sglang.srt.distributed.parallel_state as srt_parallel_state
@@ -724,6 +756,17 @@ def use_tensor_parallel_group(tp_group: GroupCoordinator):
tp_size=tp_group.world_size,
tp_rank=tp_group.rank_in_group,
tp_group=tp_group,
# Only tensor parallelism folds here, so every other dimension is
# one and the quotients come out of the shared derivation.
**derive_parallel_widths(
tp_size=tp_group.world_size,
attn_cp_size=1,
attn_dp_size=1,
moe_ep_size=1,
moe_dp_size=1,
dcp_size=1,
dcp_enabled=False,
),
):
yield
finally:
python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py
+5 -1
View File
@@ -132,7 +132,10 @@ def test_destroy_releases_sequence_parallel_subgroups_after_partial_init():
def test_srt_attention_tp_group_tracks_diffusion_tp_group():
tp_group = object()
# `world_size`, because lending the group also states the parallel widths it
# implies -- the shared `srt` vision layers ask for `attn_tp_size`, and this
# package publishes no `srt` config for that read to resolve against.
tp_group = SimpleNamespace(world_size=2)
with (
patch.object(parallel_state, "_TP", tp_group),
@@ -143,6 +146,7 @@ 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
parallel_state._clear_srt_tp_group()
python/sglang/srt/arg_groups/arg_utils.py
+28
View File
@@ -103,6 +103,34 @@ class Arg:
fallback: Any = NO_FALLBACK
@dataclasses.dataclass(frozen=True)
class Derived:
"""Metadata for a field the configuration implies, not one anyone types.
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.
Every declaration carries ``fn`` today, 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 -- an elastic scale-up
restamps ``attn_dp_size`` -- which ``ParallelContext`` answers above the
published leaf.
"""
doc: str = ""
fn: str = ""
@dataclasses.dataclass(frozen=True)
class NS:
"""Namespace-path marker for a ServerArgs field, attached alongside the
python/sglang/srt/arg_groups/fields/exec_.py
+32
View File
@@ -20,6 +20,7 @@ from typing import (
from sglang.srt.arg_groups.arg_utils import (
A,
Arg,
Derived,
)
from sglang.srt.arg_groups.choices import (
ATTENTION_BACKEND_CHOICES,
@@ -293,6 +294,27 @@ class ExecMamba:
"""Namespace ``exec.mamba``."""
_NS_PATH = "exec.mamba"
# ---- derived: whether the mamba radix cache keeps its extra state buffer.
#
# One answer, computed at publish. There used to be three spellings of this
# predicate -- a `ServerArgs` member for the resolution pipeline, a
# `runtime_context` function for readers after publish, and the shared
# helper both delegated to -- which is three places to keep saying the same
# thing. The helper stays, because resolution needs it before there is a
# bag to read; the other two are this.
#
# It reads `memory.disable_radix_cache` as well as the strategy below, so
# it spans two namespaces and could not have been a method on either bag.
enable_mamba_extra_buffer = Derived(
fn="sglang.srt.arg_groups.model_override_base.mamba_extra_buffer_of",
doc="Whether the hybrid-mamba radix cache keeps its extra state "
"buffer: the radix cache is on and the strategy asks for one.",
)
enable_mamba_extra_buffer_lazy = Derived(
fn="sglang.srt.arg_groups.overrides.mamba_extra_buffer_lazy_of",
doc="The lazy variant: the strategy is `extra_buffer_lazy` exactly.",
)
mamba_backend: A[
str,
Arg(
@@ -579,6 +601,16 @@ class ExecMoe:
"""Namespace ``exec.moe``."""
_NS_PATH = "exec.moe"
# ---- derived: computed at publish from the leaves below.
is_ep_joiner = Derived(
fn="sglang.srt.arg_groups.model_override_base.ep_joiner_of",
doc="Whether this process was launched as an elastic-EP joiner (scale or recover).",
)
is_ep_scale_joiner = Derived(
fn="sglang.srt.arg_groups.model_override_base.ep_scale_joiner_of",
doc="Whether it is a scale-up joiner specifically.",
)
enable_fused_moe_sum_all_reduce: A[
bool,
"Enable fused moe triton and sum all reduce.",
python/sglang/srt/arg_groups/fields/model.py
+7
View File
@@ -21,6 +21,7 @@ from typing import (
from sglang.srt.arg_groups.arg_utils import (
A,
Arg,
Derived,
)
from sglang.srt.arg_groups.choices import (
LOAD_FORMAT_CHOICES,
@@ -39,6 +40,12 @@ class Model:
_NS_PATH = "model"
# ---- derived: computed at publish from the leaves below.
is_startup_weight_load_overlap = Derived(
fn="sglang.srt.arg_groups.model_override_base.startup_weight_load_overlap_of",
doc="Whether weight loading overlaps startup.",
)
# -------------------------------------------------------------------------
# Model and tokenizer
# -------------------------------------------------------------------------
python/sglang/srt/arg_groups/fields/parallel.py
+42
View File
@@ -16,6 +16,7 @@ from typing import Optional
from sglang.srt.arg_groups.arg_utils import (
A,
Arg,
Derived,
)
@@ -275,3 +276,44 @@ class Parallel:
Optional[int],
"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.
attn_tp_size = Derived(
fn="sglang.srt.runtime_context.attn_tp_size_of",
doc="Attention tensor-parallel width: `tp_size` divided by the "
"attention-DP and attention-CP dimensions.",
)
attn_dp_size = Derived(
fn="sglang.srt.runtime_context.attn_dp_size_of",
doc="Attention data-parallel width: `dp_size` when DP attention is "
"on, otherwise one.",
)
attn_dcp_size = Derived(
fn="sglang.srt.runtime_context.attn_dcp_size_of",
doc="Decode context-parallel width inside the attention TP group.",
)
moe_ep_size = Derived(
fn="sglang.srt.runtime_context.moe_ep_size_of",
doc="MoE expert-parallel width, normalised from the configured value.",
)
moe_tp_size = Derived(
fn="sglang.srt.runtime_context.moe_tp_size_of",
doc="MoE tensor-parallel width: what is left of `tp_size` after the "
"expert and MoE-DP dimensions.",
)
dcp_enabled = Derived(
fn="sglang.srt.runtime_context.dcp_enabled_of",
doc="Whether decode context parallelism is in play: `dcp_size` is "
"wider than one rank, which is exactly when the group gets built.",
)
python/sglang/srt/arg_groups/model_override_base.py
+28 -6
View File
@@ -257,13 +257,35 @@ def model_config_of(server_args: Any):
return model_config
def mamba_extra_buffer_of(cfg: Any) -> bool:
"""Mid-resolution equivalent of runtime_context.mamba_extra_buffer_enabled:
reads the (possibly overlaid) strategy from a config-shaped object.
def ep_joiner_of(cfg: Any) -> bool:
"""Whether this process was launched as an elastic-EP joiner.
This is the one definition of the predicate: ``ServerArgs`` delegates its
member to it, and the runtime_context accessor is its post-publish sibling
(which cannot reuse it, because the two leaves land in different bags)."""
The one definition. After publish the answer is a bag leaf --
`get_exec().moe.is_ep_joiner` -- computed from this function by the
declaration in `arg_groups/fields/exec_.py`; resolution needs it before
there is a bag to read, which is why it is still a function.
"""
return cfg.ep_join_mode in ("scale", "recover")
def ep_scale_joiner_of(cfg: Any) -> bool:
"""The scale-up arm of :func:`ep_joiner_of`."""
return cfg.ep_join_mode == "scale"
def startup_weight_load_overlap_of(cfg: Any) -> bool:
"""Whether weight loading overlaps startup."""
return cfg.startup_weight_load_mode == "overlap"
def mamba_extra_buffer_of(cfg: Any) -> bool:
"""The predicate, read off a config-shaped object mid-resolution.
This is the one definition. After publish the answer is a bag leaf --
``get_exec().mamba.enable_mamba_extra_buffer`` -- computed from this same
function by the declaration in ``arg_groups/fields/exec_.py``. Resolution
needs it before there is a bag to read, which is why it is still a
function."""
return cfg.disable_radix_cache is False and cfg.mamba_radix_cache_strategy in (
"extra_buffer",
"extra_buffer_lazy",
python/sglang/srt/distributed/bootstrap.py
+2 -2
View File
@@ -265,8 +265,8 @@ def _init_parallel_groups(
moe_dp_size: int,
dcp_size: int,
) -> None:
is_ep_joiner = server_args.is_ep_joiner
is_scale_joiner = server_args.is_ep_scale_joiner
is_ep_joiner = get_exec().moe.is_ep_joiner
is_scale_joiner = get_exec().moe.is_ep_scale_joiner
rank_offset = get_parallel().ep_join_rank_offset if is_scale_joiner else 0
world_size = (
rank_offset + tp_size * pp_size if is_scale_joiner else tp_size * pp_size
python/sglang/srt/elastic_ep/elastic_ep.py
+1 -1
View File
@@ -110,7 +110,7 @@ class ElasticEPStateManager:
cls._on_scale = cls._on_scale_nixl
inst.ep_join_rank_offset = get_parallel().ep_join_rank_offset
if server_args.is_ep_joiner:
if get_exec().moe.is_ep_joiner:
cls._init_joiner_state(inst)
cls._instance = inst
python/sglang/srt/entrypoints/http_server.py
+1 -1
View File
@@ -2398,7 +2398,7 @@ def _wait_and_warmup(
_wait_weights_ready()
# Joiner schedulers are served through the primary after adoption.
skip_elastic_joiner_warmup = server_args.is_ep_scale_joiner
skip_elastic_joiner_warmup = get_exec().moe.is_ep_scale_joiner
if skip_elastic_joiner_warmup:
logger.debug(
"[Elastic EP] Skipping server warmup for elastic joiner (ep_join_mode=%s)",
python/sglang/srt/entrypoints/openai/serving_chat.py
+3 -9
View File
@@ -714,17 +714,11 @@ class OpenAIServingChat(OpenAIServingBase):
def _should_return_input_ids(self, request: ChatCompletionRequest) -> bool:
"""Whether prompt (input) token ids should be returned via sglext."""
return (
request.return_input_ids_in_sglext
or self.tokenizer_manager.server_args.return_input_ids
)
return request.return_input_ids_in_sglext or get_serving().return_input_ids
def _should_return_output_ids(self, request: ChatCompletionRequest) -> bool:
"""Whether sampled output token ids should be returned via sglext."""
return (
request.return_output_ids_in_sglext
or self.tokenizer_manager.server_args.return_output_ids
)
return request.return_output_ids_in_sglext or get_serving().return_output_ids
def _continuous_usage_cached_details(
self, content: Dict[str, Any]
@@ -1768,7 +1762,7 @@ class OpenAIServingChat(OpenAIServingBase):
if return_output_ids:
chunk_output_ids = content.get("output_ids")
if chunk_output_ids is not None:
if self.tokenizer_manager.server_args.incremental_streaming_output:
if get_serving().incremental_streaming_output:
accumulated = output_ids.setdefault(index, [])
if finish_reason_type == "abort":
# The abort chunk re-sends the last token plus any coalesced deltas;
python/sglang/srt/hardware_backend/mlx/model_runner_stub.py
+4 -4
View File
@@ -19,7 +19,7 @@ from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.model_executor.model_runner_components.layer_setup import (
ModelLayerInfo,
)
from sglang.srt.runtime_context import get_exec, get_memory, get_schedule
from sglang.srt.runtime_context import get_exec, get_memory, get_model, get_schedule
logger = logging.getLogger(__name__)
@@ -118,14 +118,14 @@ class MlxModelRunnerStub(ModelRunner):
return 0
@staticmethod
def validate_startup_weight_load_mode(server_args) -> None:
if server_args.is_startup_weight_load_overlap:
def validate_startup_weight_load_mode() -> None:
if get_model().is_startup_weight_load_overlap:
raise ValueError(
"--startup-weight-load-mode=overlap is not supported: CUDA only"
)
def __init__(self, *args, mlx_pool_size: int | None = None, **kwargs):
self.validate_startup_weight_load_mode(kwargs["server_args"])
self.validate_startup_weight_load_mode()
self._mlx_pool_size = mlx_pool_size
super().__init__(*args, **kwargs)
python/sglang/srt/hardware_backend/mlx/tp_worker.py
+1 -1
View File
@@ -82,7 +82,7 @@ class MlxTpModelWorker(TpModelWorker):
MlxModelRunnerStub,
)
MlxModelRunnerStub.validate_startup_weight_load_mode(self.server_args)
MlxModelRunnerStub.validate_startup_weight_load_mode()
from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+1 -1
View File
@@ -929,7 +929,7 @@ class Mamba2AttnBackend(MambaAttnBackendBase):
model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape
)
if model_runner.server_args.enable_mamba_extra_buffer():
if get_exec().mamba.enable_mamba_extra_buffer:
assert self.conv_states_shape[-1] < self.mamba_chunk_size, (
f"{self.conv_states_shape[-1]=} should be less than {self.mamba_chunk_size}"
)
python/sglang/srt/layers/dp_attention.py
+10 -1
View File
@@ -10,6 +10,10 @@ import torch
import triton
import triton.language as tl
from sglang.srt.arg_groups.model_override_base import (
ep_scale_joiner_of,
resolving_view,
)
from sglang.srt.distributed import (
GroupCoordinator,
get_attn_cp_group,
@@ -392,7 +396,12 @@ def initialize_dp_attention(
if get_exec().moe.elastic_ep_backend is not None and get_parallel().max_ep_size:
_ATTN_DP_RANK = tp_rank + get_parallel().ep_join_rank_offset
if server_args.is_ep_scale_joiner:
# Reads the resolution, not a bag: this runs under
# `initialize_dp_attention`, which the weight-cache daemon calls from
# `_init_distributed` -- and other callers reach it from processes
# whose publish is not guaranteed to have happened yet. (The daemon
# itself publishes first, at `daemon.py:284`, before `:320`.)
if ep_scale_joiner_of(resolving_view(server_args)):
dp.joiner_skip_all_gather = True
_DpGatheredBufferWrapper.set_metadata(
python/sglang/srt/managers/data_parallel_controller.py
+5 -5
View File
@@ -449,7 +449,7 @@ class DataParallelController:
Returns:
List of worker ports (same on all nodes after broadcast).
"""
is_joiner = server_args.is_ep_scale_joiner
is_joiner = get_exec().moe.is_ep_scale_joiner
if get_parallel().dist_init_addr is None or is_joiner:
na = NetworkAddress(
get_serving().host or "127.0.0.1",
@@ -561,7 +561,7 @@ class DataParallelController:
bind_host = NetworkAddress.parse(get_parallel().dist_init_addr).host
worker_ports = []
if server_args.is_ep_scale_joiner:
if get_exec().moe.is_ep_scale_joiner:
# Scale joiners connect to their pre-bound primary worker sockets.
primary = NetworkAddress.parse(get_parallel().dist_init_addr)
primary_endpoint = NetworkAddress(
@@ -626,7 +626,7 @@ class DataParallelController:
nnodes_per_tp_group = nnodes_per_pp_rank
tp_size_per_node = get_parallel().tp_size // nnodes_per_tp_group
if server_args.is_ep_scale_joiner:
if get_exec().moe.is_ep_scale_joiner:
# Scale joiners enumerate their full local TP span.
tp_rank_range = range(get_parallel().tp_size)
tp_size_per_node = get_parallel().tp_size
@@ -655,7 +655,7 @@ class DataParallelController:
rank_port_args = PortArgs.init_new(
server_args, dp_rank, worker_ports
)
if server_args.is_ep_scale_joiner:
if get_exec().moe.is_ep_scale_joiner:
# Scale-joiner outputs return through the primary tokenizer.
primary_addr = NetworkAddress.parse(
get_parallel().dist_init_addr
@@ -866,7 +866,7 @@ def run_data_parallel_controller_process(
}
)
# The primary owns routing for the expanded scheduler set.
if get_parallel().node_rank == 0 and not server_args.is_ep_scale_joiner:
if get_parallel().node_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
controller.event_loop()
for proc in controller.scheduler_procs:
proc.join()
python/sglang/srt/managers/schedule_batch.py
+6 -7
View File
@@ -4,14 +4,13 @@ from sglang.srt.dllm.config import DllmConfig
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.runtime_context import (
get_disagg,
get_exec,
get_parallel,
get_schedule,
get_serving,
get_spec,
mamba_cache_chunk_size,
mamba_checkpoint_grid,
mamba_extra_buffer_enabled,
mamba_extra_buffer_lazy_enabled,
mamba_track_grid,
)
from sglang.srt.utils.common import (
@@ -2682,7 +2681,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
req.already_computed = seq_len
req.is_retracted = False
if mamba_extra_buffer_enabled():
if get_exec().mamba.enable_mamba_extra_buffer:
track_entry = self._mamba_radix_cache_v2_req_prepare_for_extend(req)
mamba_track_mask_cpu.append(track_entry.track_mask)
mamba_track_indices_cpu.append(track_entry.track_index)
@@ -2787,7 +2786,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.extend_logprob_start_lens = extend_logprob_start_lens
self.extend_input_logprob_token_ids = extend_input_logprob_token_ids
if mamba_extra_buffer_enabled():
if get_exec().mamba.enable_mamba_extra_buffer:
self.mamba_track_indices = torch.tensor(
mamba_track_indices_cpu,
dtype=torch.int64,
@@ -2883,7 +2882,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# allocated yet; it will be allocated on demand at the track boundary
# in mamba_lazy_prealloc_at_boundary during prepare_for_decode.
req.kv.mamba_last_track_idx = req.kv.mamba_next_track_idx
if not mamba_extra_buffer_lazy_enabled():
if not get_exec().mamba.enable_mamba_extra_buffer_lazy:
req.kv.mamba_next_track_idx = (
self.req_to_token_pool.get_mamba_ping_pong_other_idx(
req.kv.mamba_next_track_idx
@@ -3398,7 +3397,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.req_pool_indices_cpu,
)
if mamba_extra_buffer_enabled():
if get_exec().mamba.enable_mamba_extra_buffer:
mamba_track_interval = mamba_track_grid(self.tree_cache.page_size)
if len(self.reqs) == 0:
@@ -3407,7 +3406,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
)
self.mamba_track_buffer_indices = []
else:
if mamba_extra_buffer_lazy_enabled():
if get_exec().mamba.enable_mamba_extra_buffer_lazy:
self.mamba_lazy_prealloc_at_boundary(mamba_track_interval)
set_mamba_track_indices_from_reqs(self)
python/sglang/srt/managers/scheduler.py
+2 -2
View File
@@ -1092,7 +1092,7 @@ class Scheduler(
def init_model_worker(self):
# Load model weights.
self.init_tp_model_worker()
if self.server_args.is_startup_weight_load_overlap:
if get_model().is_startup_weight_load_overlap:
self.tp_worker.start_startup_weight_load()
self.maybe_init_draft_worker()
@@ -1117,7 +1117,7 @@ class Scheduler(
model_runner.post_capture_resize_kv_pool()
self.kv_cache_allocation_time += time.perf_counter() - tic
if self.server_args.is_startup_weight_load_overlap:
if get_model().is_startup_weight_load_overlap:
self.tp_worker.finalize_startup_weight_load()
# Adaptive/speculative graphs and post-capture KV sizing can consume
python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+4 -4
View File
@@ -34,9 +34,9 @@ from sglang.srt.model_executor.forward_batch_info import (
)
from sglang.srt.runtime_context import (
get_disagg,
get_exec,
get_memory,
get_observability,
mamba_extra_buffer_lazy_enabled,
mamba_track_grid,
max_speculative_num_draft_tokens,
)
@@ -1131,7 +1131,7 @@ class SchedulerBatchResultProcessor:
i: int,
logits_output: LogitsProcessorOutput,
):
lazy = mamba_extra_buffer_lazy_enabled()
lazy = get_exec().mamba.enable_mamba_extra_buffer_lazy
known_mamba_boundary = None
completed_mamba_boundary = None
lookahead = 0
@@ -1206,7 +1206,7 @@ class SchedulerBatchResultProcessor:
prepare_release(req)
is_insert = (
req.mamba_lazy_is_insert
if mamba_extra_buffer_lazy_enabled()
if get_exec().mamba.enable_mamba_extra_buffer_lazy
else True
)
release_kv_cache(req, self.tree_cache, is_insert=is_insert)
@@ -1256,7 +1256,7 @@ class SchedulerBatchResultProcessor:
if req.kv.mamba_ping_pong_track_buffer is None:
return
lazy = mamba_extra_buffer_lazy_enabled()
lazy = get_exec().mamba.enable_mamba_extra_buffer_lazy
if known_boundary:
self._mamba_assert_committed_len_lookahead(req)
track_seqlen = req.kv.kv_committed_len
python/sglang/srt/managers/scheduler_components/request_receiver.py
+2 -2
View File
@@ -37,7 +37,7 @@ from sglang.srt.observability.scheduler_stage_metrics import (
SchedulerStageMetricsRecorder,
scheduler_stage_method,
)
from sglang.srt.runtime_context import get_disagg, get_parallel, is_ep_scale_joiner
from sglang.srt.runtime_context import get_disagg, get_exec, get_parallel
from sglang.srt.utils import (
broadcast_pyobj,
point_to_point_pyobj,
@@ -179,7 +179,7 @@ class SchedulerRequestReceiver:
# all-ranks gloo sync.
_local_ctrl = (
get_parallel().enable_dp_attention_local_control_broadcast
or is_ep_scale_joiner()
or get_exec().moe.is_ep_scale_joiner
)
if _local_ctrl:
control_reqs = attn_cp_tp_broadcast_pyobj(control_reqs)
python/sglang/srt/managers/tp_worker.py
+1 -1
View File
@@ -389,7 +389,7 @@ class TpModelWorker(BaseTpWorker):
# broadcast, so they reuse the target's already-broadcast seed.
if random_seed is not None:
self.random_seed = random_seed
elif server_args.is_ep_joiner:
elif get_exec().moe.is_ep_joiner:
self.random_seed = get_device().random_seed
else:
self.random_seed = broadcast_pyobj(
python/sglang/srt/mem_cache/kv_cache_builder.py
+4 -2
View File
@@ -2,6 +2,8 @@ from __future__ import annotations
import logging
from sglang.srt.runtime_context import get_exec
logger = logging.getLogger(__name__)
from dataclasses import dataclass
@@ -312,8 +314,8 @@ def build_kv_cache(
enable_metrics=enable_metrics,
enable_kv_cache_events=enable_kv_cache_events,
enable_session_radix_cache=get_memory().enable_session_radix_cache,
enable_mamba_extra_buffer=server_args.enable_mamba_extra_buffer(),
enable_mamba_extra_buffer_lazy=server_args.enable_mamba_extra_buffer_lazy(),
enable_mamba_extra_buffer=get_exec().mamba.enable_mamba_extra_buffer,
enable_mamba_extra_buffer_lazy=get_exec().mamba.enable_mamba_extra_buffer_lazy,
pp_rank=ps.pp_rank,
pp_size=ps.pp_size,
attn_cp_rank=ps.attn_cp_rank,
python/sglang/srt/mem_cache/kv_cache_configurator.py
+9 -11
View File
@@ -84,8 +84,6 @@ from sglang.srt.runtime_context import (
get_parallel,
get_schedule,
get_spec,
mamba_extra_buffer_enabled,
mamba_extra_buffer_lazy_enabled,
max_speculative_num_draft_tokens,
pre_capture_activation_reserve_mb,
)
@@ -569,7 +567,7 @@ class KVCacheConfigurator:
mamba_spec_state_size=sizes.max_running_requests,
cache_params=self.mambaish_config.mamba2_cache_params,
device=self.device,
enable_mamba_extra_buffer=mamba_extra_buffer_enabled(),
enable_mamba_extra_buffer=get_exec().mamba.enable_mamba_extra_buffer,
draft_model_idx=self.draft_model_idx,
speculative_eagle_topk=get_spec().speculative_eagle_topk,
)
@@ -691,7 +689,7 @@ class KVCacheConfigurator:
max_mamba_cache_size=get_schedule().max_mamba_cache_size,
max_num_reqs=max_num_reqs,
enable_memory_saver=get_exec().features.enable_memory_saver,
enable_mamba_extra_buffer=mamba_extra_buffer_enabled(),
enable_mamba_extra_buffer=get_exec().mamba.enable_mamba_extra_buffer,
speculative_num_draft_tokens=get_spec().speculative_num_draft_tokens,
disable_overlap_schedule=get_schedule().disable_overlap_schedule,
need_sort=get_disagg().disaggregation_mode in ("decode", "prefill"),
@@ -813,7 +811,7 @@ class KVCacheConfigurator:
extra_max_context_len=extra_max_context_len,
max_num_reqs=max_num_reqs,
enable_memory_saver=get_exec().features.enable_memory_saver,
enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(),
enable_mamba_extra_buffer=get_exec().mamba.enable_mamba_extra_buffer,
disable_overlap_schedule=get_schedule().disable_overlap_schedule,
need_sort=get_disagg().disaggregation_mode in ("decode", "prefill"),
speculative_num_draft_tokens=get_spec().speculative_num_draft_tokens,
@@ -1033,7 +1031,7 @@ class KVCacheConfigurator:
mamba_layer_ids=self._get_mamba_layer_ids_for_req_pool(),
speculative_num_draft_tokens=max_speculative_num_draft_tokens(),
speculative_eagle_topk=get_spec().speculative_eagle_topk,
enable_mamba_extra_buffer=mamba_extra_buffer_enabled(),
enable_mamba_extra_buffer=get_exec().mamba.enable_mamba_extra_buffer,
pre_alloc_size=pre_alloc_size,
enable_overlap_schedule=not get_schedule().disable_overlap_schedule,
mamba_size=get_schedule().max_mamba_cache_size,
@@ -1104,8 +1102,8 @@ class KVCacheConfigurator:
enable_memory_saver=get_exec().features.enable_memory_saver,
cache_params=self.mambaish_config.mamba2_cache_params,
mamba_layer_ids=self._get_mamba_layer_ids_for_req_pool(),
enable_mamba_extra_buffer=mamba_extra_buffer_enabled(),
enable_mamba_extra_buffer_lazy=mamba_extra_buffer_lazy_enabled(),
enable_mamba_extra_buffer=get_exec().mamba.enable_mamba_extra_buffer,
enable_mamba_extra_buffer_lazy=get_exec().mamba.enable_mamba_extra_buffer_lazy,
# A PD prefill server never runs TARGET_VERIFY, so skip the
# verify-only per-draft-token state snapshots (see the draft-head
# case above: None => the pool skips SpeculativeState).
@@ -2131,16 +2129,16 @@ class KVCacheConfigurator:
)
additional_ratio = 0
if mamba_extra_buffer_enabled():
if get_exec().mamba.enable_mamba_extra_buffer:
# ping-pong buffer size is 2 when overlap schedule is on, 1 otherwise.
# Lazy mode saves 1 slot (2 → 1) for overlap; non-overlap already uses 1.
if not get_schedule().disable_overlap_schedule:
if mamba_extra_buffer_lazy_enabled():
if get_exec().mamba.enable_mamba_extra_buffer_lazy:
additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP_LAZY
else:
additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP
else:
assert not mamba_extra_buffer_lazy_enabled(), (
assert not get_exec().mamba.enable_mamba_extra_buffer_lazy, (
"Lazy extra buffer requires overlap schedule (--disable-overlap-schedule is incompatible)"
)
additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP
python/sglang/srt/mem_cache/storage/umbp/umbp_direct_linker.py
+6 -2
View File
@@ -23,7 +23,11 @@ from sglang.srt.mem_cache.hybrid_cache.linker_pool_assembler import (
resolve_hybrid_device_pool_group,
)
from sglang.srt.mem_cache.unified_cache.unified_cache_linker import UnifiedCacheLinker
from sglang.srt.runtime_context import get_memory, get_model
from sglang.srt.runtime_context import (
get_memory,
get_model,
get_parallel,
)
from sglang.srt.utils import freeze_gc, get_device_module
logger = logging.getLogger(__name__)
@@ -278,7 +282,7 @@ class UMBPDirectLinker(UnifiedCacheLinker):
storage_config = HiCacheStorageConfig(
tp_rank=tp_rank,
tp_size=server_args.tp_size,
tp_size=get_parallel().tp_size,
pp_rank=params.pp_rank,
pp_size=params.pp_size,
attn_cp_rank=params.attn_cp_rank,
python/sglang/srt/model_executor/model_runner.py
+15 -13
View File
@@ -183,8 +183,6 @@ from sglang.srt.runtime_context import (
get_parallel,
get_schedule,
get_spec,
is_ep_joiner,
is_ep_scale_joiner,
max_speculative_num_draft_tokens,
remote_instance_transfer_engine_enabled,
set_global_dwdp_manager,
@@ -494,7 +492,10 @@ class ModelRunner:
self.graph_time_usage: dict[str, float] = {}
def _initialize_elastic_ep_joiner(self) -> None:
if not (get_exec().moe.elastic_ep_backend is not None and is_ep_scale_joiner()):
if not (
get_exec().moe.elastic_ep_backend is not None
and get_exec().moe.is_ep_scale_joiner
):
return
join_effective_ep_size = get_parallel().ep_join_rank_offset + self.ps.tp_size
@@ -703,7 +704,9 @@ class ModelRunner:
if self.is_draft_worker:
return
expert_rank = self.ps.moe_ep_rank + (
get_parallel().ep_join_rank_offset if is_ep_scale_joiner() else 0
get_parallel().ep_join_rank_offset
if get_exec().moe.is_ep_scale_joiner
else 0
)
set_global_expert_location_metadata(
compute_initial_expert_location_metadata(
@@ -1183,7 +1186,6 @@ class ModelRunner:
with self._load_format_scope(draft_load_format):
loaded = load_model_with_memory_saver(
server_args=self.server_args,
model_config=self.model_config,
load_config=self.load_config,
device=self.device,
@@ -1274,7 +1276,7 @@ class ModelRunner:
dist_barrier_after_load(
elastic_ep_backend=get_exec().moe.elastic_ep_backend,
tp_rank=self.ps.tp_rank,
is_ep_joiner=self.server_args.is_ep_joiner,
is_ep_joiner=get_exec().moe.is_ep_joiner,
)
def start_startup_weight_load(self) -> None:
@@ -1297,7 +1299,7 @@ class ModelRunner:
dist_barrier_after_load(
elastic_ep_backend=get_exec().moe.elastic_ep_backend,
tp_rank=self.ps.tp_rank,
is_ep_joiner=is_ep_joiner(),
is_ep_joiner=get_exec().moe.is_ep_joiner,
)
self.startup_weight_load = None
@@ -2022,7 +2024,7 @@ class ModelRunner:
self._rearm_eplb_after_elastic_scale()
def _report_elastic_scale_failure(self, error: str, effective_size: int) -> None:
if self.ps.tp_rank != 0 or is_ep_scale_joiner():
if self.ps.tp_rank != 0 or get_exec().moe.is_ep_scale_joiner:
return
from sglang.srt.managers.io_struct import ElasticScaleUpdateReq
@@ -2101,12 +2103,12 @@ class ModelRunner:
ElasticEPStateManager.mark_syncing_new_world()
self._elastic_scale_ready_barrier(
target_size=target_size,
log_tag="JOINER" if is_ep_scale_joiner() else "PRIMARY",
log_tag="JOINER" if get_exec().moe.is_ep_scale_joiner else "PRIMARY",
)
ElasticEPStateManager.commit_scale()
self._rearm_eplb_after_elastic_scale()
if self.ps.tp_rank == 0 and not is_ep_scale_joiner():
if self.ps.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
from sglang.srt.managers.io_struct import ElasticScaleUpdateReq
self._pending_elastic_scale_update = ElasticScaleUpdateReq(
@@ -2139,7 +2141,7 @@ class ModelRunner:
)
ElasticEPStateManager.fail_recovery(error)
self._report_elastic_scale_failure(error, effective_size)
if self.ps.tp_rank == 0 and not is_ep_scale_joiner():
if self.ps.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
logger.error("[Elastic EP] %s", error)
return
@@ -2165,7 +2167,7 @@ class ModelRunner:
ElasticEPStateManager.fail_scale(error)
self._reset_eplb_after_elastic_scale_failure()
self._report_elastic_scale_failure(error, effective_size)
if self.ps.tp_rank == 0 and not is_ep_scale_joiner():
if self.ps.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
logger.error("[Elastic EP] %s", error)
return
@@ -2181,7 +2183,7 @@ class ModelRunner:
ElasticEPStateManager.fail_scale(error)
self._reset_eplb_after_elastic_scale_failure()
self._report_elastic_scale_failure(error, effective_size)
if self.ps.tp_rank == 0 and not is_ep_scale_joiner():
if self.ps.tp_rank == 0 and not get_exec().moe.is_ep_scale_joiner:
logger.error("[Elastic EP] %s", error)
return
if not ElasticEPStateManager.begin_scale():
python/sglang/srt/model_executor/model_runner_components/load_model_utils.py
+10 -11
View File
@@ -5,7 +5,7 @@ import logging
import os
import socket
import threading
from typing import TYPE_CHECKING, Any, Optional
from typing import TYPE_CHECKING, Any
import msgspec
import torch
@@ -68,8 +68,8 @@ def maybe_precompile_model_kernels_after_loading(model, device: str) -> None:
class LoadedModel(msgspec.Struct, frozen=True, kw_only=True):
loader: Any
model: Any
remote_instance_weight_info: Optional[Any]
startup_weight_load: Optional[Any] = None
remote_instance_weight_info: Any | None
startup_weight_load: Any | None = None
def maybe_downgrade_dtype_for_legacy_gpu(*, model_config: ModelConfig) -> None:
@@ -87,7 +87,7 @@ def maybe_downgrade_dtype_for_legacy_gpu(*, model_config: ModelConfig) -> None:
def maybe_trigger_remote_instance_nccl_send_group(
*, tp_rank: int, load_format: Optional[str] = None
*, tp_rank: int, load_format: str | None = None
) -> None:
"""``load_format`` is this runner's effective format: a draft loading under
``--speculative-draft-draft-load-format`` needs its own send group, and the
@@ -136,7 +136,7 @@ def load_kv_cache_scales(*, model, kv_cache_dtype: str) -> None:
)
def resolve_sliding_window_size(model, model_config: ModelConfig) -> Optional[int]:
def resolve_sliding_window_size(model, model_config: ModelConfig) -> int | None:
# Parse other args
sliding_window_size = None
if hasattr(model, "get_attention_sliding_window_size"):
@@ -196,12 +196,12 @@ def build_load_config(
*,
server_args: ServerArgs,
tp_rank: int,
load_format: Optional[str] = None,
load_format: str | None = None,
remote_instance_weight_transporter_engine: Any,
remote_instance_weight_transporter_session_id: str,
draft_model_idx: Optional[int],
draft_model_idx: int | None,
weight_cache_mode: str,
weight_cache_socket: Optional[str],
weight_cache_socket: str | None,
) -> LoadConfig:
from sglang.srt.configs.modelopt_config import ModelOptConfig
@@ -255,7 +255,6 @@ def maybe_enable_ipc_weight_cache(
def load_model_with_memory_saver(
*,
server_args: ServerArgs,
model_config: ModelConfig,
load_config: LoadConfig,
device: str,
@@ -291,7 +290,7 @@ def load_model_with_memory_saver(
model_config=model_config,
)
device_config = DeviceConfig(device, gpu_id)
if server_args.is_startup_weight_load_overlap:
if get_model().is_startup_weight_load_overlap:
from sglang.srt.model_executor.model_runner_components.startup_weight_load import (
StartupWeightLoadManager,
)
@@ -329,7 +328,7 @@ def load_model_with_memory_saver(
def dist_barrier_after_load(
*,
elastic_ep_backend: Optional[str],
elastic_ep_backend: str | None,
tp_rank: int,
is_ep_joiner: bool = False,
) -> None:
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+1 -1
View File
@@ -398,7 +398,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
)
enable_mamba_track = (
self.model_runner.server_args.enable_mamba_extra_buffer()
get_exec().mamba.enable_mamba_extra_buffer
and self.model_runner.spec_algorithm.is_none()
)
python/sglang/srt/model_executor/runner/eager_runner.py
+3 -2
View File
@@ -58,9 +58,9 @@ from sglang.srt.model_executor.runner_utils import (
maybe_publish_prefill_shared_read_done,
)
from sglang.srt.runtime_context import (
get_exec,
get_parallel,
get_spec,
mamba_extra_buffer_enabled,
max_prefill_buffer_tokens,
max_speculative_num_draft_tokens,
)
@@ -138,7 +138,8 @@ class EagerRunner(BaseRunner):
max_num_token=max_num_token,
cache_loc_dtype=torch.int64,
enable_mamba_track=(
mamba_extra_buffer_enabled() and mr.spec_algorithm.is_none()
get_exec().mamba.enable_mamba_extra_buffer
and mr.spec_algorithm.is_none()
),
is_encoder_decoder=is_encoder_decoder,
encoder_len_fill_value=(
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+1 -1
View File
@@ -581,7 +581,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
self.raw_bs = 0
def _is_mamba_track_enabled(self) -> bool:
return self.model_runner.server_args.enable_mamba_extra_buffer() and (
return get_exec().mamba.enable_mamba_extra_buffer and (
not get_memory().disable_radix_cache
)
python/sglang/srt/models/inkling.py
+1 -2
View File
@@ -80,7 +80,6 @@ from sglang.srt.runtime_context import (
get_model,
get_parallel,
get_schedule,
mamba_extra_buffer_enabled,
)
from sglang.srt.utils import add_prefix, is_cuda, make_layers
@@ -1021,7 +1020,7 @@ class InklingForConditionalGeneration(nn.Module):
if get_disagg().disaggregation_mode != "decode":
assert not get_memory().disable_radix_cache
assert not get_schedule().disable_hybrid_swa_memory
assert mamba_extra_buffer_enabled()
assert get_exec().mamba.enable_mamba_extra_buffer
from types import SimpleNamespace
python/sglang/srt/models/inkling_common/sconv.py
+4 -2
View File
@@ -20,7 +20,6 @@ from sglang.srt.models.inkling_common.kernels.sconv import (
from sglang.srt.runtime_context import (
get_exec,
get_parallel,
mamba_extra_buffer_enabled,
)
from sglang.srt.utils import is_cuda, set_weight_attrs
@@ -271,7 +270,10 @@ class ShortConvolution(nn.Module):
draft_token_num = hidden_states.shape[1]
mamba_track_indices = getattr(forward_batch, "mamba_track_indices", None)
do_tracking = mamba_track_indices is not None and mamba_extra_buffer_enabled()
do_tracking = (
mamba_track_indices is not None
and get_exec().mamba.enable_mamba_extra_buffer
)
crossed = track_step = None
if do_tracking:
python/sglang/srt/runtime_context.py
+146 -94
View File
@@ -182,6 +182,62 @@ 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.
"""
attn_dp_size, _ = derive_attention_widths(
tp_size=cfg.tp_size,
attn_cp_size=cfg.attn_cp_size,
dp_size=cfg.dp_size,
enable_dp_attention=cfg.enable_dp_attention,
)
return derive_parallel_widths(
tp_size=cfg.tp_size,
attn_cp_size=cfg.attn_cp_size,
attn_dp_size=attn_dp_size,
moe_ep_size=cfg.ep_size,
moe_dp_size=cfg.moe_dp_size,
dcp_size=cfg.dcp_size,
dcp_enabled=cfg.dcp_size > 1,
)
def attn_tp_size_of(cfg: Any):
"""`attn_tp_size`, computed at publish. See `parallel_widths_of`."""
return parallel_widths_of(cfg)["attn_tp_size"]
def attn_dp_size_of(cfg: Any):
"""`attn_dp_size`, computed at publish. See `parallel_widths_of`."""
return parallel_widths_of(cfg)["attn_dp_size"]
def attn_dcp_size_of(cfg: Any):
"""`attn_dcp_size`, computed at publish. See `parallel_widths_of`."""
return parallel_widths_of(cfg)["attn_dcp_size"]
def moe_ep_size_of(cfg: Any):
"""`moe_ep_size`, computed at publish. See `parallel_widths_of`."""
return parallel_widths_of(cfg)["moe_ep_size"]
def moe_tp_size_of(cfg: Any):
"""`moe_tp_size`, computed at publish. See `parallel_widths_of`."""
return parallel_widths_of(cfg)["moe_tp_size"]
def dcp_enabled_of(cfg: Any):
"""`dcp_enabled`, computed at publish. See `parallel_widths_of`."""
return parallel_widths_of(cfg)["dcp_enabled"]
class ParallelContext:
"""Parallel-topology namespace: one spelling per name.
@@ -247,13 +303,17 @@ class ParallelContext:
def clear_derived_widths(self) -> None:
self._derived.clear()
def _derived_width(self, name, getter):
"""A width the leaves imply: the stamp, else the live group.
def _derived_width(self, name):
"""A width the configuration implies: override, else stamp, else the
published leaf.
The fallback keeps a process that installed groups without going
through `initialize_model_parallel` working. When neither is there,
the failure says which of the two is missing rather than surfacing a
group getter's bare assertion.
The leaf is computed at publish by `parallel_widths_of`; the stamp sits
above it because an elastic scale-up restamps `attn_dp_size` after
publish, and a scope that swaps in another TP group states the quotients
through `override`.
Nothing is recomputed on read, so overriding `tp_size` does not move
`attn_tp_size`: name the width, or publish a config.
"""
overrides = self._overrides
if name in overrides:
@@ -261,16 +321,16 @@ class ParallelContext:
derived = self._derived
if name in derived:
return derived[name]
try:
return getter()
except (AssertionError, AttributeError, RuntimeError) as exc:
raise RuntimeError(
f"derived parallel width {name!r} is not available: it is "
"computed from the configured leaves when the process groups "
"are built (initialize_model_parallel / "
"initialize_dp_attention), and neither a stamp nor a live "
"group is present"
) from exc
config = self._config
if config is not None and name in config._fields:
return getattr(config, name)
raise RuntimeError(
f"derived parallel width {name!r} is not available: it is computed "
"from the configured leaves at publish, and restamped when the "
"process groups are built. Nothing is published and nothing has "
"been stamped -- publish a parallel config, or state the width "
f"with get_parallel().override({name}=...)"
)
@contextmanager
def override(self, **kwargs):
@@ -302,12 +362,6 @@ class ParallelContext:
def pp_rank(self) -> int:
return self._v("pp_rank", _ps().get_pipeline_model_parallel_rank)
@property
def moe_ep_size(self) -> int:
return self._derived_width(
"moe_ep_size", _ps().get_moe_expert_parallel_world_size
)
@property
def moe_ep_rank(self) -> int:
return self._v("moe_ep_rank", _ps().get_moe_expert_parallel_rank)
@@ -316,22 +370,10 @@ class ParallelContext:
def moe_dp_rank(self) -> int:
return self._v("moe_dp_rank", _ps().get_moe_data_parallel_rank)
@property
def moe_tp_size(self) -> int:
return self._derived_width(
"moe_tp_size", _ps().get_moe_tensor_parallel_world_size
)
@property
def moe_tp_rank(self) -> int:
return self._v("moe_tp_rank", _ps().get_moe_tensor_parallel_rank)
@property
def attn_tp_size(self) -> int:
return self._derived_width(
"attn_tp_size", _ps().get_attn_tensor_model_parallel_world_size
)
@property
def attn_tp_rank(self) -> int:
return self._v("attn_tp_rank", _ps().get_attn_tensor_model_parallel_rank)
@@ -344,32 +386,12 @@ class ParallelContext:
def dcp_rank(self) -> int:
return self._v("dcp_rank", _ps().get_dcp_rank)
@property
def dcp_enabled(self) -> bool:
def getter():
if _ps().get_dcp_group_no_assert() is None:
return False
return _ps().get_dcp_world_size() > 1
return self._derived_width("dcp_enabled", getter)
@property
def attn_dcp_size(self) -> int:
return self._derived_width(
"attn_dcp_size",
lambda: _ps().get_dcp_world_size() if self.dcp_enabled else 1,
)
@property
def attn_dcp_rank(self) -> int:
return self._v(
"attn_dcp_rank", lambda: self.dcp_rank if self.dcp_enabled else 0
)
@property
def attn_dp_size(self) -> int:
return self._derived_width("attn_dp_size", _dp().get_attention_dp_size)
@property
def attn_dp_rank(self) -> int:
return self._v("attn_dp_rank", _dp().get_attention_dp_rank)
@@ -411,6 +433,34 @@ class ParallelContext:
return self._v("dcp_group", _ps().get_dcp_group)
def _install_derived_widths() -> None:
"""Give `ParallelContext` a property per declared quotient.
They are declared in `arg_groups/fields/parallel.py`, in the same class as
the leaves they are computed from -- unannotated, so `collect_input_fields`
leaves them off the record while they still live where the namespace does. Written here as
properties rather than answered by `__getattr__` because they are read
inside compiled model code, where an attribute load is traceable and a
dynamic lookup is not.
"""
from sglang.srt.arg_groups.arg_utils import Derived
from sglang.srt.arg_groups.fields.parallel import Parallel
for name, decl in vars(Parallel).items():
if not isinstance(decl, Derived):
continue
def getter(self, _name=name):
return self._derived_width(_name)
getter.__name__ = name
getter.__doc__ = decl.doc
setattr(ParallelContext, name, property(getter))
_install_derived_widths()
class _FlagGroupBase:
"""Shared flag-group behavior: typo-safe writes + transactional ``override()``.
@@ -845,9 +895,49 @@ def _build_config_bags(server_args: Any) -> dict:
"clashes with a subgroup of the same name"
)
bag._set(field, value)
_install_derived_leaves(tops, server_args)
return tops
def _install_derived_leaves(tops: dict, server_args: Any) -> None:
"""Compute the declared config-derived fields into their bags.
A `Derived(fn=...)` is a pure function of the published configuration, so it
is computed once, here, and stored as an ordinary leaf: readers get a plain
attribute load, and there is one answer rather than a pre-publish spelling
and a post-publish one that have to be kept saying the same thing.
The function is handed the whole resolved config, not the bag it lands in.
A derivation is free to span namespaces and they do -- the mamba
extra-buffer predicate reads `memory.disable_radix_cache` alongside its own
`exec.mamba` strategy -- which is exactly why it cannot be written as a
method on either bag.
"""
import importlib
from sglang.srt.arg_groups.arg_utils import Derived
from sglang.srt.arg_groups.overrides import resolved_view
namespaces = getattr(type(server_args), "_NAMESPACES", None)
if not namespaces:
return
view = resolved_view(server_args)
for source in namespaces:
path = getattr(source, "_NS_PATH", None)
if path is None:
continue
for name, decl in vars(source).items():
if not isinstance(decl, Derived) or not decl.fn:
continue
module, _, attr = decl.fn.rpartition(".")
bag = tops.get(path.split(".")[0])
for segment in path.split(".")[1:]:
bag = bag and getattr(bag, segment, None)
if bag is None:
continue
bag._set(name, getattr(importlib.import_module(module), attr)(view))
def _resolved_or_field(server_args: Any, name: str, default: Any) -> Any:
"""What resolution decided for `name`, falling back to the field.
@@ -1662,8 +1752,8 @@ def reset_context() -> None:
``server_args`` and install fresh ``Flags`` and ``Resources``.
``parallel`` holds the stamped derived widths, which go with the lifecycle
that stamped them: `_derived_width` prefers the stamp over the live group,
so leaving one behind lets the next test read the previous topology.
that stamped them: `_derived_width` prefers the stamp over the leaves, so
leaving one behind lets the next test read the previous topology.
"""
_CONTEXT._server_args = None
_CONTEXT._config_bags = None
@@ -1677,29 +1767,6 @@ def reset_context() -> None:
set_global_dwdp_manager(None)
def mamba_extra_buffer_enabled() -> bool:
"""Whether the mamba radix cache keeps its extra state buffer.
A predicate over two published leaves (``memory.disable_radix_cache`` and
``exec.mamba.mamba_radix_cache_strategy``), so it reads the bags rather
than the startup record — the ``ServerArgs`` member of the same name is the
pre-publish equivalent used inside the resolution pipeline.
"""
return (
get_memory().disable_radix_cache is False
and get_exec().mamba.mamba_radix_cache_strategy
in ("extra_buffer", "extra_buffer_lazy")
)
def mamba_extra_buffer_lazy_enabled() -> bool:
"""The lazy variant of :func:`mamba_extra_buffer_enabled`."""
return (
get_memory().disable_radix_cache is False
and get_exec().mamba.mamba_radix_cache_strategy == "extra_buffer_lazy"
)
def remote_instance_transfer_engine_enabled(load_format: str | None = None) -> bool:
"""Whether remote-instance weight loading runs over the transfer engine.
@@ -2034,21 +2101,6 @@ def cutedsl_moe_max_num_tokens() -> int:
return max(prefill_tokens, decode_max_bs * num_tokens_per_req)
def is_ep_joiner() -> bool:
"""True in a process launched as an elastic-EP joiner (scale or recover).
A predicate over the published ``exec.moe.ep_join_mode`` leaf, so it follows
a post-publish override; the same-named ``ServerArgs`` property is the
pre-publish equivalent.
"""
return get_exec().moe.ep_join_mode in ("scale", "recover")
def is_ep_scale_joiner() -> bool:
"""True in a process launched as an elastic-EP scale-up joiner."""
return get_exec().moe.ep_join_mode == "scale"
def describe_kv_events_publisher(server_args: Any) -> Optional[dict]:
"""Return a structured description of this server's KV-event
publisher, or `None` if publishing is disabled / misconfigured.
python/sglang/srt/server_args.py
+20 -53
View File
@@ -41,32 +41,27 @@ import logging
import tempfile
import uuid
from contextlib import contextmanager
from typing import Any, Dict, List, Optional
from typing import Any
from sglang.kernels.ops.kv_canary.consts import RealKvHashMode
from sglang.srt.arg_groups.arg_utils import NS, A, Arg, add_cli_args_from_dataclass
from sglang.srt.arg_groups.arg_utils import (
add_cli_args_from_dataclass,
)
from sglang.srt.arg_groups.argparse_actions import (
DeprecatedAction,
DeprecatedAliasStoreAction,
DeprecatedStoreConstAction,
DeprecatedStoreTrueAction,
LoRAPathAction,
)
from sglang.srt.arg_groups.model_override_base import ep_joiner_of, ep_scale_joiner_of
from sglang.srt.arg_groups.overrides import (
mamba_extra_buffer_lazy_of,
mamba_extra_buffer_of,
remote_instance_transfer_engine_of,
resolution_projection,
resolving_view,
)
from sglang.srt.environ import envs
from sglang.srt.function_call.function_call_parser import FunctionCallParser
from sglang.srt.lora.lora_registry import LoRARef
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
CudaGraphConfig,
parse_cuda_graph_config_arg,
)
from sglang.srt.model_executor.cuda_graph_config import Backend
from sglang.srt.parser.reasoning_parser import ReasoningParser
from sglang.srt.runtime_context import (
get_context,
@@ -74,13 +69,6 @@ from sglang.srt.runtime_context import (
publish,
)
from sglang.srt.speculative.decoupled_spec_io import DecoupledSpecIpcConfig
from sglang.srt.utils.common import (
LORA_TARGET_ALL_MODULES,
SUPPORTED_LORA_TARGET_MODULES,
human_readable_int,
json_list_type,
nullable_str,
)
from sglang.srt.utils.network import NetworkAddress, get_free_port, wait_port_available
logger = logging.getLogger(__name__)
@@ -290,7 +278,7 @@ class ServerArgs:
self._resolution_finished = True
@property
def launch_command(self) -> Optional[str]:
def launch_command(self) -> str | None:
"""How this record was created, verbatim.
`resolved_dict` answers with what resolution decided; this answers with
@@ -306,7 +294,7 @@ class ServerArgs:
"""
return getattr(self, "_launch_command", None)
def resolved_dict(self) -> Dict[str, Any]:
def resolved_dict(self) -> dict[str, Any]:
"""This configuration as a plain dict of resolved field values.
What the whole-object readbacks report (`/server_info` and its gRPC and
@@ -661,7 +649,7 @@ class ServerArgs:
return TokenizerWorker
def url(self, port: Optional[int] = None):
def url(self, port: int | None = None):
scheme = "https" if self.ssl_certfile else "http"
# When binding to all interfaces, use loopback for internal requests.
host = self.host
@@ -677,25 +665,6 @@ class ServerArgs:
def engine_info_bootstrap_url(self):
return self.url(port=self.engine_info_bootstrap_port)
@property
def is_ep_joiner(self) -> bool:
"""True for processes launched as elastic-EP joiners."""
cfg = resolving_view(self)
return cfg.ep_join_mode in ("scale", "recover")
@property
def is_ep_scale_joiner(self) -> bool:
cfg = resolving_view(self)
return cfg.ep_join_mode == "scale"
@property
def is_startup_weight_load_overlap(self) -> bool:
cfg = resolving_view(self)
return cfg.startup_weight_load_mode == "overlap"
def __setattr__(self, name, value):
# The record holds the operator's input. It is writable while the
# caller is still assembling it and sealed from the moment resolution
@@ -721,12 +690,6 @@ class ServerArgs:
)
object.__setattr__(self, name, value)
def enable_mamba_extra_buffer(self) -> bool:
return mamba_extra_buffer_of(resolving_view(self))
def enable_mamba_extra_buffer_lazy(self) -> bool:
return mamba_extra_buffer_lazy_of(resolving_view(self))
def check_server_args(self):
from sglang.srt.arg_groups.validation_hook import check_server_args
@@ -745,6 +708,7 @@ class ServerArgs:
# A set, not an order: `collect_input_fields` orders by `field_order.py`, which
# is what keeps the positional constructor stable.
_INPUT_NAMESPACES = [
Model,
ExecDeterministic,
@@ -774,6 +738,9 @@ ServerArgs.__annotations__ = {**_annotations, **ServerArgs.__annotations__}
# The assembled record has no base classes, so it carries the map the classes
# used to answer through their `_NS_PATH`.
ServerArgs._NS_BY_FIELD = _namespaces
# The classes themselves, so the bag projection can find the declarations
# that are not fields -- the derived half of each namespace.
ServerArgs._NAMESPACES = _INPUT_NAMESPACES
for _name, _value in _defaults.items():
setattr(ServerArgs, _name, _value)
ServerArgs = dataclasses.dataclass(ServerArgs)
@@ -892,7 +859,7 @@ def record_writable(server_args: Any):
object.__setattr__(server_args, "_input_frozen", True)
def prepare_server_args(argv: List[str]) -> ServerArgs:
def prepare_server_args(argv: list[str]) -> ServerArgs:
"""
Prepare the server arguments from the command line arguments.
@@ -962,10 +929,10 @@ class PortArgs:
metrics_ipc_name: str
# The ipc filename for MultiTokenizerRouter to receive inputs from TokenizerWorker processes (zmq)
tokenizer_worker_ipc_name: Optional[str]
tokenizer_worker_ipc_name: str | None
# The ipc endpoints between verifier scheduler and drafter scheduler
decoupled_spec_ipc_config: Optional[DecoupledSpecIpcConfig]
decoupled_spec_ipc_config: DecoupledSpecIpcConfig | None
# zmq address for load snapshot PUSH/PULL (dp-attention TCP mode only;
# empty when IPC mode derives the address from instance_id).
@@ -978,8 +945,8 @@ class PortArgs:
@staticmethod
def init_new(
server_args: ServerArgs,
dp_rank: Optional[int] = None,
worker_ports: Optional[List[int]] = None,
dp_rank: int | None = None,
worker_ports: list[int] | None = None,
) -> PortArgs:
cfg = resolving_view(server_args)
if server_args.nccl_port is None:
@@ -1046,7 +1013,7 @@ class PortArgs:
# overflow.
is_rust_server = envs.SGLANG_RUST_SERVER.get()
NUM_DERIVED_PORTS = 6 if not is_rust_server else 6 + cfg.dp_size
if server_args.is_ep_scale_joiner:
if ep_scale_joiner_of(resolving_view(server_args)):
port_base = server_args.port + ZMQ_TCP_PORT_DELTA
if port_base + NUM_DERIVED_PORTS > 65535:
port_base = server_args.port - ZMQ_TCP_PORT_DELTA
@@ -1070,7 +1037,7 @@ class PortArgs:
assert worker_ports is not None
scheduler_input_port = worker_ports[dp_rank]
is_joiner = server_args.is_ep_joiner
is_joiner = ep_joiner_of(resolving_view(server_args))
# Under SGLANG_DISTRIBUTED_INIT_METHOD_OVERRIDE, SGLang never binds
# dist_init_port / nccl_port (rendezvous uses the externally-managed
# store; see distributed/bootstrap.py:_resolve_dist_init_method), so
python/sglang/srt/speculative/spec_utils.py
+4 -5
View File
@@ -49,9 +49,8 @@ from sglang.srt.mem_cache.allocation import (
assign_req_to_token_pool_func as assign_req_to_token_pool_func,
)
from sglang.srt.runtime_context import (
get_exec,
get_spec,
mamba_extra_buffer_enabled,
mamba_extra_buffer_lazy_enabled,
mamba_track_grid,
max_speculative_num_draft_tokens,
)
@@ -778,10 +777,10 @@ def prepare_mamba_track_for_verify(batch: ScheduleBatch) -> None:
Lazy: gather the positions planned by mamba_lazy_spec_prepare. Runs
inside forward isolation, so it must not mutate req/pool state.
"""
if not mamba_extra_buffer_enabled():
if not get_exec().mamba.enable_mamba_extra_buffer:
return
track_positions = None
if mamba_extra_buffer_lazy_enabled():
if get_exec().mamba.enable_mamba_extra_buffer_lazy:
track_positions = batch.mamba_lazy_spec_track_positions_cpu
assert track_positions is not None and len(track_positions) == len(
batch.reqs
@@ -1066,7 +1065,7 @@ def spec_prepare_for_decode(batch: ScheduleBatch) -> None:
"""eagle/ngram share a stateless free function; dflash keeps stateful
prep on its draft input -- the dispatcher routes.
"""
if mamba_extra_buffer_lazy_enabled():
if get_exec().mamba.enable_mamba_extra_buffer_lazy:
# Scheduler phase (outside forward isolation).
batch.mamba_lazy_spec_prepare(
mamba_track_grid(batch.tree_cache.page_size),