diff --git a/python/sglang/srt/debug_utils/pr_fix_toggle.py b/python/sglang/srt/debug_utils/pr_fix_toggle.py index 9ff1379a6..9646bf8fe 100644 --- a/python/sglang/srt/debug_utils/pr_fix_toggle.py +++ b/python/sglang/srt/debug_utils/pr_fix_toggle.py @@ -88,13 +88,9 @@ patches: - target: sglang.srt.mem_cache.allocation_sizing.get_req_to_token_extra_context_len edits: - match: | - if ( - server_args.speculative_algorithm is not None - and server_args.page_size > 1 - and (server_args.speculative_eagle_topk or 1) > 1 - ): - extra = max(extra, get_alloc_reserve_per_decode(server_args)) - replacement: "" + extra = max(extra, get_alloc_reserve_per_decode() + page_size - 1) + replacement: | + pass """ diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 56febf102..8bb4f8b2a 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -2335,7 +2335,7 @@ class SchedulerDisaggregationDecodeMixin: # construct fake completed prefill new_batch.prepare_for_prebuilt() - new_batch.process_prebuilt(self.server_args, self.future_map) + new_batch.process_prebuilt(self.future_map) return new_batch diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 6488e2eea..c813c30cd 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -16,7 +16,6 @@ logger = logging.getLogger(__name__) if TYPE_CHECKING: from sglang.srt.managers.overlap_utils import FutureMap from sglang.srt.managers.schedule_batch import ScheduleBatch - from sglang.srt.server_args import ServerArgs class ScheduleBatchDisaggregationDecodeMixin: @@ -113,7 +112,6 @@ class ScheduleBatchDisaggregationDecodeMixin: def process_prebuilt( self: ScheduleBatch, - server_args: ServerArgs, future_map: FutureMap, ): """Assign the buffered last input id to schedule batch""" @@ -146,7 +144,6 @@ class ScheduleBatchDisaggregationDecodeMixin: spec_info = self.spec_algorithm.build_disagg_draft_input( self, - server_args, last_tokens_tensor, future_map, ) diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py index 94f4f6167..1b1b34df6 100644 --- a/python/sglang/srt/layers/attention/dsa/utils.py +++ b/python/sglang/srt/layers/attention/dsa/utils.py @@ -13,6 +13,8 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo is_in_tc_piecewise_cuda_graph, ) from sglang.srt.runtime_context import ( + get_disagg, + get_memory, get_parallel, process_model_config, ) @@ -67,27 +69,24 @@ INDEXER_K_CACHE_PRESHUFFLE_TILE = 16 if TYPE_CHECKING: from sglang.srt.model_executor.forward_batch_info import ForwardBatch - from sglang.srt.server_args import ServerArgs def compute_dsa_seqlens(original_seq_lens, dsa_index_topk: int): return original_seq_lens.clamp(max=dsa_index_topk) -def should_remap_pd_dsa_seed_to_local_slots(server_args: "ServerArgs") -> bool: +def should_remap_pd_dsa_seed_to_local_slots() -> bool: """Whether a PD seed should enter the allocator-local fused TopK domain.""" return ( is_cuda() and envs.SGLANG_DSA_FUSE_TOPK.get() - and server_args.disaggregation_mode == "decode" - and not server_args.enable_hisparse + and get_disagg().disaggregation_mode == "decode" + and not get_memory().enable_hisparse and not get_parallel().dcp_enabled ) -def should_use_dsa_fused_topk( - server_args: "ServerArgs", seed_dsa_topk_from_draft_extend: bool -) -> bool: +def should_use_dsa_fused_topk(seed_dsa_topk_from_draft_extend: bool) -> bool: """Select fused TopK for PD IndexShare. PD Prefill worker: @@ -98,10 +97,10 @@ def should_use_dsa_fused_topk( - Draft decode / target verify / draft extend: fused TopK enabled. """ pd_index_share_seed = ( - server_args.disaggregation_mode != "null" and seed_dsa_topk_from_draft_extend + get_disagg().disaggregation_mode != "null" and seed_dsa_topk_from_draft_extend ) return envs.SGLANG_DSA_FUSE_TOPK.get() and ( - not pd_index_share_seed or should_remap_pd_dsa_seed_to_local_slots(server_args) + not pd_index_share_seed or should_remap_pd_dsa_seed_to_local_slots() ) diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 46a324e44..9078dbac6 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -393,9 +393,7 @@ class DeepseekSparseAttnBackend( self.speculative_num_steps = speculative_num_steps self.speculative_num_draft_tokens = get_spec().speculative_num_draft_tokens self.speculative_step_id = speculative_step_id - self.use_fused_topk = should_use_dsa_fused_topk( - model_runner.server_args, seed_dsa_topk_from_draft_extend - ) + self.use_fused_topk = should_use_dsa_fused_topk(seed_dsa_topk_from_draft_extend) if envs.SGLANG_DSA_FUSE_TOPK.get() and not self.use_fused_topk: print_warning_once( "Disabling fused DSA top-k for IndexShare under PD disaggregation." diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 3fe9aadc9..82189ff7d 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -219,7 +219,6 @@ class FlashAttentionBackend(AttentionBackend): # Static verify width; NOTE: overwrites the config-named attr in place. self.speculative_num_draft_tokens = resolve_num_tokens_per_req( phase="target_verify", - server_args=model_runner.server_args, spec_algorithm=SpeculativeAlgorithm.from_string( get_spec().speculative_algorithm ), diff --git a/python/sglang/srt/managers/rust_server.py b/python/sglang/srt/managers/rust_server.py index 9a6c1d07e..4c596bed9 100644 --- a/python/sglang/srt/managers/rust_server.py +++ b/python/sglang/srt/managers/rust_server.py @@ -726,9 +726,7 @@ class RustServer: server_args["version"] = __version__ # Not a `server_args` field: `TokenizerManager` derives it, and the rust # ingress needs the same number for its total-token check. - server_args["num_reserved_tokens"] = compute_num_reserved_tokens( - scheduler.server_args - ) + server_args["num_reserved_tokens"] = compute_num_reserved_tokens() server_args["max_total_num_tokens"] = scheduler.max_total_num_tokens return msgspec.json.encode(server_args, enc_hook=str).decode("utf-8") diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index daca62449..46718795c 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -462,7 +462,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): self.max_req_input_len = None # Will be set later in engine.py self.enable_priority_scheduling = server_args.enable_priority_scheduling self.default_priority_value = server_args.default_priority_value - self.num_reserved_tokens = compute_num_reserved_tokens(server_args) + self.num_reserved_tokens = compute_num_reserved_tokens() self.validate_total_tokens = True def init_tokenizer_and_processor(self): diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index fe883c264..69d00c9af 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -15,12 +15,12 @@ from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers import io_struct from sglang.srt.managers.schedule_batch import Req from sglang.srt.model_executor.forward_batch_info import PPProxyTensors +from sglang.srt.runtime_context import get_spec, max_speculative_num_draft_tokens from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.state_capturer.base import TopkCaptureOutput if TYPE_CHECKING: from sglang.srt.managers.scheduler import GenerationBatchResult - from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.eagle_info import EagleDraftInput @@ -377,19 +377,23 @@ def msgpack_decode_explained(data: bytes) -> Any: raise MsgpackDecodeError(rid, msg) from e -def compute_num_reserved_tokens(server_args: ServerArgs) -> int: +def compute_num_reserved_tokens() -> int: """Output token slots reserved per request, on top of its input. The current eagle implementation stores draft tokens in the output token slots, so the context budget has to account for them; every other algorithm reserves nothing. Shared by `TokenizerManager` and the rust server's `server_args` blob (`RustServer._build_server_args`), which needs the same - number to run the total-token check in Rust. + number to run the total-token check in Rust. Both stamp the number once at + launch, so it has to cover every step an adaptive-spec run may switch to: + it reads the bags for the candidate-table ceiling and the current + `topk * steps`, not the untouched startup record. """ - algorithm = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm) + spec = get_spec() + algorithm = SpeculativeAlgorithm.from_string(spec.speculative_algorithm) if not algorithm.is_eagle(): return 0 return max( - server_args.speculative_eagle_topk * server_args.speculative_num_steps, - server_args.max_speculative_num_draft_tokens, + spec.speculative_eagle_topk * spec.speculative_num_steps, + max_speculative_num_draft_tokens(), ) diff --git a/python/sglang/srt/mem_cache/allocation_sizing.py b/python/sglang/srt/mem_cache/allocation_sizing.py index b049c0ecc..510123bdc 100644 --- a/python/sglang/srt/mem_cache/allocation_sizing.py +++ b/python/sglang/srt/mem_cache/allocation_sizing.py @@ -1,33 +1,31 @@ from __future__ import annotations -from typing import Optional - -from sglang.srt.runtime_context import get_server_args -from sglang.srt.server_args import ServerArgs +from sglang.srt.runtime_context import ( + get_schedule, + get_spec, + max_speculative_num_draft_tokens, +) -def get_alloc_len_per_decode( - server_args: ServerArgs, *, max_draft_tokens: Optional[int] = None -) -> int: - """``max_draft_tokens`` lets a caller that already resolved the draft-token - bound (the KV-cache configurator reads it off the bags) size with that same - value; the default is the handed instance's own member, never the global.""" - if server_args.speculative_algorithm is None: +def get_alloc_len_per_decode() -> int: + """KV length one request may allocate in a single decode step. + + Reads the bags: adaptive speculative decoding moves the step count and the + draft-token bound after publish, and this runs per decode batch. + """ + spec = get_spec() + if spec.speculative_algorithm is None: return 1 # Spec decoding allocates max(topk * num_steps, num_draft_tokens) per decode step. - spec_steps = server_args.speculative_num_steps or 1 - spec_topk = server_args.speculative_eagle_topk or 1 - spec_tokens = ( - max_draft_tokens - if max_draft_tokens is not None - else server_args.max_speculative_num_draft_tokens - ) - page_size = server_args.page_size + spec_steps = spec.speculative_num_steps or 1 + spec_topk = spec.speculative_eagle_topk or 1 + spec_tokens = max_speculative_num_draft_tokens() + page_size = get_schedule().page_size from sglang.srt.speculative.spec_info import SpeculativeAlgorithm - spec_algo = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm) + spec_algo = SpeculativeAlgorithm.from_string(spec.speculative_algorithm) if page_size == 1 or spec_topk == 1 or not spec_algo.has_draft_kv(): return max(spec_steps * spec_topk, spec_tokens) else: @@ -40,51 +38,29 @@ def get_alloc_len_per_decode( return max(num_new_pages_per_topk * page_size * spec_topk, spec_tokens) -def get_alloc_reserve_per_decode( - server_args: Optional[ServerArgs] = None, - *, - max_draft_tokens: Optional[int] = None, -) -> int: +def get_alloc_reserve_per_decode() -> int: """KV length reserved per request at each decode step. The 2x is a double-buffer that absorbs the kv_committed_len lag in overlap mode; see eagle_utils.eagle_prepare_for_decode. - - Callers on a request path have no config in hand, so this is the module's - single "which config" decision point: everything below it takes the - instance explicitly. """ - if server_args is None: - server_args = get_server_args() - return 2 * get_alloc_len_per_decode(server_args, max_draft_tokens=max_draft_tokens) + return 2 * get_alloc_len_per_decode() -def get_req_to_token_extra_context_len( - server_args: ServerArgs, *, max_draft_tokens: Optional[int] = None -) -> int: +def get_req_to_token_extra_context_len() -> int: """req_to_token row headroom beyond the model context length. Sized to hold the decode over-allocation; the spec v2 page>1 topk>1 holey - draft footprint can outgrow the default num_draft_tokens headroom. - - ``max_draft_tokens`` keeps this row headroom and the caller's other - draft-token-sized buffers on ONE resolved value: the KV-cache configurator - passes the bag-derived bound it also hands the pools, so the two cannot - disagree after a post-publish override. The default stays the handed - instance's member for callers sizing against a specific config object. + draft footprint can outgrow the default num_draft_tokens headroom. The row + headroom and the pools it sits next to derive from the same bag leaves, so + they cannot disagree after a post-publish override. """ - if max_draft_tokens is None: - max_draft_tokens = server_args.max_speculative_num_draft_tokens # FIXME(lsyin): temporary fix for the context length issue under spec decoding - extra = 4 + (max_draft_tokens or 0) - if server_args.speculative_algorithm is not None and server_args.page_size > 1: + extra = 4 + (max_speculative_num_draft_tokens() or 0) + page_size = get_schedule().page_size + if get_spec().speculative_algorithm is not None and page_size > 1: # kv_allocated_len is page-aligned (eagle_prepare_for_decode), so near # the context limit the aligned reserve can overshoot by page_size - 1; # without the headroom the row write silently lands in the neighbor row. - extra = max( - extra, - get_alloc_reserve_per_decode(server_args, max_draft_tokens=max_draft_tokens) - + server_args.page_size - - 1, - ) + extra = max(extra, get_alloc_reserve_per_decode() + page_size - 1) return extra diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 4d2ce9a6d..e92d34d6d 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -759,13 +759,7 @@ class KVCacheConfigurator: ) def _build_req_to_token_pool(self, *, max_num_reqs: int) -> ReqToTokenPool: - # The same bag-derived bound the pools below receive, so the row - # headroom and the speculative buffers cannot disagree after a - # post-publish override. - extra_max_context_len = get_req_to_token_extra_context_len( - self.server_args, - max_draft_tokens=max_speculative_num_draft_tokens(), - ) + extra_max_context_len = get_req_to_token_extra_context_len() if get_disagg().disaggregation_mode == "decode": # Extra slots for pre-allocated requests diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 9bcfd8085..a7a69572f 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -784,7 +784,6 @@ class ModelRunner: if self.spec_algorithm.is_speculative(): return resolve_num_tokens_per_req( phase="target_verify", - server_args=self.server_args, spec_algorithm=self.spec_algorithm, is_draft_worker=self.is_draft_worker, num_draft_tokens=num_draft_tokens, diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 8c8a318ca..8698cef49 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -43,7 +43,6 @@ from sglang.srt.runtime_context import ( get_parallel, get_schedule, get_spec, - max_speculative_num_draft_tokens, ) from sglang.srt.utils.common import ( ceil_align, @@ -633,10 +632,7 @@ class SWAChunkCapPoolConfigurator(HybridSWAPoolConfigurator): 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( - kvc.server_args, - max_draft_tokens=max_speculative_num_draft_tokens(), - ) + decode_alloc = 2 * get_alloc_len_per_decode() per_request = trailing_tokens + decode_alloc num_reqs = get_schedule().max_running_requests // kvc.ps.attn_dp_size diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index c421cc419..17b277162 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -554,7 +554,6 @@ class BaseRunner(ABC): # Speculative metadata and hidden-state capture mode. spec_info = create_dummy_verify_input( mr.spec_algorithm, - mr.server_args, buffers.custom_mask, num_tokens_per_req, mr.is_draft_worker, diff --git a/python/sglang/srt/speculative/dspark_disaggregation.py b/python/sglang/srt/speculative/dspark_disaggregation.py index 569f92c2b..0d001b6ba 100644 --- a/python/sglang/srt/speculative/dspark_disaggregation.py +++ b/python/sglang/srt/speculative/dspark_disaggregation.py @@ -10,13 +10,11 @@ from sglang.srt.speculative.dspark_components.dspark_draft import make_next_draf if TYPE_CHECKING: from sglang.srt.managers.overlap_utils import FutureMap from sglang.srt.managers.schedule_batch import ScheduleBatch - from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.spec_info import SpecInput def build_dspark_disagg_draft_input( batch: ScheduleBatch, - server_args: ServerArgs, last_tokens_tensor: torch.Tensor, future_map: FutureMap, ) -> SpecInput: diff --git a/python/sglang/srt/speculative/eagle_disaggregation.py b/python/sglang/srt/speculative/eagle_disaggregation.py index 0c7b84c1c..afec8205b 100644 --- a/python/sglang/srt/speculative/eagle_disaggregation.py +++ b/python/sglang/srt/speculative/eagle_disaggregation.py @@ -9,23 +9,25 @@ from sglang.srt.layers.attention.dsa.utils import ( ) from sglang.srt.managers.overlap_utils import RelayPayload from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode +from sglang.srt.runtime_context import get_spec from sglang.srt.speculative.eagle_info import EagleDraftInput if TYPE_CHECKING: from sglang.srt.managers.overlap_utils import FutureMap from sglang.srt.managers.schedule_batch import ScheduleBatch - from sglang.srt.server_args import ServerArgs def build_eagle_disagg_draft_input( batch: ScheduleBatch, - server_args: ServerArgs, last_tokens_tensor: torch.Tensor, future_map: FutureMap, ) -> EagleDraftInput: - num_states = server_args.speculative_eagle_topk - if server_args.enable_multi_layer_eagle: - num_states *= server_args.speculative_num_steps + # Adaptive spec moves the step count after publish, and this runs once per + # prebuilt batch. + spec = get_spec() + num_states = spec.speculative_eagle_topk + if spec.enable_multi_layer_eagle: + num_states *= spec.speculative_num_steps topk_p = torch.stack( [ @@ -58,7 +60,7 @@ def build_eagle_disagg_draft_input( dsa_indices_list = [req.output_dsa_topk_indices for req in batch.reqs] if dsa_indices_list and all(t is not None for t in dsa_indices_list): dsa_topk_indices = torch.stack(dsa_indices_list, dim=0).to(batch.device) - if should_remap_pd_dsa_seed_to_local_slots(server_args): + if should_remap_pd_dsa_seed_to_local_slots(): # PD sends request-relative positions; fused TopK consumes # decode-local physical slots. Remap once before the draft loop/graph. req_to_token = batch.req_to_token_pool.req_to_token diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index 5722c40f8..a55d18a48 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -147,9 +147,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): # Bucket sizes self.capture_bs, _ = get_batch_sizes_to_capture(model_runner) # Static capture width. - self.captured_req_width = resolve_num_tokens_per_req( - phase="draft_decode", server_args=model_runner.server_args - ) + self.captured_req_width = resolve_num_tokens_per_req(phase="draft_decode") self.max_bs = max(self.capture_bs) self.max_num_token = self.max_bs * self.captured_req_width diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 37b058772..a1f6faf70 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -130,9 +130,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): # Static capture width: full tree width (num_draft_tokens), not # num_steps + 1 -- topk > 1 draft-extend overflows the buffers. - self.captured_req_width = resolve_num_tokens_per_req( - phase="draft_extend", server_args=model_runner.server_args - ) + self.captured_req_width = resolve_num_tokens_per_req(phase="draft_extend") self.max_bs = max(self.capture_bs) self.max_num_token = self.max_bs * self.captured_req_width diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index d0d9acf7e..0712e26de 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py @@ -116,9 +116,7 @@ class FrozenKVMTPCudaGraphRunner(DecodeCudaGraphRunner): self.capture_hidden_mode = CaptureHiddenMode.LAST # Static capture width. - self.captured_req_width = resolve_num_tokens_per_req( - phase="draft_decode", server_args=model_runner.server_args - ) + self.captured_req_width = resolve_num_tokens_per_req(phase="draft_decode") self.capture_bs, _ = get_batch_sizes_to_capture( model_runner, self.captured_req_width ) diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index 97c98ad52..0db2c0e5e 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py @@ -189,10 +189,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): # tokens, which lets all steps share one buffer set. self.num_front_tokens = eagle_worker.draft_extend_num_front_tokens self.captured_req_width = ( - resolve_num_tokens_per_req( - phase="draft_extend", server_args=model_runner.server_args - ) - + self.num_front_tokens + resolve_num_tokens_per_req(phase="draft_extend") + self.num_front_tokens ) self.max_bs = max(self.capture_bs) self.max_num_token = self.max_bs * self.captured_req_width diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 440bcc0be..03967668e 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Callable, List, Optional, Tuple, Type, Union import torch +from sglang.srt.runtime_context import get_spec as get_spec_config from sglang.srt.speculative.spec_registry import ( CustomSpecAlgo, ServerArgsValidator, @@ -172,7 +173,6 @@ class SpeculativeAlgorithm(Enum): def build_disagg_draft_input( self, batch: ScheduleBatch, - server_args: ServerArgs, last_tokens_tensor: torch.Tensor, future_map: FutureMap, ) -> Optional[SpecInput]: @@ -181,16 +181,14 @@ class SpeculativeAlgorithm(Enum): build_eagle_disagg_draft_input, ) - return build_eagle_disagg_draft_input( - batch, server_args, last_tokens_tensor, future_map - ) + return build_eagle_disagg_draft_input(batch, last_tokens_tensor, future_map) if self.is_dspark(): from sglang.srt.speculative.dspark_disaggregation import ( build_dspark_disagg_draft_input, ) return build_dspark_disagg_draft_input( - batch, server_args, last_tokens_tensor, future_map + batch, last_tokens_tensor, future_map ) return None @@ -386,14 +384,20 @@ def spec_scale_global_num_tokens( def create_dummy_verify_input( spec_algorithm: SpeculativeAlgorithm, - server_args: ServerArgs, custom_mask: torch.Tensor, num_tokens_per_req: int, is_draft_worker: bool, ) -> Optional[SpecInput]: - """Dummy verify ``SpecInput`` for CUDA-graph capture (per-algorithm dispatch).""" + """Dummy verify ``SpecInput`` for CUDA-graph capture (per-algorithm dispatch). + + The tree shape comes from the bags so the dummy matches the step config + being captured; adaptive spec captures each candidate with that config's + leaves overridden. + """ from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode + spec = get_spec_config() + spec_info = None if spec_algorithm.is_eagle() or spec_algorithm.is_standalone(): from sglang.srt.speculative.eagle_info import EagleVerifyInput @@ -409,9 +413,9 @@ def create_dummy_verify_input( retrieve_next_token=None, retrieve_next_sibling=None, retrieve_cum_len=None, - spec_steps=server_args.speculative_num_steps, - topk=server_args.speculative_eagle_topk, - draft_token_num=server_args.speculative_num_draft_tokens, + spec_steps=spec.speculative_num_steps, + topk=spec.speculative_eagle_topk, + draft_token_num=spec.speculative_num_draft_tokens, capture_hidden_mode=CaptureHiddenMode.FULL, seq_lens_sum=None, seq_lens_cpu=None, @@ -423,7 +427,7 @@ def create_dummy_verify_input( spec_info = DFlashVerifyInput( draft_token=None, positions=None, - draft_token_num=server_args.speculative_num_draft_tokens, + draft_token_num=spec.speculative_num_draft_tokens, custom_mask=None, capture_hidden_mode=( CaptureHiddenMode.NULL if is_draft_worker else CaptureHiddenMode.FULL diff --git a/python/sglang/srt/speculative/spec_registry.py b/python/sglang/srt/speculative/spec_registry.py index 4ab16a2a3..1e3b84853 100644 --- a/python/sglang/srt/speculative/spec_registry.py +++ b/python/sglang/srt/speculative/spec_registry.py @@ -152,10 +152,14 @@ class CustomSpecAlgo: def build_disagg_draft_input( self, batch: ScheduleBatch, - server_args: ServerArgs, last_tokens_tensor: torch.Tensor, future_map: FutureMap, ) -> Optional[SpecInput]: + """Build the disaggregation draft input for ``batch``, or ``None``. + + The speculative config comes from ``runtime_context.get_spec()``, which + follows a runtime override where the startup record does not. + """ return None diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 19a35c66b..a2d6d91e2 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -78,7 +78,6 @@ if TYPE_CHECKING: from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo - from sglang.srt.server_args import ServerArgs if _is_cuda: @@ -98,7 +97,6 @@ logger = logging.getLogger(__name__) def resolve_num_tokens_per_req( *, phase: Literal["draft_decode", "draft_extend", "target_verify"], - server_args: ServerArgs, spec_algorithm=None, is_draft_worker: bool = False, num_draft_tokens: Optional[int] = None, @@ -107,14 +105,19 @@ def resolve_num_tokens_per_req( width (sizes capture shapes / buffers); the per-forward dynamic width lives on ``SpecInput.num_tokens_per_req``. Draft phases are EAGLE-family-only; "target_verify" is algorithm-generic via the hook. + + The widths come from the bags: adaptive spec captures each candidate step + config with that config's leaves overridden, so the buffers being sized + must follow the override rather than the startup values. """ + spec = get_spec() if phase == "draft_decode": - return server_args.speculative_eagle_topk + return spec.speculative_eagle_topk if phase == "draft_extend": - return server_args.speculative_num_draft_tokens + return spec.speculative_num_draft_tokens if phase == "target_verify": if num_draft_tokens is None: - num_draft_tokens = server_args.speculative_num_draft_tokens + num_draft_tokens = spec.speculative_num_draft_tokens return spec_algorithm.get_num_tokens_per_req_for_target_verify( num_draft_tokens, is_draft_worker ) diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 11bb34013..800c427f1 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -34,6 +34,7 @@ from sglang.srt.environ import envs from sglang.srt.layers.attention.dsa.utils import should_use_dsa_fused_topk from sglang.srt.managers.overlap_utils import FutureMap, RelayPayload from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool +from sglang.srt.runtime_context import get_context from sglang.srt.speculative.eagle_disaggregation import ( build_eagle_disagg_draft_input, ) @@ -302,17 +303,17 @@ class TestEagleDsaSeedTransfer(unittest.TestCase): device="cpu", enable_overlap=False, ) - server_args = SimpleNamespace( + # The draft-input shape comes from the spec bag. + override = get_context().override_server_args( speculative_eagle_topk=1, speculative_num_steps=5, enable_multi_layer_eagle=False, - disaggregation_mode="null", ) + override.install() + self.addCleanup(override.restore) last_tokens = torch.tensor([11, 12], dtype=torch.int64) - draft_input = build_eagle_disagg_draft_input( - batch, server_args, last_tokens, None - ) + draft_input = build_eagle_disagg_draft_input(batch, last_tokens, None) self.assertTrue(torch.equal(draft_input.dsa_topk_indices, torch.stack(seeds))) for invalid_seed in ( @@ -320,9 +321,7 @@ class TestEagleDsaSeedTransfer(unittest.TestCase): torch.full((3,), -1, dtype=torch.int32), ): batch.reqs[1].output_dsa_topk_indices = invalid_seed - draft_input = build_eagle_disagg_draft_input( - batch, server_args, last_tokens, None - ) + draft_input = build_eagle_disagg_draft_input(batch, last_tokens, None) self.assertIsNone(draft_input.dsa_topk_indices) def test_pd_decode_fused_topk_remaps_wire_positions_to_local_slots(self): @@ -347,24 +346,24 @@ class TestEagleDsaSeedTransfer(unittest.TestCase): req_to_token_pool=SimpleNamespace(req_to_token=req_to_token), seq_lens=torch.tensor([4, 4], dtype=torch.int32), ) - server_args = SimpleNamespace( + override = get_context().override_server_args( speculative_eagle_topk=1, speculative_num_steps=5, enable_multi_layer_eagle=False, disaggregation_mode="decode", enable_hisparse=False, ) + override.install() + self.addCleanup(override.restore) with envs.SGLANG_DSA_FUSE_TOPK.override(True), patch( "sglang.srt.layers.attention.dsa.utils.is_cuda", return_value=True ): self.assertTrue( - should_use_dsa_fused_topk( - server_args, seed_dsa_topk_from_draft_extend=True - ) + should_use_dsa_fused_topk(seed_dsa_topk_from_draft_extend=True) ) draft_input = build_eagle_disagg_draft_input( - batch, server_args, torch.tensor([11, 12], dtype=torch.int64), None + batch, torch.tensor([11, 12], dtype=torch.int64), None ) self.assertEqual( diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index 1bd0a16af..5d4154809 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -432,8 +432,6 @@ class TestDecodePrebuiltPriority(unittest.TestCase): scheduler.enable_overlap = False scheduler.spec_algorithm = MagicMock() scheduler.max_running_requests = 1 - # Passed whole into the (mocked) batch's process_prebuilt; never read. - scheduler.server_args = SimpleNamespace() scheduler.future_map = MagicMock() scheduler.policy = MagicMock() scheduler.policy.calc_priority.side_effect = lambda waiting_queue, _: ( diff --git a/test/registered/unit/spec/test_plugin_hook_signatures.py b/test/registered/unit/spec/test_plugin_hook_signatures.py new file mode 100644 index 000000000..44bc4505b --- /dev/null +++ b/test/registered/unit/spec/test_plugin_hook_signatures.py @@ -0,0 +1,109 @@ +"""The plugin hook takes the call the dispatch makes. + +`CustomSpecAlgo` is the out-of-tree extension point: a registered algorithm's +method is called through the same dispatch as the built-in ones, and nothing in +the tree implements it, so a drift between the two sides only ever surfaces in +somebody's plugin. It has drifted twice, both times on the disaggregation +draft-input builder: the built-in dropped a parameter the hook kept, so every +plugin call would have hit a TypeError. + +What is pinned here: + + * every method the dispatch may call on either type takes the same arguments + on both -- the set is intersected out of the two types, never listed; + * the call the dispatch actually writes binds on both types. +""" + +import ast +import inspect +import unittest +from pathlib import Path + +from sglang.srt.disaggregation import decode_schedule_batch_mixin +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.srt.speculative.spec_registry import CustomSpecAlgo +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _dispatched_methods(): + """Methods carried by both types. The dispatch calls them on whichever it + holds without knowing which, so their argument lists must agree.""" + enum_methods = { + name + for name, value in vars(SpeculativeAlgorithm).items() + if inspect.isfunction(value) + } + hook_methods = { + name + for name, value in vars(CustomSpecAlgo).items() + if inspect.isfunction(value) + } + return sorted(enum_methods & hook_methods) + + +def _dispatch_calls(): + """Every call the decode dispatch makes on the algorithm object, read from + its source: `(method, positional count, keyword names)`.""" + source = Path(decode_schedule_batch_mixin.__file__).read_text(encoding="utf-8") + calls = [] + for node in ast.walk(ast.parse(source)): + if not isinstance(node, ast.Call): + continue + func = node.func + if not ( + isinstance(func, ast.Attribute) + and isinstance(func.value, ast.Attribute) + and func.value.attr == "spec_algorithm" + ): + continue + calls.append((func.attr, len(node.args), tuple(kw.arg for kw in node.keywords))) + return calls + + +def _parameters(function): + """Parameter names, without any catch-alls.""" + return [ + name + for name, parameter in inspect.signature(function).parameters.items() + if parameter.kind + not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD) + ] + + +class TestDispatchedSignatures(CustomTestCase): + def test_the_hook_and_the_built_in_agree(self): + methods = _dispatched_methods() + self.assertNotEqual(methods, [], "the dispatched set derived to nothing") + mismatches = [] + for name in methods: + hook = _parameters(getattr(CustomSpecAlgo, name)) + builtin = _parameters(getattr(SpeculativeAlgorithm, name)) + if hook != builtin: + mismatches.append( + f"{name}: CustomSpecAlgo{tuple(hook)} vs " + f"SpeculativeAlgorithm{tuple(builtin)}" + ) + self.assertEqual( + mismatches, + [], + "a plugin implementing the hook would be called with the " + "dispatch's arguments:\n " + "\n ".join(mismatches), + ) + + def test_the_dispatch_call_binds_on_both_types(self): + calls = _dispatch_calls() + self.assertNotEqual(calls, [], "no dispatch call found to bind against") + for method, positional, keywords in calls: + self.assertIn(method, _dispatched_methods()) + for owner in (CustomSpecAlgo, SpeculativeAlgorithm): + arguments = [None] * (1 + positional) + inspect.signature(getattr(owner, method)).bind( + *arguments, **{name: None for name in keywords} + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/test_publish_precedes_bag_reads.py b/test/registered/unit/test_publish_precedes_bag_reads.py index aed5093cd..2bae967b5 100644 --- a/test/registered/unit/test_publish_precedes_bag_reads.py +++ b/test/registered/unit/test_publish_precedes_bag_reads.py @@ -111,10 +111,6 @@ _UNREAD_ENTRIES: dict = { ("srt/managers/detokenizer_manager.py", "run_detokenizer_process"): ( "DetokenizerManager reads the handed instance at this revision" ), - ("srt/managers/tokenizer_manager.py", "__init__"): ( - "the constructor and the init_* helpers it calls read the handed " - "instance at this revision" - ), } # `publish` itself and its named wrappers live here; a call inside them is the diff --git a/test/registered/unit/test_supplied_instance_exposure_ratchet.py b/test/registered/unit/test_supplied_instance_exposure_ratchet.py index f1998915d..0b77d800d 100644 --- a/test/registered/unit/test_supplied_instance_exposure_ratchet.py +++ b/test/registered/unit/test_supplied_instance_exposure_ratchet.py @@ -242,7 +242,6 @@ _EXPOSED = { ("kv_canary/capacities.py", "cuda_graph_config"), ("kv_canary/capacities.py", "speculative_num_draft_tokens"), ("kv_canary/token_oracle/install.py", "sampling_backend"), - ("layers/attention/dsa/utils.py", "disaggregation_mode"), ("layers/cp/base.py", "attn_cp_size"), ("layers/cp/base.py", "cp_strategy"), ("layers/cp/base.py", "enable_prefill_cp"), @@ -334,13 +333,6 @@ _EXPOSED = { ("managers/tp_worker.py", "random_seed"), ("managers/tp_worker.py", "speculative_algorithm"), ("managers/tp_worker.py", "tokenizer_path"), - ("managers/utils.py", "speculative_algorithm"), - ("managers/utils.py", "speculative_eagle_topk"), - ("managers/utils.py", "speculative_num_steps"), - ("mem_cache/allocation_sizing.py", "page_size"), - ("mem_cache/allocation_sizing.py", "speculative_algorithm"), - ("mem_cache/allocation_sizing.py", "speculative_eagle_topk"), - ("mem_cache/allocation_sizing.py", "speculative_num_steps"), ("mem_cache/hiradix_cache.py", "hicache_io_backend"), ("mem_cache/hiradix_cache.py", "hicache_mem_layout"), ("mem_cache/hiradix_cache.py", "served_model_name"), @@ -431,9 +423,6 @@ _EXPOSED = { "speculative/dspark_components/dspark_worker_v2.py", "speculative_num_draft_tokens", ), - ("speculative/eagle_disaggregation.py", "enable_multi_layer_eagle"), - ("speculative/eagle_disaggregation.py", "speculative_eagle_topk"), - ("speculative/eagle_disaggregation.py", "speculative_num_steps"), ("speculative/eagle_worker_v2.py", "device"), ("speculative/eagle_worker_v2.py", "enable_dp_attention"), ("speculative/eagle_worker_v2.py", "speculative_adaptive"), @@ -460,12 +449,7 @@ _EXPOSED = { ("speculative/ngram_worker.py", "speculative_num_draft_tokens"), ("speculative/ngram_worker.py", "speculative_num_steps"), ("speculative/spec_info.py", "enable_multi_layer_eagle"), - ("speculative/spec_info.py", "speculative_eagle_topk"), - ("speculative/spec_info.py", "speculative_num_draft_tokens"), - ("speculative/spec_info.py", "speculative_num_steps"), ("speculative/spec_registry.py", "disable_overlap_schedule"), - ("speculative/spec_utils.py", "speculative_eagle_topk"), - ("speculative/spec_utils.py", "speculative_num_draft_tokens"), ("speculative/standalone_worker_v2.py", "device"), ("speculative/standalone_worker_v2.py", "enable_dp_attention"), ("speculative/standalone_worker_v2.py", "speculative_algorithm"), @@ -533,8 +517,6 @@ _OVERRIDDEN_AND_READ = { ("managers/tokenizer_manager.py", "model_path"), ("managers/tokenizer_manager.py", "speculative_num_draft_tokens"), ("managers/tp_worker.py", "model_path"), - ("managers/utils.py", "speculative_num_steps"), - ("mem_cache/allocation_sizing.py", "speculative_num_steps"), ("mem_cache/hiradix_cache.py", "hicache_storage_backend"), ("mem_cache/hiradix_cache.py", "hicache_storage_backend_extra_config"), ("mem_cache/hiradix_cache.py", "hicache_storage_prefetch_policy"), @@ -561,7 +543,6 @@ _OVERRIDDEN_AND_READ = { "speculative/dspark_components/dspark_worker_v2.py", "speculative_num_draft_tokens", ), - ("speculative/eagle_disaggregation.py", "speculative_num_steps"), ("speculative/eagle_worker_v2.py", "speculative_num_draft_tokens"), ("speculative/eagle_worker_v2.py", "speculative_num_steps"), ("speculative/frozen_kv_mtp_worker_v2.py", "speculative_num_draft_tokens"), @@ -570,9 +551,6 @@ _OVERRIDDEN_AND_READ = { ("speculative/multi_layer_eagle_worker_v2.py", "speculative_num_steps"), ("speculative/ngram_worker.py", "speculative_num_draft_tokens"), ("speculative/ngram_worker.py", "speculative_num_steps"), - ("speculative/spec_info.py", "speculative_num_draft_tokens"), - ("speculative/spec_info.py", "speculative_num_steps"), - ("speculative/spec_utils.py", "speculative_num_draft_tokens"), ("speculative/standalone_worker_v2.py", "speculative_num_draft_tokens"), ("speculative/standalone_worker_v2.py", "speculative_num_steps"), ("utils/common.py", "speculative_num_draft_tokens"),