feat(pytorch): add mimo-v2-flash support

[yidingcheng0206]

Motivation

This PR adds support for Xiaomi MiMo-V2-Flash in LMDeploy's PyTorch backend. MiMo-V2-Flash uses hybrid full-attention and sliding-window-attention (SWA) layers and provides a three-layer MTP module for speculative decoding. This PR adds the required model implementation, attention and KV-cache support, MTP integration, and tool-call parsing.

Modification

• Add the MiMo-V2-Flash model configuration and target model implementation.
• Add the three-layer MiMo MTP draft model.
• Add support for the model's hybrid full-attention/SWA layer pattern.
• Add SWA state-ring cache and related CUDA attention kernels.
• Add tensor-parallel KV-head replication for full-attention and SWA layers.
• Add paged SWA KV-cache handling for speculative decoding.
• Integrate MiMo MTP with LMDeploy speculative decoding and CUDA-graph paths.
• Add the MiMo tool-call parser and register it with the parser manager.
• Add the FP8 linear and attention paths required by the model.
• Update model registration, supported-model lists, and English/Chinese documentation.

LMDeploy and Hugging Face were evaluated with aligned settings. The LMDeploy values are averages over repeated runs.

Benchmark	LMDeploy	Hugging Face	Difference
GPQA-Diamond	84.47% (8-run average)	83.70%	+0.77 pp
AIME 2025	95.10% (32-run average)	94.10%	+1.00 pp

Acceptance of MTP

The following results use TP8 + DP2 + EP8, attention TP4 per DP group, fixed 128-request ShareGPT workload, 2048 generated tokens, temperature 0, ignore_eos=true, three MTP draft positions, and three repeated runs.

Aggregate results

Engine	Mean accept length	Acceptance rate
LMDeploy	3.115	70.49% (native draft-token definition)
SGLang	3.076	76.91% (including bonus-token normalization)

Per-position results

Engine	Position 1	Position 2	Position 3	Strict 3/3 acceptance
LMDeploy	85.89%	69.38%	56.20%	56.20%
SGLang	84.86%	67.99%	54.78%	54.78%

LMDeploy's native acceptance rate is accepted_draft_tokens / drafted_tokens and excludes the bonus token. Mean accept length includes the bonus token. Position values are prefix acceptance probabilities; they must not be multiplied together. The strict 3/3 rate is the probability that all three draft positions are accepted.

Performance

The fixed-output benchmark uses the same ShareGPT inputs, TP8 + DP2 + EP8 topology, attention TP4 per DP group, concurrency 128, and three repetitions per point.

MTP throughput and TTFT

Engine	Output	Duration (s)	Request throughput (req/s)	Output throughput (tok/s)	TTFT Mean (ms)	TTFT P50 (ms)	TTFT P99 (ms)
LMDeploy MTP	2048	36.46	3.51	7195.01	934.35	909.23	1029.04
SGLang MTP	2048	35.18	3.64	7457.21	469.01	463.74	535.28
LMDeploy MTP	4096	66.54	1.93	7896.77	689.53	676.68	797.08
SGLang MTP	4096	63.66	2.01	8239.97	366.37	367.62	389.05
LMDeploy MTP	8192	121.75	1.06	8677.01	733.45	731.26	798.63
SGLang MTP	8192	119.26	1.08	8806.60	692.20	701.28	739.12

MTP TPOT and ITL

Engine	Output	TPOT Mean (ms)	TPOT P50 (ms)	TPOT P99 (ms)	ITL Mean (ms)	ITL P50 (ms)	ITL P99 (ms)
LMDeploy MTP	2048	13.21	12.93	16.96	41.75	38.39	180.02
SGLang MTP	2048	13.44	13.42	16.83	41.20	38.49	126.80
LMDeploy MTP	4096	12.08	11.84	15.74	41.19	38.62	118.37
SGLang MTP	4096	12.27	12.15	15.26	40.28	38.60	58.35
LMDeploy MTP	8192	11.32	11.20	13.57	41.12	38.96	75.59
SGLang MTP	8192	11.61	11.44	13.95	40.72	39.08	49.32

LMDeploy MTP output throughput is 96.5%, 95.8%, and 98.5% of SGLang MTP for 2048, 4096, and 8192 output tokens, respectively. MTP TPOT remains within 3% of SGLang in these tests.

Runtime validation

• TP4 and TP8 target inference and CUDA Graph serving.
• TP8 + DP2 + EP8 + DeepEP + DeepGEMM + MTP3 production topology.
• Prefix caching with MTP and high-concurrency replay.
• SWA boundary and speculative accept/reject/rollback cases.
• MiMo tool-call API round trips.

Compatibility

No intentional backward-incompatible changes are introduced. MiMo-V2-Flash support targets the PyTorch backend. MTP requires the three-layer draft weights included in the checkpoint. FA3 is optional; unsupported asymmetric head-dimension builds automatically use Triton.

Validation

• Ruff and git diff --check pass.
• Focused model, attention, cache, MTP, FP8, distributed, CUDA Graph, and parser tests pass.
• Real-checkpoint TP4/TP8 and production-topology runtime validation pass.

Assistance

Assisted with Codex, reviewed manually.

Checklist

1. Pre-commit or other linting tools are used to fix the potential lint issues.
2. The modification is covered by complete unit tests. If not, please add more unit tests to ensure the correctness.
3. If the modification has a dependency on downstream projects of a newer version, this PR should be tested with all supported versions of downstream projects.
4. The documentation has been modified accordingly, like docstring or example tutorials.

[Copilot]

Note

Copilot was unable to run its full agentic suite in this review.

Pull request overview

Adds MiMo-V2-Flash support across serving + PyTorch engine, including MiMo-specific tool parsing, MTP speculative decoding, and CUDA attention/graph improvements.

Changes:

• Introduces MiMo-V2-Flash target + MTP draft model implementations, config builder, and module registration.
• Extends speculative decoding pipeline (proposer depth transitions, DP meta handling, CUDA Graph warmup/capture safety).
• Improves attention backends/kernels (paged multi-token causal decode, FA3 capability gating, flash-attn asymmetric head tuning) and adds extensive tests/docs.

Reviewed changes

Copilot reviewed 51 out of 51 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/test_lmdeploy/serve/parsers/test_mimo_parser.py	Adds tests for MiMo tool-call parsing (complete + streaming + filtering).
tests/pytorch/test_distributed.py	Adds regression test for reduce_scatter output aliasing.
tests/pytorch/spec_decode/test_spec_agent.py	Updates spec-agent tests for new proposer depth-transition interface and DP token layouts.
tests/pytorch/spec_decode/test_cudagraph_strategy.py	Adds tests for CUDA graph buffer fill immutability and capture rollback semantics.
tests/pytorch/nn/test_linear_blocked_fp8.py	Adds tests for checkpoint-aware FP8 QKV sharding/scale behavior.
tests/pytorch/kernel/test_swa_state_ring.py	Adds CUDA kernel tests for state-ring flatten/scatter and attention equivalence.
tests/pytorch/kernel/test_paged_attention.py	Adds causal multi-token paged attention test + split-K neutrality test; refactors fixtures/classes.
tests/pytorch/kernel/test_mla_attention.py	Extends MLA builder test coverage for allow_fa3=False.
tests/pytorch/kernel/test_flash_attention.py	Adds test for asymmetric (K=192,V=128) flash attention head shape.
tests/pytorch/kernel/test_fa3_attention.py	Adds builder fallback test when FA3 wheel lacks asymmetric head instantiation.
lmdeploy/serve/parsers/tool_parser/qwen3coder_tool_parser.py	Adjusts whitespace handling so parser can preserve raw parameter text when needed.
lmdeploy/serve/parsers/tool_parser/mimo_tool_parser.py	Adds MiMo tool parser (HTML unescape + schema coercion preserving string whitespace).
lmdeploy/serve/parsers/tool_parser/init.py	Exposes MiMo tool parser in package exports.
lmdeploy/pytorch/weight_loader/model_weight_loader.py	Allows models to select checkpoint shards via select_weight_paths hook.
lmdeploy/pytorch/strategies/ar_spec/cudagraph.py	Adds MiMo MTP CUDA graph max-token sizing logic.
lmdeploy/pytorch/spec_decode/spec_agent.py	Uses proposer advance_draft_depth + proposer-provided DP token layout; improves warmup graphs.
lmdeploy/pytorch/spec_decode/proposers/mimo_mtp.py	Implements MiMo “same-position multi-depth” proposer and shared-embedding binding.
lmdeploy/pytorch/spec_decode/proposers/base.py	Introduces advance_draft_depth and DP token layout hook for draft-depth transitions.
lmdeploy/pytorch/spec_decode/proposers/init.py	Registers MiMo proposer.
lmdeploy/pytorch/nn/linear/blocked_fp8.py	Adds checkpoint-aware, block-aligned FP8 QKV sharding plan + logical/physical split_qkv.
lmdeploy/pytorch/nn/linear/init.py	Plumbs checkpoint-aware QKV sharding options into builder.
lmdeploy/pytorch/nn/attention.py	Exposes allow_fa3 and enable_paged_multi_token_decode in Attention build spec.
lmdeploy/pytorch/models/utils/cudagraph.py	Adds warmup spec API + capture-cache hook; avoids mutating caller-owned attn_metadata in fill.
lmdeploy/pytorch/models/module_map.py	Registers MiMo-V2-Flash target + MTP model entries.
lmdeploy/pytorch/models/mimo_v2_flash_mtp.py	Implements MiMo MTP draft model (3-depth predictor), loading, and CUDA graph behavior.
lmdeploy/pytorch/models/mimo_v2_flash.py	Implements MiMo-V2-Flash target model with hybrid full/SWA attention, FP8 loading, paged verification.
lmdeploy/pytorch/models/deepseek_mtp.py	Propagates spec_step_idx through generation inputs.
lmdeploy/pytorch/model_inputs.py	Adds spec_step_idx to ModelInputs/StepContext and advances it on step.
lmdeploy/pytorch/kernels/cuda/swa_state_ring.py	Adds Triton flatten/scatter kernels for SWA BF16 state rings.
lmdeploy/pytorch/kernels/cuda/pagedattention.py	Adds causal multi-token mask path for paged attention + fixes empty split-K workspace init.
lmdeploy/pytorch/kernels/cuda/flashattention.py	Fixes kernel meta selection for split-D (e.g., 192) attention heads.
lmdeploy/pytorch/engine/config_builder.py	Validates/constructs mimo_mtp spec config and MiMo-V2-Flash constraints.
lmdeploy/pytorch/distributed.py	Ensures reduce_scatter output is disjoint from input views.
lmdeploy/pytorch/configurations/mimo_v2_flash.py	Adds MiMo-V2-Flash model config builder + cache policy (paged SWA vs ring).
lmdeploy/pytorch/backends/cuda/graph_runner/runner.py	Adds model capability gating for FA3 requirement and capture rollback snapshots for mutable caches.
lmdeploy/pytorch/backends/cuda/graph_runner/full_graph.py	Threads supports_non_fa3_speculative_graph into graph meta for FA3 usage.
lmdeploy/pytorch/backends/cuda/attention/swa_state_ring.py	Adds SWA attention backend using state ring + flatten/scatter fallback for q>1.
lmdeploy/pytorch/backends/cuda/attention/fa3_capabilities.py	Adds runtime check for FA3 wheel head-dim template availability.
lmdeploy/pytorch/backends/cuda/attention/fa3.py	Advertises paged multi-token decode support on FA3 implementation.
lmdeploy/pytorch/backends/cuda/attention/default.py	Adds optional paged multi-token decode support and kernel-metadata override in forward.
lmdeploy/pytorch/backends/cuda/attention/init.py	Gates FA3 enabling on allow_fa3 + wheel capability; plumbs multi-token decode option.
lmdeploy/pytorch/backends/attention.py	Extends PagedAttentionBuildSpec with allow_fa3 and enable_paged_multi_token_decode.
lmdeploy/cli/utils.py	Adds mimo_mtp to CLI speculative algorithm choices.
docs/zh_cn/supported_models/supported_models.md	Documents MiMo-V2-Flash support and constraints (BF16 runtime, CUDA backend).
docs/zh_cn/llm/api_server_tools.md	Documents MiMo tool-call parser usage in api_server.
docs/zh_cn/advance/spec_decoding.md	Documents MiMo-V2-Flash MTP configuration and backend behavior.
docs/en/supported_models/supported_models.md	Documents MiMo-V2-Flash support and constraints (BF16 runtime, CUDA backend).
docs/en/llm/api_server_tools.md	Documents MiMo tool-call parser usage in api_server.
docs/en/advance/spec_decoding.md	Documents MiMo-V2-Flash MTP configuration and backend behavior.
README_zh-CN.md	Lists MiMo-V2-Flash as supported model.
README.md	Lists MiMo-V2-Flash as supported model.

Suppressed comments (1)

tests/pytorch/kernel/test_fa3_attention.py:1

• This test uses sys and SimpleNamespace but does not import them in the file. Add import sys and from types import SimpleNamespace (or equivalent) so the test doesn't fail with NameError.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[grimoire]

There are large change in this kernel, It could be risky to do such large change in this kernel, we better create a new kernel for swa.

[grimoire]

I am not sure if it is worth to add these restore logical. the cache might already been polluted during capture.

[grimoire]

Do we need to set/update these values in strategy/model_inputs.py?