MiMo-V2.6-Flash (MOPD) on 2× RTX PRO 6000 Blackwell: FP8 KV cache, W4A8 MoE, custom sm_120 kernels

A serving recipe that runs MiMo-V2.6-Flash-MOPD (recommended) or MiMo-V2.6-Flash-RL (309B total / 15B active, MXFP4 experts, text + image + video + audio in) on two RTX PRO 6000 Blackwell cards (sm_120, 96 GB each, PCIe, no NVLink) with vLLM at tensor-parallel 2, at 2.1× the decode speed and 1.6× the prefill speed of the stock vLLM image on the same hardware, with a 512K context window (since 2026-10-12; the benchmarks below were taken at 256K) and all four input modes enabled. The weights are Xiaomi's official checkpoint (already MXFP4 experts + FP8 dense) and are not redistributed here; this repository is the runtime quantization and the code around it:

• FP8 (E4M3) KV cache for the model's DiffKV attention layers (stock vLLM only allows bf16 there). Doubles the KV pool.
• W4A8-FP8 MoE: the MXFP4 experts with FP8 activations on Marlin (VLLM_MARLIN_INPUT_DTYPE=fp8). +15% prefill, same quality.
• Three attention fixes to vLLM's Triton DiffKV kernel (split-KV for the speculative-decode verify step, wide prefill tiles, fp8 K/V), and a purpose-built CUDA prefill attention kernel for the 9 global layers (1.34× the Triton kernel).
• A custom small-batch MoE decode kernel on b12x's FP4 layout (beats Marlin at 8–32 routed tokens; kept as an opt-in hybrid backend), the DFlash drafter at 3 draft tokens, and vLLM's CPU KV tier for multi-agent fan-out.

Everything needed to reproduce it is in recipe/ (see recipe/README.md): image build, launcher, vLLM patches, kernels, tests and the benchmark scripts. Raw benchmark logs are in results/.

Benchmarks

Two RTX PRO 6000 Blackwell Max-Q at their 300 W cap, PCIe, TP2, one 20-agent workload shape: ≈46K-token prompts at a 61:1 prefill:decode ratio. A ≈2.5 GB/GPU side process stayed resident during every run. recipe/bench/bench-sbs.py sends non-streaming requests; decode speed is the difference between a 400-token and a 16-token reply on a warm prefix cache. bench-fanout.py runs 20 agents × 3 turns (≈40K tokens each) through 8 slots. GSM8K is the first 200 test problems, greedy.

Base = the stock vllm/vllm-openai:mimo-v26-x86_64-cu130 image with the vendor recipe's settings (Marlin W4A16, DFlash with 7 draft tokens) as far as they fit on two cards (text-only, bf16 KV, 262K context). Each row adds one change.

At the default config the KV pool is 479K tokens (fp8, all encoders loaded); the 29K-token and 336K-token needle tests, tool calls, and image/video/audio probes pass. Decode is flat from 2K to 46K context and 131 tok/s at 180K.

Things that were measured and did not help on this hardware (logs in results/): b12x MXFP4 MoE alone (+15% prefill, −45% four-stream decode), MTP instead of DFlash, 8192-token prefill chunks (KV pool 467K → 293K), FlashInfer/custom all-reduce (32 MB reductions are PCIe-bound; NCCL is at the link's floor), FlashInfer's SM120 192×128 fused-attention kernels (unreachable through the 0.6.18/0.7.0 wrappers), FP8 QKᵀ in the attention kernel (−8% time, 7× the error), and 2 CTAs/SM for it (spills).

Update 2026-10-19: two tool-call parser fixes

recipe/docker/build.sh now also applies recipe/docker/pr62443-62254-mimo-parser.diff, after the vLLM #58019 parser:

• Streamed tool-call arguments match the non-streamed ones when a value contains a partial parameter opener (vLLM #62443, still open upstream, backported to this image's streaming path). A string value can contain <parameter= with no closing > (for example a doc that says "type <parameter= followed by the name"), or a spaced form like <parameter =. The streaming parser sent that text out before it could tell it was an opener. The non-streaming parse then read it as a tag, so the streamed arguments ended up with an extra key such as "b</parameter": "", or as invalid JSON. Now the parser holds back any trailing text that could still become </parameter> or <parameter= until the next characters decide it. It also drops an opener that is never closed, as the non-streaming parse does. The value itself still ends at the opener in both modes (see the limitation below); what the fix guarantees is that streaming and non-streaming agree.
• Number-like strings stay strings (vLLM #62254, applied as is; issue #62252). Under a schema that allows a string and a number, such as "type": ["string", "integer"] or an anyOf of string and number, "0017" came back as 17 and "3.50" as 3.5, both streaming and non-streaming. When the schema also allows a string, a numeric conversion now counts only if it round-trips: str(number) must give back the original text. "17" under the same schema is still 17, and an integer-only schema still turns "0017" into 17.

recipe/patches/test_mimo_parser.py checks both fixes on CPU. It uses the real tokenizer, goes through the parser entry points the server uses, and splits the text at several token and character sizes. The #58019 image fails 20 checks; with this patch all checks pass.

In a live check on our two cards (temperature 0, each request sent both streaming and non-streaming), the streamed arguments matched the non-streamed ones for a value with a partial <parameter= opener and for a value with <parameters> / <parameter-list> markup (kept verbatim); "0017" and "3.50" stayed strings; GSM8K (first 100 test problems) scored 99%. The change touches only the parser, no kernel or model code.

Known limitation (every version, including this one): a parameter value that contains the format's own closing tags (</parameter>, </function>, </tool_call>) cannot be parsed reliably, because the first </parameter> ends the value. A native <parameter=NAME> tag inside a value starts a new argument, and it can override an argument of the same name: a file whose content contains <parameter=path>…</parameter> can replace the call's real path. MiMo's tool-call format has no escaping, so no parser can tell these apart. If a tool has to carry such text, give it an encoded form (for example base64).

Update 2026-10-12: 512K context by default

serve.sh now defaults to MAX_LEN=524288 with UTIL=0.945 (was MAX_LEN=262144, UTIL=0.93). On our two Max-Q cards, next to a ~2.5 GB/GPU side process:

Speed did not change after a 488K-token request had gone through the server. A torch profile shows the same GPU time per forward pass at both settings. A needle in a 488K-token prompt was found in 3 of 4 runs (150 s of prefill each) and missed once, so retrieval near the end of the window is not perfect. Without a side process on the GPUs there is about 2.5 GB more headroom per card. MAX_LEN=262144 UTIL=0.93 restores the previous setting.

Xiaomi updated preprocessor_config.json in both checkpoints on 2026-10-09 (image min_pixels/max_pixels and normalization mean/std aligned to the training values). If you downloaded the checkpoint before that date, re-download that file; image probes pass with it.

Update 2026-10-01: three correctness fixes

• Tool-call string arguments are now passed through verbatim. The day-0 image maps --tool-call-parser mimo to vLLM's Qwen3 parser, which strips one leading and one trailing newline from every string argument. MiMo writes <parameter=…>value</parameter> with no newline wrapper, so a file-write call whose content ends in a newline lost it. recipe/docker/build.sh now applies vLLM #58019 (merged upstream 2026-09-29), which adds a MiMo parser. In a live check, 6 of 6 calls kept the trailing newline, against 0 of 6 before.
• No NaN for query rows without visible keys. In the split-KV attention reducer, a row whose segments are all empty (for example the padded rows of a CUDA-graph batch) computed exp(-inf - -inf) and produced NaN. A one-line clamp, which the upstream port of this patch (vLLM #59085) also has, makes those rows 0. recipe/patches/test_diffkv_emptyrows.py shows 16 of 16 empty rows NaN before and 0 after; normal rows are unchanged.
• Images with large uniform regions are read correctly. The day-0 image adds the vision encoder's attention sink as a bias on key 0 (sinks_bias_key0=True), which misreads uniform image regions; a solid green square came back as "a close-up of an eye". recipe/patches/mimo_v2_omni.py applies vLLM #58235 (still open upstream), which treats the sink as an extra softmax-denominator logit, plus the omni wrapper's packed_modules_mapping. Solid green, white and blue squares are now described correctly, and the image, video and audio probes pass. VIT_SINK_FIX=0 restores the stock file.

Same-session A/B of the first two fixes, every other setting unchanged (the vision fix was on in both arms; results/fixes-1001-ab.txt):

The differences are within run-to-run noise. A torch-profiler run of the fixed build (results/profiles/eco1001-cand.prof.txt) matches the earlier profile: 13.4 ms per decode step at 46K, GPU busy 94%.

Update 2026-09-28: MiMo-V2.6-Flash-MOPD

Xiaomi's MOPD checkpoint (released 2026-09-27; aimed at tool-call repetition in agent harnesses) has the same config and DFlash drafter as Flash-RL, so the recipe runs it unchanged: set MODEL to the MOPD download. Same session, same settings (results/rl-0928*.txt, results/mopd-0928*.txt):

The benchmark table above was measured on Flash-RL; MOPD is within noise of it on every row we re-ran. This repository's name predates the MOPD release; it serves both checkpoints.

Update 2026-09-24

• Exact QKV weights. The fused-FP8-QKV loader in vLLM re-quantizes each rank's rows below TP4. Measured on this checkpoint at TP2, that changed 18% of Q, ~51% of K and 73–81% of V weights in the global attention layers (1.3–2.5% relative L2). The recipe now carries a port of local-inference-lab/vllm #874 that keeps every weight on its checkpoint scale (0 weights changed; audit in recipe/patches/test_mimo_exact_qkv_audit.py). VLLM_MIMO_EXACT_QKV=0 restores the old loader.
• Sampling defaults. serve.sh now uses --generation-config auto: requests without sampling parameters get the checkpoint's temperature 1.0 / top_p 0.95 instead of vLLM's near-greedy defaults, which users report make MiMo repeat the same tool call in agent harnesses. Clients that send no sampling parameters see lower speculative acceptance at temperature 1.0 (mean accept length 2.93 vs 3.42 across our suite); send your own temperature or set GEN_CONFIG=vllm if you prefer the old behaviour.
• Same-day before/after (results/fix0924-old.txt, results/fix0924-new.txt): 46K prefill 10.1K / 10.0K tok/s, TTFT 4.6 s both; 46K decode 164 / 185–194 tok/s (noise); four streams at 46K 20.0K / 420; 20-agent fan-out 77.0 s; GSM8K-200 97.5% both; KV pool 479K → 467K tokens (the padded QKV layout). Needle, tool-call and image/video/audio probes pass. A new multi-turn tool-loop probe (recipe/bench/test-mimo-agentic.py) showed no reasoning leaks and no repeated tool calls in 8 runs with either configuration.

What the kernels do

DiffKV attention on SM120. MiMo's K and V head dims differ (192/128), so vLLM routes every target layer to its generic Triton "DiffKV" kernel (the FlashAttention path needs FA3/FA4). Profiling showed that kernel was 59% of decode time at 46K: it disables its split-KV mode whenever a request has more than one query token, and every speculative-decode verify has several, so the 9 full-context layers ran on ≈18 CTAs of a 188-SM GPU. recipe/patches fixes that (5–7× per call), adds wide prefill tiles (3×), and ports the E4M3 K/V path from the local-inference-lab fork with SM120 shared-memory fixes.

recipe/kernels/attention/prefill_attn.cu replaces the Triton kernel for prefill rows on the global layers: FA2-style, one CTA per 8 tokens × 16 GQA heads, ldmatrix + mma.m16n8k16, fp8 K/V converted in the load path. 10.5 ms vs 14.1 ms per 4096-token chunk at 26K context; Nsight shows it compute-bound (tensor pipe 55%, DRAM 0.4%).

recipe/kernels/moe/decode_moe.cu streams b12x's prepared MXFP4 layout at 92% of DRAM bandwidth for decode-sized batches and feeds decoded FP4 straight into tensor-core fragments for 8–32 tokens. It exists so that b12x (the fastest prefill MoE kernel on sm_120) and a Marlin-class decode kernel can share one weight copy; the hybrid backend is wired but not yet the default.

Credits and license

Model: Xiaomi (MIT). Serving stack: vLLM, FlashInfer, Marlin, b12x. The fp8-KV, cache-config and mixed-batch changes were ported from pull requests in local-inference-lab/vllm. The code in this repository is MIT (LICENSE); see THIRD_PARTY_NOTICES.md. Built by Diffbot for a two-card agentic-coding workstation.