Reuse draft MTP indexer topk on spec v1 EAGLE

Spec v1 EAGLE now follows the same model-config-gated MTP index reuse contract as spec v2. This lets models that declare index_share_for_mtp_iteration reuse the first draft step's NSA/DSA topk indices across internal MTP iterations while keeping the unsafe topk>1 case disabled.

Constraint: select_top_k_tokens can reorder hidden rows when topk > 1, so carried topk indices are only valid under topk == 1.

Rejected: Enable reuse unconditionally | models without the config flag may not have compatible MTP index semantics.

Rejected: Broaden to target-to-draft index reuse | separate semantic change with different correctness risks.

Confidence: high

Scope-risk: narrow

Directive: Keep spec v1 and spec v2 MTP index reuse semantics aligned, including 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.py test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py

Not-tested: full spec v1 GLM5 online throughput run

python/sglang/srt/speculative/eagle_worker.py

@@ -156,6 +156,19 @@ class EAGLEWorker(TpModelWorker):
                 memory_pool_config=target_worker.model_runner.memory_pool_config,
             )
 
+        # 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_model_runner.model_config.hf_config,
+                "index_share_for_mtp_iteration",
+                False,
+            )
+            and self.topk == 1
+        )
+
         embed, head = self.target_worker.model_runner.model.get_embed_and_head()
 
         if self.speculative_algorithm.is_eagle3():
@@ -701,6 +714,10 @@ class EAGLEWorker(TpModelWorker):
         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):
@@ -751,6 +768,10 @@ class EAGLEWorker(TpModelWorker):
             score_list, token_list, parents_list, self.speculative_num_draft_tokens
         )
 
+        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 clear_cache_pool(self):

test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py

@@ -365,6 +365,71 @@ def test_eagle_v2_draft_forward_scopes_mtp_index_reuse_to_one_draft_forward():
     assert topk_clears >= 2
 
 
+def test_eagle_v1_draft_worker_enables_mtp_index_reuse_only_from_model_config():
+    """Spec-v1 EAGLE must use the same guarded MTP index reuse contract."""
+
+    tree = _parse_module("python/sglang/srt/speculative/eagle_worker.py")
+    cls = _find_class(tree, "EAGLEWorker")
+    init = _find_method(cls, "__init__")
+
+    assert "index_share_for_mtp_iteration" in _assigned_self_attrs(init)
+
+    assign = next(
+        node
+        for node in ast.walk(init)
+        if isinstance(node, ast.Assign)
+        and any(
+            isinstance(target, ast.Attribute)
+            and isinstance(target.value, ast.Name)
+            and target.value.id == "self"
+            and target.attr == "index_share_for_mtp_iteration"
+            for target in node.targets
+        )
+    )
+
+    text = ast.unparse(assign.value)
+    assert "index_share_for_mtp_iteration" in text
+    assert "self.topk == 1" in text
+
+
+def test_eagle_v1_draft_forward_scopes_mtp_index_reuse_to_one_draft_forward():
+    """Spec-v1 EAGLE must clear transient MTP topk reuse state per call."""
+
+    tree = _parse_module("python/sglang/srt/speculative/eagle_worker.py")
+    cls = _find_class(tree, "EAGLEWorker")
+    func = _find_method(cls, "draft_forward")
+
+    reuse_assigns = []
+    topk_clears = 0
+    for node in ast.walk(func):
+        if not isinstance(node, ast.Assign):
+            continue
+        for target in node.targets:
+            if (
+                isinstance(target, ast.Attribute)
+                and isinstance(target.value, ast.Name)
+                and target.value.id == "forward_batch"
+            ):
+                if target.attr == "reuse_mtp_topk_indices":
+                    reuse_assigns.append(node.value)
+                if (
+                    target.attr == "topk_indices"
+                    and isinstance(node.value, ast.Constant)
+                    and node.value.value is None
+                ):
+                    topk_clears += 1
+
+    assert any(
+        isinstance(value, ast.Constant) and value.value is True
+        for value in reuse_assigns
+    )
+    assert any(
+        isinstance(value, ast.Constant) and value.value is False
+        for value in reuse_assigns
+    )
+    assert topk_clears >= 2
+
+
 def test_deepseek_nextn_reuses_and_updates_mtp_topk_indices_when_requested():
     """NextN must pass cached MTP topk indices into the decoder and update them.