[Spec] Restore index_share_for_mtp_iteration in EAGLE V2 draft worker (#28192)

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -188,6 +188,16 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             self.eagle_use_aux_hidden_state = eagle_config.get(
                 "use_aux_hidden_state", True
             )
+        # Reuse the first draft step's NSA/DSA indexer topk across the rest;
+        # topk == 1 only (select_top_k_tokens reorders rows, desyncing indices).
+        self.index_share_for_mtp_iteration = (
+            getattr(
+                self.draft_runner.model_config.hf_config,
+                "index_share_for_mtp_iteration",
+                False,
+            )
+            and self.topk == 1
+        )
         self.draft_tp_context = (
             draft_tp_context if server_args.enable_dp_attention else empty_context
         )
@@ -542,6 +552,9 @@ class EagleDraftWorker(EagleDraftWorkerBase):
 
         # Forward multiple steps
         scores = None
+        if self.index_share_for_mtp_iteration:
+            forward_batch.reuse_mtp_topk_indices = True
+            forward_batch.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
@@ -609,6 +622,10 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             hidden_states = logits_output.hidden_states
             forward_batch.positions.add_(1)
 
+        if self.index_share_for_mtp_iteration:
+            forward_batch.topk_indices = None
+            forward_batch.reuse_mtp_topk_indices = False
+
         # Organize the results
         if (
             self.topk == 1

python/sglang/srt/speculative/standalone_worker_v2.py

@@ -104,6 +104,15 @@ class StandaloneDraftWorker(EagleDraftWorker):
         )
         self.tree_mask_mode = TreeMaskMode.FULL_MASK
         self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
+        # draft_forward reads this (set in EagleDraftWorker.__init__, skipped here).
+        self.index_share_for_mtp_iteration = (
+            getattr(
+                self.draft_runner.model_config.hf_config,
+                "index_share_for_mtp_iteration",
+                False,
+            )
+            and self.topk == 1
+        )
 
     def alloc_memory_pool(
         self,

python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py

@@ -187,6 +187,15 @@ class _EagleDraftWorkerHarness:
         self._topk1_parents_prealloc = None
         self._topk1_score_indices_prealloc = None
         EagleDraftWorker._rebuild_topk1_chain_buffers(self)
+        # draft_forward reads this (set in EagleDraftWorker.__init__, skipped here).
+        self.index_share_for_mtp_iteration = (
+            getattr(
+                self.model_config.hf_config,
+                "index_share_for_mtp_iteration",
+                False,
+            )
+            and self.topk == 1
+        )
 
     @property
     def draft_model_runner(self):