diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index f73d9f9c3ef..ff63f9c3617 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -169,6 +169,21 @@ class DeepSeekMultiTokenPredictor(nn.Module): if mla_attn is not None and hasattr(mla_attn, "skip_topk"): mla_attn.skip_topk = skip + def compact_topk_indices(self, slot_ids: torch.Tensor): + """Gather the top-k index rows at ``slot_ids`` to the front of the buffer.""" + num_slots = slot_ids.numel() + 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, "topk_indices_buffer" + ): + topk_indices_buffer = mla_attn.topk_indices_buffer + topk_indices_buffer[:num_slots] = topk_indices_buffer[slot_ids] + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index 0efa1ac7a7e..118f27459bb 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -130,6 +130,15 @@ class DeepseekV32MultiTokenPredictor(nn.Module): if self_attn is not None and hasattr(self_attn, "skip_topk"): self_attn.skip_topk = skip + def compact_topk_indices(self, slot_ids: torch.Tensor): + """Gather the top-k index rows at ``slot_ids`` to the front of the buffer.""" + num_slots = slot_ids.numel() + for layer in self.layers.values(): + self_attn = getattr(layer.mtp_block, "self_attn", None) + if self_attn is not None and hasattr(self_attn, "topk_indices_buffer"): + topk_indices_buffer = self_attn.topk_indices_buffer + topk_indices_buffer[:num_slots] = topk_indices_buffer[slot_ids] + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 4eaf6e9e4f8..f5e4c27b4fb 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -560,6 +560,9 @@ class SpecDecodeBaseProposer: # 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) + # The topk indices were written for each query token in the multi-token + # batch. Compact the topk indices for each request's last token. + self.model.model.compact_topk_indices(token_indices_to_sample) sample_hidden_states = last_hidden_states[token_indices_to_sample]