[SpecDecoding] extend mtp support for mimo 2.5

[ZJY0516]

Purpose

support num_speculative_tokens>1 for mimo 2.5 mtp

Test Plan

vllm serve XiaomiMiMo/MiMo-V2.5 -tp 4 --trust-remote-code 
--speculative_config '{"method":"mtp","num_speculative_tokens":3}' 
--no-async-scheduling

lm_eval --model local-completions --model_args "model=XiaomiMiMo/MiMo-V2.5,base_url=http://0.0.0.0:8000/v1/completions,tokenized_requests=False,tokenizer_backend=None,num_concurrent=256,timeout=5000,max_length=40960" --tasks gsm8k --num_fewshot 5 --gen_kwargs max_gen_toks=5120

Test Result

w/o mtp

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9530	±	0.0058
		strict-match	5	exact_match	↑	0.9538	±	0.0058

mtp3

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9454	±	0.0063
		strict-match	5	exact_match	↑	0.9469	±	0.0062

mtp3 w/o async-scheduling

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9522	±	0.0059
		strict-match	5	exact_match	↑	0.9522	±	0.0059

Note: async shceduling may affect mtp accuracy

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 16e86d293a

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Load the remaining MTP layers for multi-step drafts

When num_speculative_tokens > 1, this still sets self.num_mtp_layers to 1, so every spec_step_idx maps back to model.mtp.layers.0; load_weights then ignores the checkpoint's later MTP layers. The MiMo-V2.5-Pro model card lists 3 MTP layers and its deployment example uses 3 speculative steps (https://huggingface.co/XiaomiMiMo/MiMo-V2.5-Pro/blob/main/README.md), so enabling multi-token drafting here runs all draft steps through the first-layer weights instead of the trained layer sequence, which makes the new >1 support materially incorrect/low-acceptance. Please instantiate/load the available MTP layers or keep rejecting num_speculative_tokens > 1.

Useful? React with 👍 / 👎.

[gemini-code-assist]

Code Review

This pull request removes the restriction limiting MiMo-V2 MTP to a single speculative token and implements modulo-based indexing for MTP layers. Feedback indicates that the module-level documentation needs updating to reflect this change. Additionally, a critical issue was identified where num_mtp_layers remains hardcoded to 1, which would cause all speculative steps to incorrectly reuse the first layer instead of utilizing the appropriate layers for each step.

[gemini-code-assist]

The comment still states that vLLM only supports the first MTP layer. Since this PR aims to support multiple speculative tokens, this comment should be updated to reflect that multiple layers are now supported (assuming the hardcoded layer count is also addressed).

Suggested change

	# MiMo-V2 checkpoints contain multiple MTP layers, but vLLM currently supports
	# only the first layer and only one speculative token.
	# only the first layer
	# MiMo-V2 checkpoints contain multiple MTP layers, and vLLM supports
	# multiple speculative tokens by using these layers.

[gemini-code-assist]

The variable num_mtp_layers is still hardcoded to 1. This contradicts the PR's objective of supporting num_speculative_tokens > 1. If num_mtp_layers remains 1, only the first MTP layer will be initialized and loaded. When num_speculative_tokens > 1, the forward method (line 204) will reuse this single layer for all speculative steps due to the modulo operation (spec_step_idx % 1). This is mathematically incorrect for MTP architectures like DeepSeek/MiMo where each speculative step typically requires a distinct layer trained for that specific offset. You should derive num_mtp_layers from the model configuration (e.g., config.num_nextn_predict_layers) or set it based on spec_cfg.num_speculative_tokens while ensuring it does not exceed the model's actual capacity.

Suggested change

	num_mtp_layers = 1
	num_mtp_layers = spec_cfg.num_speculative_tokens

[jeejeelee]

Can you provide the accuracy result here?

[ZJY0516]

Can you provide the accuracy result here?

Updated, PTAL

[jeejeelee]

thank you