mimo2: add MTP draft support

[tnhnyzc]

Overview

Adds MTP draft support for mimo2 architecture.

The implementation:

• loads and preserves the appended MTP tensors while keeping trunk-only GGUFs supported;
• builds the MiMo-specific MTP graph, including hidden/token fusion, fused-QKV attention, dense FFN, and shared-head fallbacks;
• creates separate MTP KV-cache memory for the appended layers;
• exposes the target model’s final hidden state to the draft context;
• uses GGUF metadata for layer counts, dimensions, head geometry, and MTP offsets.

Validation

Runtime-tested with MiMo V2.5 310B on two backends:

• IQ3_S GGUF on a MacBook Pro with M3 Max and 128 GB unified memory using Metal
• IQ2_XXS/BF16-MTP GGUF on an RTX 3090 + Tesla P40 system using CUDA

CUDA validation:

• Target-only loading with the MTP-bearing GGUF: passed
• MTP draft loading and generation: passed
• All 48 trunk and 3 NextN layers loaded correctly
• test-llama-archs: passed
• test-backend-ops: passed
• Full test suite: 56/57 passed
  • The remaining Qwen3.5 quant-selection failure reproduces on the unchanged base commit and is unrelated to this change.
• A deterministic 256-token test with unambiguous token margins produced an identical committed token sequence between target-only and MTP generation, with 127/127 draft tokens accepted.

Performance

The following measurements were collected using the IQ2_XXS/BF16-MTP GGUF on an RTX 3090 + Tesla P40 system with the CUDA backend and --spec-draft-n-max 1:

Workload	Target only	MTP	Change	Acceptance
Repeated warm prompt	~22.2 t/s	24.2–24.4 t/s	~+9–10%	84.1%
Cached chat at ~17.5k context	21.20 t/s	23.65 t/s	+11.6%	84.1%

Observed performance behavior was consistent with earlier testing of the IQ3_S GGUF using the Metal backend on an M3 Max MacBook Pro with 128 GB unified memory.

Performance remains workload-dependent.

Additional information

• I wasn't able to run and test MiMo V2.5 Pro.
• The quants that were used: AesSedai's IQ3_S, and my quants (I was only able to test with IQ3_S and IQ2_XXS quants).

Requirements

• I have read and agree with the contributing guidelines.
• AI usage disclosure: YES. GPT-5.5 and GPT-5.6 Sol were used to assist with implementation, code review, and test execution. I tested the changes and assume responsibility for the work.

[ngxson]

/bot review

[ggml-gh-bot]

Automated code review

PS: I am @ggml-gh-bot posting the review below.

I reviewed the diff (gguf-py constants, src/llama-model.cpp memory filter, and the mimo2 model/tensor/graph changes). The implementation closely follows the established Step3.5/HY-V3 MTP pattern and the trunk-graph changes are well-reasoned (they preserve the last-layer crop for normal target use and only skip it when the spec driver requests the unmasked hidden state). No blocking issues found. Details below.

Blocking

None.

Will slow the review

(point 1) src/models/mimo2.cpp — graph_mtp has no design comment. Step3.5 and HY-V3 both open their graph_mtp with a short note explaining the reference semantics being mirrored (enorm/hnorm fusion -> eh_proj -> decoder block -> head-norm -> shared LM head) and why h_nextn is the pre-head-norm state used to chain the next draft step. Reviewers will want the same here, since the pre-norm-vs-post-norm t_h_nextn chaining is the non-obvious correctness invariant. Add a brief comment block stating which reference implementation this mirrors and why layer.layer_out_norm is in the head-norm fallback chain.

(point 2) src/models/mimo2.cpp load_arch_hparams — please confirm the MTP block's attention type. ATTENTION_SLIDING_WINDOW_PATTERN is read with hparams.n_layer() (line 25), which is n_layer_all - n_layer_nextn at that point, so is_swa_impl[il] for il >= n_layer() (the MTP blocks) defaults to false and the MTP attention mask is always full-causal via build_attn_inp_kv_iswa(). If the reference MiMo MTP block actually uses sliding-window attention, the mask is silently wrong. The reported 127/127 deterministic-acceptance result implies full attention is correct for V2.5, but state that explicitly in the PR (and confirm against the converter/reference) so a future SWA-pattern GGUF doesn't break this.

Nits

(point 3) src/models/mimo2.cpp, in graph_mtp: build_inp_out_ids() never returns null (it always allocates the input tensor), so the if (inp_out_ids) { cur = ggml_get_rows(...); } guard is dead. Step3.5 and HY-V3 call ggml_get_rows(ctx0, cur, inp_out_ids) unconditionally. Drop the guard for consistency, otherwise it implies a null path that cannot occur.

(point 4) Optional follow-up: Step3.5/HY-V3 also implement an mtp_only GGUF probe (a draft-only file with no trunk tensors) via TENSOR_NOT_REQUIRED on output/trunk tensors. mimo2 only handles trunk_only. The current PR loads the full MTP-bearing GGUF as the draft, which works, but a separate lightweight draft-only MiMo MTP GGUF would not load. Not required for this PR, just flagging the parity gap.

Notes

• The trunk-graph change (crop_last_layer, guarded t_h_nextn emission) is gated on cparams.embeddings_nextn, which defaults false, so non-speculative target behavior is unchanged. The graph-params equality check already includes embeddings_nextn/embeddings_nextn_masked, so the topology change is cache-safe. Good.
• build_lora_mm handles a null scale tensor, so passing layer.nextn.eh_proj_s / model.output_s (possibly null) is fine.
• The create_memory filter addition for LLM_ARCH_MIMO2 correctly splits trunk vs MTP KV caches and matches Step3.5/HY-V3.
• AI-usage disclosure is present in the PR description. Good.

This review was generated automatically by pi coding agent using zai-org/GLM-5.2. It may contain mistakes. Maintainers make the final call.