From 8cc16c99748b8c92097d6094aad694448fb8fb0e Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sun, 10 May 2026 02:39:53 -0700 Subject: [PATCH] [Spec] Cleanup idle stub and shape-check patterns (#24881) --- python/sglang/srt/speculative/eagle_info.py | 4 ++++ python/sglang/srt/speculative/eagle_worker.py | 12 +++++------- .../sglang/srt/speculative/frozen_kv_mtp_worker.py | 4 ++-- .../srt/speculative/multi_layer_eagle_worker.py | 12 +++++------- 4 files changed, 16 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index 593e5392d..ecafdeb3b 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -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 diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index 515ef3796..7b86462af 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -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() diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py index b60380072..d73da6fc5 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py @@ -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() diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker.py b/python/sglang/srt/speculative/multi_layer_eagle_worker.py index eaca6d74d..c268168ba 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker.py @@ -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 = (