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

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

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
+                    # Install an idle EagleDraftInput so next iter's scheduler
-                    # short-circuits on None hidden_states (EagleVerifyInput
+                    # ops (merge_batch / filter_batch) see well-typed empty
-                    # has no merge_batch).
+                    # tensors instead of None.
-                    batch.spec_info = EagleDraftInput(
+                    self._draft_preprocess_idle(batch)
-                        capture_hidden_mode=CaptureHiddenMode.LAST,
-                    )
 
             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

@@ -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

@@ -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
+                    # Install an idle EagleDraftInput so next iter's scheduler
-                    # short-circuits on None hidden_states (EagleVerifyInput
+                    # ops (merge_batch / filter_batch) see well-typed empty
-                    # has no merge_batch).
+                    # tensors instead of None.
-                    batch.spec_info = EagleDraftInput(
+                    self._draft_preprocess_idle(batch)
-                        capture_hidden_mode=CaptureHiddenMode.LAST,
-                    )
 
             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 = (