From 45314a9fcb1b70a764f7d1f6c4238560fc524a25 Mon Sep 17 00:00:00 2001
From: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
Date: Mon, 29 Jun 2026 22:52:23 +0100
Subject: [PATCH] [spec] Fix index_share_for_mtp_iteration being a no-op in
 EAGLE MTP draft (#29654)

---
 python/sglang/srt/model_executor/forward_batch_info.py | 4 ++--
 python/sglang/srt/models/deepseek_nextn.py             | 4 ++--
 python/sglang/srt/speculative/eagle_info.py            | 4 ++++
 python/sglang/srt/speculative/eagle_worker_v2.py       | 4 ++--
 4 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py
index 0bc7e620f..721b1e447 100644
--- a/python/sglang/srt/model_executor/forward_batch_info.py
+++ b/python/sglang/srt/model_executor/forward_batch_info.py
@@ -430,8 +430,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     # For hidden states before normal
     return_hidden_states_before_norm: bool = False
 
-    # For NSA/DSA topk_indices reuse across forward calls (e.g., EAGLE draft)
-    topk_indices: Optional[torch.Tensor] = None
+    # Gate for reusing the first MTP draft step's indexer topk across steps;
+    # the carried topk lives on spec_info (see EagleDraftInput.mtp_topk_indices).
     reuse_mtp_topk_indices: Optional[bool] = False
 
     # === Forward-derived (built in init_new on the forward stream; FB-owned) ===
diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py
index 10ad9c3ea..74d504458 100644
--- a/python/sglang/srt/models/deepseek_nextn.py
+++ b/python/sglang/srt/models/deepseek_nextn.py
@@ -228,13 +228,13 @@ class DeepseekModelNextN(nn.Module):
                     residual,
                     zero_allocator,
                     prev_topk_indices=(
-                        forward_batch.topk_indices
+                        forward_batch.spec_info.mtp_topk_indices
                         if forward_batch.reuse_mtp_topk_indices
                         else None
                     ),
                 )
                 if forward_batch.reuse_mtp_topk_indices:
-                    forward_batch.topk_indices = topk_indices
+                    forward_batch.spec_info.mtp_topk_indices = topk_indices
 
             if not forward_batch.forward_mode.is_idle():
                 if residual is not None:
diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py
index 926886054..c1abd35b3 100644
--- a/python/sglang/srt/speculative/eagle_info.py
+++ b/python/sglang/srt/speculative/eagle_info.py
@@ -159,6 +159,10 @@ class EagleDraftInput(SpecInput):
     hidden_states: Optional[torch.Tensor] = None
     capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.FULL
 
+    # Survives across draft steps: spec_info is shared by reference across the
+    # per-step forwards (each runs on a copied ForwardBatch, dropping writebacks).
+    mtp_topk_indices: Optional[torch.Tensor] = None
+
     # Per-req bonus token (the "+1" target prediction at end of each accept
     # chain); the worker copies it here post-extend for next iter's draft.
     bonus_tokens: torch.Tensor = None
diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py
index 0d158f429..e5a1c151e 100644
--- a/python/sglang/srt/speculative/eagle_worker_v2.py
+++ b/python/sglang/srt/speculative/eagle_worker_v2.py
@@ -612,7 +612,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
         scores = None
         if self.index_share_for_mtp_iteration:
             forward_batch.reuse_mtp_topk_indices = True
-            forward_batch.topk_indices = None
+            spec_info.mtp_topk_indices = None
         for i in range(self.speculative_num_steps):
             input_ids, hidden_states, scores, tree_info = select_top_k_tokens(
                 i, topk_p, topk_index, hidden_states, scores, self.topk
@@ -688,7 +688,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             forward_batch.positions.add_(1)
 
         if self.index_share_for_mtp_iteration:
-            forward_batch.topk_indices = None
+            spec_info.mtp_topk_indices = None
             forward_batch.reuse_mtp_topk_indices = False
 
         # Organize the results