A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 01:15 UTC
[SpecDecoding] extend mtp support for mimo 2.5 (#41905)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
May 10, 2026 at 02:22 UTC
2ee8c2a56e41fbd00b4fb52f29464fb7fca48dba
1 file changed
+3
-10
vllm/model_executor/models/mimo_v2_mtp.py
+3
-10
@@ -49,7 +49,7 @@ from .mimo_v2 import MiMoV2Attention, MiMoV2MLP
49
from .utils import _merge_multimodal_embeddings, maybe_prefix
50
51
# MiMo-V2 checkpoints contain multiple MTP layers, but vLLM currently supports
52
-# only the first layer and only one speculative token.
52
+# only the first layer
53
_MIMO_V2_PRO_NUM_MTP_LAYERS = 1
54
_MIMO_V2_FLASH_NUM_MTP_LAYERS = 1
55
@@ -170,10 +170,6 @@ class MiMoV2MultiTokenPredictor(nn.Module):
170
config = vllm_config.model_config.hf_config
171
spec_cfg = vllm_config.speculative_config
172
assert spec_cfg is not None
173
- if spec_cfg.num_speculative_tokens != 1:
174
- raise ValueError(
175
- "MiMo-V2 MTP in vLLM only supports num_speculative_tokens=1."
176
- )
173
num_mtp_layers = 1
174
175
self.num_mtp_layers = num_mtp_layers
@@ -203,10 +199,10 @@ class MiMoV2MultiTokenPredictor(nn.Module):
199
inputs_embeds: torch.Tensor | None = None,
200
spec_step_idx: int = 0,
201
) -> torch.Tensor:
206
- assert spec_step_idx == 0, "MiMo-V2 MTP only supports one speculative token."
202
if inputs_embeds is None:
203
inputs_embeds = self.embed_input_ids(input_ids)
209
- return self.mtp.layers[str(spec_step_idx)](
204
+ current_step_idx = spec_step_idx % self.num_mtp_layers
205
+ return self.mtp.layers[str(current_step_idx)](
206
inputs_embeds, positions, previous_hidden_states
207
)
208
@@ -216,7 +212,6 @@ class MiMoV2MultiTokenPredictor(nn.Module):
212
lm_head: ParallelLMHead,
213
spec_step_idx: int = 0,
214
) -> torch.Tensor:
219
- assert spec_step_idx == 0, "MiMo-V2 MTP only supports one speculative token."
215
return self.logits_processor(lm_head, hidden_states)
216
217
@@ -245,7 +240,6 @@ class MiMoV2MTP(nn.Module):
240
inputs_embeds: torch.Tensor | None = None,
241
spec_step_idx: int = 0,
242
) -> torch.Tensor:
248
- assert spec_step_idx == 0, "MiMo-V2 MTP only supports one speculative token."
243
return self.model(
244
input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
245
)
@@ -255,7 +249,6 @@ class MiMoV2MTP(nn.Module):
249
hidden_states: torch.Tensor,
250
spec_step_idx: int = 0,
251
) -> torch.Tensor | None:
258
- assert spec_step_idx == 0, "MiMo-V2 MTP only supports one speculative token."
252
return self.model.compute_logits(hidden_states, self.lm_head, spec_step_idx)
253
254
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: