diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 299ad3bc4..f3557a1fa 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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 diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index e07e5b83c..3bccdd475 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/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, diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py index 1f4b8f0d9..6c77f63cb 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py +++ b/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):