MiMo-V2.5-NVFP4-TP3

lukealonso/MiMo-V2.5-NVFP4, transformed to run at tensor-parallel size 3 (e.g. on 3× 96 GB GPUs) — a geometry the original checkpoint cannot serve (64 attention heads, 4/8 KV heads, MoE intermediate 2048: none divisible by 3).

The transform is mathematically exact — no requantization, no fine-tuning, no approximation. Every real weight is moved or duplicated bit-for-bit; padded regions contribute exactly zero. Output quality is identical to the source checkpoint up to ordinary floating-point reduction-order noise (the same class of difference as running the source at TP2 vs TP4 — verified empirically, see Validation below).

Use it with the matching DFlash speculative-decoding drafter: mitomtuna/MiMo-V2.5-DFlash-TP3.

What was changed

axis	source	this repo	mechanism
full-attention heads	64 Q / 4 KV	72 Q / 9 KV	KV heads duplicated bit-identically (slot map [0,0,0,1,1,2,2,3,3]); 8 padded Q heads with zeroed q_proj rows and zeroed o_proj columns (exact-zero contribution)
sliding-window heads	64 Q / 8 KV	72 Q / 9 KV	same scheme, map [0,0,1,2,3,4,5,6,7]; per-head attention sink biases permuted alongside
MoE expert intermediate (NVFP4, g16)	2048	2304	zero rows/columns appended to the packed tensors + benign scale groups; 2048 is group-aligned, so all original quantization groups are byte-identical
dense MLP intermediate (MXFP8, g32)	16384	16512	same zero-pad scheme
everything else	—	unchanged	embeddings, lm_head, norms, routers, vision/audio towers are byte-identical copies

Because each query head attends a bit-identical copy of its original KV head, and padded heads/channels are exactly annihilated by zero output projections, per-head attention outputs are bitwise equal to the source model (proven in a standalone unit test, and end-to-end: the transform's serve-time logprob deltas vs the source are statistically indistinguishable from same-engine boot-to-boot noise).

The transform script (transform_checkpoint.py) is included in this repo for reproducibility and for adapting the scheme to other TP degrees or GQA models.

Costs vs the TP4 source (measured on 3× vs 4× RTX PRO 6000 Blackwell)

• Decode throughput ≈ 0.75–0.86× of the same engine at TP4 (bandwidth-bound MoE on one fewer GPU, +12.5% padded weights).
• Weights: ~67 GiB/GPU at TP3 (source: ~45 GiB/GPU at TP4).
• KV cache: the 9-slots-for-4-heads duplication raises per-token KV cost ~3× per rank vs the source at TP4 — expect roughly 0.5M tokens of KV pool on 3× 96 GB at gpu_memory_utilization≈0.93, i.e. practical max context ~512k (vs 1M+ at TP4).
• +12.5% attention compute (72 vs 64 head slots).
• DFlash speculative acceptance is unaffected (measured 6.7/8 at k=7 on reasoning content — identical to the source checkpoint).

Running with vLLM

The checkpoint is ModelOpt MIXED_PRECISION (MXFP8 attention/dense + NVFP4 MoE experts) and serves with the companion vLLM image (SM120/Blackwell, CUDA 13.2), which includes the required loader/model support for this format and DFlash speculative decoding:

docker pull ghcr.io/tunamitom/mimo-vllm:cu132-nvfp4-dflash

docker run --gpus '"device=0,1,2"' --ipc host --network host --init --shm-size 32g \
  -v $HF_CACHE:/root/.cache/huggingface \
  -e VLLM_USE_V2_MODEL_RUNNER=0 \
  ghcr.io/tunamitom/mimo-vllm:cu132-nvfp4-dflash \
  vllm serve mitomtuna/MiMo-V2.5-NVFP4-TP3 \
    --served-model-name mimo-v2.5 \
    --trust-remote-code \
    --tensor-parallel-size 3 \
    --kv-cache-dtype auto \
    --block-size 64 \
    --gpu-memory-utilization 0.93 \
    --max-model-len 524288 \
    --max-num-batched-tokens 4096 \
    --attention-backend TRITON_ATTN \
    --kernel-config.moe_backend flashinfer_cutlass \
    --kernel-config.linear_backend b12x \
    --mm-encoder-tp-mode data \
    --reasoning-parser mimo --tool-call-parser mimo --enable-auto-tool-choice \
    --compilation-config '{"cudagraph_mode":"PIECEWISE","custom_ops":["all"]}' \
    --async-scheduling --no-scheduler-reserve-full-isl \
    --enable-chunked-prefill --enable-prefix-caching \
    --speculative-config '{"model":"mitomtuna/MiMo-V2.5-DFlash-TP3","method":"dflash","num_speculative_tokens":7,"num_speculative_tokens_per_batch_size":[[1,4,7],[5,32,3]]}'

Notes:

• --mm-encoder-tp-mode data is required at TP3: the vision tower has 32 heads (not divisible by 3), so the multimodal encoders run replicated (data-parallel) instead of TP-sharded. Text, image, audio and video inputs all work.
• VLLM_USE_V2_MODEL_RUNNER=0 is required for the per-batch-size speculative schedule shown above; with a plain num_speculative_tokens the default runner works too.
• On PCIe-only multi-GPU boxes, add --disable-custom-all-reduce (NCCL P2P is typically the fastest correct allreduce there) and make sure NCCL P2P works on your driver configuration.
• The image also serves the original checkpoint at TP2/TP4/TP8 (lukealonso/MiMo-V2.5-NVFP4) — this TP3 repo is only needed when your GPU count is 3 (or 6).

Validation

• Unit test: per-head attention outputs bitwise-equal between source and transformed layouts (both attention types), end-to-end fp64 difference exactly 0.0.
• Serve-time logprob parity vs the source at TP4, with a same-engine restart as control: max |Δ logprob| 0.16 on dominant tokens (control: 0.25); identical structured growth into the low-probability tail. The transform adds no measurable numerics beyond boot noise.
• Greedy decoding: byte-identical outputs on 2/4 test prompts for the full reference length; the others diverge only at near-ties, as cross-boot runs of the source itself do.
• Long-context needle recall (44.8k tokens), image/audio/video understanding, and DFlash acceptance (6.74/8 @ k=7) all pass at parity with the source.

Files

• 35 model shards + model-mtp.safetensors (unused by vLLM DFlash serving; carried for completeness), model-inputscales.safetensors, amax_checkpoint.safetensors
• configuration_mimo_v2.py + auto_map in config.json (trust_remote_code config class)
• audio_tokenizer/ (used by the omni audio pipeline)
• transform_checkpoint.py — the exact transform used to produce this repo

Credits

• Base quantization: lukealonso/MiMo-V2.5-NVFP4
• Original model: XiaomiMiMo MiMo-V2.5
• DFlash drafter: XiaomiMiMo/MiMo-V2.5-DFlash