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