Reuse draft MTP indexer topk only when model config allows

EAGLE V2 draft MTP can avoid recomputing NSA/DSA indexer topk across internal draft iterations when the model declares that those indices are shareable. The port follows the upstream narrow contract: store per-forward topk on ForwardBatch, enable reuse only for topk=1, and clear the transient state after draft_forward.

Constraint: topk > 1 reorders hidden rows in select_top_k_tokens, so carried topk indices would no longer match the hidden states.

Rejected: Reuse target-side topk for draft | broader semantic change not covered by the upstream fix or current tests.

Rejected: Skip loading draft indexer weights | separate memory optimization with correctness risk for models that do not enable MTP index sharing.

Confidence: high

Scope-risk: narrow

Directive: Do not enable index_share_for_mtp_iteration without preserving the topk==1 guard and per-draft-forward cleanup.

Tested: python -m pytest test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py -q

Tested: python -m py_compile python/sglang/srt/speculative/eagle_worker_v2.py python/sglang/srt/models/deepseek_nextn.py python/sglang/srt/model_executor/forward_batch_info.py test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py

Not-tested: full GLM5 EAGLE throughput/accuracy run

python/sglang/srt/model_executor/forward_batch_info.py

@@ -401,6 +401,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     capture_hidden_mode: CaptureHiddenMode = None
     capture_draft_hidden_states: bool = False
     draft_hidden_states: Optional[torch.Tensor] = None
+    # For NSA/DSA topk_indices reuse across forward calls, e.g. EAGLE draft MTP.
+    topk_indices: Optional[torch.Tensor] = None
+    reuse_mtp_topk_indices: bool = False
 
     # For padding
     padded_static_len: int = -1  # -1 if not padded

python/sglang/srt/models/deepseek_nextn.py

@@ -530,7 +530,14 @@ class DeepseekModelNextN(nn.Module):
                 forward_batch,
                 residual,
                 zero_allocator,
+                prev_topk_indices=(
+                    forward_batch.topk_indices
+                    if forward_batch.reuse_mtp_topk_indices
+                    else None
+                ),
             )
+        if forward_batch.reuse_mtp_topk_indices:
+            forward_batch.topk_indices = topk_indices
 
         if not forward_batch.forward_mode.is_idle():
             if residual is not None:

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -152,6 +152,18 @@ class EagleDraftWorker(BaseDraftWorker):
 
         # Alias for better readability
         self.draft_runner = self.draft_worker.model_runner
+        # Reuse the first draft step's NSA/DSA indexer topk across the rest of
+        # the MTP iteration when the model config says it is safe.  The reuse is
+        # only valid for topk == 1: select_top_k_tokens reorders rows for topk
+        # > 1, which would desynchronize carried indices from hidden states.
+        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.eagle_use_aux_hidden_state = False
         if self.speculative_algorithm.is_eagle3():
             eagle_config = getattr(
@@ -441,6 +453,10 @@ class EagleDraftWorker(BaseDraftWorker):
         token_list: List[torch.Tensor] = []
         parents_list: List[torch.Tensor] = []
 
+        if self.index_share_for_mtp_iteration:
+            forward_batch.reuse_mtp_topk_indices = True
+            forward_batch.topk_indices = None
+
         # Forward multiple steps
         scores = None
         for i in range(self.speculative_num_steps):
@@ -505,6 +521,10 @@ class EagleDraftWorker(BaseDraftWorker):
             batch_size = parents_list[0].shape[0]
             parent_list = torch.empty(batch_size, 0, device=parents_list[0].device)
 
+        if self.index_share_for_mtp_iteration:
+            forward_batch.topk_indices = None
+            forward_batch.reuse_mtp_topk_indices = False
+
         return parent_list, top_scores_index, draft_tokens
 
     def draft_extend(self):