GB300 MiMo-V2.6-Pro Testing
SEPTEMBER 21, EVENING · DAY ZERO · THREE WALLS FOUND, TWO PATCHED · NO SPEED NUMBER YET
Where it stands: Xiaomi published
MiMo-V2.6-Pro-RL at 20:12 UTC today. The checkpoint is 527 GiB — twice the Station's 250.7 GiB of HBM — so this is a Grace-offload model, and a harder one than DeepSeek-V4.1-Flash: 61% of the expert bytes must live in host RAM instead of 22%. Nine boots in, the model loads with all 69 MoE layers where we put them (23 in HBM, 46 pinned in Grace) on stock vLLM nightly plus two small bind-mounted patches. The ninth boot is past every load-time wall and running its first FlashInfer autotune as this goes up. No tokens/s yet; that number lands here when it exists, not before.
What the model is (from config.json and the safetensors headers, not the card)
| MiMo-V2.6-Pro-RL | MiMo-V2.6-Flash-RL | |
|---|---|---|
| Parameters | 1.02T total / 42B active | 309B / 15B active |
| Layers | 70 (60 SWA-128 / 10 global), hidden 6144 | 48 (39 / 9), hidden 4096 |
| Experts | 384 routed, top-8, no shared expert | 256 routed, top-8 |
| Shipped precision | Experts native MXFP4 (store_dtype: mxfp4, U8 weights + U8 block scales); attention and dense FFN FP8 e4m3 blockwise; o_proj BF16 | |
| Bytes on disk | 527 GiB (experts 494.9; non-expert layers 24.6; embed/head 3.5; DFlash drafter 2.3; ViT 1.4; audio 0.5) | 161 GiB (experts 149.8) |
| Per expert / per layer | 19.1 MiB / 7.17 GiB | 12.5 MiB / 3.2 GiB |
| Context | 1M tokens; SWA layers are window-bounded so KV is dominated by the 10 (Pro) / 9 (Flash) global layers | |
| Drafter | 5-layer DFlash (block 8, 7 draft tokens), anchors on target layers 0/15/31/47/69 | |
| Modalities | text, image, video, audio in one checkpoint | |
| Revision archived | 54b10491, 155 files, byte-verified against the Hub manifest, copied to Milo-Ark | not archived yet |
Flash fits the Station's HBM outright (161 GiB + KV) and is a clean single-GPU recipe. Pro is the interesting one, and the one that does not fit. Everything below is Pro.
The memory arithmetic
HBM at --gpu-memory-utilization 0.95 is about 238 GiB. Non-expert weights are about 32 GiB. Leave ~12 GiB for CUDA graphs, activation peak and a small KV pool, and roughly 194 GiB of experts fit in HBM — 27 of 69 MoE layers. The remaining ~301 GiB (42 layers) has to be pinned in the 494 GiB Grace host and read over C2C on every decode step.
That host budget is the whole campaign. The DeepSeek-V4.1 work on this box established two things we walked in with: pinned host memory peaks at roughly 1.6× the pinned set while the loader's mmaps are live, and the box has died twice with the kernel OOM-killer at ~486 GB of shmem. Pinning 301 GiB with a 527 GiB checkpoint mmapped is exactly that shape.
Nine boots, three walls
| Boot | Config | Died at | Cause |
|---|---|---|---|
| v1 | UVA offload 305 GiB, --cpu-offload-params experts.w13_weight experts.w2_weight | Layer construction, HBM OOM at 247 GiB | Name trap. vLLM's FusedMoEFactory wraps MiMoV2's experts as mlp.experts.routed_experts.*; the segment experts.w13_weight matches nothing, the offloader offloads nothing, and every layer lands in HBM. Same trap as DSV4.1, inverted. Fix: routed_experts.w13_weight routed_experts.w2_weight. |
| v2 | v1 + correct names | Host OOM-kill, 486 GB shmem, ~60 s in | The 527 GiB checkpoint mmap plus pinning. Zero layers had been pinned yet. |
| v3 | v2 + --load-format runai_streamer, 16 GiB read buffer | Host OOM-kill, 486 GB shmem, same place | Streaming loader removed the mmap and it still died at the same number. So the mmap was not the cause. Something in pinning itself was doubling. |
Wall 2, measured: a 4-line probe inside the image. torch.empty(4.5 GiB).pin_memory() raises host Shmem by 8,192 MiB; a 4.0 GiB pin raises it by 4,168. PyTorch's CachingHostAllocator rounds every pinned block up to a power of two. MiMo-Pro's w13 tensor per layer is 384×4096×3072 B = exactly 4.5 GiB, so every one of the 42 offloaded layers pinned 8 GiB of w13 plus 4 GiB of w2 instead of 4.5 + 2.25: ~504 GiB requested for a 301 GiB set. DSV4.1 almost certainly paid the same tax (its per-layer w13 is also ~4.5 GiB) but with only 9–10 offloaded layers the overshoot fit in the host; that is very likely the “1.6× pin-time peak” we had been attributing to loader mmaps. Not re-measured tonight. cudaHostRegister on a plain pageable allocation pins the exact size; GPU reads through the UVA view measured 336 GiB/s registered vs 355 GiB/s torch-pinned on the same 4.5 GiB buffer. Patch 1 (uva.py, 30 lines): register instead of pin_memory(). | |||
| v4 | v3 + patch 1 | Host OOM-kill, ~90 s after all 308 GiB pinned cleanly | Pinning worked (host 368 GiB used, HBM at 224 GB, no spike). Then process_weights_after_loading started and HBM climbed 7 GiB per layer until the host died too. |
| v5 | v4 + per-layer memory trace | Same | Wall 3, traced: the TRT-LLM MXFP4 backend shuffles each expert layer into kernel layout at load time and installs the result with replace_parameter(), which registers a new tensor — on the GPU — and drops the UVA view. The trace printed uva=True → dHBM=+7072 MiB → uva_after=False for every offloaded layer: the offloader's work was silently undone one layer at a time. This is invisible on DSV4.1 (its MXFP4 path replaces in place) and invisible on any model that fits. |
| v6 | v5 + patch 2 | Model loaded; died wiring DFlash | Patch 2 (mxfp4.py, 40 lines): if the existing parameter is a UVA view and the shuffled tensor is the same byte size, copy the bytes into the pinned storage and re-view it, instead of registering a new HBM tensor. Trace on v6: 69/69 layers, dHBM=+144 MiB each (scratch), uva_after=True on all 42 offloaded layers. Then the DFlash drafter asked the target for the EAGLE3 aux-hidden-state interface and MiMoV2OmniForCausalLM (the class the omnimodal config resolves to) does not implement it; only the text-only MiMoV2ForCausalLM does. |
| v7 | v6 with speculation off (k=0) | Same EAGLE3 error | Launcher bug, not vLLM: ${SPEC:-dflash:7} treats an empty override as unset. Fixed to ${SPEC-...}. |
| v8 | k=0 for real, offload 305, 256K ctx | KV sizing: 9.48 GiB available, 17.23 needed for one 256K request | Model loaded (46/46 offloaded layers still UVA after the shuffle), graphs captured, then the KV check failed. Same HBM/KV pot as DSV4.1: the 194 GiB estimate for HBM experts was right, the KV budget was not. |
| v9 | offload 320, util 0.96, k=0, 256K | alive | 321.75 GiB pinned, 204.5 GiB HBM after load, load 107 s from NVMe via the streamer, CUDA graphs 2/2, KV 387,701 tokens (1.48× at 256K). Host 376 of 494 GiB used with 118 available. In FlashInfer autotune at publish time — the silent ~75-minute phase. |
Reproduce (as of v9)
IMAGE=vllm/vllm-openai:nightly-d05da62e9ccdf8e342b15bf6785d83224cc165af # 0.29.1rc1.dev422, post PR #57784
docker run -d --gpus all --ipc host --network host --ulimit memlock=-1 --cap-add IPC_LOCK \
-v /models/milo/MiMo-V2.6-Pro-RL:/model:ro \
-v $PWD/patch/uva.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/offloader/uva.py:ro \
-v $PWD/patch/mxfp4.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/quantization/mxfp4.py:ro \
-e RUNAI_STREAMER_MEMORY_LIMIT=17179869184 -e RUNAI_STREAMER_CONCURRENCY=16 \
$IMAGE --model /model --served-model-name mimo26-pro --trust-remote-code \
--tensor-parallel-size 1 --load-format runai_streamer \
--offload-backend uva --cpu-offload-gb 320 \
--cpu-offload-params routed_experts.w13_weight routed_experts.w2_weight \
--max-model-len 262144 --max-num-seqs 8 --max-num-batched-tokens 8192 \
--gpu-memory-utilization 0.96 \
--compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}' \
--tool-call-parser mimo --reasoning-parser mimo --enable-auto-tool-choice \
--generation-config vllm --port 30007The two patches will be posted with the recipe once a boot serves tokens. Both are bind-mounts over the stock image; nothing is rebuilt.
What comes next, in order
- v9 binds → protocol gates (parsed
tool_calls, reasoning split, one image and one audio request) → C1/C8/C16 knee → the first number. - Expected ceiling before any placement work: ~6.9 GiB of expert reads from Grace per token (46 of 69 layers × 8 experts × 19.1 MiB) at k=0, at ~340 GiB/s, is ~18 ms of fetch per step. Call it 35–45 tok/s single-stream if nothing else is broken. That is a prediction, not a measurement; it gets replaced.
- Routed-expert histogram on real traffic, then the pin-hot-experts hook from the DSV4.1 lane ported to this layer naming. With 61% offloaded, the hot set has to be small and skewed for it to pay.
- DFlash on the omni class, then a k-sweep. Flash-RL as the all-HBM control lane.
Two things worth saying plainly
First, both patches are bugs in stock vLLM for any UVA-offloaded MXFP4 model whose per-layer tensors are not a power of two — that is most of them. They will be filed upstream once the lane serves. Second, the DSV4.1 campaign's rule of thumb (“keep the pinned set under ~250 GiB”) was a measurement of the power-of-two rounding, not a property of the box. With exact registration the host holds 308 GiB pinned with 126 GiB to spare. That rule is retired.