diffbot/MiMo-V2.6-Flash-RL-FP8KV-W4A8-2x-RTX-PRO-6000

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

A serving recipe that runs 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 256K context 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.

Configuration	46K prefill, 1 stream (tok/s)	TTFT 46K	46K decode, 1 stream	46K, 4 streams: prefill / decode	100K, 4 streams: prefill / decode	20-agent fan-out	GSM8K-200
Base (stock image)	6,230	7.4 s	87–94	12,588 / 263	–	346 s	98.0%
+ split-KV for the verify step	6,130	7.5 s	156–164	12,340 / 335	–	344 s	98.5%
+ 3 draft tokens instead of 7	6,160	7.4 s	168–175	12,320 / 386	–	–	–
+ wide prefill tiles, CPU KV tier	8,460	5.4 s	158–175	16,873 / 385	–	97.9 s	98.5%
+ all input modes, FP8 KV cache	8,210	5.6 s	180–188	16,413 / 397	–	92.0 s	98.5%
+ one program per whole verify	8,337	5.5 s	187–191	16,608 / 412	13,334 / 348	92.3 s	99.0%
+ W4A8-FP8 MoE (Marlin)	9,532	4.8 s	184–199	18,964 / 411	14,757 / 367	80.7 s	98.5%
+ custom prefill attention kernel (this repo's default)	10,177	4.5 s	189–193	20,514 / 416	16,840 / 366	76.8 s	98.0%

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).

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.