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

[spec decoding] supports step 0 in adaptive spec decoding (updating draft kv cache without draft decoding) (#23994)

Browse Source 
Co-authored-by: shuwenn <2508695655@qq.com>
This commit is contained in:
Qiaolin Yu
2026-06-15 22:21:26 -07:00
committed by GitHub
co-authored by shuwenn
parent 800aaefc9e
commit e068355831
6 changed files with 313 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+4
View File
@@ -638,6 +638,10 @@ class Envs:
# Spec Config
SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK = EnvBool(True)
# Skip draft_extend while adaptive spec is at steps=0 (drafting disabled).
# Saves the per-step draft forward, but the draft KV goes stale: an upshift
# back to steps>0 starts from a cold draft state (low accept until it recovers).
SGLANG_SPEC_SKIP_ZERO_STEP_DRAFT_EXTEND = EnvBool(False)
# Master switch for all async-asserted invariant probes (NaN, Inf, OOB,
# page alignment). Off in prod; tests turn it on to fail-fast on
# numerical / index violations instead of getting silent NaN cascades.
python/sglang/srt/speculative/adaptive_spec_params.py
+45 -17
View File
@@ -19,7 +19,6 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
# TODO: add step=0 (nospec fallback) for BS>=8 once supported.
DEFAULT_ADAPTIVE_CONFIG: dict[str, dict] = {
"1": {
"candidate_steps": [1, 3, 7],
@@ -28,13 +27,19 @@ DEFAULT_ADAPTIVE_CONFIG: dict[str, dict] = {
"ceiling_coeff": 0,
},
"8": {
"candidate_steps": [1, 3],
"candidate_steps": [0, 1, 3],
"up_hysteresis": 0.0,
"down_hysteresis": 0.0,
"ceiling_coeff": 0,
},
"32": {
"candidate_steps": [1],
"candidate_steps": [0, 1],
"up_hysteresis": 0.0,
"down_hysteresis": 0.0,
"ceiling_coeff": 0,
},
"64": {
"candidate_steps": [0],
"up_hysteresis": 0.0,
"down_hysteresis": 0.0,
"ceiling_coeff": 0,
@@ -102,10 +107,11 @@ def _load_adaptive_config(
if (
not isinstance(steps, list)
or not steps
or not all(isinstance(s, int) and s > 0 for s in steps)
or not all(isinstance(s, int) and s >= 0 for s in steps)
):
raise ValueError(
f"BS {key}: candidate_steps must be a list of positive ints, got {steps!r}"
f"BS {key}: candidate_steps must be a list of non-negative ints, "
f"got {steps!r}"
)
bs_entries[int(key)] = entry
@@ -172,10 +178,13 @@ class AdaptiveStepSlot:
if not num_correct_drafts_per_req:
return False
batch_avg = sum(num_correct_drafts_per_req) / len(num_correct_drafts_per_req)
self.ema_accept_len = (
1 - self.ema_alpha
) * self.ema_accept_len + self.ema_alpha * batch_avg
if self.current_steps > 0:
batch_avg = sum(num_correct_drafts_per_req) / len(
num_correct_drafts_per_req
)
self.ema_accept_len = (
1 - self.ema_alpha
) * self.ema_accept_len + self.ema_alpha * batch_avg
self._batch_count += 1
if self._batch_count <= self.warmup_batches:
@@ -190,23 +199,39 @@ class AdaptiveStepSlot:
"""Recompute steps from EMA. Returns True if params changed."""
old_steps = self.current_steps
current_idx = self.candidate_steps.index(old_steps)
old_idx = current_idx
# Probe the smallest positive step after a zero-step nospec interval.
if old_steps == 0:
current_idx = min(current_idx + 1, len(self.candidate_steps) - 1)
target = self.candidate_steps[current_idx]
if target > 0 and self.ema_accept_len < 0:
# A slot initialized at steps=0 has no draft acceptance history;
# start the first positive-step probe from that step's neutral EMA.
self.ema_accept_len = float(target - 1)
return self._apply_target_steps(old_steps, target)
# TODO: Consider limiting step changes to avoid overshooting.
while current_idx > 0:
prev_step = self.candidate_steps[current_idx - 1]
drop_threshold = prev_step - 0.5 + self.down_hysteresis
# A zero-step candidate disables drafting. Treat zero accepted drafts
# as low enough to reach it when it is the floor candidate.
drop_threshold = 0.5 if prev_step == 0 else prev_step - 0.5
drop_threshold += self.down_hysteresis
if self.ema_accept_len <= drop_threshold:
current_idx -= 1
else:
break
while current_idx < len(self.candidate_steps) - 1:
current_step = self.candidate_steps[current_idx]
rise_threshold = current_step - 0.5 + self.up_hysteresis
if self.ema_accept_len > rise_threshold:
current_idx += 1
else:
break
moved_down = current_idx < old_idx
if not moved_down:
while current_idx < len(self.candidate_steps) - 1:
current_step = self.candidate_steps[current_idx]
rise_threshold = current_step - 0.5 + self.up_hysteresis
if self.ema_accept_len > rise_threshold:
current_idx += 1
else:
break
target = self.candidate_steps[current_idx]
# EMA ceiling: only caps downward — never blocks step-ups, so the
@@ -218,6 +243,9 @@ class AdaptiveStepSlot:
current_idx -= 1
target = self.candidate_steps[current_idx]
return self._apply_target_steps(old_steps, target)
def _apply_target_steps(self, old_steps: int, target: int) -> bool:
if target != old_steps:
self.current_steps = target
log_info_on_rank0(
python/sglang/srt/speculative/draft_utils.py
+2 -1
View File
@@ -37,7 +37,8 @@ class DraftBackendFactory:
return backend_map[backend_type]()
def create_decode_backend(self):
if self.speculative_num_steps == 1:
# No multi-step draft backend for steps=0 (nospec) or steps=1.
if self.speculative_num_steps <= 1:
return None
backend_map = {
python/sglang/srt/speculative/eagle_worker_v2.py
+113 -17
View File
@@ -1011,33 +1011,129 @@ class EAGLEWorkerV2(BaseSpecWorker):
topk=self.topk,
capture_hidden_mode=capture_mode,
)
with (
self.draft_worker.draft_tp_context(
self.draft_worker.draft_runner.tp_group
),
speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(),
spec_stage_span("draft"),
):
verify_input: EagleVerifyInput = self.draft_worker.draft(batch)
if self.speculative_num_steps == 0:
# Drafting disabled (high batch size). _draft_extend below still
# runs, keeping draft KV warm for when the batch shrinks.
verify_input = self._build_trivial_verify_input(batch)
else:
with (
self.draft_worker.draft_tp_context(
self.draft_worker.draft_runner.tp_group
),
speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(),
spec_stage_span("draft"),
):
verify_input: EagleVerifyInput = self.draft_worker.draft(batch)
assert verify_input.is_verify_input()
batch.spec_info = verify_input
batch_output = self.verify(batch)
# Publish before draft_extend so the fence is at verify-end.
if on_publish is not None:
on_publish(batch_output.new_seq_lens)
with (
self.draft_worker.draft_tp_context(
self.draft_worker.draft_runner.tp_group
),
speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(),
spec_stage_span("draft_extend"),
if (
self.speculative_num_steps == 0
and envs.SGLANG_SPEC_SKIP_ZERO_STEP_DRAFT_EXTEND.get()
):
self.draft_worker._draft_extend_for_decode(batch, batch_output)
self._stub_skipped_draft_extend(batch, batch_output)
else:
with (
self.draft_worker.draft_tp_context(
self.draft_worker.draft_runner.tp_group
),
speculative_moe_backend_context(),
speculative_moe_a2a_backend_context(),
spec_stage_span("draft_extend"),
):
self.draft_worker._draft_extend_for_decode(batch, batch_output)
return batch_output
def _build_trivial_verify_input(self, batch: ScheduleBatch) -> EagleVerifyInput:
"""Build a 1-node EagleVerifyInput rooted at the previous bonus token.
Used when ``speculative_num_steps == 0`` to skip drafting while still
routing through the existing TARGET_VERIFY graph captured at
``draft_token_num=1``: the kernel always accepts the root and samples
one new bonus token from target logits -- functionally a plain decode.
"""
if batch.forward_mode.is_idle():
return EagleVerifyInput.create_idle_input(
topk=self.topk, spec_steps=0, num_verify_tokens=1
)
draft_input: EagleDraftInput = batch.spec_info
bs = batch.seq_lens.shape[0]
device = self.device
retrieve_index = torch.arange(bs, dtype=torch.long, device=device).unsqueeze(1)
retrieve_next_token = torch.full((bs, 1), -1, dtype=torch.long, device=device)
retrieve_next_sibling = torch.full((bs, 1), -1, dtype=torch.long, device=device)
attn_backend = self._target_worker.model_runner.attn_backend
mask_buf, position_buf = attn_backend.get_verify_buffers_to_fill_after_draft()
if mask_buf is not None:
custom_mask = mask_buf
custom_mask.fill_(True)
else:
if batch.seq_lens_sum is not None:
seq_lens_sum = batch.seq_lens_sum
elif batch.seq_lens_cpu is not None:
seq_lens_sum = int(batch.seq_lens_cpu.sum())
else:
seq_lens_sum = bs * attn_backend.max_context_len
custom_mask = torch.ones(seq_lens_sum + bs, dtype=torch.bool, device=device)
if position_buf is not None:
positions = position_buf
positions[:bs].copy_(batch.seq_lens)
else:
positions = batch.seq_lens.to(torch.int64)
return EagleVerifyInput(
draft_token=draft_input.bonus_tokens,
custom_mask=custom_mask,
positions=positions,
retrieve_index=retrieve_index,
retrieve_next_token=retrieve_next_token,
retrieve_next_sibling=retrieve_next_sibling,
retrieve_cum_len=None,
spec_steps=0,
topk=self.topk,
draft_token_num=1,
capture_hidden_mode=CaptureHiddenMode.FULL,
seq_lens_sum=None,
seq_lens_cpu=None,
)
def _stub_skipped_draft_extend(
self, batch: ScheduleBatch, batch_output: GenerationBatchResult
) -> None:
"""Fill shape-valid stubs on next_draft_input when draft_extend is skipped.
``verify`` already set ``bonus_tokens`` (the only field the next steps=0
verify reads). The overlap FutureMap still stashes topk_p/topk_index/
hidden_states, so provide zeroed tensors of the right shape. They are never
consumed while at steps=0; an upshift to steps>0 would draft from this stale
state (cold recovery), which is the documented cost of this experimental flag.
"""
next_draft_input: EagleDraftInput = batch_output.next_draft_input
bs = batch.seq_lens.shape[0]
device = self.device
next_draft_input.topk_p = torch.zeros(
(bs, self.topk), dtype=torch.float32, device=device
)
next_draft_input.topk_index = torch.zeros(
(bs, self.topk), dtype=torch.int64, device=device
)
hidden_size = EagleDraftInput.hidden_size_for(self.draft_worker)
if hidden_size is not None:
next_draft_input.hidden_states = torch.zeros(
(bs, hidden_size),
dtype=EagleDraftInput.dtype_for(self.draft_worker),
device=device,
)
def on_verify_complete_cpu(
self, num_correct_drafts_per_req: list[int], batch_size: int = 0
) -> None: