diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index b1044dd9a..d485c922f 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/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): diff --git a/test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py b/test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py index 647159b2a..43fed3726 100644 --- a/test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py +++ b/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.