From f2ab6e306b3dc29ad155a7a7d0cdf9ff86f19a77 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Sat, 15 Aug 2026 00:39:03 -0700 Subject: [PATCH] config: the alias form of the runner-side instance read The previous batch counted `self.server_args.X` and called the runner surface done. It was not: the same read spelled through a local alias -- `server_args = model_runner.server_args` (or `sa = kvc.server_args`, `args = ...`) followed by `server_args.leaf` -- is the same process-global read wearing a local name, and the AST census counts **57 of them** across eleven files that the grep never saw. Census per function, following the alias. 52 were leaves and go to their bag (`spec` 11, `schedule` 9, `memory` 7, `exec.graph` 5, `exec.moe` 5, `parallel` 4, `disagg` 4, `model` 3, `exec.mamba` 2, `exec.overlap` 2). Five were not leaves: three derived members on the eager runner -- `max_speculative_num_draft_tokens` and `enable_mamba_extra_buffer` already had accessors, and `max_prefill_buffer_tokens` gets one (all its inputs are `schedule` leaves plus the configured PP size, so it derives from the bags and follows a post-publish override; `TestDerivedPredicatesAgreeAcrossTiers` pins it against the member over a 48-case matrix) -- plus `get_attention_backends()`, which the same commit routes through `attention_backends()`, and a dict that merely shares the name (`server_args_dict.items`). That dict is the one read left behind. `build_attention_backends` also stops resolving the pair from the record: it runs after publish, so it asks `attention_backends()` like every other consumer. The draft override on the runner still wins first. `dispatch_event_loop`'s three PP checks read the *configured* PP size, not the live topology: the MLX runner stub never initializes torch.distributed, so the live property asserts before the MLX event loop can start (a Codex catch). The configured leaf answers the same value wherever the live groups exist. `flashinfer_gdn_prefill_default`'s guard is the one read here that asks what the *operator* named rather than what the config resolved to, and the bag leaf now answers exactly that: the per-runner auto-default is stamped on the runner and deliberately never recorded process-wide, so nothing writes that leaf after launch and reading it back cannot mistake another runner's default for a flag. Three test doubles injected a `SimpleNamespace`/`MagicMock` record for exactly these reads and now publish instead (pool configurator, cache registry, GDN prefill policy) -- the fixture publishes what the case configures and hands the published instance to the whole-object contracts that still take one. The functions this sweep partially converted stop mixing sources (review catches): the flash-attention constructor's remaining seed reads (`speculative_eagle_topk`, `speculative_algorithm`, both deterministic gates) read their bags next to the leaves already converted; `_should_disable_scheduler_metadata_precompute` reads the parallel config leaves itself instead of taking the record (its alias binding was the last use); and the autotune gates (`disable_flashinfer_autotune`, deterministic, `flashinfer_autotune_skip_ops`) join the moe leaves the same function already reads from the bags. The pool-configurator fixture drops a parameter nothing published or read. --- .../attention/flashattention_backend.py | 30 ++++--- .../layers/attention/linear/gdn_backend.py | 12 +-- python/sglang/srt/managers/scheduler.py | 10 +-- .../scheduler_components/metrics_reporter.py | 8 +- python/sglang/srt/mem_cache/registry.py | 9 ++- .../srt/model_executor/cpu_graph_runner.py | 9 +-- .../attention_backend_setup.py | 14 ++-- .../srt/model_executor/pool_configurator.py | 45 +++++++---- .../runner/base_cuda_graph_runner.py | 11 ++- .../srt/model_executor/runner/eager_runner.py | 25 +++--- .../runner/flashinfer_autotune.py | 33 ++++---- python/sglang/srt/runtime_context.py | 29 +++++++ python/sglang/srt/utils/common.py | 18 +++-- .../test_gdn_prefill_backend_policy.py | 32 ++++++-- .../unit/mem_cache/test_registry.py | 70 ++++++++++------ .../model_executor/test_pool_configurator.py | 79 +++++++++++++------ .../unit/test_global_config_read_ratchet.py | 6 ++ test/registered/unit/test_runtime_context.py | 26 ++++++ 18 files changed, 319 insertions(+), 147 deletions(-) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index b9dcd6746..3fe9aadc9 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -30,7 +30,14 @@ from sglang.srt.layers.utils.cp_utils import ( from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode -from sglang.srt.runtime_context import get_schedule, get_spec +from sglang.srt.runtime_context import ( + get_exec, + get_memory, + get_model, + get_parallel, + get_schedule, + get_spec, +) from sglang.srt.speculative.ragged_verify import build_ragged_target_verify_geometry from sglang.srt.speculative.spec_info import SpecInput, SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import resolve_num_tokens_per_req @@ -49,8 +56,8 @@ from sglang.kernels.ops.attention.flash_attention import ( ) -def _should_disable_scheduler_metadata_precompute(server_args) -> bool: - return bool(server_args.enable_prefill_cp or server_args.enable_dp_attention) +def _should_disable_scheduler_metadata_precompute() -> bool: + return bool(get_parallel().enable_prefill_cp or get_parallel().enable_dp_attention) @dataclass @@ -202,7 +209,7 @@ class FlashAttentionBackend(AttentionBackend): and model_runner.token_to_kv_pool.swa_layer_nums > 0 ) - self.topk = model_runner.server_args.speculative_eagle_topk or 0 + self.topk = get_spec().speculative_eagle_topk or 0 self.speculative_num_steps = speculative_num_steps self.speculative_num_draft_tokens = get_spec().speculative_num_draft_tokens if ( @@ -214,7 +221,7 @@ class FlashAttentionBackend(AttentionBackend): phase="target_verify", server_args=model_runner.server_args, spec_algorithm=SpeculativeAlgorithm.from_string( - model_runner.server_args.speculative_algorithm + get_spec().speculative_algorithm ), is_draft_worker=True, num_draft_tokens=int(self.speculative_num_draft_tokens), @@ -283,7 +290,7 @@ class FlashAttentionBackend(AttentionBackend): self._get_fa_runtime_policy = None self._get_scheduler_metadata = None - if model_runner.server_args.enable_deterministic_inference: + if get_exec().deterministic.enable_deterministic_inference: # Must precede the first kernel compile. from sglang.kernels.ops.attention.flash_attn.cute.batch_invariance import ( set_batch_invariant, @@ -311,7 +318,7 @@ class FlashAttentionBackend(AttentionBackend): self.has_softcap = _softcapping is not None and _softcapping > 0.0 # num_splits == 0 delegates SplitKV sizing to the selected FA runtime. - deterministic = model_runner.server_args.enable_deterministic_inference + deterministic = get_exec().deterministic.enable_deterministic_inference if self._get_fa_runtime_policy is None: self.num_splits = 1 if deterministic else 0 self.decode_num_splits = self.num_splits @@ -339,11 +346,10 @@ class FlashAttentionBackend(AttentionBackend): # guard wraps both set_kv_buffer and set_mla_kv_buffer. Without this # gate, MLA + is_embedding would skip the write but still read stale # cache via get_key_buffer in the absorbed-MLA path. - server_args = model_runner.server_args self.fa_skip_kv_cache = ( - server_args.is_embedding - and server_args.chunked_prefill_size == -1 - and server_args.disable_radix_cache + get_model().is_embedding + and get_schedule().chunked_prefill_size == -1 + and get_memory().disable_radix_cache and not self.use_mla ) @@ -353,7 +359,7 @@ class FlashAttentionBackend(AttentionBackend): # combine kernel (flash_fwd_combine_launch_template.h:52). Leaving # scheduler_metadata unset uses the existing per-layer metadata path. self._disable_scheduler_metadata_precompute = ( - _should_disable_scheduler_metadata_precompute(server_args) + _should_disable_scheduler_metadata_precompute() ) def _compute_scheduler_metadata( diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 18634c283..0bc8fbd56 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -18,6 +18,7 @@ from sglang.srt.layers.radix_linear_attention import RadixLinearAttention from sglang.srt.mem_cache.memory_pool import MambaPool 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_exec, get_memory, get_schedule from sglang.srt.utils import is_cpu, is_cuda, is_hip, is_npu, is_xpu from sglang.srt.utils.common import rank0_log @@ -64,23 +65,22 @@ elif is_cpu(): def flashinfer_gdn_prefill_default(model_runner: ModelRunner) -> Optional[str]: """FlashInfer for the narrow SM100 GDN prefill domain we validated, else None.""" - args = model_runner.server_args if ( - args.linear_attn_prefill_backend is not None - or args.linear_attn_backend != "triton" - or args.enable_page_major_kv_layout + get_exec().mamba.linear_attn_prefill_backend is not None + or get_exec().mamba.linear_attn_backend != "triton" + or get_memory().enable_page_major_kv_layout or not is_cuda() or torch.cuda.get_device_capability()[0] != 10 ): return None cuda_version = torch.version.cuda - chunk_size = args.chunked_prefill_size + chunk_size = get_schedule().chunked_prefill_size config = hybrid_gdn_config(model_runner.model_config) if ( cuda_version is None or int(cuda_version.split(".", 1)[0]) < 13 - or args.enable_dynamic_chunking + or get_schedule().enable_dynamic_chunking or chunk_size is None or not 1 <= chunk_size <= 8192 or getattr(config, "linear_key_head_dim", None) != 128 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index fbe429700..9a15cb975 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -29,6 +29,7 @@ from typing import TYPE_CHECKING, Any, Deque, Dict, List, Optional, Set, Tuple, from sglang.srt.runtime_context import ( attention_backends, + configured_pp_size, get_device, get_disagg, get_exec, @@ -4900,13 +4901,12 @@ class Scheduler( def dispatch_event_loop(scheduler: Scheduler): - # Dispatch to the appropriate event loop based on the disaggregation mode - server_args = scheduler.server_args + # The live PP property asserts before torch.distributed init (MLX stub). disaggregation_mode: DisaggregationMode = scheduler.disaggregation_mode if disaggregation_mode == DisaggregationMode.NULL: if scheduler.enable_pdmux: scheduler.event_loop_pdmux() - elif server_args.pp_size > 1: + elif configured_pp_size() > 1: scheduler.event_loop_pp() elif scheduler.enable_overlap_mlx: scheduler.event_loop_overlap_mlx() @@ -4915,14 +4915,14 @@ def dispatch_event_loop(scheduler: Scheduler): else: scheduler.event_loop_normal() elif disaggregation_mode == DisaggregationMode.PREFILL: - if server_args.pp_size > 1: + if configured_pp_size() > 1: scheduler.event_loop_pp_disagg_prefill() elif scheduler.enable_overlap: scheduler.event_loop_overlap_disagg_prefill() else: scheduler.event_loop_normal_disagg_prefill() elif disaggregation_mode == DisaggregationMode.DECODE: - if server_args.pp_size > 1: + if configured_pp_size() > 1: scheduler.event_loop_pp_disagg_decode() elif scheduler.enable_overlap: scheduler.event_loop_overlap_disagg_decode() diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index 9e1f28cef..73a49a66d 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -338,15 +338,15 @@ class SchedulerMetricsReporter: "num_draft_tokens": 0, } - # Fallback to server_args if draft_worker does not have the attributes. - server_args = self.scheduler.server_args + # Fallback to the published `spec` bag if draft_worker does not have + # the attributes. num_steps = getattr( - draft_worker, "speculative_num_steps", server_args.speculative_num_steps + draft_worker, "speculative_num_steps", get_spec().speculative_num_steps ) num_draft_tokens = getattr( draft_worker, "speculative_num_draft_tokens", - server_args.speculative_num_draft_tokens, + get_spec().speculative_num_draft_tokens, ) return { diff --git a/python/sglang/srt/mem_cache/registry.py b/python/sglang/srt/mem_cache/registry.py index 63eb13065..43386f2c7 100644 --- a/python/sglang/srt/mem_cache/registry.py +++ b/python/sglang/srt/mem_cache/registry.py @@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Callable, Optional from sglang.srt.environ import envs from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.mem_cache.cache_init_params import CacheInitParams +from sglang.srt.runtime_context import get_memory from sglang.srt.utils.tensor_bridge import use_mlx if TYPE_CHECKING: @@ -128,7 +129,7 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache: ) return cache - if server_args.enable_lmcache: + if get_memory().enable_lmcache: from sglang.srt.mem_cache.storage.lmcache.lmc_radix_cache import ( LMCRadixCache, ) @@ -141,7 +142,7 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache: tp_group=ctx.tp_group, ) - if server_args.enable_flexkv: + if get_memory().enable_flexkv: # Importing the package side-effect registers the explicit # ``--radix-cache-backend=flexkv`` factory; we then call the # factory directly so --enable-flexkv stands on its own. @@ -151,8 +152,8 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache: # Honor a CLI --flexkv-config-file by forwarding it via the env # var that FlexKV's config loader actually reads. - if server_args.flexkv_config_file and not os.environ.get("FLEXKV_CONFIG_PATH"): - os.environ["FLEXKV_CONFIG_PATH"] = server_args.flexkv_config_file + if get_memory().flexkv_config_file and not os.environ.get("FLEXKV_CONFIG_PATH"): + os.environ["FLEXKV_CONFIG_PATH"] = get_memory().flexkv_config_file return _flexkv_factory(ctx) from sglang.srt.mem_cache.radix_cache import RadixCache diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index eb066cdb3..30719aa26 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/python/sglang/srt/model_executor/cpu_graph_runner.py @@ -39,7 +39,7 @@ from sglang.srt.model_executor.forward_batch_info import ( ) from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.runner_utils.capture_mode import model_capture_mode -from sglang.srt.runtime_context import get_flags, get_parallel, get_spec +from sglang.srt.runtime_context import get_exec, get_flags, get_parallel, get_spec from sglang.srt.utils import ( empty_context, log_info_on_rank0, @@ -127,14 +127,13 @@ def set_torch_compile_config(): def get_batch_sizes_to_capture(model_runner: ModelRunner): # torch compile speeds up decoding by reducing python overhead on CPU - server_args = model_runner.server_args # Reuse cuda_graph_config[decode].bs here. # Users can customize the batch sizes supported by cpu_graph, such as: # --cuda-graph-bs-decode 1 2 4 8 16 - capture_bs = server_args.cuda_graph_config.decode.bs + capture_bs = get_exec().graph.cuda_graph_config.decode.bs assert ( - max(capture_bs) <= server_args.torch_compile_max_bs - ), f"{capture_bs=}, {server_args.torch_compile_max_bs=}" + max(capture_bs) <= get_exec().graph.torch_compile_max_bs + ), f"{capture_bs=}, {get_exec().graph.torch_compile_max_bs=}" capture_bs = [bs for bs in capture_bs if bs <= model_runner.req_to_token_pool.size] capture_bs = list(sorted(set(capture_bs))) assert len(capture_bs) > 0 and capture_bs[0] > 0, f"{capture_bs=}" diff --git a/python/sglang/srt/model_executor/model_runner_components/attention_backend_setup.py b/python/sglang/srt/model_executor/model_runner_components/attention_backend_setup.py index 6d0a8af30..e60546091 100644 --- a/python/sglang/srt/model_executor/model_runner_components/attention_backend_setup.py +++ b/python/sglang/srt/model_executor/model_runner_components/attention_backend_setup.py @@ -18,6 +18,8 @@ if TYPE_CHECKING: from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.runtime_context import attention_backends, get_disagg, get_exec + logger = logging.getLogger(__name__) @@ -66,7 +68,6 @@ def configure_aux_hidden_state_capture( def build_attention_backends(*, model_runner: ModelRunner) -> AttentionBackends: """Init attention kernel backend.""" - server_args = model_runner.server_args # TODO: Refactor device-specific init branches into platform interface (separate PR). if model_runner.device in ("cuda", "musa"): @@ -81,7 +82,7 @@ def build_attention_backends(*, model_runner: ModelRunner) -> AttentionBackends: ), ) - if server_args.enable_pdmux: + if get_disagg().enable_pdmux: attn_backend = _build_resolved_backend( model_runner=model_runner, resolved=resolved, init_new_workspace=True ) @@ -91,10 +92,12 @@ def build_attention_backends(*, model_runner: ModelRunner) -> AttentionBackends: resolved=resolved, init_new_workspace=False, ) - for _ in range(server_args.sm_group_num) + for _ in range(get_disagg().sm_group_num) ] decode_attn_backend = decode_attn_backend_group[0] - elif server_args.enable_two_batch_overlap and not model_runner.is_draft_worker: + elif ( + get_exec().overlap.enable_two_batch_overlap and not model_runner.is_draft_worker + ): attn_backend = TboAttnBackend.init_new( lambda: _build_resolved_backend( model_runner=model_runner, @@ -161,7 +164,6 @@ def resolve_attention_backend_strs( target and draft coexist in one process, so it cannot come from the process-wide config. """ - server_args = model_runner.server_args is_draft_worker = model_runner.is_draft_worker draft_attn_backend = model_runner.draft_attention_backend if is_draft_worker and draft_attn_backend: @@ -172,7 +174,7 @@ def resolve_attention_backend_strs( decode=draft_attn_backend, is_draft_override=True, ) - prefill, decode = server_args.get_attention_backends() + prefill, decode = attention_backends() return ResolvedAttentionBackendStr(prefill=prefill, decode=decode) diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 9ba6ecc10..44d7f5df4 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -36,7 +36,14 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( get_compress_state_write_pad, ) from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import ( + get_disagg, + get_memory, + get_parallel, + get_schedule, + get_spec, + max_speculative_num_draft_tokens, +) from sglang.srt.utils.common import ( ceil_align, ceil_div, @@ -533,10 +540,9 @@ class SWAChunkCapPoolConfigurator(HybridSWAPoolConfigurator): super().__init__(kvc) assert self._full_layers_num > 0 - sa = kvc.server_args page_size = kvc.page_size window = kvc.sliding_window_size - draft_tokens = sa.speculative_num_draft_tokens or 1 + draft_tokens = get_spec().speculative_num_draft_tokens or 1 eviction_interval = max(1, envs.SGLANG_SWA_EVICTION_INTERVAL.get()) """ @@ -544,40 +550,47 @@ class SWAChunkCapPoolConfigurator(HybridSWAPoolConfigurator): Padding to make sure eviction point is page-aligned. """ trailing_tokens = window + eviction_interval * draft_tokens + page_size - if sa.speculative_algorithm is None: + if get_spec().speculative_algorithm is None: decode_alloc = page_size - elif sa.disable_overlap_schedule: + elif get_schedule().disable_overlap_schedule: # spec-v1: new_tokens_required_next_decode per request. - decode_alloc = spec_decode_alloc_len_per_request(sa) + decode_alloc = spec_decode_alloc_len_per_request( + page_size=page_size, + speculative_num_steps=get_spec().speculative_num_steps, + speculative_eagle_topk=get_spec().speculative_eagle_topk, + speculative_num_draft_tokens=get_spec().speculative_num_draft_tokens, + ) else: # spec-v2: the overlap allocator keeps 2 * alloc_len outstanding # (eagle_utils.eagle_prepare_for_decode: kv_committed_len + 2 * alloc_len). - decode_alloc = 2 * get_alloc_len_per_decode(sa) + decode_alloc = 2 * get_alloc_len_per_decode( + kvc.server_args, + max_draft_tokens=max_speculative_num_draft_tokens(), + ) per_request = trailing_tokens + decode_alloc - num_reqs = sa.max_running_requests // kvc.ps.attn_dp_size - if sa.disaggregation_mode == "decode": + num_reqs = get_schedule().max_running_requests // kvc.ps.attn_dp_size + if get_disagg().disaggregation_mode == "decode": self._swa_cap = ( per_request * num_reqs - + (window + page_size) * sa.disaggregation_decode_extra_slots + + (window + page_size) * get_disagg().disaggregation_decode_extra_slots ) else: - chunks_in_flight = 1 if sa.disable_overlap_schedule else 2 + chunks_in_flight = 1 if get_schedule().disable_overlap_schedule else 2 self._swa_cap = ( per_request * num_reqs - + chunks_in_flight * sa.chunked_prefill_size + + chunks_in_flight * get_schedule().chunked_prefill_size + page_size ) @staticmethod def is_applicable(kvc: KVCacheConfigurator) -> bool: """True when SWAChunkCache can be sized from explicit max requests.""" - sa = kvc.server_args - if sa.max_running_requests is None: + if get_schedule().max_running_requests is None: return False - if not sa.disable_radix_cache: + if not get_memory().disable_radix_cache: return False - if sa.chunked_prefill_size is None: + if get_schedule().chunked_prefill_size is None: return False if kvc.sliding_window_size is None: return False diff --git a/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py index 1ed5b92a1..f58717e19 100644 --- a/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py @@ -23,7 +23,10 @@ from contextlib import contextmanager from typing import TYPE_CHECKING, Any, List, Sequence, Tuple from sglang.srt.model_executor.runner.base_runner import BaseRunner -from sglang.srt.runtime_context import get_flags +from sglang.srt.runtime_context import ( + get_exec, + get_flags, +) from sglang.srt.utils import ( get_cuda_graph_batch_size_alignment, get_cuda_graph_max_batch_size, @@ -68,14 +71,14 @@ def get_batch_sizes_to_capture( """ server_args = model_runner.server_args - capture_bs = list(server_args.cuda_graph_config.decode.bs) + capture_bs = list(get_exec().graph.cuda_graph_config.decode.bs) num_max_requests = model_runner.req_to_token_pool.size mul_base = get_cuda_graph_batch_size_alignment(server_args) # TBO splits each request's rows across two micro-batches, so the # alignment constraint applies per request rather than per token row. alignment_width = captured_req_width - if server_args.enable_two_batch_overlap: + if get_exec().overlap.enable_two_batch_overlap: alignment_width = 1 # pad `num_max_requests` to avoid being filtered out @@ -92,7 +95,7 @@ def get_batch_sizes_to_capture( assert len(capture_bs) > 0 and capture_bs[0] > 0, f"{capture_bs=}" compile_bs = ( - [bs for bs in capture_bs if bs <= server_args.torch_compile_max_bs] + [bs for bs in capture_bs if bs <= get_exec().graph.torch_compile_max_bs] if get_flags().capture.enable_torch_compile else [] ) diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index 643ff7c5c..cf6d3703c 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -49,6 +49,13 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo enable_tc_piecewise_cuda_graph, set_tc_piecewise_forward_context, ) +from sglang.srt.runtime_context import ( + get_parallel, + get_spec, + mamba_extra_buffer_enabled, + max_prefill_buffer_tokens, + max_speculative_num_draft_tokens, +) from sglang.srt.utils import is_hip from sglang.srt.utils.common import ( ceil_align, @@ -76,14 +83,14 @@ class EagerRunner(BaseRunner): num_tokens_per_req = 1 if mr.spec_algorithm.is_speculative(): # speculative_adaptive can grow draft tokens at runtime; size to the max. - num_draft_tokens = sa.max_speculative_num_draft_tokens or 1 + num_draft_tokens = max_speculative_num_draft_tokens() or 1 if mr.is_draft_worker: num_tokens_per_req = max( - sa.speculative_eagle_topk or 1, + get_spec().speculative_eagle_topk or 1, num_draft_tokens, ( - 2 * (sa.speculative_num_steps or 0) - if sa.enable_multi_layer_eagle + 2 * (get_spec().speculative_num_steps or 0) + if get_spec().enable_multi_layer_eagle else 0 ), ) @@ -100,14 +107,14 @@ class EagerRunner(BaseRunner): if ( mr.is_draft_worker and mr.spec_algorithm.is_frozen_kv_mtp() - and sa.speculative_eagle_topk > 1 + and get_spec().speculative_eagle_topk > 1 ): # Frozen-KV MTP expands the draft batch by topk on the bs axis # (expand_for_topk_draft) before the eager fallback. - max_bs *= sa.speculative_eagle_topk + max_bs *= get_spec().speculative_eagle_topk # Mirror prepare_mlp_sync_batch padding so the registry holds what load_batch copies. max_bs = get_eager_max_batch_size(sa, max_bs) - prefill_ceiling = max(mr.max_total_num_tokens, sa.max_prefill_buffer_tokens()) + prefill_ceiling = max(mr.max_total_num_tokens, max_prefill_buffer_tokens()) max_num_token = max(prefill_ceiling, max_bs * num_tokens_per_req) if require_mlp_sync(sa): from sglang.srt.layers.cp.padding import get_cp_padding_align_size @@ -123,7 +130,7 @@ class EagerRunner(BaseRunner): max_num_token=max_num_token, cache_loc_dtype=torch.int64, enable_mamba_track=( - sa.enable_mamba_extra_buffer() and mr.spec_algorithm.is_none() + mamba_extra_buffer_enabled() and mr.spec_algorithm.is_none() ), is_encoder_decoder=is_encoder_decoder, encoder_len_fill_value=( @@ -134,7 +141,7 @@ class EagerRunner(BaseRunner): encoder_lens_dtype=( torch.int64 if torch.device(mr.device).type == "cpu" else torch.int32 ), - dp_size=sa.dp_size, + dp_size=get_parallel().dp_size, ) # Eager has no capture step, so warm up here (run-once via mr._kernel_warmed_up). self.warmup() diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index 34ef49ea2..fea98b385 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py @@ -25,6 +25,11 @@ import torch from sglang.srt.environ import envs from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.runtime_context import ( + get_exec, + get_model, + get_spec, +) from sglang.srt.utils import empty_context, log_info_on_rank0 if TYPE_CHECKING: @@ -37,7 +42,7 @@ FLASHINFER_AUTOTUNE_WORKAROUND_SKIPS = frozenset() def get_flashinfer_autotune_skip_ops(model_runner: ModelRunner) -> set[str]: - skip_ops = set(model_runner.server_args.flashinfer_autotune_skip_ops or ()) + skip_ops = set(get_exec().kernel.flashinfer_autotune_skip_ops or ()) skip_ops.update(FLASHINFER_AUTOTUNE_WORKAROUND_SKIPS) return skip_ops @@ -49,28 +54,29 @@ def should_run_flashinfer_autotune( mr = model_runner if mr.device != "cuda": return False - if mr.server_args.disable_flashinfer_autotune: + if get_exec().kernel.disable_flashinfer_autotune: return False - if mr.server_args.enable_deterministic_inference: + if get_exec().deterministic.enable_deterministic_inference: # Tuned configs are per problem shape, so the reduction order would follow # the batch shape. return False - server_args = mr.server_args if for_speculative_draft: backend_str = ( - server_args.speculative_moe_runner_backend or server_args.moe_runner_backend + get_spec().speculative_moe_runner_backend + or get_exec().moe.moe_runner_backend ) a2a_backend_str = ( - server_args.speculative_moe_a2a_backend or server_args.moe_a2a_backend + get_spec().speculative_moe_a2a_backend or get_exec().moe.moe_a2a_backend ) else: - backend_str = server_args.moe_runner_backend - a2a_backend_str = server_args.moe_a2a_backend + backend_str = get_exec().moe.moe_runner_backend + a2a_backend_str = get_exec().moe.moe_a2a_backend # Autotune can run before the MoE backend globals are initialized, so read - # the target or draft backend from server_args. CuteDSL v1 bypasses - # MoeRunner, and its dummy dispatch can exceed DeepEP low-latency's token limit. + # the configured backends -- the draft leaves (`get_spec()`) or the target + # leaves (`get_exec().moe`) above. CuteDSL v1 bypasses MoeRunner, and its + # dummy dispatch can exceed DeepEP low-latency's token limit. if backend_str == "flashinfer_cutedsl" and a2a_backend_str == "deepep": return False @@ -130,12 +136,11 @@ def flashinfer_autotune_cache_path(model_runner: ModelRunner) -> Path: arch = f"sm{major}{minor}" flashinfer_version = getattr(flashinfer, "__version__", "unknown") - server_args = mr.server_args model_key_parts = [ - str(server_args.model_path), + str(get_model().model_path), str(mr.dtype), - str(server_args.quantization), - str(server_args.moe_runner_backend), + str(get_model().quantization), + str(get_exec().moe.moe_runner_backend), str(mr.ps.tp_size), str(mr.ps.pp_size), str(mr.ps.attn_dp_size), diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index a71ebff22..01940988e 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -1434,6 +1434,35 @@ def remote_instance_transfer_engine_enabled(load_format: str | None = None) -> b return remote_instance_transfer_engine_of(get_model(), load_format) +def max_prefill_buffer_tokens() -> int: + """The prefill-buffer ceiling: ``chunked_prefill_size``, except PP dynamic + chunking can grow chunks toward ``max_prefill_tokens`` and probe at 1.25x. + + Every input is a published leaf (``schedule`` plus the configured PP size), + so this derives from the bags and follows a post-publish override; + ``ServerArgs.max_prefill_buffer_tokens`` is the pre-publish equivalent and + ``TestDerivedPredicatesAgreeAcrossTiers`` pins the two equal. + """ + import math + + schedule = get_schedule() + chunked = ( + schedule.chunked_prefill_size + if schedule.chunked_prefill_size and schedule.chunked_prefill_size > 0 + else 0 + ) + tokens = chunked + if ( + schedule.enable_dynamic_chunking + and _configured_parallel("pp_size") > 1 + and chunked + ): + tokens = max( + tokens, schedule.max_prefill_tokens or 0, math.ceil(chunked * 1.25) + ) + return tokens + + def pre_capture_activation_reserve_mb(gpu_mem: float | None) -> float: """The activation working-set reserve held back before cuda-graph capture. diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 811b5023e..efdd0097f 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -3903,14 +3903,20 @@ def ceil_align(x: int, y: int) -> int: return ceil_div(x, y) * y -def spec_decode_alloc_len_per_request(server_args) -> int: +def spec_decode_alloc_len_per_request( + *, + page_size, + speculative_num_steps, + speculative_eagle_topk, + speculative_num_draft_tokens, +) -> int: """Per-request KV tokens a (spec-v1) decode step allocates: the draft-decode - topk*num_steps peak vs. the verify num_draft_tokens, page-aligned. + topk*num_steps peak vs. the verify num_draft_tokens, page-aligned. A pure + function of the resolved values its one caller reads off the bags. """ - page_size = server_args.page_size - len_per_topk = server_args.speculative_num_steps or 1 - spec_topk = server_args.speculative_eagle_topk or 1 - spec_tokens = server_args.speculative_num_draft_tokens or 1 + len_per_topk = speculative_num_steps or 1 + spec_topk = speculative_eagle_topk or 1 + spec_tokens = speculative_num_draft_tokens or 1 if page_size > 1 and spec_topk > 1: # last partial page and ceil alignment diff --git a/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py b/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py index 4f5c6314b..96d44a70e 100644 --- a/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py +++ b/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py @@ -24,14 +24,29 @@ from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=5, suite="base-a-test-cpu") +def _publish(testcase, **fields): + """Install a published config for one case and restore on its cleanup -- + a failed case must not leave a partial publish for a later file in a + monolithic local run.""" + from sglang.srt.runtime_context import get_context, get_server_args + + override = get_context().override_server_args(**fields) + override.install() + testcase.addCleanup(override.restore) + return get_server_args() + + def make_runner( + testcase, *, state_dtype=torch.bfloat16, key_dim=128, value_dim=128, **arg_overrides, ): - args = SimpleNamespace( + # The policy reads the published bags, so the fixture publishes the + # configuration under test. + fields = dict( linear_attn_backend="triton", linear_attn_prefill_backend=None, uses_mamba_radix_cache=False, @@ -40,8 +55,8 @@ def make_runner( enable_dynamic_chunking=False, chunked_prefill_size=8192, ) - for name, value in arg_overrides.items(): - setattr(args, name, value) + fields.update(arg_overrides) + args = _publish(testcase, **fields) return SimpleNamespace( server_args=args, @@ -86,12 +101,13 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase): return flashinfer_gdn_prefill_default(runner) def test_selects_flashinfer_for_supported_sm100_gdn(self): - self.assertEqual(self.apply_policy(make_runner()), "flashinfer") + self.assertEqual(self.apply_policy(make_runner(self)), "flashinfer") def test_selects_flashinfer_for_radix_cache_strategies(self): for strategy in ("no_buffer", "extra_buffer", "extra_buffer_lazy"): with self.subTest(strategy=strategy): runner = make_runner( + self, uses_mamba_radix_cache=True, mamba_radix_cache_strategy=strategy, ) @@ -100,7 +116,7 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase): def test_declines_when_the_prefill_backend_is_explicit(self): for backend in ("triton", "flashinfer", "cutedsl"): with self.subTest(backend=backend): - runner = make_runner(linear_attn_prefill_backend=backend) + runner = make_runner(self, linear_attn_prefill_backend=backend) self.assertIsNone(self.apply_policy(runner)) def test_rejects_unsupported_capability(self): @@ -117,11 +133,11 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase): for name, runner_args, hardware in cases: with self.subTest(name=name): self.assertIsNone( - self.apply_policy(make_runner(**runner_args), **hardware) + self.apply_policy(make_runner(self, **runner_args), **hardware) ) def test_rejects_gdn_config_without_qwen_head_dims(self): - runner = make_runner() + runner = make_runner(self) runner.hybrid_gdn_config = SimpleNamespace() self.assertIsNone(self.apply_policy(runner)) @@ -136,7 +152,7 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase): ) for name, runner_args in cases: with self.subTest(name=name): - self.assertIsNone(self.apply_policy(make_runner(**runner_args))) + self.assertIsNone(self.apply_policy(make_runner(self, **runner_args))) def test_builds_compact_checkpoint_plan_for_packed_sequences(self): forward_batch = SimpleNamespace( diff --git a/test/registered/unit/mem_cache/test_registry.py b/test/registered/unit/mem_cache/test_registry.py index adaa68840..9e5cc31d7 100644 --- a/test/registered/unit/mem_cache/test_registry.py +++ b/test/registered/unit/mem_cache/test_registry.py @@ -19,7 +19,18 @@ from sglang.srt.mem_cache.registry import ( from sglang.test.test_utils import CustomTestCase +def _publish(testcase, **fields): + """Install a published config for one case and restore on its cleanup.""" + from sglang.srt.runtime_context import get_context, get_server_args + + override = get_context().override_server_args(**fields) + override.install() + testcase.addCleanup(override.restore) + return get_server_args() + + def _make_ctx( + testcase, *, backend=None, enable_streaming=False, @@ -32,11 +43,16 @@ def _make_ctx( effective_chunked_prefill_size=None, full_tokens_per_layer=None, ): - server_args = MagicMock() - server_args.radix_cache_backend = backend - server_args.enable_streaming_session = enable_streaming - server_args.enable_lmcache = enable_lmcache - server_args.enable_flexkv = False + # The factory reads the published bags for the cache-backend leaves, so the + # fixture publishes them; the instance stays for the whole-object contract + # `TreeCacheBuildContext` carries. + server_args = _publish( + testcase, + radix_cache_backend=backend, + enable_streaming_session=enable_streaming, + enable_lmcache=enable_lmcache, + enable_flexkv=False, + ) return TreeCacheBuildContext( server_args=server_args, params=MagicMock(), @@ -101,14 +117,14 @@ class TestCreateTreeCacheRouting(_RegistryIsolationMixin, CustomTestCase): factory = MagicMock(return_value=cache) register_radix_cache_backend("custom", factory) - result = create_tree_cache(_make_ctx(backend="custom")) + result = create_tree_cache(_make_ctx(self, backend="custom")) factory.assert_called_once() self.assertIs(result, cache) def test_unknown_backend_raises(self): with self.assertRaises(ValueError): - create_tree_cache(_make_ctx(backend="not_a_real_backend")) + create_tree_cache(_make_ctx(self, backend="not_a_real_backend")) @patch("sglang.srt.mem_cache.registry.default_radix_cache_factory") def test_unset_backend_falls_back_to_default(self, default_factory): @@ -116,7 +132,7 @@ class TestCreateTreeCacheRouting(_RegistryIsolationMixin, CustomTestCase): cache.supports_streaming_session.return_value = True default_factory.return_value = cache - result = create_tree_cache(_make_ctx(backend=None)) + result = create_tree_cache(_make_ctx(self, backend=None)) default_factory.assert_called_once() self.assertIs(result, cache) @@ -131,7 +147,7 @@ class TestCreateTreeCacheRouting(_RegistryIsolationMixin, CustomTestCase): ) as session_cls: session_cls.return_value = MagicMock(name="wrapped") result = create_tree_cache( - _make_ctx(backend="nonstreaming", enable_streaming=True) + _make_ctx(self, backend="nonstreaming", enable_streaming=True) ) session_cls.assert_called_once_with(inner) @@ -143,7 +159,7 @@ class TestCreateTreeCacheRouting(_RegistryIsolationMixin, CustomTestCase): register_radix_cache_backend("streaming", MagicMock(return_value=inner)) result = create_tree_cache( - _make_ctx(backend="streaming", enable_streaming=True) + _make_ctx(self, backend="streaming", enable_streaming=True) ) self.assertIs(result, inner) @@ -158,7 +174,9 @@ class TestDefaultRadixCacheFactory(CustomTestCase): """ def test_chunk_cache_when_chunked_prefill_and_disable_radix(self): - ctx = _make_ctx(effective_chunked_prefill_size=512, disable_radix_cache=True) + ctx = _make_ctx( + self, effective_chunked_prefill_size=512, disable_radix_cache=True + ) with patch("sglang.srt.mem_cache.chunk_cache.ChunkCache") as ChunkCache: ChunkCache.return_value = MagicMock() result = default_radix_cache_factory(ctx) @@ -167,6 +185,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase): def test_swa_chunk_cache_when_chunked_prefill_disable_and_hybrid_swa(self): ctx = _make_ctx( + self, effective_chunked_prefill_size=512, disable_radix_cache=True, is_hybrid_swa=True, @@ -179,6 +198,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase): def test_pure_swa_chunk_cache_when_chunked_prefill_disable_and_all_swa(self): ctx = _make_ctx( + self, effective_chunked_prefill_size=512, disable_radix_cache=True, is_hybrid_swa=True, @@ -193,7 +213,9 @@ class TestDefaultRadixCacheFactory(CustomTestCase): self.assertIs(result, PureSWAChunkCache.return_value) def test_cpp_radix_cache_when_env_flag_set(self): - ctx = _make_ctx() + ctx = _make_ctx( + self, + ) # `radix_cache_cpp` requires ninja + C++ extension to import, so # we inject a stand-in module rather than letting patch() trigger # the real import. @@ -215,7 +237,9 @@ class TestDefaultRadixCacheFactory(CustomTestCase): self.assertIs(result, fake_module.RadixCacheCpp.return_value) def test_unified_radix_cache_when_env_flag_set(self): - ctx = _make_ctx() + ctx = _make_ctx( + self, + ) # Shim both factory imports — each transitively loads sgl_kernel. fake_components = MagicMock() fake_radix = MagicMock() @@ -237,7 +261,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase): self.assertIs(result, fake_radix.UnifiedRadixCache.return_value) def test_hi_radix_cache_when_hierarchical(self): - ctx = _make_ctx(enable_hierarchical_cache=True) + ctx = _make_ctx(self, enable_hierarchical_cache=True) # `hiradix_cache` imports `hicache_storage` and # `memory_pool_host`, both of which transitively load # `sgl_kernel`; inject a stand-in module. @@ -254,7 +278,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase): self.assertIs(result, fake_module.HiRadixCache.return_value) def test_unified_radix_cache_when_hierarchical_and_hybrid_ssm(self): - ctx = _make_ctx(enable_hierarchical_cache=True, is_hybrid_ssm=True) + ctx = _make_ctx(self, enable_hierarchical_cache=True, is_hybrid_ssm=True) # Hybrid SSM with hierarchical cache now uses UnifiedRadixCache. fake_components = MagicMock() fake_radix = MagicMock() @@ -274,7 +298,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase): self.assertIs(result, fake_radix.UnifiedRadixCache.return_value) def test_unified_radix_cache_when_hierarchical_and_hybrid_swa(self): - ctx = _make_ctx(enable_hierarchical_cache=True, is_hybrid_swa=True) + ctx = _make_ctx(self, enable_hierarchical_cache=True, is_hybrid_swa=True) # Hybrid SWA with hierarchical cache also uses UnifiedRadixCache. fake_components = MagicMock() fake_radix = MagicMock() @@ -294,7 +318,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase): self.assertIs(result, fake_radix.UnifiedRadixCache.return_value) def test_unified_radix_cache_when_hierarchical_and_dsa(self): - ctx = _make_ctx(enable_hierarchical_cache=True, is_dsa=True) + ctx = _make_ctx(self, enable_hierarchical_cache=True, is_dsa=True) # DSA models (e.g. DeepSeek V3.2 / GLM-5.1) with hierarchical cache # use UnifiedRadixCache. fake_components = MagicMock() @@ -315,7 +339,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase): self.assertIs(result, fake_radix.UnifiedRadixCache.return_value) def test_swa_radix_cache_when_hybrid_swa(self): - ctx = _make_ctx(is_hybrid_swa=True) + ctx = _make_ctx(self, is_hybrid_swa=True) # SWA hybrid models now default to the unified radix tree. fake_components = MagicMock() fake_radix = MagicMock() @@ -331,7 +355,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase): self.assertIs(result, fake_radix.UnifiedRadixCache.return_value) def test_pure_swa_radix_cache_when_all_swa(self): - ctx = _make_ctx(is_hybrid_swa=True, full_tokens_per_layer=0) + ctx = _make_ctx(self, is_hybrid_swa=True, full_tokens_per_layer=0) with patch( "sglang.srt.mem_cache.pure_swa_radix_cache.PureSWARadixCache" ) as PureSWA: @@ -341,7 +365,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase): self.assertIs(result, PureSWA.return_value) def test_mamba_radix_cache_when_hybrid_ssm(self): - ctx = _make_ctx(is_hybrid_ssm=True) + ctx = _make_ctx(self, is_hybrid_ssm=True) # Mamba hybrid models now default to the unified radix tree. fake_components = MagicMock() fake_radix = MagicMock() @@ -357,7 +381,7 @@ class TestDefaultRadixCacheFactory(CustomTestCase): self.assertIs(result, fake_radix.UnifiedRadixCache.return_value) def test_lmc_radix_cache_when_enable_lmcache(self): - ctx = _make_ctx(enable_lmcache=True) + ctx = _make_ctx(self, enable_lmcache=True) # The lmcache backend raises at import time when the `lmcache` # package isn't installed, so inject a stand-in module instead # of letting patch() trigger the real import. @@ -377,7 +401,9 @@ class TestDefaultRadixCacheFactory(CustomTestCase): self.assertIs(result, fake_module.LMCRadixCache.return_value) def test_fallback_to_radix_cache(self): - ctx = _make_ctx() + ctx = _make_ctx( + self, + ) with patch("sglang.srt.mem_cache.radix_cache.RadixCache") as RadixCache: RadixCache.return_value = MagicMock() result = default_radix_cache_factory(ctx) diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 9d7df2df0..3d254adbb 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -11,7 +11,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock, patch from sglang.srt.distributed.parallel_state_wrapper import ParallelState -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_parallel, get_server_args from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="base-a-test-cpu") @@ -35,7 +35,21 @@ def mock_cpu_env(kv_size=2, tp_size=1, swa_eviction_interval=4): yield +def _publish_config(testcase, **fields): + """Publish the configuration a runner double describes. + + Installed per call and restored on the calling case's cleanup, so a failed + case cannot leave a partial publish for a later file in a monolithic run. + """ + from sglang.srt.runtime_context import get_context + + override = get_context().override_server_args(**fields) + override.install() + testcase.addCleanup(override.restore) + + def _make_model_runner( + testcase, *, num_kv_heads=4, head_dim=64, @@ -56,7 +70,6 @@ def _make_model_runner( disable_overlap_schedule=False, sliding_window_size=None, speculative_num_draft_tokens=None, - max_speculative_num_draft_tokens=None, speculative_algorithm=None, speculative_num_steps=None, speculative_eagle_topk=None, @@ -105,27 +118,29 @@ def _make_model_runner( mr.model_config = mc mr.kv_cache_dtype = "fake_bf16" - sa = SimpleNamespace() - sa.max_total_tokens = None - sa.swa_full_tokens_ratio = swa_full_tokens_ratio - sa.page_size = page_size - sa.disable_radix_cache = disable_radix_cache - sa.chunked_prefill_size = chunked_prefill_size - sa.disable_overlap_schedule = disable_overlap_schedule - sa.speculative_num_draft_tokens = speculative_num_draft_tokens - sa.max_speculative_num_draft_tokens = ( - max_speculative_num_draft_tokens or speculative_num_draft_tokens + # The configurator reads the published bags, so the fixture publishes the + # configuration it describes. The instance stays for the whole-object + # hand-offs the configurator still does. + _publish_config( + testcase, + max_total_tokens=None, + swa_full_tokens_ratio=swa_full_tokens_ratio, + page_size=page_size, + disable_radix_cache=disable_radix_cache, + chunked_prefill_size=chunked_prefill_size, + disable_overlap_schedule=disable_overlap_schedule, + speculative_num_draft_tokens=speculative_num_draft_tokens, + speculative_algorithm=speculative_algorithm, + speculative_num_steps=speculative_num_steps, + speculative_eagle_topk=speculative_eagle_topk, + disaggregation_mode=disaggregation_mode, + max_running_requests=max_running_requests, + disaggregation_decode_extra_slots=disaggregation_decode_extra_slots, + enable_hisparse=False, + enable_dsa_cache_layer_split=False, + kv_cache_dtype="auto", ) - sa.speculative_algorithm = speculative_algorithm - sa.speculative_num_steps = speculative_num_steps - sa.speculative_eagle_topk = speculative_eagle_topk - sa.disaggregation_mode = disaggregation_mode - sa.max_running_requests = max_running_requests - sa.disaggregation_decode_extra_slots = disaggregation_decode_extra_slots - sa.enable_hisparse = False - sa.enable_dsa_cache_layer_split = False - sa.kv_cache_dtype = "auto" - mr.server_args = sa + mr.server_args = get_server_args() spec = MagicMock() spec.is_eagle.return_value = False @@ -180,7 +195,7 @@ class TestDefaultConfigurator(unittest.TestCase): """Default (MHA): available_bytes -> tokens, memory invariant holds.""" def _run(self, available_bytes, page_size=1, **kwargs): - mr = _make_model_runner(page_size=page_size, **kwargs) + mr = _make_model_runner(self, page_size=page_size, **kwargs) with mock_cpu_env(): from sglang.srt.model_executor.pool_configurator import ( create_memory_pool_configurator, @@ -238,10 +253,12 @@ class TestDefaultConfigurator(unittest.TestCase): ): num_layers = 2 raw = _make_model_runner( + self, num_layers=num_layers, use_mla_backend=True, ) packed = _make_model_runner( + self, num_layers=num_layers, use_mla_backend=True, ) @@ -265,6 +282,7 @@ class TestHybridSWAConfigurator(unittest.TestCase): def _make_swa_runner(self, full_layers=16, swa_layers=16, ratio=0.5, page_size=1): return _make_model_runner( + self, is_hybrid_swa=True, full_attention_layer_ids=list(range(full_layers)), swa_attention_layer_ids=list(range(full_layers, full_layers + swa_layers)), @@ -353,6 +371,7 @@ class TestHybridSWAConfigurator(unittest.TestCase): def test_chunk_cache_cap_accounts_for_spec_topk_page_rounding(self): available = 1_000_000 mr = _make_model_runner( + self, is_hybrid_swa=True, full_attention_layer_ids=[0], swa_attention_layer_ids=[1], @@ -392,6 +411,7 @@ class TestHybridSWAConfigurator(unittest.TestCase): # 2*chunk(4) + page(1) = 9; cap = 41 * 2 + 9 = 91. available = 1_000_000 mr = _make_model_runner( + self, is_hybrid_swa=True, full_attention_layer_ids=[0], swa_attention_layer_ids=[1], @@ -422,6 +442,7 @@ class TestHybridSWAConfigurator(unittest.TestCase): def test_chunk_cache_cap_drops_prefill_for_disagg_decode(self): available = 1_000_000 mr = _make_model_runner( + self, is_hybrid_swa=True, full_attention_layer_ids=[0], swa_attention_layer_ids=[1], @@ -452,6 +473,7 @@ class TestHybridSWAConfigurator(unittest.TestCase): # under overlap. available = 1_000_000 mr = _make_model_runner( + self, is_hybrid_swa=True, full_attention_layer_ids=[0], swa_attention_layer_ids=[1], @@ -484,6 +506,7 @@ class TestHybridSWAConfigurator(unittest.TestCase): # request count (num_reserved_decode_tokens is a full-pool concern, not SWA). available = 2_000_000 mr = _make_model_runner( + self, is_hybrid_swa=True, full_attention_layer_ids=[0], swa_attention_layer_ids=[1], @@ -517,6 +540,7 @@ class TestAllSWAConfigurator(unittest.TestCase): def _run(self, available_bytes, ratio=0.5, page_size=1, **kwargs): mr = _make_model_runner( + self, is_hybrid_swa=True, full_attention_layer_ids=[], swa_attention_layer_ids=list(range(32)), @@ -569,7 +593,7 @@ class TestEagleConfigurator(unittest.TestCase): num_layers = 32 eagle_draft_num_layers = 4 - mr = _make_model_runner(num_layers=num_layers) + mr = _make_model_runner(self, num_layers=num_layers) mr.spec_algorithm.is_eagle.return_value = True mr.spec_algorithm.is_standalone.return_value = False mr.spec_algorithm.is_none.return_value = False @@ -591,7 +615,7 @@ class TestEagleConfigurator(unittest.TestCase): class TestFactory(unittest.TestCase): def test_default_for_non_swa(self): - mr = _make_model_runner(is_hybrid_swa=False) + mr = _make_model_runner(self, is_hybrid_swa=False) with mock_cpu_env(): from sglang.srt.model_executor.pool_configurator import ( DefaultPoolConfigurator, @@ -603,6 +627,7 @@ class TestFactory(unittest.TestCase): def test_swa_for_hybrid(self): mr = _make_model_runner( + self, is_hybrid_swa=True, full_attention_layer_ids=list(range(16)), swa_attention_layer_ids=list(range(16, 32)), @@ -621,6 +646,7 @@ class TestFactory(unittest.TestCase): # SWAChunkCapPoolConfigurator is selected only when max_running_requests is set. def _cfg(max_running_requests): mr = _make_model_runner( + self, is_hybrid_swa=True, full_attention_layer_ids=[0], swa_attention_layer_ids=[1], @@ -681,7 +707,7 @@ class TestDflashDraftKvBudget(unittest.TestCase): def test_dcp_replication_scales_draft_budget(self): """The replicated draft pool spans every DCP virtual location.""" draft_kv_per_token = 10_240 - mr = _make_model_runner() + mr = _make_model_runner(self) mr.spec_algorithm.is_dflash_family.return_value = True mr.spec_aux_config = SimpleNamespace( eagle_draft_num_layers=None, @@ -715,6 +741,7 @@ class TestDflashDraftKvBudget(unittest.TestCase): def _tokens(draft_kv_per_token): mr = _make_model_runner( + self, is_hybrid_swa=True, full_attention_layer_ids=list(range(16)), swa_attention_layer_ids=list(range(16, 32)), diff --git a/test/registered/unit/test_global_config_read_ratchet.py b/test/registered/unit/test_global_config_read_ratchet.py index 494c55f1f..d221af1c2 100644 --- a/test/registered/unit/test_global_config_read_ratchet.py +++ b/test/registered/unit/test_global_config_read_ratchet.py @@ -57,6 +57,12 @@ _CONFIGURED_SIZE_CALL_SITES = { "point, since with PP off the group is never touched, which is what lets " "the Indexer be constructed before distributed init" ), + ("srt/managers/scheduler.py", "configured_pp_size"): ( + "dispatch_event_loop picks the PP event loop; the MLX runner stub never " + "initializes torch.distributed, so the live property asserts before the " + "MLX loop can start -- the configured leaf answers the same value " + "wherever the live groups exist" + ), ("srt/mem_cache/kv_cache_configurator.py", "configured_pp_size"): ( "decides whether the token capacity needs a cross-PP all-reduce at all; " "asking the configured size keeps that decision independent of whether a " diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 86e0f90af..6e08c0dd2 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -401,6 +401,7 @@ class _FakeResolvedArgs: max_running_requests: A[int | None, Arg(help="mrr"), NS("schedule")] = None chunked_prefill_size: A[int, Arg(help="cps"), NS("schedule")] = -1 max_prefill_tokens: A[int, Arg(help="mpt"), NS("schedule")] = 16384 + enable_dynamic_chunking: A[bool, Arg(help="edc"), NS("schedule")] = False cuda_graph_config: A[object | None, Arg(help="cgc"), NS("exec.graph")] = None tp_size: A[int, Arg(help="tp"), NS("parallel")] = 1 pp_size: A[int, Arg(help="pp"), NS("parallel")] = 1 @@ -1026,6 +1027,31 @@ class TestDerivedPredicatesAgreeAcrossTiers(_IsolatedServerArgs): mamba_extra_buffer_lazy_enabled(), ) + def test_prefill_buffer_ceiling_matches_the_member(self): + from sglang.srt.runtime_context import max_prefill_buffer_tokens + + for chunked in (-1, 0, 1024, 8192): + for dynamic in (False, True): + for pp in (1, 4): + for max_prefill in (0, 2048, 16384): + with self.subTest( + chunked=chunked, + dynamic=dynamic, + pp=pp, + max_prefill=max_prefill, + ): + args = _FakeResolvedArgs( + chunked_prefill_size=chunked, + enable_dynamic_chunking=dynamic, + pp_size=pp, + max_prefill_tokens=max_prefill, + ) + get_context().set_server_args(args) + self.assertEqual( + ServerArgs.max_prefill_buffer_tokens(args), + max_prefill_buffer_tokens(), + ) + def test_activation_reserve_matches_the_member(self): from types import SimpleNamespace