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

[Spec] Remove the dead spec V1 scheduler paths (#27977)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-11 18:31:13 -07:00
committed by GitHub
parent 2e74ff192c
commit 3ffe72517f
15 changed files with 97 additions and 794 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/dsa_backend.py
+3 -3
View File
@@ -564,9 +564,9 @@ class DeepseekSparseAttnBackend(
page_table, repeats=self.speculative_num_draft_tokens, dim=0 page_table, repeats=self.speculative_num_draft_tokens, dim=0
) )
else: else:
# DRAFT_EXTEND (v1): V1 worker extends by (num_correct_drafts + 1) per request # DRAFT_EXTEND: the draft worker extends by (num_correct_drafts + 1)
# after verification. Lengths vary per request based on how many tokens # per request after verification. Lengths vary per request based on
# were accepted. # how many tokens were accepted.
page_table = torch.repeat_interleave( page_table = torch.repeat_interleave(
page_table, repeats=forward_batch.extend_seq_lens, dim=0 page_table, repeats=forward_batch.extend_seq_lens, dim=0
) )
python/sglang/srt/layers/utils/logprob.py
+2 -116
View File
@@ -2,17 +2,14 @@ from __future__ import annotations
import dataclasses import dataclasses
from enum import Enum, auto from enum import Enum, auto
from typing import TYPE_CHECKING, List, Optional, Union from typing import TYPE_CHECKING, List, Optional
import torch import torch
from sglang.srt.environ import envs from sglang.srt.environ import envs
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessorOutput from sglang.srt.layers.logits_processor import LogitsMetadata
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.speculative.eagle_info import EagleVerifyOutput
from sglang.srt.speculative.ngram_info import NgramVerifyInput
class LogprobStage(Enum): class LogprobStage(Enum):
@@ -295,117 +292,6 @@ def get_token_ids_logprobs_chunk(
return next_split_pruned_len return next_split_pruned_len
def add_output_logprobs_for_spec_v1(
batch: ScheduleBatch,
res: Union[EagleVerifyOutput, NgramVerifyInput],
logits_output: Optional[LogitsProcessorOutput] = None,
):
# Extract args
if logits_output is None:
logits_output = res.logits_output
if hasattr(res, "num_correct_drafts_per_req_cpu"):
num_correct_drafts_per_req_cpu = res.num_correct_drafts_per_req_cpu
else:
# FIXME: Get a NgramVerifyOutput class and use that instead of this hack.
num_correct_drafts_per_req_cpu = res.num_correct_drafts.tolist()
top_logprobs_nums = batch.top_logprobs_nums
token_ids_logprobs = batch.token_ids_logprobs
accept_indices = res.accept_indices
assert len(accept_indices) == len(logits_output.next_token_logits)
temperatures = batch.sampling_info.temperatures
num_draft_tokens = batch.spec_info.draft_token_num
# acceptance indices are the indices in a "flattened" batch.
# dividing it to num_draft_tokens will yield the actual batch index.
temperatures = temperatures[accept_indices // num_draft_tokens]
if envs.SGLANG_RETURN_ORIGINAL_LOGPROB.get():
logprobs = torch.nn.functional.log_softmax(
logits_output.next_token_logits, dim=-1
)
else:
logprobs = torch.nn.functional.log_softmax(
logits_output.next_token_logits / temperatures, dim=-1
)
batch_next_token_ids = res.accept_tokens
num_tokens_per_req = [accept + 1 for accept in num_correct_drafts_per_req_cpu]
# We should repeat top_logprobs_nums to match num_tokens_per_req.
top_logprobs_nums_repeat_interleaved = [
num
for num, num_tokens in zip(top_logprobs_nums, num_tokens_per_req)
for _ in range(num_tokens)
]
token_ids_logprobs_repeat_interleaved = [
token_ids
for token_ids, num_tokens in zip(token_ids_logprobs, num_tokens_per_req)
for _ in range(num_tokens)
]
# Extract logprobs
should_top_logprobs = any(x > 0 for x in top_logprobs_nums)
should_token_ids_logprobs = any(x is not None for x in token_ids_logprobs)
if should_top_logprobs:
(
logits_output.next_token_top_logprobs_val,
logits_output.next_token_top_logprobs_idx,
) = get_top_logprobs(
logprobs,
top_logprobs_nums_repeat_interleaved,
)
if should_token_ids_logprobs:
(
logits_output.next_token_token_ids_logprobs_val,
logits_output.next_token_token_ids_logprobs_idx,
) = get_token_ids_logprobs(
logprobs,
token_ids_logprobs_repeat_interleaved,
)
logits_output.next_token_logprobs = logprobs[
torch.arange(len(batch_next_token_ids), device=batch.sampling_info.device),
batch_next_token_ids,
]
# Add output logprobs to the request
pt = 0
next_token_logprobs = logits_output.next_token_logprobs.tolist()
accept_tokens_list = batch_next_token_ids.tolist()
token_top_logprobs_val = logits_output.next_token_top_logprobs_val
token_top_logprobs_idx = logits_output.next_token_top_logprobs_idx
token_ids_logprobs_val = logits_output.next_token_token_ids_logprobs_val
token_ids_logprobs_idx = logits_output.next_token_token_ids_logprobs_idx
for req, num_tokens in zip(batch.reqs, num_tokens_per_req, strict=True):
for _ in range(num_tokens):
if req.return_logprob:
req.logprob.output_token_logprobs_val.append(next_token_logprobs[pt])
req.logprob.output_token_logprobs_idx.append(accept_tokens_list[pt])
if req.logprob.top_logprobs_num > 0:
assert (
should_top_logprobs
), "Inconsistent state: should_top_logprobs is False"
req.logprob.output_top_logprobs_val.append(
token_top_logprobs_val[pt]
)
req.logprob.output_top_logprobs_idx.append(
token_top_logprobs_idx[pt]
)
if (
req.logprob.token_ids_logprob is not None
and should_token_ids_logprobs
):
req.logprob.output_token_ids_logprobs_val.append(
token_ids_logprobs_val[pt]
)
req.logprob.output_token_ids_logprobs_idx.append(
token_ids_logprobs_idx[pt]
)
pt += 1
def compute_spec_v2_logprobs( def compute_spec_v2_logprobs(
batch, batch,
logits_output, logits_output,
python/sglang/srt/managers/overlap_utils.py
+3 -3
View File
@@ -95,7 +95,7 @@ def resolve_forward_inputs(batch: ScheduleBatch, future_map: FutureMap) -> None:
# Only the overlap path relays spec extras through the future_map; the # Only the overlap path relays spec extras through the future_map; the
# synchronous (non-overlap) V2 path installs next_draft_input directly. # synchronous (non-overlap) V2 path installs next_draft_input directly.
if batch.enable_overlap and batch.is_spec_v2: if batch.enable_overlap and not batch.spec_algorithm.is_none():
future_map._resolve_spec_extras(batch) future_map._resolve_spec_extras(batch)
@@ -326,8 +326,8 @@ class FutureMap:
if indices.shape[0] == 0: if indices.shape[0] == 0:
# DP idle: payload is empty stub; lazy-init shape peek would IndexError. # DP idle: payload is empty stub; lazy-init shape peek would IndexError.
return return
# Dispatch by payload type, not spec_algo: spec_v1 (non-overlap spec) # Dispatch by payload type, not spec_algo: non-spec decode passes a
# also passes a token Tensor here. # token Tensor here.
# FIXME(lsyin): unify this relay path with a dataclass instead of the # FIXME(lsyin): unify this relay path with a dataclass instead of the
# Tensor / EagleDraftInput type switch. # Tensor / EagleDraftInput type switch.
if isinstance(payload, torch.Tensor): if isinstance(payload, torch.Tensor):
python/sglang/srt/managers/schedule_batch.py
+5 -39
View File
@@ -2355,25 +2355,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
new_pages = sum(1 for r in requests if r.kv_committed_len % page_size == 0) new_pages = sum(1 for r in requests if r.kv_committed_len % page_size == 0)
return new_pages * page_size return new_pages * page_size
if self.is_spec_v2: return self._new_tokens_required_next_decode_spec_v2(requests, page_size)
return self._new_tokens_required_next_decode_spec_v2(requests, page_size)
server_args = get_global_server_args()
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
if page_size > 1 and spec_topk > 1:
# last partial page and ceil alignment
len_per_topk = ceil_align(len_per_topk + page_size, page_size)
spec_tokens = ceil_align(spec_tokens, page_size)
elif page_size > 1:
# only page alignment
len_per_topk = ceil_align(len_per_topk, page_size)
spec_tokens = ceil_align(spec_tokens, page_size)
num_tokens = max(len_per_topk * spec_topk, spec_tokens) * len(requests)
return num_tokens
def _new_tokens_required_next_decode_spec_v2(self, requests, page_size): def _new_tokens_required_next_decode_spec_v2(self, requests, page_size):
"""Tight estimate matching eagle_info_v2.prepare_for_decode allocation.""" """Tight estimate matching eagle_info_v2.prepare_for_decode allocation."""
@@ -2498,12 +2480,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.model_config.vocab_size, self.model_config.vocab_size,
) )
@property
def is_spec_v2(self):
# Whether the V2 worker/schema is used. Independent of overlap: the
# non-overlap path also drives the V2 worker, just synchronously.
return self.spec_algorithm.supports_spec_v2()
def mamba_lazy_prealloc_at_boundary(self, mamba_track_interval: int): def mamba_lazy_prealloc_at_boundary(self, mamba_track_interval: int):
"""Allocate a temporary second ping-pong slot for reqs at a track boundary. """Allocate a temporary second ping-pong slot for reqs at a track boundary.
@@ -2547,14 +2523,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
if hasattr(self, "attn_cp_metadata") and self.attn_cp_metadata is not None: if hasattr(self, "attn_cp_metadata") and self.attn_cp_metadata is not None:
self.attn_cp_metadata = None self.attn_cp_metadata = None
if self.is_spec_v2: if not self.spec_algorithm.is_none():
# TODO(spec-v2): all spec v2 should go through this path # Spec decoding: the draft input owns decode preparation
# (allocation, pre-claim, seq-lens bookkeeping).
draft_input: EagleDraftInput = self.spec_info draft_input: EagleDraftInput = self.spec_info
draft_input.prepare_for_decode(self) draft_input.prepare_for_decode(self)
if not self.spec_algorithm.is_none():
# if spec decoding is used, the decode batch is prepared inside
# `forward_batch_speculative_generation` after running draft models.
return return
if self.sampling_info.penalizer_orchestrator.is_required: if self.sampling_info.penalizer_orchestrator.is_required:
@@ -2638,8 +2611,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self, self,
chunked_req_to_exclude: Optional[Union[Req, List[Req]]] = None, chunked_req_to_exclude: Optional[Union[Req, List[Req]]] = None,
keep_indices: Optional[List[int]] = None, keep_indices: Optional[List[int]] = None,
# FIXME(lsyin): deprecate this API after spec v1 is deprecated
v1_spec_info_filtered: Optional[bool] = False,
): ):
if keep_indices is None: if keep_indices is None:
if isinstance(chunked_req_to_exclude, Req): if isinstance(chunked_req_to_exclude, Req):
@@ -2706,15 +2677,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.has_grammar = any(req.grammar for req in self.reqs) self.has_grammar = any(req.grammar for req in self.reqs)
self.sampling_info.filter_batch(keep_indices, keep_indices_device) self.sampling_info.filter_batch(keep_indices, keep_indices_device)
# NOTE: spec_info filtered before batch filtering only happens in:
# - Spec v1's verify phase
# - Only for decode batch (running_batch)
has_been_filtered = v1_spec_info_filtered and not self.is_spec_v2
if self.spec_info: if self.spec_info:
self.spec_info.filter_batch( self.spec_info.filter_batch(
new_indices=keep_indices_device, new_indices=keep_indices_device,
has_been_filtered=has_been_filtered, has_been_filtered=False,
) )
def merge_batch(self, other: ScheduleBatch): def merge_batch(self, other: ScheduleBatch):
python/sglang/srt/managers/scheduler.py
+17 -24
View File
@@ -1377,8 +1377,7 @@ class Scheduler(
) )
def _abort_on_running_timeout(self): def _abort_on_running_timeout(self):
# NOTE: this should be called before a batch is launched, # NOTE: this should be called before a batch is launched.
# as current spec-v1 still filters batch inside verify stage.
timeout_s = envs.SGLANG_REQ_RUNNING_TIMEOUT.get() timeout_s = envs.SGLANG_REQ_RUNNING_TIMEOUT.get()
if timeout_s <= 0: if timeout_s <= 0:
return return
@@ -1541,7 +1540,7 @@ class Scheduler(
# TODO(lsyin): support overlap + spec + grammar # TODO(lsyin): support overlap + spec + grammar
need_grammar_sync = ( need_grammar_sync = (
batch batch
and batch.is_spec_v2 and not batch.spec_algorithm.is_none()
and batch.has_grammar and batch.has_grammar
and batch.forward_mode.is_decode() and batch.forward_mode.is_decode()
and len(self.result_queue) > 0 and len(self.result_queue) > 0
@@ -2822,7 +2821,7 @@ class Scheduler(
and new_batch.input_embeds is None and new_batch.input_embeds is None
): ):
# TODO (lianmin): support return_logprob + mixed chunked prefill # TODO (lianmin): support return_logprob + mixed chunked prefill
self.running_batch.filter_batch(v1_spec_info_filtered=True) self.running_batch.filter_batch()
if not self.running_batch.is_empty(): if not self.running_batch.is_empty():
self.running_batch.prepare_for_decode() self.running_batch.prepare_for_decode()
new_batch.mix_with_running(self.running_batch) new_batch.mix_with_running(self.running_batch)
@@ -2867,7 +2866,7 @@ class Scheduler(
"""Update the current running decoding batch.""" """Update the current running decoding batch."""
initial_bs = batch.batch_size() initial_bs = batch.batch_size()
batch.filter_batch(v1_spec_info_filtered=True) batch.filter_batch()
if batch.is_empty(): if batch.is_empty():
batch.batch_is_full = False batch.batch_is_full = False
return batch return batch
@@ -2987,7 +2986,7 @@ class Scheduler(
passes overlap=False. passes overlap=False.
""" """
# 1. snapshot # 1. snapshot
snapshot_v2_full = batch.is_spec_v2 snapshot_v2_full = not batch.spec_algorithm.is_none()
sched_snapshot = ( sched_snapshot = (
{f.name: getattr(batch, f.name) for f in dataclasses.fields(batch)} {f.name: getattr(batch, f.name) for f in dataclasses.fields(batch)}
if snapshot_v2_full if snapshot_v2_full
@@ -3061,7 +3060,7 @@ class Scheduler(
self.future_map.publish, future_indices self.future_map.publish, future_indices
) )
} }
if batch.is_spec_v2 if not batch.spec_algorithm.is_none()
else {} else {}
) )
@@ -3069,7 +3068,7 @@ class Scheduler(
batch_result = self.model_worker.forward_batch_generation( batch_result = self.model_worker.forward_batch_generation(
batch, **fwd_kwargs batch, **fwd_kwargs
) )
if not batch.is_spec_v2: if batch.spec_algorithm.is_none():
self.future_map.publish(future_indices, batch.seq_lens + 1) self.future_map.publish(future_indices, batch.seq_lens + 1)
# Park any refs the worker wants kept alive 2 iters # Park any refs the worker wants kept alive 2 iters
# (cross-stream tensor lifetime; pinned in the same # (cross-stream tensor lifetime; pinned in the same
@@ -3083,7 +3082,7 @@ class Scheduler(
if batch_result.delay_sample_func is None: if batch_result.delay_sample_func is None:
stash_payload = ( stash_payload = (
batch_result.next_draft_input batch_result.next_draft_input
if batch.is_spec_v2 if not batch.spec_algorithm.is_none()
else batch_result.next_token_ids else batch_result.next_token_ids
) )
self.future_map.stash(future_indices, stash_payload) self.future_map.stash(future_indices, stash_payload)
@@ -3097,7 +3096,7 @@ class Scheduler(
# Next-iter input_ids relayed via future_map. # Next-iter input_ids relayed via future_map.
batch.input_ids = None batch.input_ids = None
if batch.is_spec_v2: if not batch.spec_algorithm.is_none():
batch.spec_info = batch_result.next_draft_input batch.spec_info = batch_result.next_draft_input
batch.spec_info.future_indices = future_indices batch.spec_info.future_indices = future_indices
elif self.enable_pdmux and batch.forward_mode.is_split_prefill(): elif self.enable_pdmux and batch.forward_mode.is_split_prefill():
@@ -3108,8 +3107,8 @@ class Scheduler(
batch.req_pool_indices, batch_result.next_token_ids batch.req_pool_indices, batch_result.next_token_ids
) )
batch.input_ids = None batch.input_ids = None
elif batch.is_spec_v2: elif not batch.spec_algorithm.is_none():
# Non-overlap V2: drive the V2 worker synchronously (no # Non-overlap: drive the V2 worker synchronously (no
# future_map relay / on_publish). # future_map relay / on_publish).
resolve_forward_inputs(batch, self.future_map) resolve_forward_inputs(batch, self.future_map)
with self._forward_isolation(batch, overlap=False): with self._forward_isolation(batch, overlap=False):
@@ -3141,17 +3140,11 @@ class Scheduler(
batch, **kwargs batch, **kwargs
) )
if isinstance(batch_result.next_token_ids, torch.Tensor): if isinstance(batch_result.next_token_ids, torch.Tensor):
if self.spec_algorithm.is_none(): # Non-spec: relay via future_map, gathered next iter.
# Non-spec: relay via future_map, gathered next iter. self.future_map.stash(
self.future_map.stash( batch.req_pool_indices, batch_result.next_token_ids
batch.req_pool_indices, batch_result.next_token_ids )
) batch.input_ids = None
batch.input_ids = None
else:
# Spec_v1 (NGRAM / DFLASH, non-overlap): worker shape
# doesn't match req_pool_indices; relay is unused (worker
# rebuilds input_ids inside verify).
batch.input_ids = batch_result.next_token_ids.to(torch.int64)
self.update_cache_from_scheduler(batch, batch_result) self.update_cache_from_scheduler(batch, batch_result)
# These 2 values are needed for processing the output, but the values can be # These 2 values are needed for processing the output, but the values can be
@@ -3796,7 +3789,7 @@ class Scheduler(
self.cur_batch = None self.cur_batch = None
if recv_req.mode == "retract" and not self.running_batch.is_empty(): if recv_req.mode == "retract" and not self.running_batch.is_empty():
self.running_batch.filter_batch(v1_spec_info_filtered=True) self.running_batch.filter_batch()
if len(self.running_batch.reqs) != 0: if len(self.running_batch.reqs) != 0:
retracted_reqs = self.running_batch.retract_all(self.server_args) retracted_reqs = self.running_batch.retract_all(self.server_args)
for req in retracted_reqs: for req in retracted_reqs:
python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+22 -44
View File
@@ -631,10 +631,6 @@ class SchedulerBatchResultProcessor:
self.token_to_kv_pool_allocator.free_group_begin() self.token_to_kv_pool_allocator.free_group_begin()
# Spec V1 handles output_ids, update_finish_state, grammar, and reasoning tokens
# in the verify phase. Non-spec and V2 handle them here in post-processing.
is_spec_v1 = not batch.spec_algorithm.is_none() and not batch.is_spec_v2
for i, req in enumerate(batch.reqs): for i, req in enumerate(batch.reqs):
req: Req req: Req
@@ -645,19 +641,6 @@ class SchedulerBatchResultProcessor:
# And all the over-allocated tokens will be freed in `release_kv_cache`. # And all the over-allocated tokens will be freed in `release_kv_cache`.
continue continue
if is_spec_v1:
req.time_stats.set_last_decode_finish_time()
self._handle_finish_state_updated_req(
req, batch, result, i, logits_output
)
if req.return_hidden_states and logits_output.hidden_states is not None:
req.hidden_states.append(
logits_output.hidden_states[i].cpu().clone().tolist()
)
if req.grammar is not None:
req.grammar.finished = req.finished()
continue
# Non-spec and V2: full post-processing # Non-spec and V2: full post-processing
next_token_id = next_token_ids[i] next_token_id = next_token_ids[i]
new_accepted_len = 1 new_accepted_len = 1
@@ -715,31 +698,27 @@ class SchedulerBatchResultProcessor:
next_token_ids: Union[torch.Tensor, List[int]], next_token_ids: Union[torch.Tensor, List[int]],
) -> Tuple[Union[List[int], List[List[int]]], Optional[List[float]]]: ) -> Tuple[Union[List[int], List[List[int]]], Optional[List[float]]]:
next_token_logprobs = None next_token_logprobs = None
if batch.spec_algorithm.is_none() or batch.is_spec_v2: if not batch.spec_algorithm.is_none():
if batch.is_spec_v2: next_token_ids = self._resolve_spec_v2_tokens(result, batch)
next_token_ids = self._resolve_spec_v2_tokens(result, batch) elif isinstance(next_token_ids, list):
elif isinstance(next_token_ids, list): pass # MLX path: already a list[int], skip torch round-trip
pass # MLX path: already a list[int], skip torch round-trip else:
else: next_token_ids = next_token_ids.tolist()
next_token_ids = next_token_ids.tolist()
if batch.return_logprob: if batch.return_logprob:
next_token_logprobs = logits_output.next_token_logprobs.tolist() next_token_logprobs = logits_output.next_token_logprobs.tolist()
if logits_output.next_token_top_logprobs_val: if logits_output.next_token_top_logprobs_val:
logits_output.next_token_top_logprobs_val = [ logits_output.next_token_top_logprobs_val = [
v.tolist() for v in logits_output.next_token_top_logprobs_val v.tolist() for v in logits_output.next_token_top_logprobs_val
] ]
logits_output.next_token_top_logprobs_idx = [ logits_output.next_token_top_logprobs_idx = [
x.tolist() for x in logits_output.next_token_top_logprobs_idx x.tolist() for x in logits_output.next_token_top_logprobs_idx
] ]
if logits_output.next_token_token_ids_logprobs_val: if logits_output.next_token_token_ids_logprobs_val:
logits_output.next_token_token_ids_logprobs_val = [ logits_output.next_token_token_ids_logprobs_val = [
v.tolist() v.tolist() for v in logits_output.next_token_token_ids_logprobs_val
for v in logits_output.next_token_token_ids_logprobs_val ]
]
# else: Spec V1 — output_ids, update_finish_state, grammar, and reasoning tokens
# are already handled in the verify phase (eagle_info.py / ngram_info.py).
return next_token_ids, next_token_logprobs return next_token_ids, next_token_logprobs
def _apply_decode_logprobs( def _apply_decode_logprobs(
@@ -752,9 +731,8 @@ class SchedulerBatchResultProcessor:
next_token_logprobs: list, next_token_logprobs: list,
logits_output: LogitsProcessorOutput, logits_output: LogitsProcessorOutput,
) -> None: ) -> None:
# Spec v1 handles logprobs inside its own worker. # Normalize: non-spec has 1 token, spec decoding has multiple.
# Normalize: non-spec has 1 token, spec v2 has multiple. if not batch.spec_algorithm.is_none():
if batch.is_spec_v2:
accepted_logprobs = next_token_logprobs[i] accepted_logprobs = next_token_logprobs[i]
accepted_ids = next_token_id accepted_ids = next_token_id
max_accept = len(accepted_logprobs) max_accept = len(accepted_logprobs)
@@ -795,7 +773,7 @@ class SchedulerBatchResultProcessor:
if batch.spec_algorithm.is_none(): if batch.spec_algorithm.is_none():
# Normal decode: single token # Normal decode: single token
req.grammar.accept_token(next_token_id) req.grammar.accept_token(next_token_id)
elif batch.is_spec_v2: else:
# Speculative decode: next_token_id is a list of accepted tokens # Speculative decode: next_token_id is a list of accepted tokens
for token_id in next_token_id: for token_id in next_token_id:
req.grammar.accept_token(token_id) req.grammar.accept_token(token_id)
python/sglang/srt/mem_cache/common.py
+2 -4
View File
@@ -159,10 +159,8 @@ def get_alloc_len_per_decode(server_args: Optional[ServerArgs] = None) -> int:
if server_args.speculative_algorithm is None: if server_args.speculative_algorithm is None:
return 1 return 1
# Spec v1: # Spec decoding allocates max(topk * num_steps, num_draft_tokens) per
# 1) alloc topk * num_steps when draft decoding and then restore the allocation # decode step (draft chain and verify block share the reservation).
# 2) alloc num_draft_tokens when verifying the drafts
# Sepc v2: allocate max(topk * num_steps, num_draft_tokens)
spec_steps = server_args.speculative_num_steps or 1 spec_steps = server_args.speculative_num_steps or 1
spec_topk = server_args.speculative_eagle_topk or 1 spec_topk = server_args.speculative_eagle_topk or 1
python/sglang/srt/mem_cache/kv_cache_builder.py
+6 -13
View File
@@ -54,19 +54,12 @@ def get_draft_kv_pool(
if draft_worker is None or spec_algorithm.is_ngram(): if draft_worker is None or spec_algorithm.is_ngram():
return None, None return None, None
# V2 (EAGLE family) nests the runner under `.draft_worker`; DFLASH / # V2 workers nest the draft runner under `.draft_worker`.
# FROZEN_KV_MTP expose `.model_runner` directly. if server_args.enable_multi_layer_eagle:
if spec_algorithm.supports_spec_v2(): draft_runner = draft_worker.draft_worker.draft_runner_list[0]
if server_args.enable_multi_layer_eagle: else:
draft_runner = draft_worker.draft_worker.draft_runner_list[0] draft_runner = draft_worker.draft_worker.draft_runner
else: return draft_runner.token_to_kv_pool, draft_runner.model_config
draft_runner = draft_worker.draft_worker.draft_runner
return draft_runner.token_to_kv_pool, draft_runner.model_config
return (
draft_worker.model_runner.token_to_kv_pool,
draft_worker.model_config,
)
def maybe_register_hicache_draft( def maybe_register_hicache_draft(
python/sglang/srt/speculative/eagle_info.py
+2 -525
View File
@@ -1,23 +1,13 @@
import copy
import logging import logging
from dataclasses import dataclass from dataclasses import dataclass
from typing import List, Optional, Tuple from typing import List, Optional, Tuple
import torch import torch
import torch.nn.functional as F
from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
from sglang.srt.distributed import get_tp_group
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
from sglang.srt.layers.dp_attention import (
get_attention_tp_group,
is_dp_attention_enabled,
)
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.sampler import apply_custom_logit_processor
from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.common import ( from sglang.srt.mem_cache.common import (
alloc_paged_token_slots_extend, alloc_paged_token_slots_extend,
alloc_token_slots, alloc_token_slots,
@@ -29,29 +19,13 @@ from sglang.srt.speculative.eagle_info_v2 import (
EagleDraftInputV2Mixin, EagleDraftInputV2Mixin,
EagleVerifyInputV2Mixin, EagleVerifyInputV2Mixin,
) )
from sglang.srt.speculative.eagle_utils import verify_tree_greedy_func
from sglang.srt.speculative.spec_info import SpecInput, SpecInputType from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
from sglang.srt.speculative.spec_utils import ( from sglang.srt.speculative.spec_utils import (
SIMULATE_ACC_LEN,
TREE_SPEC_KERNEL_AVAILABLE,
align_evict_mask_to_page_size,
assign_req_to_token_pool_func, assign_req_to_token_pool_func,
create_extend_after_decode_spec_info, create_extend_after_decode_spec_info,
create_num_accept_tokens_filter,
filter_finished_cache_loc_kernel,
generate_simulated_accept_index,
get_src_tgt_cache_loc,
get_target_cache_loc,
) )
from sglang.srt.utils import is_cuda, is_musa, next_power_of_2 from sglang.srt.utils import next_power_of_2
from sglang.srt.utils.async_probe import maybe_detect_nan, maybe_detect_oob from sglang.srt.utils.async_probe import maybe_detect_oob
if is_cuda() or is_musa():
from sgl_kernel import (
top_k_renorm_prob,
top_p_renorm_prob,
tree_speculative_sampling_target_only,
)
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -239,464 +213,6 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
return kv_indices, cum_kv_seq_len, qo_indptr, self.custom_mask return kv_indices, cum_kv_seq_len, qo_indptr, self.custom_mask
def verify(
self,
batch: ScheduleBatch,
logits_output: LogitsProcessorOutput,
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
page_size: int,
vocab_mask: Optional[torch.Tensor] = None, # For grammar
) -> torch.Tensor:
"""
Verify and find accepted tokens based on logits output and batch
(which contains spec decoding information).
WARNING: This API in-place modifies the states of logits_output
This API updates values inside logits_output based on the accepted
tokens. I.e., logits_output.next_token_logits only contains
accepted token logits.
"""
if batch.forward_mode.is_idle():
# hidden_size=None: worker fixup in forward_draft_extend_after_decode
# rebuilds via EagleDraftExtendInput.hidden_size_for(worker)
# (single source incl. EAGLE-3 aux widening).
draft_extend_input = EagleDraftExtendInput.create_idle_input(
device=batch.device,
hidden_size=None,
dtype=None,
capture_hidden_mode=CaptureHiddenMode.LAST,
)
return EagleVerifyOutput.create_idle(
draft_extend_input=draft_extend_input,
logits_output=logits_output,
device=batch.device,
spec_steps=self.spec_steps,
)
bs = self.retrieve_index.shape[0]
candidates = self.draft_token.reshape(bs, self.draft_token_num)
sampling_info = batch.sampling_info
predict_shape = list(logits_output.next_token_logits.shape)[:-1]
predict_shape[-1] += 1
predict = torch.empty(predict_shape, dtype=torch.int32, device=batch.device)
accept_index = torch.full(
(bs, self.spec_steps + 1), -1, dtype=torch.int32, device=batch.device
)
num_correct_drafts = torch.empty((bs,), dtype=torch.int32, device=batch.device)
if bs != len(sampling_info):
sampling_info = copy.deepcopy(sampling_info)
# NOTE: retrieve_index are the indices of the requests that are kept.
sampling_info.filter_batch(
self.retrieve_index.tolist(), self.retrieve_index
)
# Apply the custom logit processors if registered in the sampling info.
if sampling_info.has_custom_logit_processor:
apply_custom_logit_processor(
logits_output.next_token_logits,
sampling_info,
num_tokens_in_batch=self.draft_token_num,
)
# Apply penalty
if (
sampling_info.penalizer_orchestrator.is_required
or sampling_info.logit_bias is not None
):
# This is a relaxed version of penalties for speculative decoding.
sampling_info.penalizer_orchestrator.apply(
logits_output.next_token_logits, repeat=self.draft_token_num
)
if sampling_info.logit_bias is not None:
logits_output.next_token_logits.add_(
torch.repeat_interleave(
sampling_info.logit_bias, self.draft_token_num, dim=0
)
)
# Apply grammar mask
if vocab_mask is not None:
assert self.grammar is not None
self.grammar.apply_vocab_mask(
logits=logits_output.next_token_logits, vocab_mask=vocab_mask
)
# Sample tokens. Force greedy sampling on AMD
is_all_greedy = sampling_info.is_all_greedy
if (not is_all_greedy) and (not TREE_SPEC_KERNEL_AVAILABLE):
logger.warning(
"Tree speculative sampling kernel unavailable (likely AMD/HIP build). "
"Falling back to greedy verification."
)
if is_all_greedy or not TREE_SPEC_KERNEL_AVAILABLE:
target_predict = torch.argmax(logits_output.next_token_logits, dim=-1)
target_predict = target_predict.reshape(bs, self.draft_token_num)
predict, accept_index, num_correct_drafts = verify_tree_greedy_func(
predicts=predict, # mutable
accept_index=accept_index, # mutable
accept_token_num=num_correct_drafts, # mutable
candidates=candidates,
retrieve_index=self.retrieve_index,
retrieve_next_token=self.retrieve_next_token,
retrieve_next_sibling=self.retrieve_next_sibling,
target_predict=target_predict,
topk=self.topk,
)
else:
# apply temperature and get target probs
expanded_temperature = torch.repeat_interleave(
sampling_info.temperatures, self.draft_token_num, dim=0
) # (bs * draft_token_num, 1)
target_probs = F.softmax(
logits_output.next_token_logits / expanded_temperature, dim=-1
) # (bs * draft_token_num, vocab_size)
maybe_detect_nan(target_probs, "verify: target_probs after softmax")
target_probs = top_k_renorm_prob(
target_probs,
torch.repeat_interleave(
sampling_info.top_ks, self.draft_token_num, dim=0
),
) # (bs * draft_token_num, vocab_size)
maybe_detect_nan(target_probs, "verify: target_probs after top_k_renorm")
if sampling_info.need_top_p_sampling:
target_probs = top_p_renorm_prob(
target_probs,
torch.repeat_interleave(
sampling_info.top_ps, self.draft_token_num, dim=0
),
)
maybe_detect_nan(
target_probs, "verify: target_probs after top_p_renorm"
)
target_probs = target_probs.reshape(bs, self.draft_token_num, -1)
draft_probs = torch.zeros(
target_probs.shape, dtype=torch.float32, device=batch.device
)
# coins for rejection sampling
coins = torch.rand_like(
candidates, dtype=torch.float32, device=batch.device
)
# coins for final sampling
coins_for_final_sampling = torch.rand(
(bs,), dtype=torch.float32, device=batch.device
)
tree_speculative_sampling_target_only(
predicts=predict, # mutable
accept_index=accept_index, # mutable
accept_token_num=num_correct_drafts, # mutable
candidates=candidates,
# kwarg LHS retained as `retrive_*` to match sgl_kernel op schema.
retrive_index=self.retrieve_index,
retrive_next_token=self.retrieve_next_token,
retrive_next_sibling=self.retrieve_next_sibling,
uniform_samples=coins,
uniform_samples_for_final_sampling=coins_for_final_sampling,
target_probs=target_probs,
draft_probs=draft_probs,
threshold_single=get_global_server_args().speculative_accept_threshold_single,
threshold_acc=get_global_server_args().speculative_accept_threshold_acc,
deterministic=True,
)
# Sync sampling results across TP ranks: different GPUs may
# produce slightly different target_probs due to floating-point
# non-determinism in softmax/top_k/top_p, causing different
# sampled tokens. Broadcast from rank 0 to ensure consistency.
tp_group = (
get_attention_tp_group()
if is_dp_attention_enabled()
else get_tp_group()
)
if tp_group.world_size > 1:
tp_group.broadcast(predict, src=0)
tp_group.broadcast(accept_index, src=0)
tp_group.broadcast(num_correct_drafts, src=0)
if SIMULATE_ACC_LEN > 0.0:
# Do simulation
accept_index = generate_simulated_accept_index(
accept_index=accept_index,
predict=predict, # mutable
num_correct_drafts=num_correct_drafts, # mutable
bs=bs,
spec_steps=self.spec_steps,
)
# accept_index values index batch.out_cache_loc (size = bs * draft_token_num);
# -1 is the reject sentinel.
maybe_detect_oob(
accept_index,
-1,
bs * self.draft_token_num,
"eagle verify accept_index post-sampling",
)
maybe_detect_oob(
num_correct_drafts,
0,
self.draft_token_num + 1,
"eagle verify num_correct_drafts post-sampling",
)
unfinished_index = []
unfinished_accept_index = []
accept_index_cpu = accept_index.tolist()
predict_cpu = predict.tolist()
has_finished = False
think_end_id = batch.model_config.think_end_id
# Iterate every accepted token and check if req has finished after append the token
# should be checked BEFORE free kv cache slots
for i, (req, accept_index_row) in enumerate(zip(batch.reqs, accept_index_cpu)):
num_accept_tokens = 0
for j, idx in enumerate(accept_index_row):
if idx == -1:
break
num_accept_tokens += 1
id = predict_cpu[idx]
req.output_ids.append(id)
if req.require_reasoning and think_end_id is not None:
req.update_reasoning_tokens(id, think_end_id)
req.update_finish_state()
if not req.finished() and req.grammar is not None:
try:
req.grammar.accept_token(id)
except ValueError as e:
logger.info(
f"{i=}, {req=}\n" f"{accept_index=}\n" f"{predict=}\n"
)
raise e
req.update_finish_state()
if req.finished():
has_finished = True
# set all tokens after finished token to -1 and break
accept_index[i, j + 1 :] = -1
break
# Update KV cache tracking for the accepted tokens
req.kv_committed_len += num_accept_tokens
req.kv_allocated_len = req.kv_committed_len
if not req.finished():
unfinished_index.append(i)
if idx == -1:
unfinished_accept_index.append(accept_index[i, :j])
else:
unfinished_accept_index.append(accept_index[i])
req.spec_verify_ct += 1
num_correct_drafts_this_req = (
sum(1 for idx in accept_index_row if idx != -1) - 1
)
req.spec_num_correct_drafts += num_correct_drafts_this_req
req.update_spec_correct_drafts_histogram(num_correct_drafts_this_req)
if has_finished:
num_correct_drafts = (accept_index != -1).sum(dim=1) - 1
# Free the KV cache for unaccepted tokens
# TODO: fuse them
accept_index = accept_index[accept_index != -1]
accept_tokens = predict[accept_index]
maybe_detect_oob(
accept_tokens,
0,
batch.model_config.vocab_size,
"eagle verify accept_tokens",
)
evict_mask = torch.full_like(self.draft_token, True, dtype=torch.bool)
evict_mask[accept_index] = False
num_correct_drafts_cpu = num_correct_drafts.cpu()
num_accept_tokens_cpu = num_correct_drafts_cpu + 1
# FIXME: this `tolist()` fixes the numerical calculation consistency
# try to unify the tensor representation and list representation
num_correct_drafts_list = num_correct_drafts_cpu.tolist()
num_accept_tokens_list = num_accept_tokens_cpu.tolist()
if page_size == 1:
# TODO: boolean array index leads to a device sync. Remove it.
token_to_kv_pool_allocator.free(batch.out_cache_loc[evict_mask])
else:
if self.topk == 1:
# Only evict full empty page. Do not evict partial empty page
align_evict_mask_to_page_size[len(batch.seq_lens),](
batch.seq_lens,
evict_mask,
page_size,
self.draft_token_num,
next_power_of_2(self.draft_token_num),
)
token_to_kv_pool_allocator.free(batch.out_cache_loc[evict_mask])
else:
# Shift the accepted tokens to the beginning.
# Only evict the last part
src_cache_loc, tgt_cache_loc, to_free_num_slots = get_src_tgt_cache_loc(
batch.seq_lens,
batch.out_cache_loc,
accept_index,
num_correct_drafts,
self.draft_token_num,
page_size,
)
to_free_slots = torch.empty(
(to_free_num_slots.sum().item(),),
dtype=torch.int64,
device=to_free_num_slots.device,
)
# out_cache_loc: [0 1 2, 3 4 5, 6 7 8]
# accept_index: [0 -1 2, 3 4 -1, 6 -1 -1]
# tgt_cache_loc: [0 1 , 3 4 , 6 ]
# to_free_slots: [ 2, 5, 7 8]
# to_free_slots also needs to be page-aligned without the first partial page
#
# split each row of out_cache_loc into two parts.
# 1. the first part goes to tgt_cache_loc. length = num_correct_drafts[i] + 1
# 2. the second part goes to to_free_slots.
get_target_cache_loc[(bs,)](
tgt_cache_loc,
to_free_slots,
num_correct_drafts,
to_free_num_slots,
batch.out_cache_loc,
self.draft_token_num,
next_power_of_2(self.draft_token_num),
next_power_of_2(bs),
)
# Free the kv cache
token_to_kv_pool_allocator.free(to_free_slots)
# Copy the kv cache
batch.token_to_kv_pool_allocator.get_kvcache().move_kv_cache(
tgt_cache_loc, src_cache_loc
)
# Construct EagleVerifyOutput
if not has_finished:
if page_size == 1 or self.topk == 1:
batch.out_cache_loc = batch.out_cache_loc[accept_index]
assign_req_to_token_pool_func(
batch.req_pool_indices,
batch.req_to_token_pool.req_to_token,
batch.seq_lens,
batch.seq_lens + num_correct_drafts + 1,
batch.out_cache_loc,
bs,
)
else:
batch.out_cache_loc = tgt_cache_loc
batch.seq_lens.add_(num_correct_drafts + 1)
batch.seq_lens_cpu.add_(num_accept_tokens_cpu)
draft_extend_input = EagleDraftExtendInput(
hidden_states=(
batch.spec_info.hidden_states[accept_index]
if batch.spec_info.hidden_states is not None
else None
),
num_correct_drafts=num_correct_drafts,
num_accept_tokens=num_correct_drafts + 1,
num_accept_tokens_cpu=num_accept_tokens_list,
input_ids=accept_tokens,
seq_lens=batch.seq_lens,
seq_lens_cpu=batch.seq_lens_cpu,
req_pool_indices=batch.req_pool_indices,
)
return EagleVerifyOutput(
draft_extend_input=draft_extend_input,
logits_output=logits_output,
accept_tokens=accept_tokens,
num_correct_drafts_per_req_cpu=num_correct_drafts_list,
accept_indices=accept_index,
)
else:
if page_size == 1 or self.topk == 1:
assign_req_to_token_pool_func(
batch.req_pool_indices,
batch.req_to_token_pool.req_to_token,
batch.seq_lens,
batch.seq_lens + num_correct_drafts + 1,
batch.out_cache_loc[accept_index],
bs,
)
batch.seq_lens.add_(num_correct_drafts + 1)
batch.seq_lens_cpu.add_(num_accept_tokens_cpu)
if len(unfinished_accept_index) > 0:
unfinished_accept_index = torch.cat(unfinished_accept_index)
unfinished_index_device = torch.tensor(
unfinished_index, dtype=torch.int64, device=predict.device
)
draft_input_num_correct_drafts_cpu = [
num_correct_drafts_list[i] for i in unfinished_index
]
draft_input_num_accept_tokens_cpu = [
num_accept_tokens_list[i] for i in unfinished_index
]
if page_size == 1 or self.topk == 1:
batch.out_cache_loc = batch.out_cache_loc[unfinished_accept_index]
else:
batch.out_cache_loc = torch.empty(
len(unfinished_index) + sum(draft_input_num_correct_drafts_cpu),
dtype=torch.int64,
device=predict.device,
)
num_accept_tokens_filter = create_num_accept_tokens_filter(
num_correct_drafts,
unfinished_index_device,
batch.seq_lens,
)
batch.seq_lens_cpu.add_(num_accept_tokens_cpu)
filter_finished_cache_loc_kernel[(bs,)](
batch.out_cache_loc,
tgt_cache_loc,
num_correct_drafts,
num_accept_tokens_filter,
next_power_of_2(bs),
next_power_of_2(self.draft_token_num),
)
unfinished_num_correct_drafts = num_correct_drafts[
unfinished_index_device
]
draft_extend_input = EagleDraftExtendInput(
hidden_states=(
batch.spec_info.hidden_states[unfinished_accept_index]
if batch.spec_info.hidden_states is not None
else None
),
num_accept_tokens_cpu=draft_input_num_accept_tokens_cpu,
num_correct_drafts=unfinished_num_correct_drafts,
num_accept_tokens=unfinished_num_correct_drafts + 1,
input_ids=predict[unfinished_accept_index],
seq_lens=batch.seq_lens[unfinished_index_device],
seq_lens_cpu=batch.seq_lens_cpu[unfinished_index],
req_pool_indices=batch.req_pool_indices[unfinished_index_device],
)
else:
# hidden_size=None: worker fixup rebuilds via
# EagleDraftExtendInput.hidden_size_for(worker) (single source).
draft_extend_input = EagleDraftExtendInput.create_idle_input(
device=batch.device,
hidden_size=None,
dtype=None,
capture_hidden_mode=CaptureHiddenMode.LAST,
)
return EagleVerifyOutput(
draft_extend_input=draft_extend_input,
logits_output=logits_output,
accept_tokens=accept_tokens,
num_correct_drafts_per_req_cpu=num_correct_drafts_list,
accept_indices=accept_index,
)
@dataclass @dataclass
class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
@@ -1007,42 +523,3 @@ class EagleDraftExtendInput(SpecInput):
req_to_token.size(1), req_to_token.size(1),
) )
return kv_indices, cum_kv_seq_len, qo_indptr, None return kv_indices, cum_kv_seq_len, qo_indptr, None
@dataclass
class EagleVerifyOutput:
# Next iter's draft-extend input, installed as `batch.spec_info` for the
# draft-extend forward.
draft_extend_input: EagleDraftExtendInput
# Logit outputs from target worker.
logits_output: LogitsProcessorOutput
# All accepted tokens flat across all reqs incl. those that finished this
# step. Includes the bonus token. Used for output processing.
accept_tokens: torch.Tensor
# Accepted token length per sequence in a batch in CPU (full set).
num_correct_drafts_per_req_cpu: List[int]
# Accepted indices from logits_output.next_token_logits
accept_indices: torch.Tensor
# Whether the target verify forward ran a captured cuda graph. Set by
# the worker after `EagleVerifyInput.sample` returns; default kept so
# idle / direct constructions don't have to pass it.
can_run_cuda_graph: bool = False
@classmethod
def create_idle(
cls,
*,
draft_extend_input: EagleDraftExtendInput,
logits_output: LogitsProcessorOutput,
device: torch.device,
spec_steps: int,
) -> "EagleVerifyOutput":
return cls(
draft_extend_input=draft_extend_input,
logits_output=logits_output,
accept_tokens=torch.empty(0, dtype=torch.long, device=device),
num_correct_drafts_per_req_cpu=[],
accept_indices=torch.full(
(0, spec_steps + 1), -1, dtype=torch.int32, device=device
),
)
python/sglang/srt/speculative/spec_info.py
-8
View File
@@ -151,14 +151,6 @@ class SpeculativeAlgorithm(Enum):
) )
return None return None
def supports_spec_v2(self) -> bool:
return (
self.is_eagle()
or self.is_standalone()
or self.is_ngram()
or self.is_dflash()
)
def need_topk(self) -> bool: def need_topk(self) -> bool:
return self.is_eagle() or self.is_standalone() return self.is_eagle() or self.is_standalone()
python/sglang/srt/speculative/spec_registry.py
+20 -5
View File
@@ -4,6 +4,7 @@ should use that classmethod API; do not import from this module directly.
from __future__ import annotations from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Callable, Dict, Optional, Type from typing import TYPE_CHECKING, Callable, Dict, Optional, Type
import torch import torch
@@ -17,6 +18,8 @@ if TYPE_CHECKING:
WorkerFactory = Callable[["ServerArgs"], Type] WorkerFactory = Callable[["ServerArgs"], Type]
ServerArgsValidator = Callable[["ServerArgs"], None] ServerArgsValidator = Callable[["ServerArgs"], None]
logger = logging.getLogger(__name__)
class CustomSpecAlgo: class CustomSpecAlgo:
"""A plugin-registered speculative algorithm. Duck-types """A plugin-registered speculative algorithm. Duck-types
@@ -28,8 +31,12 @@ class CustomSpecAlgo:
branches like ``if spec_algorithm.is_eagle():`` in scheduler / branches like ``if spec_algorithm.is_eagle():`` in scheduler /
model_runner). Pass the subclass via ``spec_class=...`` at registration. model_runner). Pass the subclass via ``spec_class=...`` at registration.
Defaults: all ``is_*()`` return ``False`` except ``is_speculative``; Defaults: all ``is_*()`` return ``False`` except ``is_speculative``.
``supports_spec_v2`` follows ``supports_overlap``.
``supports_overlap=False`` is deprecated: the spec V1 worker path has been
removed, so such algorithms run on the V2 scheduler schema with overlap
disabled (synchronous). Migrate plugin workers to the V2 schema and
overlap scheduling.
""" """
def __init__( def __init__(
@@ -82,14 +89,22 @@ class CustomSpecAlgo:
# Conservative default: the larger KV reserve. # Conservative default: the larger KV reserve.
return True return True
def supports_spec_v2(self) -> bool:
return self.supports_overlap
def create_worker(self, server_args: ServerArgs) -> Type: def create_worker(self, server_args: ServerArgs) -> Type:
if not server_args.disable_overlap_schedule and not self.supports_overlap: if not server_args.disable_overlap_schedule and not self.supports_overlap:
raise ValueError( raise ValueError(
f"Speculative algorithm {self.name} does not support overlap scheduling." f"Speculative algorithm {self.name} does not support overlap scheduling."
) )
if not self.supports_overlap:
# Reached only when overlap is disabled, so the algorithm really
# does run synchronously on the V2 schema below.
logger.warning(
"Speculative algorithm %s is registered with "
"supports_overlap=False, which is deprecated: the spec V1 "
"worker path has been removed, and the algorithm now runs on "
"the V2 scheduler schema with overlap disabled (synchronous). "
"Migrate the plugin worker to support overlap scheduling.",
self.name,
)
return self.factory(server_args) return self.factory(server_args)
def get_num_tokens_per_bs_for_target_verify( def get_num_tokens_per_bs_for_target_verify(
test/registered/spec/eagle/test_spec_eagle_topk.py
+1 -1
View File
@@ -27,7 +27,7 @@ class TestEagle3Topk16(Eagle3Base, SpecCorrectnessKit, SpecAccuracyKit, SpecLogp
spec_topk = 16 spec_topk = 16
spec_tokens = 64 spec_tokens = 64
disable_overlap = True # topk>1 -> spec v1 disable_overlap = True # synchronous baseline; SpecV2 subclass flips overlap on
cuda_graph_max_bs = 5 cuda_graph_max_bs = 5
acc_length_thres = 3.1 acc_length_thres = 3.1
batch_accept_len_thres = 1.75 batch_accept_len_thres = 1.75
test/registered/unit/hardware_backend/mlx/test_attention_patching.py
+4 -1
View File
@@ -1125,6 +1125,8 @@ class TestMlxOverlapScheduler(unittest.TestCase):
# (deferred input materialization) before launching the forward. # (deferred input materialization) before launching the forward.
# Without resolve_forward_inputs in _launch_fresh, input_ids stays # Without resolve_forward_inputs in _launch_fresh, input_ids stays
# None and async_forward_batch_generation_mlx dereferences a None. # None and async_forward_batch_generation_mlx dereferences a None.
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
class _StopLoop(Exception): class _StopLoop(Exception):
pass pass
@@ -1151,7 +1153,8 @@ class TestMlxOverlapScheduler(unittest.TestCase):
prefill_input_ids_cpu=torch.tensor([1, 2, 3], dtype=torch.int64), prefill_input_ids_cpu=torch.tensor([1, 2, 3], dtype=torch.int64),
input_ids=None, input_ids=None,
mix_running_indices=None, mix_running_indices=None,
is_spec_v2=False, enable_overlap=True,
spec_algorithm=SpeculativeAlgorithm.NONE,
device="cpu", device="cpu",
) )
scheduler.get_next_batch_to_run = lambda: batch scheduler.get_next_batch_to_run = lambda: batch
test/registered/unit/spec/test_decode_bookkeeping_ownership.py
-4
View File
@@ -69,10 +69,6 @@ _OWNER_SITES = {
"DFlashDraftInputV2.prepare_for_decode", "DFlashDraftInputV2.prepare_for_decode",
"kv_allocated_len", "kv_allocated_len",
): 1, ): 1,
# spec v1: each verify path owns its own settlement
("speculative/eagle_info.py", "EagleVerifyInput.verify", "kv_committed_len"): 1,
("speculative/eagle_info.py", "EagleVerifyInput.verify", "kv_allocated_len"): 1,
("speculative/eagle_info.py", "EagleVerifyInput.verify", "spec_verify_ct"): 1,
# disaggregation decode prealloc # disaggregation decode prealloc
( (
"disaggregation/decode.py", "disaggregation/decode.py",
test/registered/unit/spec/test_spec_registry.py
+10 -4
View File
@@ -162,16 +162,22 @@ class TestCustomSpecAlgoInterface(_RegistryIsolated):
self.assertEqual(self.algo.is_some(), not self.algo.is_none()) self.assertEqual(self.algo.is_some(), not self.algo.is_none())
self.assertEqual(SpeculativeAlgorithm.EAGLE.is_some(), self.algo.is_some()) self.assertEqual(SpeculativeAlgorithm.EAGLE.is_some(), self.algo.is_some())
def test_supports_spec_v2_follows_supports_overlap(self): def test_supports_overlap_false_warns_deprecation(self):
# Plugin registered with supports_overlap=False -> not spec_v2. # supports_overlap=False plugins run the V2 schema synchronously; the
self.assertFalse(self.algo.supports_spec_v2()) # removed V1 path is surfaced as a deprecation warning at create time.
server_args = MagicMock()
server_args.disable_overlap_schedule = True
with self.assertLogs("sglang.srt.speculative.spec_registry", "WARNING") as logs:
self.algo.create_worker(server_args)
self.assertTrue(any("deprecated" in line for line in logs.output))
@SpeculativeAlgorithm.register("MY_V2", supports_overlap=True) @SpeculativeAlgorithm.register("MY_V2", supports_overlap=True)
def _factory(server_args): def _factory(server_args):
return MagicMock return MagicMock
v2 = SpeculativeAlgorithm.from_string("MY_V2") v2 = SpeculativeAlgorithm.from_string("MY_V2")
self.assertTrue(v2.supports_spec_v2()) server_args.disable_overlap_schedule = False
self.assertIs(v2.create_worker(server_args), MagicMock)
def test_create_worker_calls_factory(self): def test_create_worker_calls_factory(self):
server_args = MagicMock() server_args = MagicMock()