[BugFix][Spec Decode] Compact shared topk indices buffer after first MTP draft step (#47238)
Giancarlo Delfin committed
Jul 1, 2026 at 23:38 UTC
d63c8e944481e057d00dfee20bc49544d291e521
3 files changed
+27
vllm/model_executor/models/deepseek_mtp.py
+15
@@ -169,6 +169,21 @@ class DeepSeekMultiTokenPredictor(nn.Module):
169
if mla_attn is not None and hasattr(mla_attn, "skip_topk"):
170
mla_attn.skip_topk = skip
171
172
+ def compact_topk_indices(self, slot_ids: torch.Tensor):
173
+ """Gather the top-k index rows at ``slot_ids`` to the front of the buffer."""
174
+ num_slots = slot_ids.numel()
175
+ for layer in self.layers.values():
176
+ mtp_block = getattr(layer, "mtp_block", None)
177
+ if mtp_block is not None:
178
+ self_attn = getattr(mtp_block, "self_attn", None)
179
+ if self_attn is not None:
180
+ mla_attn = getattr(self_attn, "mla_attn", None)
181
+ if mla_attn is not None and hasattr(
182
+ mla_attn, "topk_indices_buffer"
183
+ ):
184
+ topk_indices_buffer = mla_attn.topk_indices_buffer
185
+ topk_indices_buffer[:num_slots] = topk_indices_buffer[slot_ids]
186
+
187
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
188
return self.embed_tokens(input_ids)
189
vllm/models/deepseek_v32/nvidia/mtp.py
+9
@@ -130,6 +130,15 @@ class DeepseekV32MultiTokenPredictor(nn.Module):
130
if self_attn is not None and hasattr(self_attn, "skip_topk"):
131
self_attn.skip_topk = skip
132
133
+ def compact_topk_indices(self, slot_ids: torch.Tensor):
134
+ """Gather the top-k index rows at ``slot_ids`` to the front of the buffer."""
135
+ num_slots = slot_ids.numel()
136
+ for layer in self.layers.values():
137
+ self_attn = getattr(layer.mtp_block, "self_attn", None)
138
+ if self_attn is not None and hasattr(self_attn, "topk_indices_buffer"):
139
+ topk_indices_buffer = self_attn.topk_indices_buffer
140
+ topk_indices_buffer[:num_slots] = topk_indices_buffer[slot_ids]
141
+
142
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
143
return self.embed_tokens(input_ids)
144
vllm/v1/spec_decode/llm_base_proposer.py
+3
@@ -560,6 +560,9 @@ class SpecDecodeBaseProposer:
560
# and read the indices that step 0 just wrote into the shared buffer.
561
if self._share_mtp_indices and hasattr(self.model.model, "set_skip_topk"):
562
self.model.model.set_skip_topk(True)
563
+ # The topk indices were written for each query token in the multi-token
564
+ # batch. Compact the topk indices for each request's last token.
565
+ self.model.model.compact_topk_indices(token_indices_to_sample)
566
567
sample_hidden_states = last_hidden_states[token_indices_to_sample]
568