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

[Spec] Cleanup idle stub and shape-check patterns (#24881)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-05-10 02:39:53 -07:00
committed by GitHub
parent c7f674e427
commit 8cc16c9974
4 changed files with 16 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/speculative/eagle_info.py
+4
View File
@@ -644,6 +644,10 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
@dataclass
class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
# For idle stubs use `create_idle_input`, not the bare ctor: `filter_batch`
# / `merge_batch` slice / cat `topk_p` / `topk_index` / `hidden_states` /
# `bonus_tokens` unconditionally.
# shape: (b, topk)
topk_p: torch.Tensor = None
topk_index: torch.Tensor = None
python/sglang/srt/speculative/eagle_worker.py
+5 -7
View File
@@ -515,12 +515,10 @@ class EAGLEWorker(TpModelWorker):
batch.spec_info = next_draft_input
else:
# All reqs finished and dp_attention isn't forcing extend.
# Stash an empty EagleDraftInput so next iter's merge_batch
# short-circuits on None hidden_states (EagleVerifyInput
# has no merge_batch).
batch.spec_info = EagleDraftInput(
capture_hidden_mode=CaptureHiddenMode.LAST,
)
# Install an idle EagleDraftInput so next iter's scheduler
# ops (merge_batch / filter_batch) see well-typed empty
# tensors instead of None.
self._draft_preprocess_idle(batch)
set_time_batch(
batch.reqs, "set_spec_draft_extend_end_time", trace_only=True
@@ -1128,7 +1126,7 @@ class EAGLEWorker(TpModelWorker):
input_is_idle = batch.forward_mode.is_idle()
if not input_is_idle and draft_extend_input.input_ids.numel() == 0:
if not input_is_idle and draft_extend_input.input_ids.shape[0] == 0:
# All reqs finished this verify; swap to an idle ExtendInput.
batch = batch.copy()
batch.prepare_for_idle()
python/sglang/srt/speculative/frozen_kv_mtp_worker.py
+2 -2
View File
@@ -462,7 +462,7 @@ class FrozenKVMTPWorker(TpModelWorker):
draft_extend_input = verify_output.draft_extend_input
if (
self.server_args.enable_dp_attention
or draft_extend_input.input_ids.numel() > 0
or draft_extend_input.input_ids.shape[0] > 0
):
# Stash for the seed step; _run_assistant_seed_step swaps in
# a fresh FrozenKVMTPDraftInput for next iter.
@@ -512,7 +512,7 @@ class FrozenKVMTPWorker(TpModelWorker):
draft_extend_input: FrozenKVMTPDraftExtendInput = batch.spec_info
input_is_idle = batch.forward_mode.is_idle()
if not input_is_idle and draft_extend_input.input_ids.numel() == 0:
if not input_is_idle and draft_extend_input.input_ids.shape[0] == 0:
# All reqs finished; stash an idle FrozenKVMTPDraftInput so the
# next-iter draft sees a valid spec_info.
batch = batch.copy()
python/sglang/srt/speculative/multi_layer_eagle_worker.py
+5 -7
View File
@@ -293,12 +293,10 @@ class MultiLayerEagleWorker(TpModelWorker):
batch.spec_info = next_draft_input
else:
# All reqs finished and dp_attention isn't forcing extend.
# Stash an empty EagleDraftInput so next iter's merge_batch
# short-circuits on None hidden_states (EagleVerifyInput
# has no merge_batch).
batch.spec_info = EagleDraftInput(
capture_hidden_mode=CaptureHiddenMode.LAST,
)
# Install an idle EagleDraftInput so next iter's scheduler
# ops (merge_batch / filter_batch) see well-typed empty
# tensors instead of None.
self._draft_preprocess_idle(batch)
return GenerationBatchResult(
logits_output=logits_output,
@@ -676,7 +674,7 @@ class MultiLayerEagleWorker(TpModelWorker):
input_is_idle = batch.forward_mode.is_idle()
if not input_is_idle and draft_extend_input.input_ids.numel() == 0:
if not input_is_idle and draft_extend_input.input_ids.shape[0] == 0:
batch = batch.copy()
batch.prepare_for_idle()
hidden_size = (