From aaf9a957631326f62b25309cb8a11a246f564ed1 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Sun, 6 Sep 2026 21:44:24 -0700 Subject: [PATCH] [Config] Round 6.5: a namespace declares what it derives, next to what it derives it from (#38113) Fifth of five; stacked on #38049. The split gave every namespace a file, but only for the half an operator types. This is the other half. ## The parallel quotients are declared, not written out `attn_tp_size` and its five siblings were sixty lines of near-identical properties in the runtime context, a file away from the leaves they are quotients of, so reading `parallel.py` told you what you could set and nothing about what that decides. They are declared in `Parallel` now, in the same class as those leaves. They carry no annotation, so they are not dataclass fields and `collect_input_fields` never puts them on the record -- the same mechanism that already keeps `_NS_PATH` off it. That is the right exclusion: a quotient has no operator input to preserve, and the record is what crosses a process boundary, where a stamped width is one an elastic scale-up will not refresh. ## A quotient is a value in the bag, like every other derived one `_derived_width` answered from a stamp or, failing that, a live process group. The group read could never disagree with the stamp: - `initialize_model_parallel` stamps all six as its last statement, unconditionally; - an elastic scale-up restamps `attn_dp_size` through `update_dp_attention_post_scale` -- the comment claiming it does *not* was wrong; - no hardware backend builds groups of its own; - `multimodal_gen`, which has its own `initialize_model_parallel` and does not stamp, never reads a quotient. So a built group was always already stamped, and the group read goes -- and with it the last reason for a quotient to be resolved on every read. Every input to `derive_parallel_widths` is a record field. `dcp_enabled` is `decode_context_parallel_size > 1`, not a fact about a built group; it was spelled `_DCP is not None`, which is a longer way to say the same thing. So the six are fixed once the configuration is fixed -- the same test every other `Derived(fn=...)` in this PR passes. They are declared the same way and computed the same way: once, at publish, into ordinary bag leaves. What remains is override -> stamp -> published leaf. The stamp stays above the leaf because an elastic scale-up restamps `attn_dp_size`; the override stays on top because that is how a test names a width. ## One answer for the config-derived predicates `enable_mamba_extra_buffer` and its lazy variant, `is_ep_joiner`, `is_ep_scale_joiner`, `is_startup_weight_load_overlap`: each existed as a `ServerArgs` member for the resolution pipeline and, for most of them, again as a `runtime_context` function for readers after publish. Three places to keep saying the same thing. A `Derived(fn=...)` is a pure function of the published configuration, so `publish` computes it once and stores it as an ordinary bag leaf -- a plain attribute load, which is what a read inside compiled model code needs. The function is handed the whole resolved config rather than the bag it lands in, because a derivation is free to span namespaces and the mamba one does: it reads `memory.disable_radix_cache` alongside its own `exec.mamba` strategy, which is why it could never have been a method on either bag. The pre-publish helpers stay -- resolution needs the predicate before there is a bag to read -- and three readers keep them, because they run before their own process publishes: `initialize_dp_attention`, which the weight-cache daemon calls while building its groups thirty lines before its `publish`, and `PortArgs.init_new`, a factory handed the record that already reads eighteen other fields off it. ## Notes for a reviewer **Overriding a leaf does not move its quotient.** `override(tp_size=2)` leaves `attn_tp_size` where the published config put it, because nothing is recomputed on read. A test states a topology by publishing a config -- which is what a real process does -- or by naming the width it wants, `override(attn_tp_size=2)`. Six tests say it that way now. This is the price of having one answer computed once, and it is the same price every other derived value in the config already carries. A caller that reads a quotient without publishing or overriding now gets an explicit error naming the field, instead of a default that an uninitialised group happened to supply. One fixture was in that state -- `TestMlaWriteDoorsUnderDcp` built a bare pool and asked whether DCP was on -- and it publishes a config now, which is what the process it stands in for does. Eighteen sites read these predicates without calling them. That is correct -- they are properties -- but it is worth saying they were checked, because a census that assumes otherwise reports eighteen always-true conditions. ## The skill that documents this subsystem is updated with it `.claude/rules/modify-component-must-read.md` points at `.claude/skills/sglang-runtime-context/SKILL.md` before anyone touches these files, so a stale sentence there is a wrong instruction rather than a stale note. Four of its load-bearing statements stopped being true across this series and are corrected here: `NS(...)` is no longer how a field states its namespace (the declaring class is); the DCP degrade rule is gone, because the quotients are not live reads; `mamba_extra_buffer_enabled()` and the other predicate functions it named as the shape to copy no longer exist; and the namespace-coverage ratchet is described in terms of the marker. The docstring of `test_server_args_namespaces.py` said the same thing and is fixed too. The consequence a test author actually trips over is stated there as well: overriding a leaf no longer moves its quotient, so a topology is stated by publishing a config or by naming the width. ## Verification A full registered-unit sweep (648 files) against this stack's merge-base: 19 failures on both sides, the same 19 -- AMD `gfx950`, `modelopt`, `cuda_vmm`, `weight_checker` and friends, none of them config. The narrower 139-file config sweep used earlier in this series does not contain the files this change reaches -- `test_kv_index_translator` never names `get_parallel()`, it constructs an object that does -- which is why the baseline differential over everything is what is quoted here. --- .../skills/sglang-runtime-context/SKILL.md | 50 ++-- python/sglang/benchmark/one_batch.py | 11 +- .../runtime/distributed/parallel_state.py | 51 +++- ...est_component_accuracy_parallel_runtime.py | 6 +- python/sglang/srt/arg_groups/arg_utils.py | 28 ++ python/sglang/srt/arg_groups/fields/exec_.py | 32 +++ python/sglang/srt/arg_groups/fields/model.py | 7 + .../sglang/srt/arg_groups/fields/parallel.py | 42 +++ .../srt/arg_groups/model_override_base.py | 34 ++- python/sglang/srt/distributed/bootstrap.py | 4 +- python/sglang/srt/elastic_ep/elastic_ep.py | 2 +- python/sglang/srt/entrypoints/http_server.py | 2 +- .../srt/entrypoints/openai/serving_chat.py | 12 +- .../hardware_backend/mlx/model_runner_stub.py | 8 +- .../srt/hardware_backend/mlx/tp_worker.py | 2 +- .../attention/hybrid_linear_attn_backend.py | 2 +- python/sglang/srt/layers/dp_attention.py | 11 +- .../srt/managers/data_parallel_controller.py | 10 +- python/sglang/srt/managers/schedule_batch.py | 13 +- python/sglang/srt/managers/scheduler.py | 4 +- .../batch_result_processor.py | 8 +- .../scheduler_components/request_receiver.py | 4 +- python/sglang/srt/managers/tp_worker.py | 2 +- .../sglang/srt/mem_cache/kv_cache_builder.py | 6 +- .../srt/mem_cache/kv_cache_configurator.py | 20 +- .../storage/umbp/umbp_direct_linker.py | 8 +- .../sglang/srt/model_executor/model_runner.py | 28 +- .../load_model_utils.py | 21 +- .../runner/decode_cuda_graph_runner.py | 2 +- .../srt/model_executor/runner/eager_runner.py | 5 +- .../runner/prefill_cuda_graph_runner.py | 2 +- python/sglang/srt/models/inkling.py | 3 +- .../sglang/srt/models/inkling_common/sconv.py | 6 +- python/sglang/srt/runtime_context.py | 240 ++++++++++------- python/sglang/srt/server_args.py | 73 ++--- python/sglang/srt/speculative/spec_utils.py | 9 +- .../unit/constrained/test_grammar_manager.py | 82 +++++- .../entrypoints/openai/test_serving_chat.py | 36 ++- .../mlx/test_tp_worker_routing.py | 14 +- .../moe/test_qwen35_flashinfer_fusion.py | 23 +- .../unit/layers/test_minicpm_sparse_cache.py | 12 + .../test_generation_auxiliary_output.py | 18 +- .../mem_cache/test_decode_radix_lock_ref.py | 24 +- .../unit/mem_cache/test_full_loc_fast_path.py | 11 + .../mem_cache/test_kv_index_translator.py | 44 +++ .../mem_cache/test_multi_ended_allocator.py | 69 ++++- .../unit/mem_cache/test_registry.py | 17 +- .../mem_cache/test_unified_byte_accounting.py | 20 ++ .../unit/mem_cache/test_unified_mha_views.py | 12 + .../mem_cache/test_unified_mla_gpu_parity.py | 8 + .../test_startup_weight_load.py | 50 +++- .../test_resolution_reads_the_declarations.py | 5 +- .../test_eagle_draft_cuda_graph_runner.py | 8 + test/registered/unit/test_runtime_context.py | 254 ++++++++++++++---- .../unit/test_server_args_migration.py | 9 +- .../unit/test_server_args_namespaces.py | 13 +- 56 files changed, 1118 insertions(+), 379 deletions(-) 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