contrib: add MiMo-V2.5 (FP8 on Trn2)

[whn09]

Description

Adds contrib/models/MiMo-V2.5/ — a NeuronX Distributed Inference port of XiaomiMiMo/MiMo-V2.5 with an FP8 serving recipe on trn2.48xlarge.

MiMo-V2.5 is the successor to the earlier MiMo-V2-Flash release. The decoder-only MoE architecture is identical; the main delta on the NxDI side is that V2.5 ships the attention projection fused on disk (attention_projection_layout="fused_qkv") in a per-layer interleaved-group layout, so the preprocess script has to reconstruct per-proj q/k/v tensors. The multimodal (vision + audio) heads in the HF checkpoint are not used by the NxDI language path.

Model Information

Model Name: MiMo-V2.5

Model Architecture: Decoder-only MoE transformer with hybrid attention (9 full + 39 sliding-window layers), 256 routed experts × 8 top-k, asymmetric Q/K (192) vs V (128) head dims, partial RoPE (rotary_dim=64 of head_dim=192), sigmoid + noaux_tc router.

Purpose: Text generation.

Checklist

Required Components

• Accuracy Test (test/integration/test_model.py)

  • Imports MiMoV2InferenceConfig and NeuronMiMoV2ForCausalLM and asserts the required-attributes contract.
  • An FP8 end-to-end smoke test (perf_test/smoke_compile_mimo_v2_5.py + smoke_generate_mimo_v2_5.py) runs on Trn2 and produces a coherent MiMo self-introduction.

• README.md sections:

  • Usage Example — full Python snippet building MoENeuronConfig and running generate().
  • Compatibility Matrix — trn2.48xlarge / SDK 2.29+ / PyTorch 2.9.
  • Example Checkpoints — XiaomiMiMo/MiMo-V2.5 on HF Hub.
  • Testing Instructions — pytest contrib/models/MiMo-V2.5/test/integration/test_model.py -v.

• Source Code (src/)

  • src/modeling_mimo_v2.py — NxDI-compatible modeling (NeuronMiMoV2Attention, NeuronMiMoV2ForCausalLM, MiMoV2InferenceConfig).
  • src/conversion_script/preprocess_mimo_v2_5_fp8.py — streaming OCP-FP8 → Neuron-FP8 rescale with V2.5-specific fused-qkv splitting.

Optional Components

• perf_test/ — start_vllm_server.sh, sanity_check.sh, run_bench_single.sh, bench_mimo_v2_5.sh, 0_setup.sh for vLLM serving + benchmarking; smoke_compile_mimo_v2_5.py / smoke_generate_mimo_v2_5.py for direct NxDI.

Folder Structure

contrib/models/MiMo-V2.5/
  README.md
  src/
    __init__.py
    modeling_mimo_v2.py
    conversion_script/
      preprocess_mimo_v2_5_fp8.py
  perf_test/
    0_setup.sh
    start_vllm_server.sh
    sanity_check.sh
    run_bench_single.sh
    bench_mimo_v2_5.sh
    smoke_compile_mimo_v2_5.py
    smoke_generate_mimo_v2_5.py
    vllm-neuron-patch.patch
  test/
    unit/__init__.py
    integration/
      __init__.py
      test_model.py

Testing

How did you test this change?

• Downloaded XiaomiMiMo/MiMo-V2.5 from HF Hub (~295 GB FP8 blockwise).
• Ran preprocess_mimo_v2_5_fp8.py to produce Neuron-FP8 checkpoint (~311 GB, ~16 min, ~15 GB peak RAM).
• Ran smoke_compile_mimo_v2_5.py STAGE=all on trn2.48xlarge: compile OK (NEFF cached in neuronx-cc cache, shard_checkpoint for 64 ranks took ~29 min and populated weights/tp{0..63}_sharded_checkpoint.safetensors). Removing redundant keys from checkpoint: [] — no state_dict drops.
• Ran smoke_generate_mimo_v2_5.py with apply_chat_template: produced fluent MiMo self-introduction ("Hi there! I'm MiMo, a large language model developed by Xiaomi's LLM team...").
• Ran bench_mimo_v2_5.sh end-to-end (server launch + sanity + 3 bench runs at c=1/16/32). Results in README "Performance" section.

Test Results:

vLLM serving on trn2.48xlarge, FP8, BS=32, TP=64 / moe_ep=64, continuous batching + bucketing, 900/90 random I/O:

Concurrency	Output throughput (tok/s)	TPOT median (ms)	TTFT median (ms)
1	15.88	58.28	485
16	113.92	130.85	863
32	147.39	190.48	1798

Median inter-token latency stays at ~58 ms across all concurrencies (the cost of one BS=32 TKG NEFF forward), which matches expectations for this fixed-shape graph.

Compatibility

Tested with:

• Neuron SDK Version(s): 2.29
• Instance Type(s): trn2.48xlarge
• PyTorch Version: 2.9
• Python Version: 3.12

Additional Information

Known constraints on the FP8 serving path (detailed in README.md#fp8-configuration-notes):

• moe_tp_degree=1, moe_ep_degree=64 is the only supported MoE ratio. moe_tp_degree=64 collapses the per-rank blockwise FP8 scale to a singleton because intermediate=2048 / 64 = 32 < block_size=128; NxDI's _setup_for_scale then drops per-channel granularity and the model produces repetition collapse after ~30 decode tokens.
• batch_size >= 32 required by NxDI's TKG path: batch_size >= num_experts / top_k = 256/8 = 32 when Expert Parallelism is enabled. Single-stream BS=1 FP8 latency demos are not currently possible on V2.5.
• Outer ep_degree = 1. The MoE-internal EP factor is controlled only by moe_ep_degree; setting the outer ep_degree>1 multiplies world_size past the 64-NC cap.

Checkpoint preparation has two V2.5-specific oddities handled by the preprocess script:

1. Fused qkv split. The HF checkpoint stores self_attn.qkv_proj.weight as 4 num_kv-groups × [16 Q heads, 1-2 K heads, 1-2 V heads] per group (full layers get 1 KV/group, SWA layers get 2). The preprocess slices this back out into per-proj q_proj / k_proj / v_proj tensors to match NxDI's hard-coded ColumnParallelLinear layout. Naive [Q|K|V] concat slicing produces garbled output — we verified empirically by probing per-group scale magnitudes.
2. Inconsistent model.safetensors.index.json. V2.5's index references legacy model_N-00001-of-00002.safetensors filenames that don't match the model_pp0_epN_shardM.safetensors LFS objects actually on the Hub. LazyWeightMap rebuilds weight_map directly from the on-disk shards at startup rather than trusting the index.

Related Issues

None.

vLLM Integration

• This model is intended for use with vLLM
• Documentation includes vLLM registration instructions (see README.md#vllm-integration and perf_test/vllm-neuron-patch.patch)

The vLLM piggybacks on upstream vLLM's builtin MiMoV2FlashForCausalLM arch support (Xiaomi's upstream PR) — preprocess rewrites the checkpoint's architectures to the Flash name so vLLM 0.16's pydantic arch validator accepts it without requiring a vLLM-side PR to add a new arch. auto_map still points at the V2.5 configuration / modeling modules and trust_remote_code=True loads V2.5 classes.

---

By submitting this PR, I confirm that:

• I have read and followed the contributing guidelines
• This is a community contribution and may have limited testing compared to officially-supported models
• The code follows best practices and is well-documented
• All required components listed above are included