1849 lines
73 KiB
Python
1849 lines
73 KiB
Python
# Copyright 2023-2026 SGLang Team
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""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_server_args()`` returns the process-wide ``ServerArgs``. This is the
|
|
user's raw input, kept **read-only** for debug and reproduction; what
|
|
resolution decided lives in the declarations (``resolution_result``) and, for
|
|
business code, in the namespace bags below -- never on this object's fields. The context owns the storage:
|
|
publishing goes through ``RuntimeContext.set_server_args`` (the legacy
|
|
``set_global_server_args_for_scheduler`` / ``get_global_server_args`` are thin
|
|
shims over this slot).
|
|
|
|
``get_exec()`` / ``get_memory()`` / ``get_schedule()`` / ``get_device()`` /
|
|
``get_model()`` / ``get_spec()`` / ``get_lora()`` / ``get_mm()`` /
|
|
``get_disagg()`` / ``get_serving()`` / ``get_observability()`` return the
|
|
resolved **config namespace bags** — the single source of truth for config,
|
|
snapshotted from ``server_args`` at publish and driven by the ``NS(...)``
|
|
metadata on each field (multi-level under ``exec.*``). Reads are attribute
|
|
chains (``get_exec().moe.moe_runner_backend``); bags are read-only by bare
|
|
assignment (written via ``override``).
|
|
|
|
``get_flags()`` returns the runtime-flags tier: state that is **not** a pure
|
|
function of config (the capture lifecycle, ACTIVE MoE backend, DP runtime) —
|
|
never a mirror of config. Flags live in typed dataclass groups; reads and
|
|
writes are plain attribute access, and each group offers a transactional,
|
|
test-only ``override(**kw)``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
import functools
|
|
import logging
|
|
import math
|
|
import os
|
|
import sys
|
|
from contextlib import contextmanager
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.server_args import ServerArgs
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Imported lazily so this module has no import-time dependencies: any module can
|
|
# import get_parallel at module level without risking an import cycle.
|
|
def _ps():
|
|
from sglang.srt.distributed import parallel_state
|
|
|
|
return parallel_state
|
|
|
|
|
|
def _dp():
|
|
from sglang.srt.layers import dp_attention
|
|
|
|
return dp_attention
|
|
|
|
|
|
@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.
|
|
"""
|
|
from sglang.srt.arg_groups.arg_utils import namespace_of
|
|
from sglang.srt.server_args import ServerArgs
|
|
|
|
return frozenset(
|
|
field
|
|
for field, path in namespace_of(ServerArgs).items()
|
|
if path.split(".")[0] == "parallel"
|
|
)
|
|
|
|
|
|
_PARALLEL_FIELDS = frozenset(
|
|
{
|
|
"world_size",
|
|
"world_rank",
|
|
"tp_size",
|
|
"tp_rank",
|
|
"pp_size",
|
|
"pp_rank",
|
|
"moe_ep_size",
|
|
"moe_ep_rank",
|
|
"moe_dp_size",
|
|
"moe_dp_rank",
|
|
"moe_tp_size",
|
|
"moe_tp_rank",
|
|
"attn_tp_size",
|
|
"attn_tp_rank",
|
|
"attn_cp_size",
|
|
"attn_cp_rank",
|
|
"dcp_enabled",
|
|
"dcp_size",
|
|
"dcp_rank",
|
|
"attn_dcp_size",
|
|
"attn_dcp_rank",
|
|
"attn_dp_size",
|
|
"attn_dp_rank",
|
|
"world_group",
|
|
"tp_group",
|
|
"pp_group",
|
|
"moe_ep_group",
|
|
"moe_dp_group",
|
|
"moe_tp_group",
|
|
"attn_tp_group",
|
|
"attn_cp_group",
|
|
"dcp_group",
|
|
}
|
|
)
|
|
|
|
|
|
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.
|
|
"""
|
|
attn_dp_size = dp_size if enable_dp_attention else 1
|
|
return attn_dp_size, tp_size // attn_dp_size // attn_cp_size
|
|
|
|
|
|
def derive_parallel_widths(
|
|
*,
|
|
tp_size: int,
|
|
attn_cp_size: int,
|
|
attn_dp_size: int,
|
|
moe_ep_size: int,
|
|
moe_dp_size: int,
|
|
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.
|
|
|
|
`world_size` is not among them: it is not a quotient, and `get_world_size()`
|
|
answers with the live WORLD group, which stays right through an elastic
|
|
scale-up that a stamp taken at group build would not survive.
|
|
"""
|
|
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_tp_size": derive_attention_widths(
|
|
tp_size=tp_size,
|
|
attn_cp_size=attn_cp_size,
|
|
dp_size=attn_dp_size,
|
|
enable_dp_attention=True,
|
|
)[1],
|
|
"moe_ep_size": moe_ep_size,
|
|
"moe_tp_size": tp_size // moe_ep_size // moe_dp_size,
|
|
"dcp_enabled": dcp_enabled,
|
|
"attn_dcp_size": dcp_size if dcp_enabled else 1,
|
|
}
|
|
|
|
|
|
class ParallelContext:
|
|
"""Parallel-topology namespace: one spelling per name.
|
|
|
|
Ranks and group handles are read-through ``@property`` over the canonical
|
|
getters, so they answer with the **live** process groups and raise before
|
|
distributed init. Every other name — ``tp_size`` and its size siblings
|
|
included, alongside config-only leaves such as ``nccl_port`` — is answered
|
|
from the published ``parallel`` bag, in any process at any point after
|
|
publish.
|
|
|
|
A size is read from the configuration because the groups are built at
|
|
exactly the configured widths. Two things do not follow that rule and are
|
|
asked of the group itself: ``initialize_model_parallel`` aliases ``_MOE_DP``
|
|
to ``_ATTN_CP`` when ``attn_cp_size > moe_dp_size``, so a reader that means
|
|
the MoE communicator's width calls ``get_moe_cp_size()``; and
|
|
``patch_tensor_parallel_group`` runs a scope under a different TP group,
|
|
which it declares by overriding ``tp_size``, ``tp_rank`` and ``tp_group``
|
|
for its duration. Elastic EP is a third case, and it needs no rule here: it
|
|
scales ``ep_size`` / ``dp_size`` on the published bag while the group
|
|
coordinators keep the width they were constructed with, so the two are
|
|
different names rather than two answers to one name.
|
|
"""
|
|
|
|
__slots__ = ("_overrides", "_config", "_derived")
|
|
|
|
def __init__(self):
|
|
self._overrides = {}
|
|
self._config = None # parallel config bag, wired at publish
|
|
self._derived = {} # widths stamped when the groups are built
|
|
|
|
def __getattr__(self, name):
|
|
if name.startswith("_"):
|
|
# This also breaks the recursion when the ``_config`` slot itself is
|
|
# still unset (pickle/copy protocols probe attributes before
|
|
# __init__ runs).
|
|
raise AttributeError(name)
|
|
overrides = self._overrides
|
|
if name in overrides:
|
|
return overrides[name]
|
|
config = self._config
|
|
if config is not None:
|
|
if name in config._fields:
|
|
return getattr(config, name)
|
|
elif name in _parallel_config_leaves():
|
|
raise ValueError("config namespace 'parallel' not published")
|
|
raise AttributeError(f"ParallelContext has no {name!r}")
|
|
|
|
def _v(self, name, getter):
|
|
overrides = self._overrides
|
|
return overrides[name] if name in overrides else getter()
|
|
|
|
def stamp_derived_widths(self, **widths) -> None:
|
|
"""Record the widths derived from the leaves, as the groups are built.
|
|
|
|
`initialize_model_parallel` computes the set through
|
|
`derive_parallel_widths` and hands it here; `initialize_dp_attention`
|
|
stamps `attn_dp_size` again once it knows the effective width, and
|
|
elastic EP restamps it where it already updates the live one. A stamped
|
|
width is what the readers answer with.
|
|
"""
|
|
self._derived.update(widths)
|
|
|
|
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.
|
|
|
|
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.
|
|
"""
|
|
overrides = self._overrides
|
|
if name in overrides:
|
|
return overrides[name]
|
|
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
|
|
|
|
@contextmanager
|
|
def override(self, **kwargs):
|
|
"""Temporarily force parallel values, restoring on exit. Validates keys and
|
|
supports nesting."""
|
|
unknown = set(kwargs) - _PARALLEL_FIELDS
|
|
if unknown:
|
|
raise ValueError(f"unknown parallel field(s): {sorted(unknown)}")
|
|
saved = dict(self._overrides)
|
|
self._overrides.update(kwargs)
|
|
try:
|
|
yield self
|
|
finally:
|
|
self._overrides = saved
|
|
|
|
@property
|
|
def world_size(self) -> int:
|
|
return self._v("world_size", _ps().get_world_size)
|
|
|
|
@property
|
|
def world_rank(self) -> int:
|
|
return self._v("world_rank", _ps().get_world_rank)
|
|
|
|
@property
|
|
def tp_rank(self) -> int:
|
|
return self._v("tp_rank", _ps().get_tensor_model_parallel_rank)
|
|
|
|
@property
|
|
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)
|
|
|
|
@property
|
|
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)
|
|
|
|
@property
|
|
def attn_cp_rank(self) -> int:
|
|
return self._v("attn_cp_rank", _ps().get_attn_context_model_parallel_rank)
|
|
|
|
@property
|
|
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)
|
|
|
|
@property
|
|
def world_group(self) -> Any:
|
|
return self._v("world_group", _ps().get_world_group)
|
|
|
|
@property
|
|
def tp_group(self) -> Any:
|
|
return self._v("tp_group", _ps().get_tp_group)
|
|
|
|
@property
|
|
def pp_group(self) -> Any:
|
|
return self._v("pp_group", _ps().get_pp_group)
|
|
|
|
@property
|
|
def moe_ep_group(self) -> Any:
|
|
return self._v("moe_ep_group", _ps().get_moe_ep_group)
|
|
|
|
@property
|
|
def moe_dp_group(self) -> Any:
|
|
return self._v("moe_dp_group", _ps().get_moe_dp_group)
|
|
|
|
@property
|
|
def moe_tp_group(self) -> Any:
|
|
return self._v("moe_tp_group", _ps().get_moe_tp_group)
|
|
|
|
@property
|
|
def attn_tp_group(self) -> Any:
|
|
return self._v("attn_tp_group", _ps().get_attn_tp_group)
|
|
|
|
@property
|
|
def attn_cp_group(self) -> Any:
|
|
return self._v("attn_cp_group", _ps().get_attn_cp_group)
|
|
|
|
@property
|
|
def dcp_group(self) -> Any:
|
|
return self._v("dcp_group", _ps().get_dcp_group)
|
|
|
|
|
|
class _FlagGroupBase:
|
|
"""Shared flag-group behavior: typo-safe writes + transactional ``override()``.
|
|
|
|
Groups are plain dataclasses; ``__dataclass_fields__`` is the single source
|
|
of truth for which leaves exist, so a mistyped name fails loudly instead of
|
|
creating a stray attribute.
|
|
"""
|
|
|
|
def __setattr__(self, name: str, value: Any) -> None:
|
|
if name not in type(self).__dataclass_fields__:
|
|
raise AttributeError(
|
|
f"{type(self).__name__} has no flag '{name}' (leaves are "
|
|
"declared as dataclass fields; check for typos)"
|
|
)
|
|
object.__setattr__(self, name, value)
|
|
|
|
@contextmanager
|
|
def override(self, **kwargs):
|
|
"""Temporarily force flag values, restoring on exit. Transactional
|
|
(keys validated before any write) — the test-only injection
|
|
primitive."""
|
|
fields = type(self).__dataclass_fields__
|
|
unknown = set(kwargs) - set(fields)
|
|
if unknown:
|
|
raise ValueError(
|
|
f"unknown flag(s) for {type(self).__name__}: {sorted(unknown)}"
|
|
)
|
|
saved = {name: getattr(self, name) for name in kwargs}
|
|
for name, value in kwargs.items():
|
|
object.__setattr__(self, name, value)
|
|
try:
|
|
yield self
|
|
finally:
|
|
for name, value in saved.items():
|
|
object.__setattr__(self, name, value)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class CaptureFlags(_FlagGroupBase):
|
|
"""Capture-time flags; never frozen (written during cuda-graph capture)."""
|
|
|
|
# Seeded from server_args at publish; a model whose _can_torch_compile is
|
|
# False clears it during warmup (the only post-publish writer).
|
|
enable_torch_compile: bool = False
|
|
|
|
# Set for the duration of decode/spec graph capture (model_capture_mode).
|
|
# While set, dispose_tensor() is a no-op so deep_gemm's pre-permute does not
|
|
# free hidden_states that the dual-stream MoE shared expert reads afterward.
|
|
disable_dispose_tensor: bool = False
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class MoeFlags(_FlagGroupBase):
|
|
"""MoE runtime flags, materialized by ``initialize_moe_config`` (scheduler
|
|
init, after distributed setup). ``a2a_backend`` / ``runner_backend`` /
|
|
``disable_fp4_allgather`` are the ACTIVE values: the speculative contexts
|
|
in ``layers.moe.utils`` swap them around draft-model forwards. Values are
|
|
the parsed enums from ``layers.moe.utils``; ``None`` means "not
|
|
initialized yet" and the accessors fall back lazily.
|
|
"""
|
|
|
|
a2a_backend: Any = None
|
|
runner_backend: Any = None
|
|
speculative_runner_backend: Any = None
|
|
speculative_a2a_backend: Any = None
|
|
deepep_mode: Any = None
|
|
deepep_config: str | None = None
|
|
tbo_enabled: bool | None = None
|
|
sbo_enabled: bool | None = None
|
|
tbo_token_distribution_threshold: float | None = None
|
|
disable_fp4_allgather: bool | None = None
|
|
quantization: str | None = None
|
|
# The shared-experts-fusion decision, per runner — the runner_backend /
|
|
# speculative_runner_backend shape. Both leaves are seeded from the config
|
|
# intent by ``initialize_moe_config``; each MoE model's gate
|
|
# (determine_num_fused_shared_experts) refines the ACTIVE leaf, both ways,
|
|
# before its layers build and read it. ``speculative_moe_backend_context``
|
|
# brackets a draft's build: on exit the draft's effective decision is
|
|
# persisted onto the speculative leaf (inspectable afterwards) and the
|
|
# target's ACTIVE value returns.
|
|
disable_shared_experts_fusion: bool | None = None
|
|
speculative_disable_shared_experts_fusion: bool | None = None
|
|
# Lifecycle marker (the capture.disable_dispose_tensor family): set while
|
|
# speculative_moe_backend_context is active, so a draft gate's write also
|
|
# lands on the speculative leaf.
|
|
in_speculative_scope: bool = False
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class DpFlags(_FlagGroupBase):
|
|
"""DP-attention runtime flags, materialized by ``initialize_dp_attention``
|
|
(after distributed setup; reads the model config). Topology values
|
|
(sizes/ranks) stay on ``layers.dp_attention`` until the parallel vertical
|
|
migrates them."""
|
|
|
|
enabled: bool = False
|
|
use_world_group_for_gather: bool = False
|
|
joiner_skip_all_gather: bool = False
|
|
# Hybrid-SSM models materialize idle ranks via the MAX_LEN fabricated-row
|
|
# conversion (set when hf_config has hybrid_override_pattern).
|
|
max_len_with_idle: bool = False
|
|
# DP gathered-buffer allocation metadata (model hidden size / dtype /
|
|
# device), set by initialize_dp_attention alongside the flags above.
|
|
buffer_hidden_size: Any = None
|
|
buffer_dtype: Any = None
|
|
buffer_device: Any = None
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Flags(_FlagGroupBase):
|
|
"""Root of the runtime-flags tier.
|
|
|
|
Resolved configuration lives in the config bags below (projected from the
|
|
declarations at publish) — this tier only carries genuine runtime
|
|
state whose value is not a function of the configuration alone, grouped
|
|
by lifecycle (``capture``) or subsystem (``moe`` / ``dp``).
|
|
"""
|
|
|
|
capture: CaptureFlags = dataclasses.field(default_factory=CaptureFlags)
|
|
moe: MoeFlags = dataclasses.field(default_factory=MoeFlags)
|
|
dp: DpFlags = dataclasses.field(default_factory=DpFlags)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Resources(_FlagGroupBase):
|
|
"""Process-level resource handles: named slots with one reset lifecycle,
|
|
scoped test injection via ``override()``, and the creation/publish
|
|
semantics kept in the owning modules' accessors (which are thin shims
|
|
over these slots)."""
|
|
|
|
# CUDA graph memory pool shared across the prefill and decode graph
|
|
# backends (created lazily by model_executor.runner_utils.pool).
|
|
graph_memory_pool: Any = None
|
|
# EPLB: per-process recorder and the publish-once location metadata
|
|
# (owning accessors live in sglang.srt.eplb).
|
|
expert_distribution_recorder: Any = None
|
|
expert_location_metadata: Any = None
|
|
# LPLB: layer_id -> solver.
|
|
lplb_solvers: dict = dataclasses.field(default_factory=dict)
|
|
# Named side streams (see RuntimeContext.get_stream): name -> stream.
|
|
streams: dict = dataclasses.field(default_factory=dict)
|
|
# Named persistent buffers (see RuntimeContext.get_buffer): name -> tensor.
|
|
# Accessors with bespoke semantics (grow-only, per-device keys) manage
|
|
# their entries directly.
|
|
buffers: dict = dataclasses.field(default_factory=dict)
|
|
# Persistent reusable CUDA events for non-EP DP TBO, keyed by
|
|
# (kind, subbatch) — see dp_attention._tbo_event for why reuse matters.
|
|
tbo_event_pool: dict = dataclasses.field(default_factory=dict)
|
|
# State capturers (installed by their subsystems when capture is on).
|
|
indexer_capturer: Any = None
|
|
experts_capturer: Any = None
|
|
# The shared TCPStore created during distributed initialization.
|
|
tcp_store: Any = None
|
|
# Trace verbosity; the accessor seeds it lazily from SGLANG_TRACE_LEVEL.
|
|
trace_level: Any = None
|
|
|
|
|
|
class ForwardFlags:
|
|
"""Per-forward runtime flags with one API and two backings.
|
|
|
|
Flags read only from eager Python are backed by context variables, so
|
|
nested scopes and threads stay isolated (a new thread sees the defaults).
|
|
Flags that are read or written *inside torch.compile-traced model code*
|
|
(``_GRAPH_VISIBLE``) are backed by plain dict slots instead: dynamo
|
|
cannot trace ``ContextVar.get``/``set``, while plain reads it guards on
|
|
— the storage form these flags had before joining the tier. Their
|
|
writers and readers are single-threaded per process (TBO interleaves
|
|
ubatches on one thread; attention-TP input scattering excludes TBO), so
|
|
context isolation is not needed for correctness.
|
|
|
|
``scoped(**kw)`` — the one regular write path — restores on exit for
|
|
both backings. ``set()`` exists for the legacy unscoped setters' shims.
|
|
"""
|
|
|
|
_DEFAULTS = {
|
|
"multi_stream": False,
|
|
"moe_output_buffer": None,
|
|
# Attention-TP input-scattering (set per forward by
|
|
# AttnTpContext.maybe_input_scattered / set_attn_inputs).
|
|
"attn_input_scattered": False,
|
|
"attn_inputs": None,
|
|
# Sticky across forwards: every ForwardBatch construction writes it;
|
|
# graph runners force False around capture.
|
|
"is_extend_in_batch": False,
|
|
# Per-layer MLP collective control (set by decoder via scoped()
|
|
# around the MLP / MoE / hybrid mixer call).
|
|
# fuse_mlp_allreduce: next residual+LN absorbs the post-MLP all-reduce.
|
|
# mlp_reduce_scatter: postprocess will reduce-scatter (skip MLP AR).
|
|
# flashinfer_trtllm_bypass: deepseek dual-stream graph topk bypass.
|
|
"fuse_mlp_allreduce": False,
|
|
"mlp_reduce_scatter": False,
|
|
"flashinfer_trtllm_bypass": False,
|
|
}
|
|
|
|
# Read/written inside compiled graphs (vocab embedding, communicator,
|
|
# EP dispatch, DP gather/scatter, MLP/MoE skip-AR): plain-slot backed.
|
|
# Before moving a flag out of this set, prove no read/write site sits
|
|
# under torch.compile.
|
|
_GRAPH_VISIBLE = frozenset(
|
|
{
|
|
"attn_input_scattered",
|
|
"attn_inputs",
|
|
"is_extend_in_batch",
|
|
"fuse_mlp_allreduce",
|
|
"mlp_reduce_scatter",
|
|
"flashinfer_trtllm_bypass",
|
|
}
|
|
)
|
|
|
|
__slots__ = ("_vars", "_plain")
|
|
|
|
def __init__(self):
|
|
import contextvars
|
|
|
|
object.__setattr__(
|
|
self,
|
|
"_plain",
|
|
{
|
|
name: default
|
|
for name, default in self._DEFAULTS.items()
|
|
if name in self._GRAPH_VISIBLE
|
|
},
|
|
)
|
|
object.__setattr__(
|
|
self,
|
|
"_vars",
|
|
{
|
|
name: contextvars.ContextVar(f"forward.{name}", default=default)
|
|
for name, default in self._DEFAULTS.items()
|
|
if name not in self._GRAPH_VISIBLE
|
|
},
|
|
)
|
|
|
|
def __getattr__(self, name: str) -> Any:
|
|
plain = self._plain
|
|
if name in plain:
|
|
return plain[name]
|
|
try:
|
|
return self._vars[name].get()
|
|
except KeyError:
|
|
raise AttributeError(
|
|
f"ForwardFlags has no flag '{name}' (flags are declared in "
|
|
"ForwardFlags._DEFAULTS; check for typos)"
|
|
) from None
|
|
|
|
def __setattr__(self, name: str, value: Any) -> None:
|
|
raise AttributeError(
|
|
"ForwardFlags is written through scoped(**kw) (or the legacy "
|
|
"set() shim), never by attribute assignment"
|
|
)
|
|
|
|
def set(self, name: str, value: Any) -> None:
|
|
"""Unscoped write for legacy setter shims; persists until the next
|
|
write (current context only, for contextvar-backed flags)."""
|
|
if name in self._plain:
|
|
self._plain[name] = value
|
|
else:
|
|
self._vars[name].set(value)
|
|
|
|
@contextmanager
|
|
def scoped(self, **kwargs):
|
|
"""Set flags for the current scope, restoring on exit. Transactional
|
|
(keys validated before any write) and exception-safe."""
|
|
unknown = set(kwargs) - set(self._DEFAULTS)
|
|
if unknown:
|
|
raise ValueError(f"unknown forward flag(s): {sorted(unknown)}")
|
|
plain_saved = [
|
|
(name, self._plain[name]) for name in kwargs if name in self._plain
|
|
]
|
|
tokens = []
|
|
for name, value in kwargs.items():
|
|
if name in self._plain:
|
|
self._plain[name] = value
|
|
else:
|
|
tokens.append((self._vars[name], self._vars[name].set(value)))
|
|
try:
|
|
yield self
|
|
finally:
|
|
for var, token in reversed(tokens):
|
|
var.reset(token)
|
|
for name, value in reversed(plain_saved):
|
|
self._plain[name] = value
|
|
|
|
|
|
class _ConfigBag:
|
|
"""A resolved-config namespace bag.
|
|
|
|
Values are snapshotted from ``server_args`` at ``publish`` and this bag is
|
|
the **single source of truth** for its fields thereafter. Read is plain
|
|
attribute access; the bag is read-only by bare assignment. The sanctioned
|
|
writers are ``get_context().override(source, ...)`` (permanent) and
|
|
the scoped ``.override(**kw)`` context manager (tests). Sub-namespaces
|
|
(e.g. ``exec.moe``) are nested ``_ConfigBag`` instances reached by attribute.
|
|
|
|
Leaves and sub-bags are stored as **real instance attributes** (in
|
|
``__dict__``), so ``bag.leaf`` / ``bag.sub`` is a plain attribute load that
|
|
``torch.compile`` / dynamo can trace — config reads inside a compiled model
|
|
forward (e.g. ``get_exec().comm.enable_symm_mem`` in the embedding layer)
|
|
must not graph-break. ``_fields`` / ``_subs`` keep the authoritative
|
|
name→value maps used for override routing, membership, and scoped restore;
|
|
``__getattr__`` is only a fallback for genuinely absent names. (Deliberately
|
|
no ``__slots__``: leaves are dynamic, and the ``__dict__`` is what makes the
|
|
reads traceable.)
|
|
"""
|
|
|
|
def __init__(self, path: str):
|
|
object.__setattr__(self, "_path", path)
|
|
object.__setattr__(self, "_fields", {}) # {leaf: value}
|
|
object.__setattr__(self, "_subs", {}) # {subname: _ConfigBag}
|
|
|
|
def __getattr__(self, name: str) -> Any:
|
|
# Fallback only: real leaves/sub-bags resolve via __dict__ before this
|
|
# runs. Uses object.__getattribute__ (not self._fields) to stay safe if
|
|
# invoked before __init__ populates the bookkeeping dicts.
|
|
fields = object.__getattribute__(self, "_fields")
|
|
if name in fields:
|
|
return fields[name]
|
|
subs = object.__getattribute__(self, "_subs")
|
|
if name in subs:
|
|
return subs[name]
|
|
path = object.__getattribute__(self, "_path")
|
|
raise AttributeError(f"config namespace {path!r} has no leaf/subgroup {name!r}")
|
|
|
|
def __setattr__(self, name: str, value: Any) -> None:
|
|
raise AttributeError(
|
|
f"config namespace {self._path!r} is read-only; write via "
|
|
"get_context().override(source, ...) or the scoped .override(**kw)"
|
|
)
|
|
|
|
def _set(self, name: str, value: Any) -> None:
|
|
"""Internal write (publish + override) that bypasses the read-only guard.
|
|
Updates both the bookkeeping map and the real attribute (traceable read)."""
|
|
object.__getattribute__(self, "_fields")[name] = value
|
|
object.__setattr__(self, name, value)
|
|
|
|
def _set_sub(self, name: str, sub: _ConfigBag) -> None:
|
|
"""Register a nested bag as both a bookkeeping entry and a real
|
|
attribute (so ``bag.sub`` is a plain, traceable attribute load)."""
|
|
object.__getattribute__(self, "_subs")[name] = sub
|
|
object.__setattr__(self, name, sub)
|
|
|
|
def __contains__(self, name: str) -> bool:
|
|
return name in object.__getattribute__(self, "_fields")
|
|
|
|
@contextmanager
|
|
def override(self, **kwargs):
|
|
"""Scoped, transactional override of this bag's own leaves (keys
|
|
validated before any write; restored on exit).
|
|
|
|
For a window where one runner's value differs from the process's — a
|
|
draft model loading under ``--speculative-draft-load-format`` while the
|
|
target keeps ``--load-format`` — and for tests forcing a code path.
|
|
A permanent change goes through ``get_context().override``."""
|
|
fields = object.__getattribute__(self, "_fields")
|
|
unknown = set(kwargs) - set(fields)
|
|
if unknown:
|
|
path = object.__getattribute__(self, "_path")
|
|
raise ValueError(f"unknown config leaf for {path!r}: {sorted(unknown)}")
|
|
saved = {name: fields[name] for name in kwargs}
|
|
for name, value in kwargs.items():
|
|
self._set(name, value)
|
|
try:
|
|
yield self
|
|
finally:
|
|
for name, value in saved.items():
|
|
self._set(name, value)
|
|
|
|
|
|
def _build_config_bags(server_args: Any) -> dict:
|
|
"""Snapshot the resolution result into the namespace bag tree, driven by
|
|
the ``NS(...)`` metadata on the dataclass fields. Each leaf comes from
|
|
``resolution_result`` -- the declaration if resolution made one, else what
|
|
the caller supplied. Returns
|
|
``{top_level_name: _ConfigBag}``, arbitrarily nested (``exec.moe.eplb.…``).
|
|
Only dataclass fields carry ``NS`` markers, so derived properties/methods are
|
|
naturally excluded (they stay on the bag). A name used as both a leaf and a
|
|
subgroup at the same level is a hard error — no silent shadowing."""
|
|
from sglang.srt.arg_groups.arg_utils import namespace_of
|
|
from sglang.srt.arg_groups.overrides import resolution_result
|
|
|
|
_MISSING = object()
|
|
tops: dict = {}
|
|
for field, path in namespace_of(type(server_args)).items():
|
|
value = resolution_result(server_args, field, _MISSING)
|
|
if value is _MISSING:
|
|
# Every NS-declared field is a dataclass field, so a resolved config
|
|
# always carries it; a miss means a malformed/partial config object
|
|
# was published. Fail loud here rather than silently omitting the
|
|
# leaf (which surfaces later as a confusing "not a published leaf").
|
|
raise AttributeError(
|
|
f"config field {field!r} is declared NS({path!r}) but absent from "
|
|
f"the published {type(server_args).__name__}; cannot project its bag leaf"
|
|
)
|
|
parts = path.split(".")
|
|
bag = tops.get(parts[0])
|
|
if bag is None:
|
|
bag = tops[parts[0]] = _ConfigBag(parts[0])
|
|
for depth in range(1, len(parts)):
|
|
name = parts[depth]
|
|
if name in object.__getattribute__(bag, "_fields"):
|
|
raise ValueError(
|
|
f"config namespace collision: {'.'.join(parts[: depth + 1])!r} "
|
|
"is declared as both a leaf and a subgroup"
|
|
)
|
|
subs = object.__getattribute__(bag, "_subs")
|
|
child = subs.get(name)
|
|
if child is None:
|
|
child = _ConfigBag(".".join(parts[: depth + 1]))
|
|
bag._set_sub(name, child)
|
|
bag = child
|
|
if field in object.__getattribute__(bag, "_subs"):
|
|
raise ValueError(
|
|
f"config namespace collision: leaf {field!r} under {path!r} "
|
|
"clashes with a subgroup of the same name"
|
|
)
|
|
bag._set(field, value)
|
|
return tops
|
|
|
|
|
|
class RuntimeContext:
|
|
"""Container for the structured runtime accessors; exposes ``parallel``,
|
|
``server_args``, the resolved config namespace bags, ``flags``,
|
|
``resources``, and ``forward``."""
|
|
|
|
__slots__ = (
|
|
"parallel",
|
|
"_server_args",
|
|
"_config_bags",
|
|
"_overrides_log",
|
|
"_publish_role",
|
|
"flags",
|
|
"resources",
|
|
"forward",
|
|
)
|
|
|
|
def __init__(self, parallel: ParallelContext):
|
|
self.parallel = parallel
|
|
self._server_args: ServerArgs | None = None
|
|
self._config_bags: dict | None = None
|
|
self._overrides_log: list = []
|
|
self._publish_role: str | None = None
|
|
self.flags = Flags()
|
|
self.resources = Resources()
|
|
self.forward = ForwardFlags()
|
|
|
|
def get_stream(self, name: str) -> Any:
|
|
"""Named process-level side stream: get-or-create, shared by
|
|
name (the keyed-lazy pattern of the persistent buffers). Creation is
|
|
a driver call that must stay outside cuda-graph capture — call sites
|
|
lease their stream at init/warmup time."""
|
|
stream = self.resources.streams.get(name)
|
|
if stream is None:
|
|
import torch
|
|
|
|
from sglang.srt.arg_groups.overrides import resolution_result
|
|
|
|
device = (
|
|
resolution_result(self._server_args, "device")
|
|
if self._server_args
|
|
else "cuda"
|
|
)
|
|
stream = torch.get_device_module(device).Stream()
|
|
self.resources.streams[name] = stream
|
|
return stream
|
|
|
|
def set_stream(self, name: str, stream: Any) -> Any:
|
|
"""Install (or replace) the named stream — explicit injection for
|
|
tests and backends that bring their own stream."""
|
|
self.resources.streams[name] = stream
|
|
return stream
|
|
|
|
def get_buffer(self, name: str, factory: Any) -> Any:
|
|
"""Named process-level persistent buffer: get-or-create via
|
|
``factory()``, shared by name (the keyed-lazy pattern of the
|
|
persistent buffers / named streams)."""
|
|
buf = self.resources.buffers.get(name)
|
|
if buf is None:
|
|
buf = factory()
|
|
self.resources.buffers[name] = buf
|
|
return buf
|
|
|
|
@property
|
|
def server_args(self) -> ServerArgs:
|
|
"""The process-wide ``ServerArgs`` (context-owned slot)."""
|
|
server_args = self._server_args
|
|
if server_args is None:
|
|
# Verbatim legacy message: tests and user scripts may match on it.
|
|
raise ValueError("Global server args is not set yet!")
|
|
return server_args
|
|
|
|
def set_server_args(self, server_args: ServerArgs) -> None:
|
|
"""Publish the process-wide ``ServerArgs`` into the context-owned slot.
|
|
|
|
Overwrite-allowed: a re-publish replaces the slot (test kits re-publish
|
|
per test; production ordering discipline lives at the call-sites, e.g.
|
|
the draft-worker guard in ``ModelRunner.__init__``). The published
|
|
object is the raw input; the resolution it carries is its declaration
|
|
stash, which is what the bags are projected from.
|
|
"""
|
|
# Seed the capture tier for the new lifecycle (defaults for sentinel
|
|
# and mock publishes, which carry no config).
|
|
self.flags.capture.enable_torch_compile = getattr(
|
|
server_args, "enable_torch_compile", False
|
|
)
|
|
self._server_args = server_args
|
|
# The adaptive draft-token bound memoizes on the config *path*, so a new
|
|
# publication that reuses the path must not inherit the bound computed
|
|
# from the file's previous contents.
|
|
_adaptive_draft_token_bound.cache_clear()
|
|
# Snapshot resolved config into the namespace bags (the single source of
|
|
# truth for config reads). Driven by NS(...) metadata; a mock/partial
|
|
# config with no NS markers yields an empty tree (no bags projected).
|
|
self._config_bags = _build_config_bags(server_args)
|
|
# Wire the published `parallel` bag onto the live wrapper: it is the slot
|
|
# the `config` property reads, which is how config-only leaves like
|
|
# pp_max_micro_batch_size are spelled.
|
|
self.parallel._config = self._config_bags.get("parallel")
|
|
# A direct install is roleless; ``publish`` assigns the role afterwards.
|
|
self._overrides_log = []
|
|
self._publish_role = None
|
|
|
|
def config_bag(self, name: str) -> _ConfigBag:
|
|
"""Return the top-level config namespace bag (``device`` / ``model`` /
|
|
``exec`` / ``schedule`` / ``memory`` / ``spec`` / ``lora`` / ``mm`` /
|
|
``disagg`` / ``serving`` / ``observability``). Fails closed until
|
|
``publish`` / ``set_server_args`` has projected it."""
|
|
bags = self._config_bags
|
|
if not bags or name not in bags:
|
|
raise ValueError(f"config namespace {name!r} not published")
|
|
if _ROLE_NS_MODE != "off":
|
|
self._check_role_namespace(name)
|
|
return bags[name]
|
|
|
|
def is_config_namespace_published(self, name: str) -> bool:
|
|
"""Return whether a config namespace exists in the current context."""
|
|
bags = self._config_bags
|
|
return bags is not None and name in bags
|
|
|
|
def _check_role_namespace(self, name: str) -> None:
|
|
# Out of line so the mode gate above stays one dead-branch-prunable
|
|
# check under dynamo in the default "off" mode (config_bag runs inside
|
|
# compiled model forwards).
|
|
role = self._publish_role
|
|
if _ROLE_NS_MODE == "record":
|
|
if not _is_compiling():
|
|
_record_namespace_read(role, name)
|
|
elif _ROLE_NS_MODE == "enforce" and role is not None:
|
|
if role not in ROLE_NAMESPACE_SETS:
|
|
raise ValueError(
|
|
f"publish role {role!r} has no ROLE_NAMESPACE_SETS entry; "
|
|
"declare its namespace set (None for the full tree)."
|
|
)
|
|
allowed = ROLE_NAMESPACE_SETS[role]
|
|
if allowed is not None and name not in allowed:
|
|
raise ValueError(
|
|
f"config namespace {name!r} is outside the declared set "
|
|
f"for publish role {role!r} ({sorted(allowed)}). If this "
|
|
"read is legitimate for the process type, extend "
|
|
"ROLE_NAMESPACE_SETS; if not, the read belongs in a "
|
|
"different process or behind a per-instance boundary."
|
|
)
|
|
|
|
def override(self, source: str, **fields) -> None:
|
|
"""The business mutation entry: write resolved config
|
|
leaves onto the namespace bags — the single source of truth. It does
|
|
**not** touch ``server_args`` (the pristine startup record) and there is
|
|
no write-through, so the old "wrote one store, read another" desync class
|
|
cannot occur.
|
|
|
|
Each flat field name is routed to its bag by the ``NS`` metadata (flat
|
|
names are unique across namespaces). Validation is all-or-nothing: an
|
|
unknown / unprojected field aborts before any write. ``source`` is
|
|
recorded for provenance / reproduction.
|
|
"""
|
|
if not fields:
|
|
return
|
|
bags = self._config_bags
|
|
if bags is None:
|
|
raise ValueError("config not published; cannot override")
|
|
from sglang.srt.arg_groups.arg_utils import namespace_of
|
|
|
|
nsmap = namespace_of(type(self._server_args))
|
|
targets = [] # (bag, leaf, value) — resolved before any write
|
|
for name, value in fields.items():
|
|
path = nsmap.get(name)
|
|
if path is None:
|
|
raise ValueError(
|
|
f"override: unknown config field {name!r} (no NS namespace) — "
|
|
"not a resolved config leaf"
|
|
)
|
|
parts = path.split(".")
|
|
bag = bags.get(parts[0])
|
|
if bag is None:
|
|
raise ValueError(f"override: namespace {parts[0]!r} not published")
|
|
for seg in parts[1:]:
|
|
bag = object.__getattribute__(bag, "_subs").get(seg)
|
|
if bag is None:
|
|
raise ValueError(
|
|
f"override: subgroup {seg!r} missing under {path!r}"
|
|
)
|
|
if name not in bag:
|
|
raise ValueError(f"override: field {name!r} not projected on {path!r}")
|
|
targets.append((bag, name, value))
|
|
for bag, name, value in targets:
|
|
bag._set(name, value)
|
|
self._overrides_log.append((source, dict(fields)))
|
|
|
|
def config_leaf(self, name: str):
|
|
"""One resolved config leaf by field name — the read side of ``override``.
|
|
|
|
Callers that hold a field name rather than a namespace (a readback
|
|
endpoint, a control-plane handler) would otherwise have to know which
|
|
bag it lives in.
|
|
"""
|
|
bags = self._config_bags
|
|
if bags is None:
|
|
raise ValueError("config not published; cannot read a config leaf")
|
|
from sglang.srt.arg_groups.arg_utils import namespace_of
|
|
|
|
path = namespace_of(type(self._server_args)).get(name)
|
|
if path is None:
|
|
raise ValueError(f"{name!r} is not a config leaf (no NS namespace)")
|
|
parts = path.split(".")
|
|
bag = self.config_bag(parts[0])
|
|
for seg in parts[1:]:
|
|
bag = object.__getattribute__(bag, "_subs").get(seg)
|
|
if bag is None:
|
|
raise ValueError(f"subgroup {seg!r} missing under {path!r}")
|
|
return getattr(bag, name)
|
|
|
|
def overrides_log(self) -> list:
|
|
"""Provenance of post-publish ``override`` calls: ``[(source, {field: value})]``.
|
|
|
|
Returns deep-ish copies (source, dict(fields)) so callers inspecting the
|
|
log cannot mutate the recorded provenance in place."""
|
|
return [(source, dict(fields)) for source, fields in self._overrides_log]
|
|
|
|
def resolved_server_args_dict(self, base: dict | None = None) -> dict:
|
|
"""Serialize the *resolved* config: the pristine ``server_args`` fields
|
|
with every post-publish ``override`` overlaid.
|
|
|
|
``get_internal_state`` reports this, and ``/server_info`` carries it in
|
|
the ``internal_states`` block, so scheduler-side runtime changes show up
|
|
in a readback: HiCache attach/detach, the generated forward-pass-metrics
|
|
endpoint, tunables set via ``/set_internal_state``.
|
|
|
|
``base`` defaults to ``server_args.resolved_dict()`` -- the record's
|
|
fields as resolution decided them, nested dataclasses expanded. (It used
|
|
to be ``dict(vars(server_args))``, which carried the private resolution
|
|
bookkeeping and the ``model_config`` memo into the readback.) Override
|
|
leaves are flat ``ServerArgs`` field names, so overlaying them onto the
|
|
top level of the base is exact.
|
|
|
|
The log is per process: it carries what *this* process overrode. A
|
|
weight reload records ``model_path`` and ``load_format`` from the
|
|
scheduler process (``ModelRunner.update_model_fields``); the tokenizer
|
|
process records only ``load_format`` and keeps ``model_path`` /
|
|
``served_model_name`` as ``TokenizerManager`` attributes, which
|
|
``TokenizerManager.resolved_config_dict`` overlays on top of this dump.
|
|
The top-level ``/server_info`` fields are the startup record, not this
|
|
dump.
|
|
"""
|
|
d = self.server_args.resolved_dict() if base is None else dict(base)
|
|
for _source, fields in self._overrides_log:
|
|
d.update(fields)
|
|
return d
|
|
|
|
def override_server_args(self, **fields) -> _ServerArgsOverride:
|
|
"""Test-only scoped override for the config tier — the sibling of
|
|
``get_parallel().override()`` and the flag groups' ``override()``:
|
|
tests force execution paths by overriding the context instead of
|
|
hand-building config objects.
|
|
|
|
``install()`` (or entering it as a context manager) publishes a fresh
|
|
dummy-boundary ``ServerArgs`` carrying ``fields`` and returns it;
|
|
``restore()`` (or exiting) reinstates whatever the slot held before.
|
|
|
|
This is the sanctioned way for a test to get a published context, and
|
|
it stays. The transitional reason it was introduced for — production
|
|
code branching on raw ``server_args`` fields at runtime — is gone (the
|
|
read ratchet pins business reads at zero), but a test that exercises
|
|
bag readers still needs bags, and the bag tree is projected *from an
|
|
instance*: something has to publish one. Prefer the finer-grained
|
|
scoped overrides (``get_exec().override(...)``, the flag groups'
|
|
``override``) on top of a published context when a test only needs to
|
|
force one leaf.
|
|
"""
|
|
return _ServerArgsOverride(self, fields)
|
|
|
|
|
|
class _ServerArgsOverride:
|
|
"""Scoped config override (see ``RuntimeContext.override_server_args``).
|
|
|
|
Deliberately a plain class rather than a generator context manager:
|
|
fixtures that live for a whole test case install the override without a
|
|
``with`` block, and a suspended generator would run its restore whenever
|
|
the garbage collector closes it — un-publishing the active config at a
|
|
nondeterministic point.
|
|
"""
|
|
|
|
__slots__ = (
|
|
"_context",
|
|
"_fields",
|
|
"_prev_server_args",
|
|
"_prev_bags",
|
|
"_prev_overrides_log",
|
|
"_prev_publish_role",
|
|
"_prev_parallel_config",
|
|
"_prev_capture",
|
|
"_installed",
|
|
)
|
|
|
|
def __init__(self, context: RuntimeContext, fields: dict):
|
|
self._context = context
|
|
self._fields = fields
|
|
self._installed = False
|
|
|
|
def install(self) -> ServerArgs:
|
|
"""Publish a fresh dummy-boundary ``ServerArgs`` carrying the
|
|
overrides; returns the published instance."""
|
|
from sglang.srt.server_args import ServerArgs
|
|
|
|
assert not self._installed, "override_server_args already installed"
|
|
ctx = self._context
|
|
self._prev_server_args = ctx._server_args
|
|
self._prev_bags = ctx._config_bags
|
|
self._prev_overrides_log = ctx._overrides_log
|
|
self._prev_publish_role = ctx._publish_role
|
|
self._prev_parallel_config = ctx.parallel._config
|
|
self._prev_capture = ctx.flags.capture.enable_torch_compile
|
|
from sglang.srt.arg_groups.overrides import (
|
|
_apply_fields,
|
|
declare_late_resolution,
|
|
)
|
|
|
|
server_args = ServerArgs(model_path="dummy")
|
|
server_args.resolve_once()
|
|
# Underscore names seed private property caches (the strict guard
|
|
# exempts them); everything else must be a real config field.
|
|
unknown = {name for name in self._fields if not name.startswith("_")} - set(
|
|
type(server_args).__dataclass_fields__
|
|
)
|
|
if unknown:
|
|
raise ValueError(
|
|
f"override_server_args: unknown ServerArgs field(s): {sorted(unknown)}"
|
|
)
|
|
# Declared so the projection sees it; late, because the record is
|
|
# resolved already and not yet published.
|
|
# Underscore names are not fields at all (they seed private property
|
|
# caches), so they stay a direct write.
|
|
declared = {
|
|
name: value for name, value in self._fields.items() if name[0] != "_"
|
|
}
|
|
if declared:
|
|
declare_late_resolution(server_args, "override_server_args", **declared)
|
|
# This hook stands in for a launch: the caller's values are both what
|
|
# the operator passed and what resolution decided, so they go on the
|
|
# record as well as into the stash. Production late resolution declares
|
|
# only -- there the record stays the operator's input.
|
|
_apply_fields(server_args, self._fields)
|
|
ctx.set_server_args(server_args)
|
|
self._installed = True
|
|
return server_args
|
|
|
|
def restore(self) -> None:
|
|
"""Reinstate the exact pre-install lifecycle state (or the empty slot)."""
|
|
if not self._installed:
|
|
return
|
|
self._installed = False
|
|
ctx = self._context
|
|
ctx._server_args = self._prev_server_args
|
|
ctx._config_bags = self._prev_bags
|
|
ctx._overrides_log = self._prev_overrides_log
|
|
ctx._publish_role = self._prev_publish_role
|
|
ctx.parallel._config = self._prev_parallel_config
|
|
ctx.flags.capture.enable_torch_compile = self._prev_capture
|
|
self._prev_server_args = None
|
|
self._prev_bags = None
|
|
self._prev_overrides_log = None
|
|
self._prev_parallel_config = None
|
|
|
|
def __enter__(self) -> ServerArgs:
|
|
return self.install()
|
|
|
|
def __exit__(self, *exc) -> None:
|
|
self.restore()
|
|
|
|
|
|
_PARALLEL = ParallelContext()
|
|
_CONTEXT = RuntimeContext(parallel=_PARALLEL)
|
|
|
|
|
|
def get_context() -> RuntimeContext:
|
|
return _CONTEXT
|
|
|
|
|
|
def get_parallel() -> ParallelContext:
|
|
return _PARALLEL
|
|
|
|
|
|
def get_server_args() -> ServerArgs:
|
|
return _CONTEXT.server_args
|
|
|
|
|
|
def get_flags() -> Flags:
|
|
return _CONTEXT.flags
|
|
|
|
|
|
def get_resources() -> Resources:
|
|
return _CONTEXT.resources
|
|
|
|
|
|
def get_forward() -> ForwardFlags:
|
|
return _CONTEXT.forward
|
|
|
|
|
|
# --- Resolved config namespaces -------------------------
|
|
# Each returns the top-level snapshot bag; reads are `get_exec().moe.field` etc.
|
|
# All fail with ValueError("... not published") until publish has projected them.
|
|
# ``parallel`` has no bag getter: ``get_parallel()`` answers its leaves
|
|
# directly, alongside the live topology they belong to.
|
|
def get_device() -> _ConfigBag:
|
|
return _CONTEXT.config_bag("device")
|
|
|
|
|
|
def get_model() -> _ConfigBag:
|
|
return _CONTEXT.config_bag("model")
|
|
|
|
|
|
def get_exec() -> _ConfigBag:
|
|
return _CONTEXT.config_bag("exec")
|
|
|
|
|
|
def get_schedule() -> _ConfigBag:
|
|
return _CONTEXT.config_bag("schedule")
|
|
|
|
|
|
def get_memory() -> _ConfigBag:
|
|
return _CONTEXT.config_bag("memory")
|
|
|
|
|
|
def get_spec() -> _ConfigBag:
|
|
return _CONTEXT.config_bag("spec")
|
|
|
|
|
|
def get_lora() -> _ConfigBag:
|
|
return _CONTEXT.config_bag("lora")
|
|
|
|
|
|
def get_mm() -> _ConfigBag:
|
|
return _CONTEXT.config_bag("mm")
|
|
|
|
|
|
def get_disagg() -> _ConfigBag:
|
|
return _CONTEXT.config_bag("disagg")
|
|
|
|
|
|
def get_serving() -> _ConfigBag:
|
|
return _CONTEXT.config_bag("serving")
|
|
|
|
|
|
def get_observability() -> _ConfigBag:
|
|
return _CONTEXT.config_bag("observability")
|
|
|
|
|
|
# --- Per-role namespace sets (2c) -------------------------------------------
|
|
#
|
|
# ``publish(role=...)`` records which process type installed the config; this
|
|
# table declares which top-level config namespaces each role reads. ``None``
|
|
# means the full tree — either the role genuinely needs everything (scheduler)
|
|
# or its deployment shape has not been audited yet (restrict only what smoke
|
|
# coverage can verify). ``parallel`` is served by ``get_parallel()`` and
|
|
# every process legitimately reads topology config, so it is not in this table.
|
|
#
|
|
# ``SGLANG_ROLE_NAMESPACES`` selects the mode (read once at import):
|
|
# off (default) no bookkeeping, zero overhead;
|
|
# record audit mode — collect (role, namespace) reads per process and dump
|
|
# them at exit (the data that seeds this table). Reads made inside
|
|
# torch.compile-traced code are NOT observed (recording is pruned
|
|
# under tracing to keep capture legal) — run audits with
|
|
# compilation disabled before restricting a role.
|
|
# enforce fail closed — a bag read outside the role's declared set raises.
|
|
ROLE_NAMESPACE_SETS: dict[str, frozenset[str] | None] = {
|
|
# Reads (almost) everything by design — the model-executing process.
|
|
"scheduler": None,
|
|
"test": None,
|
|
# The DP controller's static read set, checked against the module: the
|
|
# elastic-EP gate, the load-balance method, the watchdog timeout, and the
|
|
# disaggregation mode.
|
|
"dp_controller": frozenset({"exec", "parallel", "device", "disagg"}),
|
|
# Record-mode audit (2026-08-06, text model, /generate + /get_server_info +
|
|
# /v1/models): reads exactly {"serving"} — the per-instance managers read
|
|
# self.server_args by design. Still declared full, because that run did not
|
|
# exercise the multimodal processors, LoRA/score endpoints, the disagg
|
|
# roles, or the gRPC bridge; narrowing needs those shapes audited too, and
|
|
# a wrong set fails a request rather than a test.
|
|
"tokenizer": None,
|
|
# Deployment shapes not exercised locally; audit before restricting.
|
|
"detokenizer": None,
|
|
"encoder": None,
|
|
"expert_backup": None,
|
|
"weight_cache_daemon": None,
|
|
# The diffusion GPU worker runs a model and publishes a placeholder so
|
|
# shared SRT reads do not fail closed; declared full for that reason.
|
|
"diffusion_gpu_worker": None,
|
|
}
|
|
|
|
|
|
def _validated_role_ns_mode(value: str) -> str:
|
|
mode = value.strip().lower()
|
|
if mode not in ("off", "record", "enforce"):
|
|
raise ValueError(
|
|
f"SGLANG_ROLE_NAMESPACES={value!r} is not one of off / record / "
|
|
"enforce — refusing to guess (a typo here would silently disable "
|
|
"enforcement)."
|
|
)
|
|
return mode
|
|
|
|
|
|
def _role_ns_mode_from_env() -> str:
|
|
# Resolved once at import so the config_bag gate stays a dynamo-prunable
|
|
# constant; validated fail-loud here (EnvField's warn-and-default parse
|
|
# would silently turn a typo into "off").
|
|
from sglang.srt.environ import envs
|
|
|
|
return _validated_role_ns_mode(envs.SGLANG_ROLE_NAMESPACES.get())
|
|
|
|
|
|
_ROLE_NS_MODE = _role_ns_mode_from_env()
|
|
_RECORDED_NS_READS: set[tuple[str | None, str]] = set()
|
|
_RECORD_DUMP_REGISTERED = False
|
|
|
|
|
|
def _is_compiling() -> bool:
|
|
# Recording has Python side effects (set mutation, file I/O, atexit) that
|
|
# must never run under tracing; torch.compiler.is_compiling() is dynamo's
|
|
# sanctioned probe. The function-level import keeps this module
|
|
# import-light; a sys.modules lookup here breaks fullgraph tracing (dynamo
|
|
# enumerates the dict, which other imports mutate mid-trace).
|
|
import torch
|
|
|
|
return torch.compiler.is_compiling()
|
|
|
|
|
|
def _ensure_record_dump_registered() -> None:
|
|
global _RECORD_DUMP_REGISTERED
|
|
if not _RECORD_DUMP_REGISTERED:
|
|
_RECORD_DUMP_REGISTERED = True
|
|
import atexit
|
|
|
|
atexit.register(_dump_recorded_namespace_reads)
|
|
|
|
|
|
def _append_role_ns_out(role: str | None, name: str) -> None:
|
|
# Persist immediately: worker processes are routinely torn down with
|
|
# signals that skip atexit, and the audit must survive that.
|
|
from sglang.srt.environ import envs
|
|
|
|
out = envs.SGLANG_ROLE_NAMESPACES_OUT.get()
|
|
if not out:
|
|
return
|
|
try:
|
|
with open(out, "a") as f:
|
|
f.write(f"{role} {name}\n")
|
|
except OSError as e:
|
|
# The entry stays in the in-memory set; the exit summary still covers it.
|
|
print(
|
|
f"[role-namespaces] pid={os.getpid()} failed to append "
|
|
f"({role}, {name}) to {out!r}: {e}",
|
|
file=sys.stderr,
|
|
flush=True,
|
|
)
|
|
|
|
|
|
def _record_namespace_read(role: str | None, name: str) -> None:
|
|
if (role, name) in _RECORDED_NS_READS:
|
|
return
|
|
_RECORDED_NS_READS.add((role, name))
|
|
_append_role_ns_out(role, name)
|
|
_ensure_record_dump_registered()
|
|
|
|
|
|
def _dump_recorded_namespace_reads() -> None:
|
|
"""Emit the record-mode audit: one line per role with the namespaces its
|
|
process actually read (multi-process runs dump once per process). The
|
|
process's own publish role is always included, so a zero-read role emits
|
|
an (empty) line rather than being indistinguishable from a process where
|
|
recording never ran."""
|
|
by_role: dict = {}
|
|
own_role = _CONTEXT._publish_role
|
|
if own_role is not None:
|
|
by_role.setdefault(own_role, set())
|
|
for role, name in _RECORDED_NS_READS:
|
|
if name == "-": # publish-time marker, not a namespace read
|
|
by_role.setdefault(role, set())
|
|
continue
|
|
by_role.setdefault(role, set()).add(name)
|
|
for role in sorted(by_role, key=str):
|
|
print(
|
|
f"[role-namespaces] pid={os.getpid()} role={role} "
|
|
f"read={','.join(sorted(by_role[role]))}",
|
|
file=sys.stderr,
|
|
flush=True,
|
|
)
|
|
|
|
|
|
def publish(server_args, *, role: str, hf_config: Any = None) -> RuntimeContext:
|
|
"""Install process-wide config for this OS process.
|
|
|
|
Records the process ``role`` — one of the ``ROLE_NAMESPACE_SETS`` keys,
|
|
which is the one place the roles are enumerated — and
|
|
projects the config bags. Draft workers skip publish (they must not clobber
|
|
the target). ``role`` is provenance, and — when ``SGLANG_ROLE_NAMESPACES``
|
|
is ``enforce`` — the key into ``ROLE_NAMESPACE_SETS`` for fail-closed
|
|
namespace-read enforcement (``record`` audits the reads instead).
|
|
``hf_config`` is accepted for forward-compat and currently unused.
|
|
|
|
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**
|
|
(bags re-projected, provenance reset, role overwritten), which is what
|
|
lets one process rebuild an engine after shutting the previous one down.
|
|
"""
|
|
if _ROLE_NS_MODE == "enforce" and role not in ROLE_NAMESPACE_SETS:
|
|
# Fail closed at publish time, not at the first stray read.
|
|
raise ValueError(
|
|
f"publish role {role!r} has no ROLE_NAMESPACE_SETS entry; declare "
|
|
"its namespace set (None for the full tree)."
|
|
)
|
|
server_args.resolve_once()
|
|
discarded = _CONTEXT.overrides_log()
|
|
_CONTEXT.set_server_args(server_args)
|
|
if discarded:
|
|
logger.warning(
|
|
"publish(role=%s) re-projected the config bags and dropped %d "
|
|
"override(s) taken since the last publish: %s",
|
|
role,
|
|
len(discarded),
|
|
", ".join(
|
|
f"{source}({', '.join(sorted(fields))})" for source, fields in discarded
|
|
),
|
|
)
|
|
_CONTEXT._publish_role = role
|
|
if _ROLE_NS_MODE == "record":
|
|
# The '-' marker distinguishes a zero-read role from a process where
|
|
# recording never ran (signal teardown skips atexit).
|
|
_record_namespace_read(role, "-")
|
|
print(
|
|
f"[role-namespaces] pid={os.getpid()} role={role} recording; note: "
|
|
"reads inside torch.compile-traced code are not observed — audit "
|
|
"with compilation disabled before restricting a role.",
|
|
file=sys.stderr,
|
|
flush=True,
|
|
)
|
|
return _CONTEXT
|
|
|
|
|
|
def assert_published(server_args, *, role: str) -> RuntimeContext:
|
|
"""This record, under this role, is already published -- or fail loud.
|
|
|
|
Publishing is the process entry's job: `run_scheduler_process`,
|
|
`init_multi_tokenizer`, a spawned encoder worker, the benchmark work
|
|
functions. A constructor arriving here unpublished means one of those
|
|
entries is missing.
|
|
|
|
A `publish` at this point re-projects the bags over a live process,
|
|
discarding every `override()` taken since and the provenance log with it,
|
|
so this raises.
|
|
"""
|
|
if _CONTEXT._server_args is server_args and _CONTEXT._publish_role == role:
|
|
return _CONTEXT
|
|
if _CONTEXT._server_args is None:
|
|
detail = "nothing is published in this process"
|
|
elif _CONTEXT._server_args is not server_args:
|
|
detail = (
|
|
"a different record is published "
|
|
f"(role={_CONTEXT._publish_role!r}); this constructor was handed "
|
|
"one the process never published"
|
|
)
|
|
else:
|
|
detail = (
|
|
f"this record is published under role "
|
|
f"{_CONTEXT._publish_role!r}, not {role!r}"
|
|
)
|
|
raise RuntimeError(
|
|
f"config not published for role {role!r}: {detail}. The process entry "
|
|
"publishes -- add publish(server_args, role=...) there rather than "
|
|
"publishing from a constructor."
|
|
)
|
|
|
|
|
|
def publish_role() -> str | None:
|
|
"""The role recorded by the last ``publish`` (None for a legacy set)."""
|
|
return _CONTEXT._publish_role
|
|
|
|
|
|
def get_stream(name: str) -> Any:
|
|
return _CONTEXT.get_stream(name)
|
|
|
|
|
|
def set_stream(name: str, stream: Any) -> Any:
|
|
return _CONTEXT.set_stream(name, stream)
|
|
|
|
|
|
def get_buffer(name: str, factory: Any) -> Any:
|
|
return _CONTEXT.get_buffer(name, factory)
|
|
|
|
|
|
_GLOBAL_DWDP_MANAGER: Any = None
|
|
|
|
|
|
def get_global_dwdp_manager() -> Any:
|
|
return _GLOBAL_DWDP_MANAGER
|
|
|
|
|
|
def set_global_dwdp_manager(manager: Any) -> None:
|
|
global _GLOBAL_DWDP_MANAGER
|
|
_GLOBAL_DWDP_MANAGER = manager
|
|
|
|
|
|
def _group_leaves(group: _FlagGroupBase) -> dict[str, Any]:
|
|
"""The leaf values of a flag group, recursively."""
|
|
leaves: dict[str, Any] = {}
|
|
for name in type(group).__dataclass_fields__:
|
|
value = getattr(group, name)
|
|
if isinstance(value, _FlagGroupBase):
|
|
leaves[name] = _group_leaves(value)
|
|
elif isinstance(value, (dict, list)):
|
|
leaves[name] = type(value)(value)
|
|
else:
|
|
leaves[name] = value
|
|
return leaves
|
|
|
|
|
|
def _restore_leaves(group: _FlagGroupBase, leaves: dict[str, Any]) -> None:
|
|
for name, value in leaves.items():
|
|
current = getattr(group, name)
|
|
if isinstance(current, _FlagGroupBase):
|
|
_restore_leaves(current, value)
|
|
elif isinstance(current, dict):
|
|
current.clear()
|
|
current.update(value)
|
|
elif isinstance(current, list):
|
|
current[:] = value
|
|
else:
|
|
setattr(group, name, value)
|
|
|
|
|
|
def snapshot_context() -> dict[str, Any]:
|
|
"""Everything a publish replaces, so a failed launch can put it back.
|
|
|
|
Enumerated from ``__slots__`` rather than listed by hand: a hand-picked copy
|
|
of context state is one field behind the day a slot is added, and the copy
|
|
that silently drops one is worse than none. Flag groups are snapshotted by
|
|
leaf, not by reference: publish writes *into* the same ``Flags`` object
|
|
(``capture.enable_torch_compile``), so a reference held here would already
|
|
carry the failed launch's value by the time it is put back.
|
|
"""
|
|
state: dict[str, Any] = {}
|
|
for name in RuntimeContext.__slots__:
|
|
if name == "parallel":
|
|
continue
|
|
value = getattr(_CONTEXT, name)
|
|
if isinstance(value, _FlagGroupBase):
|
|
state[name] = (value, _group_leaves(value))
|
|
elif isinstance(value, list):
|
|
state[name] = list(value)
|
|
else:
|
|
state[name] = value
|
|
state["__parallel__"] = {
|
|
name: getattr(_CONTEXT.parallel, name)
|
|
for name in type(_CONTEXT.parallel).__slots__
|
|
}
|
|
state["__dwdp__"] = get_global_dwdp_manager()
|
|
return state
|
|
|
|
|
|
def restore_context(state: dict[str, Any]) -> None:
|
|
"""Put back what ``snapshot_context`` captured."""
|
|
for name in RuntimeContext.__slots__:
|
|
if name == "parallel":
|
|
continue
|
|
value = state[name]
|
|
if isinstance(value, tuple) and isinstance(value[0], _FlagGroupBase):
|
|
group, leaves = value
|
|
setattr(_CONTEXT, name, group)
|
|
_restore_leaves(group, leaves)
|
|
else:
|
|
setattr(_CONTEXT, name, value)
|
|
for name, value in state["__parallel__"].items():
|
|
setattr(_CONTEXT.parallel, name, value)
|
|
_adaptive_draft_token_bound.cache_clear()
|
|
set_global_dwdp_manager(state["__dwdp__"])
|
|
|
|
|
|
def reset_context() -> None:
|
|
"""Clear the context-owned store (unit-test teardown): drop the published
|
|
``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.
|
|
"""
|
|
_CONTEXT._server_args = None
|
|
_CONTEXT._config_bags = None
|
|
_adaptive_draft_token_bound.cache_clear()
|
|
_CONTEXT._overrides_log = []
|
|
_CONTEXT._publish_role = None
|
|
_CONTEXT.parallel._config = None
|
|
_CONTEXT.parallel.clear_derived_widths()
|
|
_CONTEXT.flags = Flags()
|
|
_CONTEXT.resources = Resources()
|
|
_CONTEXT.forward = ForwardFlags()
|
|
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.
|
|
|
|
Every input is a ``model`` leaf, so this derives from the bags and follows a
|
|
post-publish override; ``ServerArgs.remote_instance_weight_loader_use_transfer_engine``
|
|
is the pre-publish equivalent, and both go through the same helper.
|
|
``load_format`` is the caller's own (a draft runner loading under
|
|
``--speculative-draft-load-format`` has one the process record does not).
|
|
"""
|
|
from sglang.srt.arg_groups.overrides import remote_instance_transfer_engine_of
|
|
|
|
return remote_instance_transfer_engine_of(get_model(), load_format)
|
|
|
|
|
|
def max_prefill_buffer_tokens() -> int:
|
|
"""The prefill-buffer ceiling: ``chunked_prefill_size``, except PP dynamic
|
|
chunking can grow chunks toward ``max_prefill_tokens`` and probe at 1.25x.
|
|
|
|
Every input is a published leaf (``schedule`` plus the configured PP size),
|
|
so this derives from the bags and follows a post-publish override;
|
|
``ServerArgs.max_prefill_buffer_tokens`` is the pre-publish equivalent and
|
|
``TestDerivedPredicatesAgreeAcrossTiers`` pins the two equal.
|
|
"""
|
|
import math
|
|
|
|
schedule = get_schedule()
|
|
chunked = (
|
|
schedule.chunked_prefill_size
|
|
if schedule.chunked_prefill_size and schedule.chunked_prefill_size > 0
|
|
else 0
|
|
)
|
|
tokens = chunked
|
|
if schedule.enable_dynamic_chunking and get_parallel().pp_size > 1 and chunked:
|
|
tokens = max(
|
|
tokens, schedule.max_prefill_tokens or 0, math.ceil(chunked * 1.25)
|
|
)
|
|
return tokens
|
|
|
|
|
|
def pre_capture_activation_reserve_mb(gpu_mem: float | None) -> float:
|
|
"""The activation working-set reserve held back before cuda-graph capture.
|
|
|
|
Derived from published leaves across four bags (``disagg`` / ``schedule`` /
|
|
``exec.graph`` / ``spec``) plus the configured parallel sizes, so it follows
|
|
a post-publish override; ``ServerArgs.pre_capture_activation_reserve_mb`` is
|
|
the pre-publish equivalent and
|
|
``TestDerivedPredicatesAgreeAcrossTiers`` pins the two equal.
|
|
"""
|
|
schedule = get_schedule()
|
|
if get_disagg().disaggregation_mode == "decode":
|
|
running_requests = (
|
|
schedule.max_running_requests
|
|
or get_exec().graph.cuda_graph_config.decode.max_bs
|
|
or 1
|
|
)
|
|
activation_tokens = max(
|
|
running_requests * (get_spec().speculative_num_draft_tokens or 1), 2048
|
|
)
|
|
elif schedule.chunked_prefill_size > 0:
|
|
activation_tokens = max(schedule.chunked_prefill_size, 2048)
|
|
else:
|
|
activation_tokens = max(schedule.max_prefill_tokens, 2048)
|
|
parallel = get_parallel()
|
|
reserved_mem = (
|
|
512 + activation_tokens * 1.5 + parallel.tp_size * parallel.pp_size / 8 * 1024
|
|
)
|
|
if gpu_mem is not None and gpu_mem > 60 * 1024:
|
|
reserved_mem = max(reserved_mem, 10 * 1024)
|
|
return reserved_mem
|
|
|
|
|
|
# --- Derived config accessors ------------------------------------------------
|
|
#
|
|
# A few values are computed from several config fields plus the HF config, so
|
|
# they are ``ServerArgs`` members rather than namespace leaves. Business code
|
|
# must not reach for the startup record to get them: these accessors are the
|
|
# named home, and this module — which owns the slot — is the only place that
|
|
# reads it. Each one keeps the member's exact semantics, including which model
|
|
# config it derives from (always the process's, i.e. the target's).
|
|
|
|
|
|
def mamba_cache_chunk_size() -> int:
|
|
"""The caching point granularity for mamba state: ``max(the model's mamba
|
|
chunk size, page_size)``. Cached on the config after the first call."""
|
|
return get_server_args().mamba_cache_chunk_size
|
|
|
|
|
|
def mamba_checkpoint_grid(tree_page: int) -> int:
|
|
"""The granularity a donated mamba checkpoint's depth must land on so the
|
|
radix tree can name it. Pass the page the tree actually allocates on: DCP
|
|
widens it past ``mamba_cache_chunk_size``, and deriving that here would be a
|
|
second copy of a predicate that already lives in the cache builder."""
|
|
return math.lcm(mamba_cache_chunk_size(), tree_page)
|
|
|
|
|
|
def mamba_track_grid(tree_page: int) -> int:
|
|
"""The same granularity for a decode-donated checkpoint, which additionally
|
|
has to land on the requested ``mamba_track_interval``."""
|
|
return math.lcm(
|
|
mamba_checkpoint_grid(tree_page), get_exec().mamba.mamba_track_interval
|
|
)
|
|
|
|
|
|
def max_speculative_num_draft_tokens() -> int | None:
|
|
"""The largest draft-token count speculative decoding may use.
|
|
|
|
All three inputs are ``spec`` leaves, so this derives from the bags and
|
|
follows a post-publish override; ``ServerArgs.max_speculative_num_draft_tokens``
|
|
is the pre-publish equivalent. Adaptive spec resolves the count from its
|
|
candidate-step table instead of the flat field.
|
|
"""
|
|
spec = get_spec()
|
|
if spec.speculative_num_draft_tokens is None:
|
|
return None
|
|
if not spec.speculative_adaptive:
|
|
return spec.speculative_num_draft_tokens
|
|
# The adaptive branch parses a JSON config, and this is called per decode
|
|
# batch (`spec_prepare_for_decode`), so memoize on the inputs -- keyed, not
|
|
# cached once, so a post-publish override still recomputes.
|
|
return _adaptive_draft_token_bound(spec.speculative_adaptive_config)
|
|
|
|
|
|
@functools.lru_cache(maxsize=8)
|
|
def _adaptive_draft_token_bound(cfg_path: str | None) -> int:
|
|
from sglang.srt.speculative.adaptive_spec_params import (
|
|
resolve_candidate_steps_from_config,
|
|
)
|
|
|
|
candidate_steps = resolve_candidate_steps_from_config(cfg_path=cfg_path)
|
|
# Adaptive spec requires topk=1 today, so each runtime state needs
|
|
# steps + 1 draft-token slots (mirrors the ServerArgs member).
|
|
return max(candidate_steps) + 1
|
|
|
|
|
|
def uses_mla_backend() -> bool:
|
|
"""Whether this process's model runs the MLA attention path."""
|
|
return get_server_args().use_mla_backend()
|
|
|
|
|
|
def attention_backends() -> tuple:
|
|
"""The configured ``(prefill, decode)`` backend pair, split fields falling
|
|
back to ``attention_backend``.
|
|
|
|
All three inputs are ``exec.kernel`` leaves, so this derives from the bags
|
|
and follows a post-publish override; ``ServerArgs.get_attention_backends``
|
|
is the pre-publish equivalent the resolution pipeline uses. A built runner
|
|
stamps its own resolved pair (``ModelRunner.prefill_attention_backend_str``);
|
|
read that when there is a runner in hand.
|
|
"""
|
|
from sglang.srt.arg_groups.overrides import attention_backends_of
|
|
|
|
# All three leaves live in the same bag, so the resolution pipeline's own
|
|
# helper applies directly -- one definition of the fallback rule.
|
|
return attention_backends_of(get_exec().kernel)
|
|
|
|
|
|
def process_model_config():
|
|
"""The process's ``ModelConfig`` (built once from the published config)."""
|
|
return get_server_args().get_model_config()
|
|
|
|
|
|
def cutedsl_moe_max_num_tokens() -> int:
|
|
"""The CuteDSL A2A per-rank token budget.
|
|
|
|
Every input is a published leaf (``spec``, ``schedule``, ``exec.graph``), so
|
|
this derives from the bags and follows a post-publish override;
|
|
``ServerArgs.cutedsl_moe_max_num_tokens`` is the pre-publish equivalent the
|
|
resolution pipeline uses. Max over the prefill bound, the piecewise-prefill
|
|
capture, and the decode/verify bound.
|
|
"""
|
|
from sglang.srt.model_executor.cuda_graph_config import Backend
|
|
|
|
spec = get_spec()
|
|
num_tokens_per_req = (
|
|
(spec.speculative_num_draft_tokens or 1) if spec.speculative_algorithm else 1
|
|
)
|
|
prefill_tokens = get_schedule().max_prefill_tokens
|
|
cg_config = get_exec().graph.cuda_graph_config
|
|
if cg_config is not None and cg_config.prefill.backend == Backend.TC_PIECEWISE:
|
|
prefill_tokens = max(prefill_tokens, cg_config.prefill.max_bs or 0)
|
|
decode_max_bs = (cg_config.decode.max_bs if cg_config is not None else 0) or 0
|
|
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"
|