[feature] add index share feature for DSA MTP (#44420)

Signed-off-by: JaredforReal <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

vllm/model_executor/models/deepseek_mtp.py

@@ -115,7 +115,9 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
         )
 
         hidden_states, residual = self.mtp_block(
-            positions=positions, hidden_states=hidden_states, residual=None
+            positions=positions,
+            hidden_states=hidden_states,
+            residual=None,
         )
         hidden_states = residual + hidden_states
         return hidden_states
@@ -147,6 +149,22 @@ class DeepSeekMultiTokenPredictor(nn.Module):
         )
         self.logits_processor = LogitsProcessor(config.vocab_size)
 
+    def set_skip_topk(self, skip: bool):
+        """Toggle skip_topk on all MTP layers with sparse attention.
+
+        Called by the proposer to implement index_share_for_mtp_iteration:
+        step 0 sets skip=False (compute own indices), steps 1+ set skip=True
+        (reuse step 0's indices).
+        """
+        for layer in self.layers.values():
+            mtp_block = getattr(layer, "mtp_block", None)
+            if mtp_block is not None:
+                self_attn = getattr(mtp_block, "self_attn", None)
+                if self_attn is not None:
+                    mla_attn = getattr(self_attn, "mla_attn", None)
+                    if mla_attn is not None and hasattr(mla_attn, "skip_topk"):
+                        mla_attn.skip_topk = skip
+
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.embed_tokens(input_ids)
 
@@ -225,7 +243,11 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
         spec_step_idx: int = 0,
     ) -> torch.Tensor:
         hidden_states = self.model(
-            input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
+            input_ids,
+            positions,
+            hidden_states,
+            inputs_embeds,
+            spec_step_idx,
         )
         return hidden_states
 

vllm/model_executor/models/deepseek_v2.py

@@ -1018,19 +1018,20 @@ class DeepseekV2MLAAttention(nn.Module):
                 is_inplace_rope=self.indexer_rope_emb.enabled(),
             )
 
-            # Enable IndexCache for DeepSeek models to reduce redundant top-k
-            # token selection computations in sparse attention.
-            use_index_cache = getattr(config, "use_index_cache", False)
-            if use_index_cache:
-                # IndexCache config
-                # Refer: https://arxiv.org/abs/2603.12201 for more details.
-                _index_topk_freq = getattr(config, "index_topk_freq", 1)
-                _index_topk_pattern = getattr(config, "index_topk_pattern", None)
-                layer_id = extract_layer_index(prefix)
-                if _index_topk_pattern is None:
-                    _skip_topk = max(layer_id - 1, 0) % _index_topk_freq != 0
-                elif 0 <= layer_id < len(_index_topk_pattern):
-                    _skip_topk = _index_topk_pattern[layer_id] == "S"
+            # IndexCache config
+            # Refer: https://arxiv.org/abs/2603.12201 for more details.
+            _index_topk_freq = getattr(config, "index_topk_freq", 1)
+            _index_topk_pattern = getattr(config, "index_topk_pattern", None)
+            _index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2)
+            layer_id = extract_layer_index(prefix)
+
+            if _index_topk_pattern is None:
+                _skip_topk = (
+                    max(layer_id - _index_skip_topk_offset + 1, 0) % _index_topk_freq
+                    != 0
+                )
+            elif 0 <= layer_id < len(_index_topk_pattern):
+                _skip_topk = _index_topk_pattern[layer_id] == "S"
 
         else:
             self.indexer_rope_emb = None
@@ -1252,8 +1253,8 @@ class DeepseekV2Model(nn.Module):
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers,
             lambda prefix: DeepseekV2DecoderLayer(
-                vllm_config,
-                prefix,
+                vllm_config=vllm_config,
+                prefix=prefix,
                 topk_indices_buffer=topk_indices_buffer,
             ),
             prefix=f"{prefix}.layers",

vllm/transformers_utils/model_arch_config_convertor.py

@@ -50,7 +50,7 @@ class ModelArchConfigConvertorBase:
             # special case for deepseek_v4
             if hasattr(self.hf_text_config, "compress_ratios"):
                 return self.hf_text_config.head_dim
-            qk_rope_head_dim = getattr(self.hf_text_config, "qk_rope_head_dim", 0)
+            qk_rope_head_dim = self._get_qk_rope_head_dim()
             if not envs.VLLM_MLA_DISABLE:
                 return self.hf_text_config.kv_lora_rank + qk_rope_head_dim
             else:
@@ -71,6 +71,38 @@ class ModelArchConfigConvertorBase:
         # FIXME(woosuk): This may not be true for all models.
         return self.get_hidden_size() // total_num_attention_heads
 
+    def _get_qk_rope_head_dim(self) -> int:
+        """Get qk_rope_head_dim, fixing the transformers v5.4+ attribute_map bug."""
+        cfg = self.hf_text_config
+        qk_rope_head_dim = getattr(cfg, "qk_rope_head_dim", 0)
+        qk_nope_head_dim = getattr(cfg, "qk_nope_head_dim", 0)
+
+        # In valid MLA configs, qk_rope_head_dim != qk_nope_head_dim.
+        if qk_rope_head_dim == 0 or qk_rope_head_dim != qk_nope_head_dim:
+            return qk_rope_head_dim  # not corrupted
+
+        # Read the correct value from raw config.json.
+        from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
+
+        model_path = self.hf_config.name_or_path
+        if not model_path:
+            return qk_rope_head_dim
+        raw = get_hf_file_to_dict("config.json", model_path)
+        if raw and "qk_rope_head_dim" in raw:
+            correct = raw["qk_rope_head_dim"]
+            if correct != qk_rope_head_dim:
+                logger.info(
+                    "Fixing qk_rope_head_dim: %d -> %d "
+                    "(transformers v5.4+ attribute_map bug)",
+                    qk_rope_head_dim,
+                    correct,
+                )
+                # Patch the config so downstream model layers also get
+                # the correct value.
+                cfg.qk_rope_head_dim = correct
+                return correct
+        return qk_rope_head_dim
+
     def get_total_num_kv_heads(self) -> int:
         attributes = [
             # For Falcon:

vllm/v1/spec_decode/llm_base_proposer.py

@@ -70,6 +70,7 @@ class SpecDecodeBaseProposer:
         self.draft_model_config = self.speculative_config.draft_model_config
         self.method = self.speculative_config.method
         self.pass_hidden_states_to_model = pass_hidden_states_to_model
+        self._share_mtp_indices = False
 
         self.device = device
         self.dtype = vllm_config.model_config.dtype
@@ -490,6 +491,11 @@ class SpecDecodeBaseProposer:
         model_kwargs, slot_mapping_size = self.build_model_inputs_first_pass(
             num_tokens, num_input_tokens, mm_embed_inputs
         )
+        # Step 0 of index_share_for_mtp_iteration: let the MTP layer
+        # compute its own indices (skip_topk=False) so subsequent steps
+        # can reuse them.
+        if self._share_mtp_indices and hasattr(self.model.model, "set_skip_topk"):
+            self.model.model.set_skip_topk(False)
 
         with set_forward_context(
             per_layer_attn_metadata,
@@ -508,6 +514,11 @@ class SpecDecodeBaseProposer:
             else:
                 last_hidden_states, hidden_states = ret_hidden_states
 
+        # After step 0: switch to reuse mode so steps 1+ skip the indexer
+        # and read the indices that step 0 just wrote into the shared buffer.
+        if self._share_mtp_indices and hasattr(self.model.model, "set_skip_topk"):
+            self.model.model.set_skip_topk(True)
+
         sample_hidden_states = last_hidden_states[token_indices_to_sample]
 
         # Early exit if there is only one draft token to be generated.
@@ -1418,16 +1429,34 @@ class SpecDecodeBaseProposer:
                         )
 
         if hasattr(target_language_model.model, "topk_indices_buffer"):
+            target_buffer = target_language_model.model.topk_indices_buffer
             if hasattr(self.model.model, "topk_indices_buffer"):
                 del self.model.model.topk_indices_buffer
-            self.model.model.topk_indices_buffer = (
-                target_language_model.model.topk_indices_buffer
-            )
+            self.model.model.topk_indices_buffer = target_buffer
+            # Also share at per-module level so that the indexer and
+            # sparse-attention backends in each MTP layer read from
+            # the target model's buffer.
+            for _, module in self.model.model.named_modules():
+                if hasattr(module, "topk_indices_buffer"):
+                    module.topk_indices_buffer = target_buffer
             logger.info(
                 "Detected MTP model with topk_indices_buffer. "
                 "Sharing target model topk_indices_buffer with the draft model."
             )
 
+        # Detect index_share_for_mtp_iteration: when True, the proposer
+        # toggles skip_topk so step 0 computes MTP's own indices and
+        # steps 1+ reuse them.
+        spec_config = self.vllm_config.speculative_config
+        draft_hf_config = (
+            spec_config.draft_model_config.hf_config
+            if spec_config is not None
+            else None
+        )
+        self._share_mtp_indices = getattr(
+            draft_hf_config, "index_share_for_mtp_iteration", False
+        )
+
         if self.use_local_argmax_reduction:
             if not hasattr(self.model, "get_top_tokens"):
                 raise ValueError(

vllm/v1/worker/gpu/spec_decode/eagle/utils.py

@@ -76,10 +76,15 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod
                     del sh.head
                     sh.head = target_lm_head
 
-    # MTP also shares a topk_indices_buffer between target and draft.
+    # MTP shares topk_indices_buffer with the target model. We update
+    # every module in the draft that holds a buffer reference so that
+    # the per-layer indexer and sparse-attention backends all point to
+    # the target's buffer.
     if hasattr(target_inner, "topk_indices_buffer"):
-        if hasattr(draft_inner, "topk_indices_buffer"):
-            del draft_inner.topk_indices_buffer
-        draft_inner.topk_indices_buffer = target_inner.topk_indices_buffer
+        target_buffer = target_inner.topk_indices_buffer
+        if target_buffer is not None:
+            for _, module in draft_inner.named_modules():
+                if hasattr(module, "topk_indices_buffer"):
+                    module.topk_indices_buffer = target_buffer
 
     return eagle_model