diff --git a/.claude/skills/sglang-runtime-context/SKILL.md b/.claude/skills/sglang-runtime-context/SKILL.md index 16de40f65..05193dcf8 100644 --- a/.claude/skills/sglang-runtime-context/SKILL.md +++ b/.claude/skills/sglang-runtime-context/SKILL.md @@ -47,8 +47,11 @@ with what the operator typed, not with what resolution decided.** compilation disabled before restricting a role), and `=enforce` fails closed on bag reads outside the role's `ROLE_NAMESPACE_SETS` entry (`None` = full tree; only audited roles are restricted). -- Bag membership is metadata on the dataclass: every `ServerArgs` field carries - `NS("path")` (e.g. `NS("exec.moe")`); coverage is linted two-way +- Bag membership is **where the field is declared**: one class per namespace under + `arg_groups/fields/`, each carrying the `_NS_PATH` it stands for, and `ServerArgs` + is assembled from them (`collect_input_fields`). The per-field `NS("path")` marker + survives only for a class that cannot express this — an ad-hoc dataclass spanning + namespaces, which is what the config-bag tests build. Coverage is linted two-way (`test_server_args_namespaces.py`, `test_runtime_context_config_bags.py`). - **Reading config**: `get_()[.sub].field` — e.g. `get_exec().moe.moe_a2a_backend`, `get_schedule().max_running_requests`. Bag leaves @@ -253,10 +256,18 @@ A process-global seed field-read of one of these sizes (`get_server_args().tp_size`, or an alias of it) is a read-ratchet failure. A `server_args` the object was *handed* is a different thing and not a ratchet matter — see "Reads that legitimately stay on a ServerArgs instance". -Fail-loud is narrower: before dist init, a live size/group read raises — except -the DCP pair, which degrades instead (`dcp_enabled` → `False`, -`attn_dcp_size` → `1` when no group is installed; -`test_attn_dcp_defaults_when_group_is_uninitialized` pins this). After init, +Fail-loud is narrower: before dist init, a live *rank/group* read raises. The six +parallel quotients are not live reads at all — `attn_tp_size`, `attn_dp_size`, +`attn_dcp_size`, `moe_ep_size`, `moe_tp_size`, `dcp_enabled` are a function of the +configured leaves, computed once at publish into bag leaves, and answered +override → stamp → published leaf. So `dcp_enabled` means "the launch configured +DCP" (`dcp_size > 1`), not "a DCP group is installed here"; in a scheduler the +stamp makes the two identical, in a process that publishes without dist init they +differ. `test_a_topology_is_stated_by_naming_the_width` and its neighbours in +`test_runtime_context.py` pin this; they replaced +`test_attn_dcp_defaults_when_group_is_uninitialized`. One consequence for tests: +overriding a leaf no longer moves its quotient — state a topology by publishing a +config, or by naming the width. After init, only the DCP group is optional (`_DCP` exists only when `dcp_size > 1`; attn-CP and moe-DP always install, as size-1 aliases if unused). The `config` hop is deliberately dynamo-traceable (a plain property over a slot, no @@ -283,13 +294,21 @@ where an object was handed one; it is not a global accessor. raises on a non-leaf. A call site that knows its field reads the bag leaf. - **the live topology** → `get_parallel()` (bare names). - **a value derived from published leaves** → an accessor in `runtime_context` that - derives it *from the bags*: `mamba_extra_buffer_enabled()` / - `mamba_extra_buffer_lazy_enabled()` read `get_memory()` and `get_exec()`, so - they see post-publish overrides. Prefer this shape whenever the inputs are - leaves; the same-named `ServerArgs` members are the pre-publish equivalents the - resolution pipeline uses, and wrapping one of those instead would quietly cost - you override visibility. `is_ep_joiner()` / `is_ep_scale_joiner()` are the same - shape over `exec.moe.ep_join_mode`, `attention_backends()` derives the + derives it *from the bags*. The strongest form of this is a `Derived(fn=...)` + declared beside the leaves it is computed from, in the namespace's own + `arg_groups/fields/` class: `publish` computes it once and stores it as an + ordinary bag leaf, so the read is a plain attribute load and it sees + post-publish overrides. `enable_mamba_extra_buffer`, `is_ep_joiner`, + `is_ep_scale_joiner` and `is_startup_weight_load_overlap` are declared that way + now — read them where they are declared: + `get_exec().mamba.enable_mamba_extra_buffer`, `get_exec().moe.is_ep_joiner`, + `get_model().is_startup_weight_load_overlap`. (The namespace is the class that + declares the field, not the namespaces its `fn` happens to read: the mamba one + spans `exec.mamba` and `memory`, which is exactly why it could not be a method + on either bag.) The + old `mamba_extra_buffer_enabled()` / `is_ep_joiner()` functions and the + same-named `ServerArgs` members are gone. The pre-publish helpers that remain + exist for resolution, which has no bag to read yet. `attention_backends()` derives the `(prefill, decode)` pair from the three `exec.kernel` leaves, and `max_speculative_num_draft_tokens()` / `cutedsl_moe_max_num_tokens()` derive theirs from `spec` / `schedule` / `exec.graph`. @@ -548,8 +567,9 @@ ONE thread — do not design for TBO threads that don't exist. flag-owning layers are pinned by name. A new module-level runtime global belongs on a flags group / resources slot instead; migrating a pinned survivor must shrink the pin. 7. **Namespace coverage** (`test_server_args_namespaces.py`, - `test_runtime_context_config_bags.py`): every `ServerArgs` field carries `NS(...)` - metadata and the projected bags must cover the fields exactly (two-way). + `test_runtime_context_config_bags.py`): every `ServerArgs` field resolves to a + namespace — from the `arg_groups/fields/` class that declares it — and the + projected bags must cover the fields exactly (two-way). Never module-skip a test "until the migration settles" — seed the context instead (the deferral ratchet that once pinned this is retired; the rule stands). diff --git a/python/sglang/benchmark/one_batch.py b/python/sglang/benchmark/one_batch.py index f21ed2faf..2d891af3f 100644 --- a/python/sglang/benchmark/one_batch.py +++ b/python/sglang/benchmark/one_batch.py @@ -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( diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index a22b54dfc..b6584805e 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py @@ -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: diff --git a/python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py b/python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py index a33f862b4..a7e6d0ac5 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_accuracy_parallel_runtime.py @@ -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() diff --git a/python/sglang/srt/arg_groups/arg_utils.py b/python/sglang/srt/arg_groups/arg_utils.py index 93dde06e6..fe535ff40 100644 --- a/python/sglang/srt/arg_groups/arg_utils.py +++ b/python/sglang/srt/arg_groups/arg_utils.py @@ -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 diff --git a/python/sglang/srt/arg_groups/fields/exec_.py b/python/sglang/srt/arg_groups/fields/exec_.py index ab70e7d12..bb1c3bbed 100644 --- a/python/sglang/srt/arg_groups/fields/exec_.py +++ b/python/sglang/srt/arg_groups/fields/exec_.py @@ -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.", diff --git a/python/sglang/srt/arg_groups/fields/model.py b/python/sglang/srt/arg_groups/fields/model.py index 62f925433..088139e95 100644 --- a/python/sglang/srt/arg_groups/fields/model.py +++ b/python/sglang/srt/arg_groups/fields/model.py @@ -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 # ------------------------------------------------------------------------- diff --git a/python/sglang/srt/arg_groups/fields/parallel.py b/python/sglang/srt/arg_groups/fields/parallel.py index bb3457f6b..6afad0463 100644 --- a/python/sglang/srt/arg_groups/fields/parallel.py +++ b/python/sglang/srt/arg_groups/fields/parallel.py @@ -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.", + ) diff --git a/python/sglang/srt/arg_groups/model_override_base.py b/python/sglang/srt/arg_groups/model_override_base.py index 58658476f..2d764449d 100644 --- a/python/sglang/srt/arg_groups/model_override_base.py +++ b/python/sglang/srt/arg_groups/model_override_base.py @@ -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", diff --git a/python/sglang/srt/distributed/bootstrap.py b/python/sglang/srt/distributed/bootstrap.py index 595eb7eba..32b985c0e 100644 --- a/python/sglang/srt/distributed/bootstrap.py +++ b/python/sglang/srt/distributed/bootstrap.py @@ -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 diff --git a/python/sglang/srt/elastic_ep/elastic_ep.py b/python/sglang/srt/elastic_ep/elastic_ep.py index 298dfa453..7894e11c8 100644 --- a/python/sglang/srt/elastic_ep/elastic_ep.py +++ b/python/sglang/srt/elastic_ep/elastic_ep.py @@ -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 diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 56e2946a1..4d9f7ec08 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -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)", diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 021f926ab..75d5f631f 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -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; diff --git a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py b/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py index 52ffce041..d69e0511b 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py +++ b/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py @@ -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) diff --git a/python/sglang/srt/hardware_backend/mlx/tp_worker.py b/python/sglang/srt/hardware_backend/mlx/tp_worker.py index 1b1b0e277..6d2e6803f 100644 --- a/python/sglang/srt/hardware_backend/mlx/tp_worker.py +++ b/python/sglang/srt/hardware_backend/mlx/tp_worker.py @@ -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 diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 1f0500e64..fe494bafe 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -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}" ) diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index 7340070c9..93fa29e7e 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -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( diff --git a/python/sglang/srt/managers/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py index 2959a1ede..c4b078b90 100644 --- a/python/sglang/srt/managers/data_parallel_controller.py +++ b/python/sglang/srt/managers/data_parallel_controller.py @@ -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() diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 204485c10..d5e85b986 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 212a66be7..85b9c203c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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 diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index eb6cf4a8f..94c710388 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -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 diff --git a/python/sglang/srt/managers/scheduler_components/request_receiver.py b/python/sglang/srt/managers/scheduler_components/request_receiver.py index 74f039949..e99e86f35 100644 --- a/python/sglang/srt/managers/scheduler_components/request_receiver.py +++ b/python/sglang/srt/managers/scheduler_components/request_receiver.py @@ -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) diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 816b7d2d3..12b32503f 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -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( diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 92adb20dc..a0ae0d8ab 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -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, diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 0c293e90d..d3c615a03 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -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 diff --git a/python/sglang/srt/mem_cache/storage/umbp/umbp_direct_linker.py b/python/sglang/srt/mem_cache/storage/umbp/umbp_direct_linker.py index 39c70a5d1..41767ef2e 100644 --- a/python/sglang/srt/mem_cache/storage/umbp/umbp_direct_linker.py +++ b/python/sglang/srt/mem_cache/storage/umbp/umbp_direct_linker.py @@ -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, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index c577d56ab..ea7f75da7 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -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(): diff --git a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py index 86de114cc..54850895e 100644 --- a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py +++ b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py @@ -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: diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 0aba7db01..b4cbebedf 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -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() ) diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index 1e6a37c07..690c6e423 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -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=( diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index a4a233ebb..d04a57f23 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -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 ) diff --git a/python/sglang/srt/models/inkling.py b/python/sglang/srt/models/inkling.py index a419de4b4..ebad2504f 100644 --- a/python/sglang/srt/models/inkling.py +++ b/python/sglang/srt/models/inkling.py @@ -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 diff --git a/python/sglang/srt/models/inkling_common/sconv.py b/python/sglang/srt/models/inkling_common/sconv.py index a846da550..ebfe6986e 100644 --- a/python/sglang/srt/models/inkling_common/sconv.py +++ b/python/sglang/srt/models/inkling_common/sconv.py @@ -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: diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index a26b09513..8249119fe 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -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. diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 4b96ceb26..652bfead8 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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 diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 941a94535..be1419c66 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -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), diff --git a/test/registered/unit/constrained/test_grammar_manager.py b/test/registered/unit/constrained/test_grammar_manager.py index 0d70b4008..43dc421c9 100644 --- a/test/registered/unit/constrained/test_grammar_manager.py +++ b/test/registered/unit/constrained/test_grammar_manager.py @@ -26,10 +26,13 @@ from sglang.srt.constrained.base_grammar_backend import ( from sglang.srt.constrained.grammar_manager import GrammarManager from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject from sglang.srt.distributed.communication_tags import P2PTag +from sglang.srt.runtime_context import get_context, publish, reset_context from sglang.srt.sampling.sampling_params import ( REQUEST_REASONING_END_TOKEN_IDS_KEY, ) +from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import enter_override register_cpu_ci(2.0, "base-a-test-cpu") @@ -38,13 +41,24 @@ register_cpu_ci(est_time=5, suite="stage-b-test-cpu-intel") def _make_scheduler(grammar_backend_name="none", skip_tokenizer=False): - """Create a mock scheduler with necessary attributes.""" + """Create a mock scheduler with necessary attributes. + + The grammar manager reads its config from the bags, so the settings that + used to be hung off the mock are published instead. The caller resets the + context; every test here goes through `_GrammarFixture`. + """ + reset_context() + server_args = ServerArgs( + model_path="dummy", + grammar_backend=grammar_backend_name, + skip_tokenizer_init=skip_tokenizer, + reasoning_parser=None, + constrained_json_whitespace_pattern=None, + constrained_json_disable_any_whitespace=False, + ) + publish(server_args, role="scheduler") scheduler = MagicMock() - scheduler.server_args.grammar_backend = grammar_backend_name - scheduler.server_args.skip_tokenizer_init = skip_tokenizer - scheduler.server_args.reasoning_parser = None - scheduler.server_args.constrained_json_whitespace_pattern = None - scheduler.server_args.constrained_json_disable_any_whitespace = False + scheduler.server_args = server_args scheduler.model_config.request_selectable_think_end_id_sequences = None # Distributed group mocks @@ -84,13 +98,19 @@ def _make_req( class TestGrammarManagerInit(unittest.TestCase): + def setUp(self): + reset_context() + self.addCleanup(reset_context) + """Test GrammarManager initialization.""" @patch("sglang.srt.constrained.grammar_manager.create_grammar_backend") def test_init_with_backend(self, mock_create): mock_create.return_value = MagicMock(spec=BaseGrammarBackend) scheduler = _make_scheduler("xgrammar") - scheduler.server_args.skip_tokenizer_init = False + enter_override( + self, get_context().override_server_args(skip_tokenizer_init=False) + ) mgr = GrammarManager(scheduler) self.assertIsNotNone(mgr.grammar_backend) @@ -114,7 +134,9 @@ class TestGrammarManagerInit(unittest.TestCase): mock_backend = MagicMock(spec=BaseGrammarBackend) mock_create.return_value = mock_backend scheduler = _make_scheduler() - scheduler.server_args.skip_tokenizer_init = False + enter_override( + self, get_context().override_server_args(skip_tokenizer_init=False) + ) mgr = GrammarManager(scheduler) mgr.clear() @@ -129,11 +151,17 @@ class TestGrammarManagerInit(unittest.TestCase): class TestProcessReqWithGrammar(unittest.TestCase): + def setUp(self): + reset_context() + self.addCleanup(reset_context) + """Test process_req_with_grammar dispatch and caching.""" def _make_mgr(self): scheduler = _make_scheduler() - scheduler.server_args.skip_tokenizer_init = True + enter_override( + self, get_context().override_server_args(skip_tokenizer_init=True) + ) mgr = GrammarManager(scheduler) mgr.grammar_backend = MagicMock(spec=BaseGrammarBackend) return mgr @@ -240,7 +268,9 @@ class TestProcessReqWithGrammar(unittest.TestCase): def test_no_backend_aborts(self): """No grammar backend should abort request.""" scheduler = _make_scheduler() - scheduler.server_args.skip_tokenizer_init = True + enter_override( + self, get_context().override_server_args(skip_tokenizer_init=True) + ) mgr = GrammarManager(scheduler) mgr.grammar_backend = None @@ -357,11 +387,17 @@ class TestProcessReqWithGrammar(unittest.TestCase): class TestAbortRequests(unittest.TestCase): + def setUp(self): + reset_context() + self.addCleanup(reset_context) + """Test abort_requests handling.""" def _make_mgr_with_queue(self): scheduler = _make_scheduler() - scheduler.server_args.skip_tokenizer_init = True + enter_override( + self, get_context().override_server_args(skip_tokenizer_init=True) + ) mgr = GrammarManager(scheduler) mgr.grammar_backend = MagicMock(spec=BaseGrammarBackend) return mgr @@ -435,11 +471,17 @@ class TestAbortRequests(unittest.TestCase): class TestGetReadyGrammarRequests(unittest.TestCase): + def setUp(self): + reset_context() + self.addCleanup(reset_context) + """Test get_ready_grammar_requests polling and result handling.""" def _make_mgr(self): scheduler = _make_scheduler() - scheduler.server_args.skip_tokenizer_init = True + enter_override( + self, get_context().override_server_args(skip_tokenizer_init=True) + ) mgr = GrammarManager(scheduler) mgr.grammar_backend = MagicMock(spec=BaseGrammarBackend) # Use very short poll interval for tests @@ -710,11 +752,17 @@ class _FakePPGroup: class TestGrammarManagerPPSync(unittest.TestCase): + def setUp(self): + reset_context() + self.addCleanup(reset_context) + """Test PP synchronization of grammar ready/failed indexes.""" def _make_mgr_for_pp(self, pp_rank, pp_size, pp_group): scheduler = _make_scheduler() - scheduler.server_args.skip_tokenizer_init = True + enter_override( + self, get_context().override_server_args(skip_tokenizer_init=True) + ) scheduler.ps.pp_rank = pp_rank scheduler.ps.pp_size = pp_size scheduler.pp_group = pp_group @@ -772,11 +820,17 @@ class TestGrammarManagerPPSync(unittest.TestCase): class TestStrictReasoningPaths(unittest.TestCase): + def setUp(self): + reset_context() + self.addCleanup(reset_context) + """Test _enable_strict_thinking code paths in GrammarManager.""" def _make_mgr(self): scheduler = _make_scheduler() - scheduler.server_args.skip_tokenizer_init = True + enter_override( + self, get_context().override_server_args(skip_tokenizer_init=True) + ) mgr = GrammarManager(scheduler) mgr.grammar_backend = MagicMock(spec=BaseGrammarBackend) mgr._enable_strict_thinking = True diff --git a/test/registered/unit/entrypoints/openai/test_serving_chat.py b/test/registered/unit/entrypoints/openai/test_serving_chat.py index 70533a7aa..1cac9012d 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_chat.py +++ b/test/registered/unit/entrypoints/openai/test_serving_chat.py @@ -6,7 +6,7 @@ or python -m unittest discover -s tests -p "test_*unit.py" -v """ -from sglang.test.test_utils import maybe_stub_sgl_kernel +from sglang.test.test_utils import enter_override, maybe_stub_sgl_kernel maybe_stub_sgl_kernel() # must precede any import that pulls in sgl_kernel @@ -3318,8 +3318,12 @@ class ServingChatTestCase(unittest.TestCase): self.assertIsNone(response.sglext) def test_non_streaming_ids_server_default_enables_flag(self): - self.tm.server_args.return_input_ids = True - self.tm.server_args.return_output_ids = True + enter_override( + self, + get_context().override_server_args( + return_input_ids=True, return_output_ids=True + ), + ) req = ChatCompletionRequest( model="x", messages=[{"role": "user", "content": "Hi?"}] ) @@ -3371,7 +3375,12 @@ class ServingChatTestCase(unittest.TestCase): ): """Stream chunks with incremental or cumulative output_ids; return parsed sglext chunks, or raw SSE strings when return_raw.""" - self.tm.server_args.incremental_streaming_output = incremental + enter_override( + self, + get_context().override_server_args( + incremental_streaming_output=incremental + ), + ) if framed: self.fastapi_request.headers["x-sglext-ids-framed"] = "1" @@ -3495,7 +3504,12 @@ class ServingChatTestCase(unittest.TestCase): self, normal_chunks, abort_output_ids, incremental, abort_completion_tokens ): """Stream normal chunks then a graceful-abort chunk; return parsed sglext chunks.""" - self.tm.server_args.incremental_streaming_output = incremental + enter_override( + self, + get_context().override_server_args( + incremental_streaming_output=incremental + ), + ) async def _mock_generate(): generated = 0 @@ -3661,14 +3675,18 @@ class ServingChatTestCase(unittest.TestCase): def test_continuous_usage_reports_cached_tokens(self): """continuous_usage_stats chunks include cached tokens when cache reporting is on.""" - self.enterContext(get_context().override_server_args(enable_cache_report=True)) + enter_override( + self, get_context().override_server_args(enable_cache_report=True) + ) usages = self._collect_continuous_usage(cached_tokens=6) self.assertTrue(usages, "continuous_usage_stats attached no usage") self.assertEqual(usages[0]["prompt_tokens_details"]["cached_tokens"], 6) def test_continuous_usage_omits_cached_tokens_when_report_disabled(self): """With cache reporting off, continuous_usage_stats must not leak cached tokens.""" - self.enterContext(get_context().override_server_args(enable_cache_report=False)) + enter_override( + self, get_context().override_server_args(enable_cache_report=False) + ) usages = self._collect_continuous_usage(cached_tokens=6) self.assertTrue(usages, "continuous_usage_stats attached no usage") self.assertIsNone(usages[0].get("prompt_tokens_details")) @@ -3684,8 +3702,8 @@ class ServingChatTestCase(unittest.TestCase): Regression test for https://github.com/sgl-project/sglang/issues/22510. """ # Enable incremental_streaming_output on the mock - self.enterContext( - get_context().override_server_args(incremental_streaming_output=True) + enter_override( + self, get_context().override_server_args(incremental_streaming_output=True) ) # Simulate incremental streaming: each yield has ONLY the new text (delta), diff --git a/test/registered/unit/hardware_backend/mlx/test_tp_worker_routing.py b/test/registered/unit/hardware_backend/mlx/test_tp_worker_routing.py index e419c7a52..790982bb9 100644 --- a/test/registered/unit/hardware_backend/mlx/test_tp_worker_routing.py +++ b/test/registered/unit/hardware_backend/mlx/test_tp_worker_routing.py @@ -226,14 +226,18 @@ class TestMlxExtendRouting(CustomTestCase): MlxModelRunnerStub, ) from sglang.srt.hardware_backend.mlx.tp_worker import MlxTpModelWorker + from sglang.srt.runtime_context import get_context worker = MlxTpModelWorker.__new__(MlxTpModelWorker) - worker.server_args = SimpleNamespace(is_startup_weight_load_overlap=True) - with self.assertRaisesRegex(ValueError, "CUDA only"): - MlxModelRunnerStub.validate_startup_weight_load_mode(worker.server_args) - with self.assertRaisesRegex(ValueError, "CUDA only"): - worker._init_model_runner() + # The guard reads `get_model().is_startup_weight_load_overlap`, which is + # derived from `startup_weight_load_mode`. Stating it on a `server_args` + # of the worker's own no longer reaches it. + with get_context().override_server_args(startup_weight_load_mode="overlap"): + with self.assertRaisesRegex(ValueError, "CUDA only"): + MlxModelRunnerStub.validate_startup_weight_load_mode() + with self.assertRaisesRegex(ValueError, "CUDA only"): + worker._init_model_runner() # ---------- the shared decision helper ---------- # The helper takes no seq_len: length cannot distinguish a 1-token diff --git a/test/registered/unit/layers/moe/test_qwen35_flashinfer_fusion.py b/test/registered/unit/layers/moe/test_qwen35_flashinfer_fusion.py index 62c5a27e1..6288c41bf 100644 --- a/test/registered/unit/layers/moe/test_qwen35_flashinfer_fusion.py +++ b/test/registered/unit/layers/moe/test_qwen35_flashinfer_fusion.py @@ -14,8 +14,14 @@ from sglang.srt.layers.moe.qwen35_flashinfer_fusion import ( is_supported_forward_mode, resolve_max_m, ) +from sglang.srt.model_executor.cuda_graph_config import ( + CudaGraphConfig, + PhaseConfig, +) from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.models.qwen3_5_text import Qwen3_5ForCausalLM +from sglang.srt.runtime_context import publish, reset_context +from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=9, suite="base-a-test-cpu") @@ -65,12 +71,19 @@ def test_supported_forward_modes(forward_mode, expected): return_value=8192, ) def test_framework_capacity_is_maximum_of_all_sources(_cutedsl_moe_max_num_tokens): - graph = SimpleNamespace( - decode=SimpleNamespace(max_bs=512, bs=[1, 64, 256]), - prefill=SimpleNamespace(max_bs=4096, bs=[1024, 2048, 4096]), + # The graph bounds are a bag leaf, so the test states them by publishing. + reset_context() + publish( + ServerArgs( + model_path="dummy", + cuda_graph_config=CudaGraphConfig( + decode=PhaseConfig(max_bs=512, bs=[1, 64, 256]), + prefill=PhaseConfig(max_bs=4096, bs=[1024, 2048, 4096]), + ), + ), + role="test", ) - server_args = SimpleNamespace(cuda_graph_config=graph) - runner = SimpleNamespace(server_args=server_args, max_running_requests=2048) + runner = SimpleNamespace(server_args=SimpleNamespace(), max_running_requests=2048) assert resolve_max_m(runner) == 8192 diff --git a/test/registered/unit/layers/test_minicpm_sparse_cache.py b/test/registered/unit/layers/test_minicpm_sparse_cache.py index 8db2c22cf..5369270a1 100644 --- a/test/registered/unit/layers/test_minicpm_sparse_cache.py +++ b/test/registered/unit/layers/test_minicpm_sparse_cache.py @@ -16,6 +16,8 @@ from sglang.srt.managers.scheduler_components.pool_stats_observer import ( ) from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +from sglang.srt.runtime_context import publish, reset_context +from sglang.srt.server_args import ServerArgs from sglang.srt.session.streaming_session import SessionSlot, StreamingSession from sglang.test.ci.ci_register import register_cpu_ci @@ -91,6 +93,16 @@ def alloc_extend(pool, req_pool_idx: int, seq_len: int): ) +def setup_function(_): + # The cache reads its parallel topology from the context. + reset_context() + publish(ServerArgs(model_path="dummy"), role="test") + + +def teardown_function(_): + reset_context() + + def test_extend_allocates_at_sparse_boundaries(): pool, _, req_pool_idx, allocator = make_pool_and_req() cache = pool._aux_cache diff --git a/test/registered/unit/managers/test_generation_auxiliary_output.py b/test/registered/unit/managers/test_generation_auxiliary_output.py index 723a82fce..e35154f18 100644 --- a/test/registered/unit/managers/test_generation_auxiliary_output.py +++ b/test/registered/unit/managers/test_generation_auxiliary_output.py @@ -15,6 +15,8 @@ from sglang.srt.managers.scheduler_pp_mixin import PPBatchMetadata from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.model_executor.forward_batch_info import PPProxyTensors from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.runtime_context import publish, reset_context +from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.test.ci.ci_register import register_cpu_ci @@ -69,13 +71,27 @@ def _model_runner_for_sampling_path( spec_algorithm=SpeculativeAlgorithm.NONE, dllm_algorithm=None, ): + # `supports_sampling_observer` reads the dLLM algorithm from the bags, so + # the path is stated by publishing it rather than by standing one in. + reset_context() + publish(ServerArgs(model_path="dummy", dllm_algorithm=dllm_algorithm), role="test") runner = object.__new__(ModelRunner) - runner.server_args = SimpleNamespace(dllm_algorithm=dllm_algorithm) + runner.server_args = SimpleNamespace() runner.spec_algorithm = spec_algorithm runner._sampling_observer = None return runner +def setup_function(_): + # The code under test reads its config from the bags. + reset_context() + publish(ServerArgs(model_path="dummy"), role="test") + + +def teardown_function(_): + reset_context() + + def test_auxiliary_output_releases_device_holder_after_copy(): device_output = DeviceOutput(torch.tensor([1.0, 2.0])) logits_output = LogitsProcessorOutput( diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index d36d01152..fcf354709 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -20,7 +20,10 @@ Usage: python -m pytest test/registered/unit/mem_cache/test_decode_radix_lock_ref.py -v """ +from sglang.srt.runtime_context import get_context, publish, reset_context +from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import enter_override register_cpu_ci(est_time=11, suite="base-a-test-cpu") @@ -97,14 +100,26 @@ def _make_req(fill_ids, req_pool_idx=0, cache_protected_len=0, last_node=None): class TestDecodeLockRefScenarios(unittest.TestCase): + def setUp(self): + # The decode queue reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="scheduler") + """Test lock_ref balance across decode transfer scenarios.""" def test_swa_tail_len_keeps_page_aligned_matchable_window(self): + enter_override( + self, + get_context().override_server_args( + disaggregation_decode_enable_radix_cache=True + ), + ) queue = DecodePreallocQueue.__new__(DecodePreallocQueue) queue._uses_swa_tail_prealloc = MagicMock(return_value=True) queue.scheduler = SimpleNamespace( sliding_window_size=127, - server_args=SimpleNamespace(disaggregation_decode_enable_radix_cache=True), + server_args=SimpleNamespace(), ) queue.token_to_kv_pool_allocator = MagicMock(page_size=64) @@ -420,7 +435,12 @@ class TestDecodeLockRefScenarios(unittest.TestCase): running_batch = MagicMock() running_batch.reqs = [] server_args = MagicMock() - server_args.disaggregation_decode_enable_radix_cache = True + enter_override( + self, + get_context().override_server_args( + disaggregation_decode_enable_radix_cache=True + ), + ) scheduler = MagicMock() scheduler.running_batch = running_batch scheduler.server_args = server_args diff --git a/test/registered/unit/mem_cache/test_full_loc_fast_path.py b/test/registered/unit/mem_cache/test_full_loc_fast_path.py index 54ab0ab0e..8c301633e 100644 --- a/test/registered/unit/mem_cache/test_full_loc_fast_path.py +++ b/test/registered/unit/mem_cache/test_full_loc_fast_path.py @@ -280,6 +280,17 @@ class TestMlaWriteDoorsUnderDcp(unittest.TestCase): already collapsed -- so there is no single correct translation and refusing is the contract.""" + def setUp(self): + # `set_kv_buffer` asks the parallel context whether DCP is in play, and + # that is a value the configuration decides at publish -- so a case + # here states its topology the way a process does, by publishing one. + from sglang.srt.runtime_context import publish, reset_context + from sglang.srt.server_args import ServerArgs + + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="test") + def _bare_mla_pool(self): from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool diff --git a/test/registered/unit/mem_cache/test_kv_index_translator.py b/test/registered/unit/mem_cache/test_kv_index_translator.py index 02fdb834e..a5e818e24 100644 --- a/test/registered/unit/mem_cache/test_kv_index_translator.py +++ b/test/registered/unit/mem_cache/test_kv_index_translator.py @@ -17,6 +17,8 @@ CPU-only: these exercise the builder's pure-torch reference path, not the Triton kernel. """ +from sglang.srt.runtime_context import publish, reset_context +from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="base-a-test-cpu") @@ -115,6 +117,12 @@ def _reference_table(req_to_token, req_pool_indices, seq_lens, v2p, mult, ps, wi class TestPassthrough(unittest.TestCase): + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def test_non_unified_returns_same_objects(self): """Strict passthrough: no copy, no branch. Any tensor op on the non-unified path breaks byte-identity for every static-pool server.""" @@ -160,6 +168,12 @@ def _alloc_and_fill(allocator, ps, lens): class TestReadTableBuild(unittest.TestCase): + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def test_read_table_matches_reference_across_multipliers(self): """Both read tables must equal the independent per-element derivation, across page sizes and both multiplier regimes (MLA=1, MHA=2L); the swa @@ -274,6 +288,12 @@ class TestBuildInto(unittest.TestCase): FULL-side read-table entries -- the trtllm_mla / flashmla consumption route (their rows ARE the read table's rows).""" + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def test_prefix_filled_tail_sentinel_preserved_width_capped(self): """The -1 tail sentinel belongs to the backend, and a table padded WIDER than the req_to_token page span (trtllm's LCM alignment) must be @@ -339,6 +359,12 @@ class TestPoolOwnership(unittest.TestCase): to address a buffer with only num_slots rows. """ + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def test_real_factory_bundle_satisfies_the_ownership_identity(self): """The guard rests on `allocator.get_kvcache() is token_to_kv_pool`, so a factory returning a pool the allocator does not hold would silently @@ -420,6 +446,12 @@ class TestPoolOwnership(unittest.TestCase): class TestCaptureContract(unittest.TestCase): + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def test_caller_owned_table_is_returned_whole_and_filled_prefix_only(self): ps = 4 allocator = _build_composite(ps) @@ -503,6 +535,12 @@ class TestViewMemo(unittest.TestCase): identity, so per-batch state stays out of the ForwardBatch while one metadata build's many consumers still share one table build.""" + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def _fb(self, allocator, ps, lens): req_to_token, rows, seq_lens = _alloc_and_fill(allocator, ps, lens=lens) fb = _FakeForwardBatch( @@ -562,6 +600,12 @@ class TestWriteLoc(unittest.TestCase): POINTWISE from the full-side values -- pads, slices, and fresh copies included -- with no handover and no stored per-forward state.""" + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def _built(self, ps=1, n=4): allocator = _build_composite(ps) req_to_token, rows, seq_lens = _alloc_and_fill(allocator, ps, lens=[max(n, 1)]) diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index 9e7c6e868..d7abd02a2 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -44,7 +44,8 @@ from sglang.srt.mem_cache.unified_memory_pool import ( MLASubPoolSpec, UnifiedKVPool, ) -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_parallel, publish, reset_context +from sglang.srt.server_args import ServerArgs _DEV = "cpu" @@ -103,6 +104,12 @@ class _RejectScalarIndexTensor: class TestUnifiedKVPoolViews(unittest.TestCase): + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def test_min_slot_index_and_disjoint_bytes(self): full = _make_mha_spec("full", "up", layer_num=4) mamba = _make_mamba_spec("mamba", "down", layer_num=2) @@ -173,6 +180,12 @@ class TestUnifiedKVPoolViews(unittest.TestCase): class TestMultiEndedAllocator(unittest.TestCase): + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def _build_pair(self, n_full_slots=64, n_mamba_slots=16): full = _make_mha_spec("full", "up", layer_num=2) mamba = _make_mamba_spec("mamba", "down", layer_num=2) @@ -526,6 +539,12 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase): """The SWA composite: joint byte-budget, slot-conservation, `free_swa` tombstone semantics, and divergent compaction of the two sub-pools.""" + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def _build( self, n_full_slots=32, @@ -891,6 +910,12 @@ class TestPagedMultiEndedAllocator(unittest.TestCase): """`MultiEndedAllocator(page_size=8)`: free-list, v2p/p2v and compaction are page-granular, while the external API stays in token ids as at page_size 1.""" + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + PAGE_SIZE = 8 def _build(self, n_full_pages=16, n_swa_pages=8, full_layer_num=2, swa_layer_num=2): @@ -1693,6 +1718,12 @@ class TestLazyCompaction(unittest.TestCase): `_flush` only waits on `forward_stream` when one is passed, so these allocators stay CPU-only.""" + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def _make_full(self, *, lazy: bool, n_full_slots=64, n_mamba_slots=16): full = _make_mha_spec("full", "up", layer_num=2) mamba = _make_mamba_spec("mamba", "down", layer_num=2) @@ -1969,6 +2000,12 @@ class TestO3FusedAllocBind(unittest.TestCase): """Fused take_physical_pages + bind_pages via `_alloc_bind_fast_or_slow`. GPU-only: the fused kernel is Triton.""" + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + @classmethod def setUpClass(cls): if not torch.cuda.is_available(): @@ -2218,6 +2255,12 @@ class TestSWACompositeKernelIdSurface(unittest.TestCase): `translate_kv_loc_for_kernel` / `full_v2p_page_table`, and every id must follow `kernel_id(t) = v2p[t // ps] * (ps * mult) + t % ps`.""" + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + PS = 4 FULL_L = 4 SWA_L = 2 @@ -2328,6 +2371,12 @@ class TestPs64MLACompositeFeasibility(unittest.TestCase): pages stress the sink-page floor, the per-layer-view tail pad and the page-granular alloc at once.""" + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + PS = 64 LAYERS = 3 @@ -2424,6 +2473,12 @@ class TestChainFrontierWalk(unittest.TestCase): """N-pool chain walk: 2-pool byte-identity with the single-peer formulas, transparent-middle skipping, and growth-side-neighbor credit routing.""" + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def _build_pair(self): full = _make_mha_spec("full", "up", layer_num=2) mamba = _make_mamba_spec("mamba", "down", layer_num=2) @@ -2552,6 +2607,12 @@ class TestFloatMultiEndedAllocator(unittest.TestCase): larger-gap extension, boundary absorption with park-on-empty transparency, and the on-demand movers `make_room` and `compact_holes`.""" + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def _build_tri(self, n_state=8, n_float=32, n_full=32): state = _make_mamba_spec("state", "up", layer_num=2) fl = _make_mha_spec("swa", "float", layer_num=2) @@ -2844,6 +2905,12 @@ class TestDcpWidening(unittest.TestCase): """`dcp_size > 1`: the alloc surface speaks a widened virtual id space while the pool keeps storing one row per `dcp_size` logical ids.""" + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + @contextlib.contextmanager def _dcp(self, dcp_size, dcp_rank=0): """The width comes from the parallel context, not a constructor diff --git a/test/registered/unit/mem_cache/test_registry.py b/test/registered/unit/mem_cache/test_registry.py index f2a741cc9..47bb5671c 100644 --- a/test/registered/unit/mem_cache/test_registry.py +++ b/test/registered/unit/mem_cache/test_registry.py @@ -1,5 +1,6 @@ """Unit tests for the radix-cache registry, routing, and selection chain.""" +from sglang.srt.runtime_context import get_context from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=11, suite="base-a-test-cpu") @@ -16,7 +17,7 @@ from sglang.srt.mem_cache.registry import ( register_radix_cache_backend, registered_radix_cache_backends, ) -from sglang.test.test_utils import CustomTestCase +from sglang.test.test_utils import CustomTestCase, enter_override def _publish(testcase, **fields): @@ -340,14 +341,14 @@ class TestDefaultRadixCacheFactory(CustomTestCase): from sglang.srt.mem_cache.storage.umbp import umbp_direct_linker ctx = _make_ctx(self) - object.__setattr__( - ctx.server_args, "enable_unified_cache_external_linker", True + # The factory reads the linker settings from the bags. + enter_override( + self, + get_context().override_server_args( + enable_unified_cache_external_linker=True, + unified_cache_external_linker_backend="mori", + ), ) - object.__setattr__( - ctx.server_args, "unified_cache_external_linker_backend", "mori" - ) - self.assertTrue(ctx.server_args.enable_unified_cache_external_linker) - self.assertEqual(ctx.server_args.unified_cache_external_linker_backend, "mori") fake_components = MagicMock() fake_components.ComponentType.FULL = "full" fake_radix = MagicMock() diff --git a/test/registered/unit/mem_cache/test_unified_byte_accounting.py b/test/registered/unit/mem_cache/test_unified_byte_accounting.py index e808d81ac..ee7a7e74a 100644 --- a/test/registered/unit/mem_cache/test_unified_byte_accounting.py +++ b/test/registered/unit/mem_cache/test_unified_byte_accounting.py @@ -32,6 +32,8 @@ from test_multi_ended_allocator import ( ) from sglang.srt.mem_cache.allocator import unified_sub_pool as mea +from sglang.srt.runtime_context import publish, reset_context +from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="base-a-test-cpu") @@ -51,6 +53,12 @@ def _paged_pair(lazy: bool): class TestHealthyLifecycleReportsClean(unittest.TestCase): + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def test_lazy_end_pool_clean_through_free_and_flush(self): full, _swa = _paged_pair(lazy=True) self.assertEqual(full._byte_accounting_violations(), []) @@ -68,6 +76,12 @@ class TestDriftReportsLoudly(unittest.TestCase): passes every other test -- the pool still 'works', it just lies about capacity.""" + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def _lazy_full(self): full, _swa = _paged_pair(lazy=True) v = full.alloc(full.page_size * 4) @@ -112,6 +126,12 @@ class TestDriftReportsLoudly(unittest.TestCase): class TestChainFrontierOrder(unittest.TestCase): + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def test_overlapping_frontiers_report(self): """Both bands hold pages, then the up member's watermark is pushed past the down member's LIVE low frontier. Both sides must be populated diff --git a/test/registered/unit/mem_cache/test_unified_mha_views.py b/test/registered/unit/mem_cache/test_unified_mha_views.py index 9fb4ce2b9..b97de3d34 100644 --- a/test/registered/unit/mem_cache/test_unified_mha_views.py +++ b/test/registered/unit/mem_cache/test_unified_mha_views.py @@ -367,6 +367,18 @@ class TestUnifiedMHATokenToKVPool(unittest.TestCase): class TestFactoryViews(unittest.TestCase): + def setUp(self): + # `KVIndexTranslator.__init__` asks the parallel context for + # `attn_dcp_size`, which is a quotient of the configured leaves and is + # computed at publish. A bare process has none, so state one the way a + # real process does. + from sglang.srt.runtime_context import publish, reset_context + from sglang.srt.server_args import ServerArgs + + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="test") + """Over the real SWA factory: matching kernel-facing multipliers in the composite allocator, and a rebind that emits both write locs.""" diff --git a/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py b/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py index 9b6d1f278..9cb799690 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py +++ b/test/registered/unit/mem_cache/test_unified_mla_gpu_parity.py @@ -26,6 +26,8 @@ import unittest import torch +from sglang.srt.runtime_context import publish, reset_context +from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small") @@ -107,6 +109,12 @@ def _rand_locs(max_tokens: int, ps: int, n: int) -> torch.Tensor: @unittest.skipUnless(_HAS_CUDA, "requires CUDA") class TestUnifiedMLAPoolGPUParity(unittest.TestCase): + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def _assert_parity(self, unified, ref, locs, ps, layers=range(_L)): for l in layers: got = unified.get_key_buffer(l)[_kernel_id(locs, ps)] diff --git a/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py b/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py index 99c7e61e8..3816716af 100644 --- a/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py +++ b/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py @@ -158,6 +158,9 @@ class _TiedWeightModel(nn.Module): class TestStartupWeightLoadSelector(CustomTestCase): def setUp(self): + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") self.load_config = LoadConfig(load_format=LoadFormat.SAFETENSORS) self.loader = DefaultModelLoader(self.load_config) self.device_config = DeviceConfig("cuda", 0) @@ -324,6 +327,12 @@ class TestStartupWeightLoadSelector(CustomTestCase): class TestStartupWeightLoadManager(CustomTestCase): + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def _manager(self, loader): return StartupWeightLoadManager( loader=loader, @@ -512,6 +521,12 @@ class TestStartupWeightLoadManager(CustomTestCase): class TestModelStorageManifest(CustomTestCase): + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def test_in_place_updates_preserve_the_manifest(self): model = _TiedWeightModel() manifest = ModelStorageManifest.capture(model) @@ -554,6 +569,12 @@ class TestModelStorageManifest(CustomTestCase): class TestCaptureSafeWeightInitialization(CustomTestCase): + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def test_only_parameters_are_filled(self): model = _TiedWeightModel() @@ -576,6 +597,12 @@ class _LifecycleRunner: class TestStartupWeightLoadFanout(CustomTestCase): + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def test_primary_and_multi_runner_extras_are_started_once(self): trace = [] primary = _LifecycleRunner("primary", trace) @@ -629,6 +656,12 @@ class _RunnerStartupManager: class TestModelRunnerStartupWeightLoadOwnership(CustomTestCase): + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + @staticmethod def _runner(manager): runner = ModelRunner.__new__(ModelRunner) @@ -689,14 +722,21 @@ class _SchedulerWorker: class TestStartupWeightLoadSchedulerRouting(CustomTestCase): - @staticmethod - def _scheduler(worker, trace, *, mode, draft_worker=None): + def setUp(self): + reset_context() + self.addCleanup(reset_context) + + def _scheduler(self, worker, trace, *, mode, draft_worker=None): from sglang.srt.managers.scheduler import Scheduler - scheduler = Scheduler.__new__(Scheduler) - scheduler.server_args = SimpleNamespace( - is_startup_weight_load_overlap=mode == "overlap" + # The schedule reads the mode from the bags, so the test states it by + # publishing a record rather than by standing one in. + reset_context() + publish( + ServerArgs(model_path="dummy", startup_weight_load_mode=mode), + role="scheduler", ) + scheduler = Scheduler.__new__(Scheduler) scheduler.init_tp_model_worker = lambda: setattr(scheduler, "tp_worker", worker) scheduler.maybe_init_draft_worker = lambda: setattr( scheduler, "draft_worker", draft_worker diff --git a/test/registered/unit/server_args/test_resolution_reads_the_declarations.py b/test/registered/unit/server_args/test_resolution_reads_the_declarations.py index 3f3c78d64..acdc926d4 100644 --- a/test/registered/unit/server_args/test_resolution_reads_the_declarations.py +++ b/test/registered/unit/server_args/test_resolution_reads_the_declarations.py @@ -472,8 +472,9 @@ class TestResolutionReadsTheDeclarations(CustomTestCase): ) members = _record_members() # The floor is here to catch the scan collapsing, not to pin the - # class's size. - self.assertGreater(len(members), 15, f"only {len(members)} members were found") + # class's size -- it drops as derived members move to their namespaces + # and become declarations rather than methods on the record. + self.assertGreater(len(members), 10, f"only {len(members)} members were found") offenders = [] for name, fn in sorted(members.items()): holders = _holders(fn) | {"self"} diff --git a/test/registered/unit/spec/test_eagle_draft_cuda_graph_runner.py b/test/registered/unit/spec/test_eagle_draft_cuda_graph_runner.py index 827128217..2ab325d8c 100644 --- a/test/registered/unit/spec/test_eagle_draft_cuda_graph_runner.py +++ b/test/registered/unit/spec/test_eagle_draft_cuda_graph_runner.py @@ -19,6 +19,8 @@ from types import SimpleNamespace import torch +from sglang.srt.runtime_context import publish, reset_context +from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.eagle_draft_cuda_graph_runner import ( EAGLEDraftCudaGraphRunner, ) @@ -59,6 +61,12 @@ class _RecordingDraftBackend: class TestEagleDraftCudaGraphRunner(CustomTestCase): + def setUp(self): + # The code under test reads its config from the bags. + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy"), role="tokenizer") + def _build_runner(self, backend): runner = EAGLEDraftCudaGraphRunner.__new__(EAGLEDraftCudaGraphRunner) runner.deepep_adapter = SimpleNamespace(replay=lambda: None) diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 5d2a6ba9d..eaddcd9a9 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -53,21 +53,22 @@ _SRT = _pathlib.Path(next(iter(_sglang.__path__))).resolve() / "srt" _PS = "sglang.srt.distributed.parallel_state" _DP = "sglang.srt.layers.dp_attention" +# Ranks and the world size read the live group: they are not implied by +# anything, so there is nothing to derive them from. The quotients used to be +# in this table and are not any more -- `attn_tp_size` and its siblings are +# functions of the configured leaves, and `TestDerivedWidthsComeFromTheLeaves` +# is what pins them. SIZE_RANK_DELEGATIONS = [ ("world_size", f"{_PS}.get_world_size"), ("world_rank", f"{_PS}.get_world_rank"), ("tp_rank", f"{_PS}.get_tensor_model_parallel_rank"), ("dcp_rank", f"{_PS}.get_dcp_rank"), ("pp_rank", f"{_PS}.get_pipeline_model_parallel_rank"), - ("moe_ep_size", f"{_PS}.get_moe_expert_parallel_world_size"), ("moe_ep_rank", f"{_PS}.get_moe_expert_parallel_rank"), ("moe_dp_rank", f"{_PS}.get_moe_data_parallel_rank"), - ("moe_tp_size", f"{_PS}.get_moe_tensor_parallel_world_size"), ("moe_tp_rank", f"{_PS}.get_moe_tensor_parallel_rank"), - ("attn_tp_size", f"{_PS}.get_attn_tensor_model_parallel_world_size"), ("attn_tp_rank", f"{_PS}.get_attn_tensor_model_parallel_rank"), ("attn_cp_rank", f"{_PS}.get_attn_context_model_parallel_rank"), - ("attn_dp_size", f"{_DP}.get_attention_dp_size"), ("attn_dp_rank", f"{_DP}.get_attention_dp_rank"), ] @@ -176,36 +177,49 @@ class TestParallelOverride(_IsolatedOverrides): class TestParallelDCP(_IsolatedOverrides): - def test_attn_dcp_defaults_when_group_is_uninitialized(self): + """The DCP width is a quotient; the DCP rank is a live reading. + + They used to be tested the same way, by mocking the group getters, because + the width read the group too. It does not: `attn_dcp_size` is + `dcp_size if dcp_enabled else 1`, so the way to state it is to state the + leaves. + """ + + def _published(self, **fields): + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy", **fields), role="test") + return get_parallel() + + def test_attn_dcp_is_one_when_dcp_is_off(self): + parallel = self._published(tp_size=8, dcp_size=1) + self.assertFalse(parallel.dcp_enabled) + self.assertEqual(parallel.attn_dcp_size, 1) + + def test_attn_dcp_is_the_configured_width_when_on(self): + parallel = self._published(tp_size=8, dcp_size=8) + self.assertTrue(parallel.dcp_enabled) + self.assertEqual(parallel.attn_dcp_size, 8) + + def test_the_dcp_rank_still_reads_the_group(self): + """A rank is not implied by the configuration, so it reads the group -- + gated on a width that is.""" with ( - patch(f"{_PS}.get_dcp_group_no_assert", return_value=None), - patch(f"{_PS}.get_dcp_world_size", side_effect=AssertionError), + get_parallel().override(tp_size=8, dcp_size=8, dcp_enabled=False), patch(f"{_PS}.get_dcp_rank", side_effect=AssertionError), ): - self.assertFalse(get_parallel().dcp_enabled) - self.assertEqual(get_parallel().attn_dcp_size, 1) self.assertEqual(get_parallel().attn_dcp_rank, 0) - - def test_attn_dcp_delegates_when_enabled(self): with ( - patch(f"{_PS}.get_dcp_group_no_assert", return_value=object()), - patch(f"{_PS}.get_dcp_world_size", return_value=8), + get_parallel().override(tp_size=8, dcp_size=8, dcp_enabled=True), patch(f"{_PS}.get_dcp_rank", return_value=3), ): - self.assertTrue(get_parallel().dcp_enabled) - self.assertEqual(get_parallel().attn_dcp_size, 8) self.assertEqual(get_parallel().attn_dcp_rank, 3) - def test_dcp_enablement_is_platform_agnostic(self): - with ( - patch(f"{_PS}.get_dcp_group_no_assert", return_value=object()), - patch("sglang.srt.utils.is_cuda", return_value=False) as is_cuda, - patch(f"{_PS}.get_dcp_world_size", return_value=8), - patch(f"{_PS}.get_dcp_rank", return_value=3), - ): - self.assertTrue(get_parallel().dcp_enabled) - self.assertEqual(get_parallel().attn_dcp_size, 8) - self.assertEqual(get_parallel().attn_dcp_rank, 3) + def test_the_width_does_not_consult_the_platform(self): + with patch("sglang.srt.utils.is_cuda", return_value=False) as is_cuda: + parallel = self._published(tp_size=8, dcp_size=8) + self.assertTrue(parallel.dcp_enabled) + self.assertEqual(parallel.attn_dcp_size, 8) is_cuda.assert_not_called() @@ -1150,27 +1164,32 @@ class TestDerivedPredicatesAgreeAcrossTiers(_IsolatedServerArgs): _STRATEGIES = ("auto", "no_buffer", "extra_buffer", "extra_buffer_lazy") - def test_mamba_extra_buffer_matches_the_member(self): - from sglang.srt.runtime_context import ( - mamba_extra_buffer_enabled, - mamba_extra_buffer_lazy_enabled, - ) - + def test_the_mamba_extra_buffer_predicate_has_one_answer(self): + """It used to be asserted that two spellings agreed. There is one now: + the declaration computes it at publish, and the bag carries it.""" for disable_radix_cache in (False, True): for strategy in self._STRATEGIES: with self.subTest(radix=disable_radix_cache, strategy=strategy): - args = _FakeResolvedArgs( - disable_radix_cache=disable_radix_cache, - mamba_radix_cache_strategy=strategy, + reset_context() + publish( + ServerArgs( + model_path="dummy", + disable_radix_cache=disable_radix_cache, + mamba_radix_cache_strategy=strategy, + ), + role="test", ) - get_context().set_server_args(args) - self.assertEqual( - ServerArgs.enable_mamba_extra_buffer(args), - mamba_extra_buffer_enabled(), + expected = disable_radix_cache is False and strategy in ( + "extra_buffer", + "extra_buffer_lazy", ) self.assertEqual( - ServerArgs.enable_mamba_extra_buffer_lazy(args), - mamba_extra_buffer_lazy_enabled(), + get_exec().mamba.enable_mamba_extra_buffer, expected + ) + self.assertEqual( + get_exec().mamba.enable_mamba_extra_buffer_lazy, + disable_radix_cache is False + and strategy == "extra_buffer_lazy", ) def test_prefill_buffer_ceiling_matches_the_member(self): @@ -1446,6 +1465,62 @@ class TestDerivedWidths(_IsolatedOverrides): ) ) + def test_the_published_configuration_decides_the_widths(self): + """The quotients are computed once, at publish, from the leaves. + + Every input is a record field, so there is nothing to recompute on a + read: `publish` fills the bag and the bag is the answer. + """ + reset_context() + self.addCleanup(reset_context) + publish( + ServerArgs( + model_path="dummy", tp_size=8, dp_size=2, enable_dp_attention=True + ), + role="test", + ) + self.assertEqual(get_parallel().attn_tp_size, 4) + self.assertEqual(get_parallel().attn_dp_size, 2) + self.assertEqual(get_parallel().moe_tp_size, 8) + + reset_context() + publish( + ServerArgs(model_path="dummy", tp_size=8, ep_size=4, moe_dp_size=2), + role="test", + ) + self.assertEqual(get_parallel().moe_tp_size, 1) + + def test_a_topology_is_stated_by_naming_the_width(self): + """Overriding a leaf does not move the quotient -- the quotient is not + recomputed on read. Naming it is how a test states one.""" + reset_context() + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy", tp_size=8), role="test") + self.assertEqual(get_parallel().attn_tp_size, 8) + with get_parallel().override(tp_size=2): + self.assertEqual(get_parallel().attn_tp_size, 8) + with get_parallel().override(attn_tp_size=4): + self.assertEqual(get_parallel().attn_tp_size, 4) + + def test_an_unstated_topology_still_fails(self): + """Neutral leaves are for the dimensions a caller is not using, not for + a caller that stated nothing: every width would come back 1, which is a + plausible-looking number invented out of nothing.""" + with self.assertRaises(RuntimeError) as caught: + get_parallel().attn_tp_size + self.assertIn("not available", str(caught.exception)) + + def test_a_stamp_and_a_live_group_both_win_over_the_leaves(self): + """Order is stamp, then live group, then the leaves. Where a group + exists it is the truth -- elastic scale-up moves the group without + restamping -- so the leaf derivation only answers where there is none. + """ + parallel = get_parallel() + parallel.stamp_derived_widths(attn_tp_size=7) + self.addCleanup(parallel.clear_derived_widths) + with parallel.override(tp_size=8, attn_dp_size=2): + self.assertEqual(parallel.attn_tp_size, 7) + def test_the_quotients_come_from_the_leaves(self): widths = derive_parallel_widths( tp_size=8, @@ -1497,10 +1572,24 @@ class TestDerivedWidths(_IsolatedOverrides): self.assertEqual(parallel.attn_tp_size, 1) self.assertEqual(parallel.attn_tp_size, 4) - def test_without_a_stamp_the_live_group_still_answers(self): - """A process that installed groups by hand keeps working.""" - with patch(f"{_PS}.get_attn_tensor_model_parallel_world_size", return_value=2): - self.assertEqual(get_parallel().attn_tp_size, 2) + def test_the_group_is_never_consulted(self): + """There is no third source. A quotient comes from an override, a stamp + or the published leaf -- never from a group coordinator, which could + only ever agree, since `initialize_model_parallel` stamps as its last + statement.""" + reset_context() + self.addCleanup(reset_context) + with patch( + f"{_PS}.get_attn_tensor_model_parallel_world_size", + side_effect=AssertionError("the group must not be consulted"), + ): + publish( + ServerArgs( + model_path="dummy", tp_size=8, dp_size=2, enable_dp_attention=True + ), + role="test", + ) + self.assertEqual(get_parallel().attn_tp_size, 4) def test_with_neither_the_failure_names_the_cause(self): with patch( @@ -1533,24 +1622,23 @@ class TestDerivedWidths(_IsolatedOverrides): parallel.stamp_derived_widths(attn_dp_size=4) self.assertEqual(parallel.attn_dp_size, 4) parallel.clear_derived_widths() - with patch(f"{_DP}.get_attention_dp_size", return_value=1): + with parallel.override(tp_size=8, attn_dp_size=1): self.assertEqual(parallel.attn_dp_size, 1) def test_reset_context_drops_the_stamp(self): """The stamp belongs to the lifecycle that made it. - `_derived_width` prefers the stamp over the live group, so a stamp that - outlived `reset_context()` would let the next test read the previous - topology. + `_derived_width` prefers the stamp over the published leaf, so a stamp + that outlived `reset_context()` would let the next test read the + previous topology. """ - from sglang.srt.runtime_context import reset_context - parallel = get_parallel() parallel.stamp_derived_widths(attn_tp_size=4) self.assertEqual(parallel.attn_tp_size, 4) reset_context() - with patch(f"{_PS}.get_attn_tensor_model_parallel_world_size", return_value=1): - self.assertEqual(get_parallel().attn_tp_size, 1) + self.addCleanup(reset_context) + publish(ServerArgs(model_path="dummy", tp_size=1), role="test") + self.assertEqual(get_parallel().attn_tp_size, 1) def test_the_arithmetic_has_one_home(self): """`parallel_state` builds its groups from the same dict it stamps, and @@ -1589,5 +1677,67 @@ class TestDerivedWidths(_IsolatedOverrides): self.assertEqual(attn_dp_size, widths["attn_dp_size"]) +class TestTheDerivedHalfIsDeclared(CustomTestCase): + """The quotients are declared beside the leaves, in the same class. + + A namespace is one file and one class. `Parallel` says both what an + operator can set and what that decides; the quotients are unannotated, so + they are not dataclass fields and never reach the record. + `ParallelContext` installs a property per declaration rather than carrying + its own list, so the two cannot drift. + """ + + def test_every_declared_quotient_has_a_property(self): + from sglang.srt.arg_groups.arg_utils import Derived + from sglang.srt.arg_groups.fields.parallel import Parallel + + declared = { + name for name, value in vars(Parallel).items() if isinstance(value, Derived) + } + self.assertTrue(declared, "the derived half is empty") + for name in declared: + self.assertIsInstance( + getattr(type(get_context().parallel), name, None), + property, + f"{name} is declared but no property was installed", + ) + + def test_the_declared_set_is_what_derive_parallel_widths_produces(self): + """The declaration is not a second list to keep in step: it names + exactly the quotients the derivation returns.""" + from sglang.srt.arg_groups.arg_utils import Derived + from sglang.srt.arg_groups.fields.parallel import Parallel + + declared = { + name for name, value in vars(Parallel).items() if isinstance(value, Derived) + } + produced = set( + derive_parallel_widths( + tp_size=8, + attn_cp_size=1, + attn_dp_size=2, + moe_ep_size=1, + moe_dp_size=1, + dcp_size=1, + dcp_enabled=False, + ) + ) + self.assertEqual(declared, produced) + + def test_a_declared_quotient_is_not_a_record_field(self): + """It has no operator input to preserve, and the record is what crosses + a process boundary.""" + import dataclasses + + from sglang.srt.arg_groups.arg_utils import Derived + from sglang.srt.arg_groups.fields.parallel import Parallel + from sglang.srt.server_args import ServerArgs + + fields = {f.name for f in dataclasses.fields(ServerArgs)} + for name, value in vars(Parallel).items(): + if isinstance(value, Derived): + self.assertNotIn(name, fields) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/test_server_args_migration.py b/test/registered/unit/test_server_args_migration.py index b1a94065d..2a3d16c72 100644 --- a/test/registered/unit/test_server_args_migration.py +++ b/test/registered/unit/test_server_args_migration.py @@ -8,6 +8,7 @@ import argparse import unittest from sglang.srt.arg_groups.overrides import resolution_result +from sglang.srt.runtime_context import get_model, publish, reset_context from sglang.srt.server_args import ServerArgs from sglang.srt.utils.common import configure_media_url_security from sglang.test.ci.ci_register import register_cpu_ci @@ -132,9 +133,13 @@ class TestServerArgsAnnotatedCli(CustomTestCase): serial = self._parse([]) overlap = self._parse(["--startup-weight-load-mode", "overlap"]) self.assertEqual(serial.startup_weight_load_mode, "serial") - self.assertFalse(serial.is_startup_weight_load_overlap) self.assertEqual(overlap.startup_weight_load_mode, "overlap") - self.assertTrue(overlap.is_startup_weight_load_overlap) + # The predicate over that leaf is a bag leaf now, computed at publish. + for record, expected in ((serial, False), (overlap, True)): + reset_context() + self.addCleanup(reset_context) + publish(record, role="test") + self.assertIs(get_model().is_startup_weight_load_overlap, expected) with self.assertRaises(SystemExit): self.parser.parse_args( diff --git a/test/registered/unit/test_server_args_namespaces.py b/test/registered/unit/test_server_args_namespaces.py index 2e575eada..dac4df82f 100644 --- a/test/registered/unit/test_server_args_namespaces.py +++ b/test/registered/unit/test_server_args_namespaces.py @@ -1,9 +1,14 @@ """Coverage lint for the ServerArgs -> RuntimeContext namespace split. -Every ServerArgs field must carry an ``NS("")`` marker in its ``Annotated`` -metadata, and every path must be one of the known domains. This is the guardrail -that fails when an upstream PR adds a ServerArgs field without assigning it a -namespace (the property that retires the old hand-maintained mirror file). +Every ServerArgs field must resolve to a namespace, and every path must be one of +the known domains. A field gets its namespace from the ``arg_groups/fields/`` +class that declares it -- each carries the ``_NS_PATH`` it stands for, so the +module a declaration lives in *is* the answer, and there is no per-field marker +to forget. (``NS("")`` survives for the one shape a class cannot express: +an ad-hoc dataclass spanning namespaces, which the config-bag tests build.) + +This is the guardrail that fails when an upstream PR adds a field to a namespace +class that has no ``_NS_PATH``, or adds one outside the taxonomy below. """ import dataclasses