Gaudi Cookbook / Recipes
/

XiaomiMiMo/MiMo-V2.6-Pro-RL

1.02T MoE (42B active), 128-token sliding windows with sinks, MXFP4 experts on the MXFP4 decode kernels, 256K context on one 8-card node.

Verified on Gaudi 2 Hugging Face Full guide How it was ported Port code
moe1.02T / 42B active262,144 ctx servedmxfp4vllm-gaudi v0.26.0PortText
Hardware
8 cards · TP 8 · EP
Image
vllm-gaudi v0.26.0 @ b059f9c · vLLM @ 568afb3a · Gaudi SW 1.24.1
Variant
The public checkpoint (566 GB: FP8 block attention, MXFP4 routed experts). The experts stay MXFP4 on the device, about 67 GiB of weights per card on eight cards.
Features
Fixes
profile v0.26.0vllm-long-prefill-share
13 fixes this configuration depends on; opt-in fixes are named.
Launch
Shell 
export CHECKPOINT_DIR=/data/mimo-v2.6-pro-rl-hpu
export PT_HPU_LAZY_MODE=0
export VLLM_T_COMPILE_DYNAMIC_SHAPES=false
export PT_HPU_ENABLE_LAZY_COLLECTIVES=true
export VLLM_SKIP_WARMUP=false
export MIMO_EXPERT_DTYPE=mxfp4
export MIMO_MXK=1
export MIMO_SWA_KV_GROUPS=1
export MIMO_GA_TILE=1
export VLLM_LONG_PREFILL_SHARE=0.5
export VLLM_LONG_PREFILL_SHARE_ASYNC=1
export VLLM_LONG_PREFILL_MIN_TOKENS=65536
export VLLM_MOE_CHUNK=256
export VLLM_MOE_TOKEN_BOUNDARY=2048
export VLLM_ENGINE_ITERATION_TIMEOUT_S=3600
export VLLM_RPC_TIMEOUT=3600000
export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=3600

python3 "$COOKBOOK/ports/mimo-v2-6-flash/convert/stage_checkpoint.py" XiaomiMiMo/MiMo-V2.6-Pro-RL "$CHECKPOINT_DIR" --revision 73875d00b30a89ef8cc353a0b60b0e9f9561952d || exit 1

# Fixes: the v0.26.0 profile plus this recipe's opt-in fixes, in the pack's order.
# Stops with a non-zero exit code if the image does not match.
python3 "$COOKBOOK/fixes/apply_all.py" --profile v0.26.0 vllm-long-prefill-share || exit 1

vllm serve "$CHECKPOINT_DIR" \
  --served-model-name XiaomiMiMo/MiMo-V2.6-Pro-RL \
  --tensor-parallel-size 8 \
  --max-model-len 262144 \
  --max-num-seqs 8 \
  --long-prefill-token-threshold 2048 \
  --max-num-batched-tokens 8192 \
  --block-size 128 \
  --gpu-memory-utilization 0.8 \
  --enable-prompt-tokens-details \
  --trust-remote-code \
  --generation-config vllm \
  --override-generation-config '{"temperature": 1.0, "top_p": 0.95}' \
  --chat-template "$COOKBOOK/ports/mimo-v2-6-flash/serving/chat_template.jinja" \
  --reasoning-parser mimo \
  --enable-auto-tool-choice \
  --tool-call-parser mimo
Run inside a vllm-gaudi container. $COOKBOOK is a checkout of this repository there.
At a glance

Checks that passed

  • install.sh: profile, runner patches, plugin, kernel build --variant glm --pt2; library gate 3,253 GUIDs and both op gates pass
  • mxk_ops.py on one card: 28/28 decode cases within the parity bound, kernel worst rel 1.71e-3 vs 2.3e-3 to 4.8e-3 for the stock op; compiled = eager
  • decode 40.3 / 27.5 / 15.1 tok/s per stream at 1 / 4 / 8 streams (floor 12)
  • tool calls 6/6, including tool_choice required and a call after a 32.9K-token document
  • needles at 240K x 3 depths 3/3; four 240K-token prompts at once 4/4 (899 s); 128K 2/3 (one refusal: the model called the needle inserted text)
  • window-sized KV vs the stock layout: not completed (pod preempted); identical on single requests on Flash with the same port code

Notes

  • Needs the port image (ports/mimo-v2-6-flash, which serves both MiMo-V2.6 models). Its install.sh applies the v0.26.0 profile and vllm-long-prefill-share before the port's runner patches, and builds the MXFP4 MoE kernels (glm variant with the torch.compile op library).
  • TP8 only: window-sized KV needs the global and sliding layers to hold the same KV heads per rank.
  • --gpu-memory-utilization 0.8: in development runs 0.9 ran out of memory at a short prefill and 0.85 failed on a prefix-cached prefill.
Guide

Xiaomi's 1.02T mixture-of-experts model (42B active) at 256K tokens per request on one 8-card Gaudi 2 node. It has the same architecture as MiMo-V2.6-Flash, at a larger size: 60 of its 70 layers attend to a 128-token window with a learned sink. It runs in the same MiMo-V2.6 port image. The difference from Flash is the experts. Pro keeps them in MXFP4, because FP8 experts would not fit on one node, and runs decode on the MXFP4 MoE kernels.
Model XiaomiMiMo/MiMo-V2.6-Pro-RL (MIT), revision 73875d0
Hardware 8x Gaudi 2, one node (TP 8, experts split over the 8 ranks: 48 of 384 per card)
Image the port image, built from vllm-gaudi v0.26.0 (b059f9c), vLLM 568afb3a, Gaudi SW 1.24.1
Precision bf16 attention (FP8 block weights dequantised at load), MXFP4 routed experts (the checkpoint's bytes), bf16 KV
Kernels MXFP4 MoE decode kernels, glm variant with the torch.compile op library (MIMO_MXK=1)
Context 262,144 per request (the checkpoint allows 1,048,576)
Concurrency --max-num-seqs 8
Parsers tools mimo, reasoning mimo
Recipe recipes/XiaomiMiMo/MiMo-V2.6-Pro-RL/recipe.yaml

Build the image

The port installs into a stock vllm-gaudi v0.26.0 image (Getting started). It is the same image as for Flash. From a checkout of this repository:

docker build -f ports/mimo-v2-6-flash/Dockerfile --build-arg BASE=<your-vllm-gaudi-v0.26.0-image> \
  -t <your-mimo-v2-6-image> .

install.sh runs these steps and stops at the first one that does not apply:

  1. check the release commits;
  2. apply the v0.26.0 fix profile and vllm-long-prefill-share;
  3. apply the port's runner patches and install the plugin;
  4. build the MXFP4 MoE kernels with kernels/mxfp4-moe/build.sh --variant glm --pt2 into /opt/gaudi-cookbook-mimo/mxk.

The kernel build needs tpc-clang, which is in the Habana base image, and runs its library gate and op gates. See ports/mimo-v2-6-flash/README.md.

Required fixes

All from Fixes, the same set as Flash. install.sh applies them when the image is built; the launch command runs the same line again and only confirms them. The opt-in one is named on that line; the rest come from --profile v0.26.0, which also applies the release's other generic fixes.

Fix Why this model needs it
prompt-order prefix caching can read another request's KV in the same step
sampling-rows-refresh stale device sampling rows after a batch reorder
structured-output-rows structured-output steps sample each row with another request's parameters
bad-words-index bad_words raises IndexError in the engine core
logitsprocs-rows logit_bias, min_p and min_tokens are silently ignored
async-inputs a decode the scheduler leaves out of a step (behind long-prompt chunks) reads a stale token and repeats it
async-discard-counter AttributeError on the first preemption under async scheduling
resumed-prefill-logits a request resumed after preemption repeats its earlier output
gather-logprobs-int32 top_logprobs ids come back corrupted
logprobs-decode-guard a bad logprob id is replaced instead of stopping the API server
swa-decode-window the exact decode window (one extra block, the oldest block head-masked); the port's runner patch turns it on for mimo_v2
rank-exception-fatal an exception on a non-zero TP rank stops the engine instead of leaving ranks out of step
vllm-long-prefill-share (opt-in) a long prompt mid-prefill no longer holds every other stream to a crawl; measured on Flash, not on Pro

rope-cpu-cache is in the profile; the port also rebuilds MiMo's RoPE tables on the CPU in float64 itself (see the port write-up).

Verified

Run on one 8-card Gaudi 2 node from the stock vllm-gaudi v0.26.0 image (b059f9c, vLLM 568afb3a), with install.sh run as the Dockerfile runs it and this page's Launch block as written.

  • install.sh: commit check, the v0.26.0 profile plus vllm-long-prefill-share, both runner patches, the plugin install, and the kernel build --variant glm --pt2. The library gate (3,253 GUIDs: the 3,250 stock ones plus mxk_w13x, mxk_w2x, mxk_deq) passed, and the op gate passed on both op libraries.
  • tests/hpu/mxk_ops.py on one card, in compile mode with the torch.compile op library, on layers 1 and 69 for TP8 ranks 0 and 7: 28 of 28 decode cases (1 to 16 tokens) within the parity bound. The kernels' worst relative error against an fp32 CPU reference was 1.71e-3, against 2.3e-3 to 4.8e-3 for the stock op. Compiled output equals eager output bit for bit.
  • kernels/mxfp4-moe/tests/run_parity.sh --variant glm --e-local 48 --inter 2048 (lazy-mode op library, random weights of Pro's per-rank shapes): PARITY PASS, PREFILL PASS, SIG_DEVICE PASS.
  • tests/hpu/device_ops.py: sliding-window attention with sinks within bound in bf16 at Pro's head layout (16 query heads, 1 KV head per rank; worst relative error 0.0036).
  • stage_checkpoint.py on the public checkpoint (revision 73875d0): 130 shards, 60 sliding and 10 global layers.

Serving gates, PROFILE=pro tests/hpu/run_gates.sh against this page's command:

Gate Result
Decode, 256 greedy tokens (floor 12 tok/s per stream; the stock MoE op stays under it at 4 and 8 streams) 40.3 / 27.5 / 15.1 tok/s per stream at 1 / 4 / 8 streams
Tool calls: one call with arguments, choosing between tools, no call, using a result, tool_choice: "required", a call after a 32.9K-token document 6/6
Needles at 240K tokens x depths 0.1 / 0.5 / 0.9 (greedy) 3/3
Four 240K-token prompts at once 4/4, 899 s
Needles at 128K tokens x depths 0.1 / 0.5 / 0.9 (greedy) 2/3. At depth 0.9 the model called the line about the access code inserted text and declined to repeat it (see Known issues)
Window-sized KV against the stock layout not completed: the test pod was preempted while the stock-layout server was loading. The same comparison on MiMo-V2.6-Flash, with the same port code, was token-identical on single requests

On a test deployment of the same model code and command line (installed from its own copies, kernel library built from the same glm kernel code), the same server measured with the stock MoE op and then with the kernels:

Gate Stock MoE op With the kernels
Prompt NLL, 255-token text (fp32 CPU oracle 0.6708) 0.6797 (perplexity 1.973) 0.6797 (perplexity 1.973)
Decode vs teacher-forced scoring, 200 greedy tokens, mean abs difference 0.035 nats 0.031 nats
GSM8K 5-shot, flexible extract, 200 samples not run 0.925

Performance

Per MoE call on one card (mxk_ops.py, 48 experts of layer 1 or 69, compiled, warm): 53 to 56 us with the kernels against 353 to 357 us for the stock op at 1 token, and 247 to 248 us against 2,612 to 2,613 us at 8 tokens.

End to end, from the decode gate above (one run, warm server, 256 greedy tokens, TP 8):

Streams Tokens/s per stream Aggregate
1 40.3 40
4 27.5 110
8 15.1 121

On the test deployment, the stock MoE op ran 18.1 / 8.8 / 6.2 tok/s per stream at 1 / 4 / 8 streams, against 39.7 / 23.5 / 16.2 with the kernels (same server, same node).

  • The KV pool held about 2.0M tokens, 7.6 requests at the full 262,144-token window.
  • A needle request at about 240K tokens took 3.8 to 6.4 minutes from submission to answer, most of it prefill.
  • Starting from storage that was not in the page cache took about 35 minutes to load and 45 minutes to warm up on the test deployment.

Known issues

  • Pro sometimes reads an inserted fact as prompt injection and declines to repeat it, even when the fact is phrased plainly ("For the record, the access code ... is ..."). In the gates above, the needle at 128K and depth 0.9 was refused this way; on the test deployment, one of three needles at 240K on the stock-op server. The needle gates count a refusal as a miss.
  • In development runs of the port (earlier code, before window-sized KV), --gpu-memory-utilization 0.9 ran out of memory at a 1.5K-token prefill, and 0.85 failed on a prefix-cached prefill with "synLaunch Sections' validation failure". 0.8 is the value the test deployment ran.
  • Compile mode computes a new batch shape slightly differently the first time it runs. Compare outputs only on warm runs.
  • The kernels sum in a different order than the stock op, so greedy outputs with MIMO_MXK=1 and MIMO_MXK=0 can part ways at a near-tie. Compare quality (perplexity, needles, GSM8K), not tokens.
  • Text only. The checkpoint's vision tower, audio encoder, MTP layers and draft model are not loaded.

What not to do

  • Do not set MIMO_EXPERT_DTYPE=fp8 for Pro: the experts would need twice the memory and do not fit on one node.
  • Do not use another kernel variant, and do not turn MIMO_MXK=1 on in lazy mode. The port only loads a glm build with its torch.compile op library, and otherwise keeps the stock op.
  • Do not run window-sized KV at a TP other than 8: the global and sliding layers then hold different KV heads per rank, and the port refuses to start.
  • Do not turn prefix caching off without VLLM_CONTIGUOUS_PA=false: the contiguous KV layout bypasses the per-group block tables, and the port refuses to start.
  • Do not add speculative decoding, or switch to lazy mode with HPU graphs, with window-sized KV on: the per-layer metadata swap needs regionally compiled layers, and the port refuses both.
  • Do not drop MIMO_GA_TILE=1. Without it, the first long prefill with context builds an attention bias that warmup never profiled: 4.3 GB of bf16 for an 8,192-token chunk against 256K of context.