minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

spec: size the speculative buffers from the bags, not the startup record (#35024)

Browse Source
This commit is contained in:
Cheng Wan
2026-08-17 16:16:56 -07:00
committed by GitHub
parent 3d7ec00179
commit d2bc697396
28 changed files with 216 additions and 179 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/debug_utils/pr_fix_toggle.py
+3 -7
View File
@@ -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
"""
python/sglang/srt/disaggregation/decode.py
+1 -1
View File
@@ -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
python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py
-3
View File
@@ -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,
)
python/sglang/srt/layers/attention/dsa/utils.py
+8 -9
View File
@@ -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()
)
python/sglang/srt/layers/attention/dsa_backend.py
+1 -3
View File
@@ -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."
python/sglang/srt/layers/attention/flashattention_backend.py
-1
View File
@@ -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
),
python/sglang/srt/managers/rust_server.py
+1 -3
View File
@@ -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")
python/sglang/srt/managers/tokenizer_manager.py
+1 -1
View File
@@ -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):
python/sglang/srt/managers/utils.py
+10 -6
View File
@@ -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(),
)
python/sglang/srt/mem_cache/allocation_sizing.py
+28 -52
View File
@@ -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
python/sglang/srt/mem_cache/kv_cache_configurator.py
+1 -7
View File
@@ -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
python/sglang/srt/model_executor/model_runner.py
-1
View File
@@ -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,
python/sglang/srt/model_executor/pool_configurator.py
+1 -5
View File
@@ -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
python/sglang/srt/model_executor/runner/base_runner.py
-1
View File
@@ -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,
python/sglang/srt/speculative/dspark_disaggregation.py
-2
View File
@@ -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:
python/sglang/srt/speculative/eagle_disaggregation.py
+8 -6
View File
@@ -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
python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
+1 -3
View File
@@ -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
python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
+1 -3
View File
@@ -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
python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py
+1 -3
View File
@@ -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
)
python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py
+1 -4
View File
@@ -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
python/sglang/srt/speculative/spec_info.py
+15 -11
View File
@@ -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
python/sglang/srt/speculative/spec_registry.py
+5 -1
View File
@@ -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
python/sglang/srt/speculative/spec_utils.py
+8 -5
View File
@@ -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
)
test/registered/unit/disaggregation/test_disaggregation_wire.py
+12 -13
View File
@@ -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(
test/registered/unit/managers/test_priority_scheduling_disaggregation.py
-2
View File
@@ -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, _: (
test/registered/unit/spec/test_plugin_hook_signatures.py
+109
View File
@@ -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()
test/registered/unit/test_publish_precedes_bag_reads.py
-4
View File
@@ -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
test/registered/unit/test_supplied_instance_exposure_ratchet.py
-22
View File
@@ -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"),