[contrib] Add MiMo-V2-Flash (Xiaomi, TP=64, EP=64 MoE)

[whn09]

Description

Adds MiMo-V2-Flash (Xiaomi) to contrib/. All code lives under contrib/models/MiMo-V2-Flash/ — git diff upstream/main..HEAD -- src/ is empty.

Two inference paths are supported, picked via neuron_config.quantized:

• FP8 (recommended) — the HF FP8 checkpoint is preprocessed directly into Neuron-FP8 (OCP ±448 → Neuron ±240) with per-row attention scales and blockwise (128×128) MoE scales. Requires moe_tp_degree=1, moe_ep_degree=64, batch_size>=32 on Trn2 (see "FP8 Configuration Notes" below for why other configs collapse).
• BF16 (fallback) — the HF FP8 checkpoint is preprocessed to BF16 once, then loaded normally. Uses the MoE use_torch_block_wise=true fallback because Neuron SDK 2.29 does not yet ship the NKI kernel that the default blockwise path pulls in. ~2× slower than FP8 but serves as a known-good reference.

Model Information

Model Name: MiMo-V2-Flash

Model Architecture: 48 decoder layers, 256 MoE experts (top-8) with sigmoid routing and topk_method=noaux_tc (e_score_correction_bias folded into top-k selection), no shared experts, hybrid attention (full + sliding window per-layer pattern), asymmetric Q/K/V head dims (Q/K=192, V=128), partial RoPE (34%), QK RMSNorm before reshape, attention sink bias on SWA layers, attention_value_scale=0.707 applied to V right after v_proj.

Purpose: Text generation.

HuggingFace Checkpoint: XiaomiMiMo/MiMo-V2-Flash

The original checkpoint uses blockwise FP8 (E4M3 OCP ±448) which is incompatible with Neuron FP8 (E4M3 IEEE-754 ±240). Two preprocess scripts are provided:

• src/conversion_script/preprocess_mimo_v2_flash_fp8.py — FP8 → Neuron-FP8 for the native FP8 inference path. Streaming (per-layer) to cap peak RAM at ~24 GB and finishes in ~20 minutes on trn2.48xlarge; produces a ~310 GB model_layer{0..47}.safetensors set plus model_extras.safetensors. o_proj is kept BF16 per HF's quantization_config.ignored_layers.
• src/conversion_script/preprocess_mimo_v2_fp8.py — FP8 → BF16 for the BF16 inference path.

Quick Start (FP8)

End-to-end recipe — README has detailed steps, timings, and sanity-check curl:

# 1. Download HF FP8
huggingface-cli download XiaomiMiMo/MiMo-V2-Flash --local-dir /opt/dlami/nvme/models/MiMo-V2-Flash

# 2. Preprocess (HF FP8 -> Neuron FP8, ~20 min)
source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate
python contrib/models/MiMo-V2-Flash/src/conversion_script/preprocess_mimo_v2_flash_fp8.py \
    --hf_model_path /opt/dlami/nvme/models/MiMo-V2-Flash \
    --save_path     /opt/dlami/nvme/models/MiMo-V2-Flash-Neuron-FP8 \
    --tp_degree 64

# 3. (Optional) smoke-verify without vLLM
source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate
python contrib/models/MiMo-V2-Flash/perf_test/smoke_compile_mimo_v2_flash.py
python contrib/models/MiMo-V2-Flash/perf_test/smoke_generate_mimo_v2_flash.py

# 4. Install vllm-neuron + contrib patch
bash contrib/models/MiMo-V2-Flash/perf_test/0_setup.sh

# 5. Start vLLM + bench (BS=32/moe_ep=64, BS=128/moe_ep=64)
bash contrib/models/MiMo-V2-Flash/perf_test/bench_mimo_v2_flash.sh

FP8 Configuration Notes

Three non-obvious constraints on Trn2; each took a repro to find:

1. moe_tp_degree=1, moe_ep_degree=64 is the only working FP8 ratio. At moe_tp=64 each rank's intermediate slice is 32 rows (<128 blockwise block), and NxDI's _setup_for_scale collapses the per-rank scale to a singleton — losing per-channel FP8 scale granularity. The resulting drift compounds across Flash's 47 MoE layers into output collapse ("helpful helpful helpful...") after ~30 decode tokens. moe_tp=32/ep=2 and moe_tp=16/ep=4 have both been empirically tested and still collapse. Only moe_tp=1/ep=64 keeps each expert's weight + blockwise scale intact on a single rank and produces correct output.

2. batch_size >= 32 on the FP8 path. NxDI's TKG path refuses Expert Parallelism when batch_size < num_experts / top_k = 256 / 8 = 32; BS=1 latency demos on FP8 are not possible. For single-stream latency use the BF16 checkpoint with moe_tp=64, moe_ep=1, batch_size=1.

3. Keep outer ep_degree=1. MoENeuronConfig.ep_degree is the full-model expert-parallel factor and multiplies world_size to tp_degree * ep_degree. At world_size > 64 on a 64-NC Trn2, sharded-checkpoint size grows linearly, ranks beyond 63 have no backing hardware, and load fails. MoE EP is controlled exclusively via moe_ep_degree.

Checklist

Required Components

• Accuracy Test — contrib/models/MiMo-V2-Flash/test/integration/test_model.py: import, required-attribute, and state-dict-converter tests.
• README.md
  • Usage Example — flat-import bootstrap (sys.path.insert + from modeling_mimo_v2 import ...), mirroring upstream contrib/Qwen2-Audio-7B
  • Compatibility Matrix — Trn2 (trn2.48xlarge), Neuron SDK 2.29, PyTorch 2.9
  • Example Checkpoints — HuggingFace link + preprocessing instructions (BF16 and FP8)
  • Testing Instructions — pytest contrib/models/MiMo-V2-Flash/test/integration/test_model.py
  • Quick Start — end-to-end reproduction steps with approximate timings
• Source Code under contrib/models/MiMo-V2-Flash/src/:
  • modeling_mimo_v2.py — full modeling, with FP8-only runtime patches gated on quantized=True
  • conversion_script/preprocess_mimo_v2_fp8.py — FP8 → BF16 preprocessor
  • conversion_script/preprocess_mimo_v2_flash_fp8.py — FP8 → Neuron-FP8 preprocessor (streaming)

Optional Components

• Performance benchmarks — perf_test/0_setup.sh (install vllm-neuron, fetch weights), perf_test/bench_mimo_v2_flash.sh (BS=32 and BS=128 with moe_tp=1 / moe_ep=64; save_sharded_checkpoint=true to skip re-sharding on restart), perf_test/vllm-neuron-patch.patch (maps MiMo architecture to Qwen2 loader in vllm-neuron, plumbs hf_config through).
• FP8 smoke tests — perf_test/smoke_compile_mimo_v2_flash.py and perf_test/smoke_generate_mimo_v2_flash.py. These bypass vLLM so FP8 bring-up can iterate without paying vllm-neuron startup cost. STAGE={instantiate,compile,load,all}, DRY_RUN=1, SKIP_WARMUP=1, MOE_TP/MOE_EP knobs for quick debug loops.

Folder Structure

contrib/models/MiMo-V2-Flash/
  README.md
  src/
    __init__.py
    modeling_mimo_v2.py
    conversion_script/
      preprocess_mimo_v2_fp8.py          (FP8 -> BF16)
      preprocess_mimo_v2_flash_fp8.py    (FP8 -> Neuron FP8, streaming)
  test/
    integration/
      test_model.py
    unit/
  perf_test/
    0_setup.sh
    bench_mimo_v2_flash.sh
    sanity_check.sh
    run_bench_single.sh
    smoke_compile_mimo_v2_flash.py
    smoke_generate_mimo_v2_flash.py
    vllm-neuron-patch.patch

Testing

How did you test this change?

Validated on trn2.48xlarge with Neuron SDK 2.29 / PyTorch 2.9 / Python 3.12:

1. Imports and config — integration tests pass on the NxDI venv (import, required attributes, state-dict converter, MoENeuronConfig class resolution).
2. BF16 end-to-end serving — vllm-neuron with the included patch boots the model at TP=64/EP=64 and responds to chat-completion requests.
3. BF16 throughput benchmarking — historical numbers preserved in README.
4. FP8 smoke (direct NxDI) — TP=64 / moe_tp=1 / moe_ep=64 / BS=32 / SEQ=1024 produces coherent multi-sentence output on Chinese and English prompts, including 40-token chat-template prompts with <|im_start|>system ... <|im_end|> frames. No repetition collapse through 100 tokens of generation.
5. FP8 end-to-end vLLM serving — same recipe as smoke, served via /v1/chat/completions. Example 40-token Chinese prompt "你好，介绍一下你自己。比如年龄，性别，职业，爱好是什么。" produces a fluent multi-paragraph Markdown response (verified directly against a H200 sglang reference serving the same HF checkpoint).

FP8 debugging history (kept in the branch history for future reference):

• V-pad (symmetric K/V head_dim via preprocess-side padding): attempted, did not help, reverted.
• attention_value_scale=0.707: real HF behaviour; was previously overridden to 1.0, now matches reference.
• Router bias init (arange + bf16): guards against XLA constant-folding the + bias op when the init is uniform; matches Jim Burtoft's MiniMax-M2 fix.
• outer ep_degree = 1: catches a latent world-size overflow the moment anyone sets MOE_EP > 1.
• MoE scale expansion uses moe_tp_degree: catches a shape mismatch the moment moe_tp_degree != tp_degree.

Compatibility

Tested with:

• Neuron SDK Version(s): 2.29
• Instance Type(s): Trn2 (trn2.48xlarge, logical_nc_config=2 → 64 logical cores)
• PyTorch Version: 2.9
• Python Version: 3.12

Additional Information

Import convention follows upstream contrib/Qwen2-Audio-7B: tests and examples add this model's src/ directory to sys.path and import modeling files by their flat module name. This keeps the contrib package self-contained (no registration in utils/constants.py) and allows the upstream package to remain untouched.

FP8 runtime pieces (all installed as monkey-patches, gated on neuron_config.quantized):

• _apply_ep_scale_fix — don't EP-shard singleton [1,1,W] scales.
• _apply_blockwise_scale_stride_fix — force partition_stride=1 for BLOCKWISE_SYMMETRIC to avoid strided-split failures when per-rank weight is smaller than a 128-wide scale block.
• _apply_2d_per_channel_fix — 2D attention/dense-MLP weights use per-row (out, 1) scales; flip their from_float q_config from BLOCKWISE_SYMMETRIC to PER_CHANNEL_SYMMETRIC at construction.
• _apply_router_noaux_tc_fix — register e_score_correction_bias on RouterTopK (init as torch.arange(num_experts, dtype=torch.bfloat16) so XLA cannot constant-fold the subsequent + bias op and the checkpoint values actually bind) and fold the bias into top-k selection. (Flash's topk_method=noaux_tc.)
• save_quantized_state_dict override — skip the HF-side re-quantize path (requires CUDA; materializes a ~600 GB BF16 copy) when the preprocess-produced Neuron-FP8 index is already on disk.
• convert_mimo_v2_hf_to_neuron_state_dict — replicate per-row K/V .scale in lockstep with existing CONVERT_TO_MHA weight replication; expand MoE blockwise gate_up_proj / down_proj .scale along the TP-partitioned dim so per_partition_size == 1 after sharding (expansion uses moe_tp_degree, not tp_degree).

Related Issues

Part of a cleanup of my earlier whn09:contrib/llm-models branch, which originally combined MiMo-V2-Flash and MiniMax-M2 and also touched src/neuronx_distributed_inference/utils/constants.py. The branch has been split into two zero-invasion PRs (this one, and a companion MiniMax-M2 PR).

vLLM Integration

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

---

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