[Spec] Install EagleDraftExtendInput as the V2 draft-extend spec_info (#24860)

python/sglang/srt/kv_canary/plan_input.py

@@ -118,7 +118,7 @@ def _extract_prefix_lens_and_extend_seq_lens(
         out_prefix_lens.copy_(forward_batch.seq_lens[:bs].to(torch.int64))
         out_extend_seq_lens.fill_(int(spec_info.draft_token_num))
     elif forward_mode.is_draft_extend_v2():
-        # Evidence: EagleDraftInputV2Mixin.prepare_for_extend_to_fill_draft_kvcache bumps
+        # Evidence: EagleDraftExtendInputV2Mixin.prepare_for_extend_to_fill_draft_kvcache bumps
         # seq_lens by num_draft_tokens. FlashAttentionBackend.init_forward_metadata reads the
         # draft-extend-v2 query length from spec_info.extend_seq_lens_tensor when available.
         # CUDA-graph replay passes extend_seq_lens but omits extend_prefix_lens, so derive the

python/sglang/srt/speculative/eagle_info.py

@@ -9,6 +9,7 @@ from sglang.srt.environ import envs
 from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
 from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
 from sglang.srt.speculative.eagle_info_v2 import (
+    EagleDraftExtendInputV2Mixin,
     EagleDraftInputV2Mixin,
     EagleVerifyInputV2Mixin,
 )
@@ -168,10 +169,6 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
 
     # V2 overlap worker only: req_pool_indices used as buf slot keys.
     future_indices: Optional[torch.Tensor] = None
-    # V2 reuses `EagleDraftInput` across phases (V1 has a separate
-    # `EagleDraftExtendInput` for these). Set during V2's draft-extend.
-    num_correct_drafts: Optional[torch.Tensor] = None
-    num_accept_tokens: Optional[torch.Tensor] = None
 
     def __post_init__(self):
         super().__init__(SpecInputType.EAGLE_DRAFT)
@@ -276,17 +273,18 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
 
 
 @dataclass
-class EagleDraftExtendInput(SpecInput):
+class EagleDraftExtendInput(SpecInput, EagleDraftExtendInputV2Mixin):
-    """Inputs to the draft-extend forward (the per-accepted-token pass after verify).
+    """Inputs to the draft-extend forward (the fill-draft-kvcache pass after
+    target prefill / verify).
 
-    Produced by `EagleVerifyInput.verify`, installed on `batch.spec_info` for
+    Installed on `batch.spec_info` by the worker's `_draft_extend_for_*`
-    the draft-extend forward, then replaced with a fresh `EagleDraftInput` for
+    (and synthetically by draft-extend cuda-graph capture), then replaced
-    the next iter's draft.
+    with a fresh `EagleDraftInput` for the next iter's draft.
     """
 
-    # shape: (total_accepted, hidden_size). Sliced from verify-time hidden_states
+    # Target-model hidden states for the draft-extend forward; None when the
-    # by accept_index; consumed by the draft-extend forward. None when the spec
+    # draft doesn't read hidden_states (e.g., STANDALONE). Shape: decode
-    # algorithm's draft doesn't read hidden_states (e.g., STANDALONE).
+    # (bs * num_draft_tokens, hidden), prefill (extend_num_tokens, hidden).
     hidden_states: Optional[torch.Tensor] = None
 
     # Per-req accept counts. `num_accept_tokens = num_correct_drafts + 1`.

python/sglang/srt/speculative/eagle_info_v2.py

@@ -60,7 +60,11 @@ if TYPE_CHECKING:
     from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
         EAGLEDraftCudaGraphRunner,
     )
-    from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
+    from sglang.srt.speculative.eagle_info import (
+        EagleDraftExtendInput,
+        EagleDraftInput,
+        EagleVerifyInput,
+    )
 
 if is_cuda() or is_musa():
     from sgl_kernel import (
@@ -302,8 +306,10 @@ class EagleDraftInputV2Mixin:
         can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run(forward_batch)
         return forward_batch, can_cuda_graph
 
+
+class EagleDraftExtendInputV2Mixin:
     def prepare_for_extend_to_fill_draft_kvcache(
-        self,
+        self: EagleDraftExtendInput,
         batch: ScheduleBatch,
         predict: torch.Tensor,
         num_draft_tokens: int,

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -55,7 +55,11 @@ from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
 from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
     EAGLEDraftExtendCudaGraphRunner,
 )
-from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
+from sglang.srt.speculative.eagle_info import (
+    EagleDraftExtendInput,
+    EagleDraftInput,
+    EagleVerifyInput,
+)
 from sglang.srt.speculative.eagle_info_v2 import fill_bonus_tokens
 from sglang.srt.speculative.eagle_utils import (
     TreeMaskMode,
@@ -655,17 +659,15 @@ class EagleDraftWorker(BaseDraftWorker):
                 )
                 pt += extend_len
 
-        # Construct spec_info
+        # Draft-extend spec_info for the extend forward; carries only
-        next_draft_input = EagleDraftInput(
+        # hidden_states + shape info.
+        batch.spec_info = EagleDraftExtendInput(
             hidden_states=target_hidden_states,
-            bonus_tokens=next_token_ids,
             # draft mode is same with decode mode, only 1 token per req
             num_tokens_per_req=1,
             num_tokens_for_logprob_per_req=1,
         )
 
-        batch.spec_info = next_draft_input
-
         # Run forward (LAST mode: only the final hidden state per request,
         # to feed the next draft step which expects [bs, hidden_dim]).
         # STANDALONE skips hidden states end-to-end.
@@ -696,20 +698,27 @@ class EagleDraftWorker(BaseDraftWorker):
         maybe_detect_nan(logits_output.next_token_logits, "draft_extend_for_prefill")
         maybe_detect_inf(logits_output.next_token_logits, "draft_extend_for_prefill")
 
-        # Update spec_info for the next draft step
+        # Assemble the next-iter draft spec_info from the extend output.
         probs = torch.softmax(logits_output.next_token_logits, dim=-1)
-        next_draft_input.topk_p, next_draft_input.topk_index = fast_topk(
+        topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
-            probs, self.topk, dim=-1
+        return EagleDraftInput(
+            topk_p=topk_p,
+            topk_index=topk_index,
+            hidden_states=logits_output.hidden_states,
+            bonus_tokens=next_token_ids,
+            num_tokens_per_req=1,
+            num_tokens_for_logprob_per_req=1,
         )
-        next_draft_input.hidden_states = logits_output.hidden_states
-        return next_draft_input
 
     def _draft_extend_for_decode(
         self, batch: ScheduleBatch, batch_result: GenerationBatchResult
     ):
         # Batch 2: Draft extend
-        draft_input = EagleDraftInput(
+        draft_extend_input = EagleDraftExtendInput(
             hidden_states=batch_result.logits_output.hidden_states,
+            # accept_lens includes the bonus token; correct drafts exclude it.
+            num_correct_drafts=batch_result.accept_lens - 1,
+            num_accept_tokens=batch_result.accept_lens,
             # Draft-extend fills the whole tree width (num_draft_tokens) per req,
             # not num_steps + 1, so DP MLP-sync padding stays consistent for topk > 1.
             num_tokens_per_req=self.speculative_num_draft_tokens,
@@ -724,7 +733,7 @@ class EagleDraftWorker(BaseDraftWorker):
 
         # Prepare for draft extend in a separate stream
         with self.plan_stream_ctx:
-            forward_batch = draft_input.prepare_for_extend_to_fill_draft_kvcache(
+            forward_batch = draft_extend_input.prepare_for_extend_to_fill_draft_kvcache(
                 batch,
                 batch_result.next_token_ids,
                 self.speculative_num_draft_tokens,
@@ -737,12 +746,6 @@ class EagleDraftWorker(BaseDraftWorker):
                 self.plan_stream
             )
 
-        if forward_batch.spec_info.num_correct_drafts is None:
-            # `batch_result.accept_lens` already includes the bonus token, so use it
-            # directly for `num_accept_tokens` and subtract 1 for `num_correct_drafts`.
-            forward_batch.spec_info.num_correct_drafts = batch_result.accept_lens - 1
-            forward_batch.spec_info.num_accept_tokens = batch_result.accept_lens
-
         # Run draft extend batch in the main compute stream
         can_cuda_graph = (
             self.cuda_graph_runner_for_draft_extend

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -43,7 +43,11 @@ from sglang.srt.model_executor.forward_batch_info import (
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.speculative.base_spec_worker import BaseDraftWorker, BaseSpecWorker
 from sglang.srt.speculative.draft_utils import DraftBackendFactory
-from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
+from sglang.srt.speculative.eagle_info import (
+    EagleDraftExtendInput,
+    EagleDraftInput,
+    EagleVerifyInput,
+)
 from sglang.srt.speculative.eagle_info_v2 import fill_bonus_tokens
 from sglang.srt.speculative.eagle_utils import TreeMaskMode, build_tree_kernel_efficient
 from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import (
@@ -414,16 +418,15 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
             "draft_extend_for_prefill: next_token_ids before draft embed",
         )
 
-        # Construct spec_info
+        # Draft-extend spec_info for the extend forward; carries only
-        next_draft_input = EagleDraftInput(
+        # hidden_states + shape info.
+        extend_input = EagleDraftExtendInput(
             hidden_states=target_hidden_states,
-            bonus_tokens=next_token_ids,
             # draft mode is same with decode mode, only 1 token per req
             num_tokens_per_req=1,
             num_tokens_for_logprob_per_req=1,
         )
-
+        batch.spec_info = extend_input
-        batch.spec_info = next_draft_input
 
         # Chain-style MTP needs FULL to get all-token hidden states;
         # non-chain only needs LAST (the target model's hidden states).
@@ -484,8 +487,16 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
                     forward_batch.extend_seq_lens,
                     topk_index,
                 )
-        next_draft_input.topk_p = torch.cat(topk_p_list, dim=1)
+        next_draft_input = EagleDraftInput(
-        next_draft_input.topk_index = torch.cat(topk_index_list, dim=1)
+            topk_p=torch.cat(topk_p_list, dim=1),
+            topk_index=torch.cat(topk_index_list, dim=1),
+            # Chain-style left the last step's hidden_states on the extend
+            # input; non-chain keeps the target hidden states.
+            hidden_states=extend_input.hidden_states,
+            bonus_tokens=next_token_ids,
+            num_tokens_per_req=1,
+            num_tokens_for_logprob_per_req=1,
+        )
 
         # Update req_to_hidden_states_pool for KV Cache reversion
         if forward_batch.extend_seq_lens is not None:
@@ -504,7 +515,7 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
         self, batch: ScheduleBatch, batch_result: GenerationBatchResult
     ):
         # Batch 2: Draft extend
-        draft_input = EagleDraftInput(
+        draft_extend_input = EagleDraftExtendInput(
             hidden_states=batch_result.logits_output.hidden_states,
             num_tokens_per_req=self.speculative_num_steps + 1,
             num_tokens_for_logprob_per_req=1,
@@ -513,7 +524,7 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
         # Prepare for draft extend in a separate stream
         # Notice that here we use batch_result.next_token_ids as the input ids
         with self.plan_stream_ctx:
-            forward_batch = draft_input.prepare_for_extend_to_fill_draft_kvcache(
+            forward_batch = draft_extend_input.prepare_for_extend_to_fill_draft_kvcache(
                 batch,
                 batch_result.next_token_ids,
                 self.speculative_num_draft_tokens,