diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index b8987a99872..d46eb67c5ea 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/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 diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index e80d00437c7..80d518dacbd 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/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", diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index d706b505742..250aee50378 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/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: diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 38db5483937..cf9b70a7ca6 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/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( diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index fcbfc5569ef..c292c2b6cb4 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/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