[vllmellm]

Overview

MiMo-V2.6 has asymmetric K/V head dims (192 / 128), which today pins it to the DiffKV attention backends: any other --attention-backend is silently replaced by the DiffKV auto-pick, and the native KV-cache write would copy V with K's head width. This PR makes --attention-backend ROCM_AITER_UNIFIED_ATTN usable for the model (opt-in, defaults unchanged) and fixes the write kernel.

Claims

• --attention-backend ROCM_AITER_UNIFIED_ATTN now selects unified attention for every MiMo-V2.6 layer instead of being dropped; verified on Flash (TP2, TP4, TP8) and on Pro (TP8, gsm8k n=100 parity).
• reshape_and_cache_flash writes V with its own head size and head stride. Symmetric callers run the same code as before: the new checks live in the asymmetric branch, and the existing kernel is byte-identical to main.
• Accuracy and speculative-decoding parity with the default backend (gsm8k at n=1319, MTP x3 acceptance).
• At long context prefill is 1.43x faster at 128k and 1.21x at 32k, with decode 1.17x faster per token (TPOT p50 2.94 vs 3.45 ms) and output throughput up 21-45%. At 8k with concurrency 8 the two are at parity. Measured with --block-size 16 on both arms, which tested better than the backend's default for this shape. On Pro at TP8 the same benchmark gives 1.13x prefill and 1.66x lower inter-token latency.

Validation

Commands

# serve (TP4 shown; TP2/TP8 differ only in --tensor-parallel-size; the control arm drops
# --attention-backend and gets the DiffKV auto-pick)
vllm serve XiaomiMiMo/MiMo-V2.6-Flash-MOPD \
  --tensor-parallel-size 4 --max-model-len 270336 --gpu-memory-utilization 0.82 \
  --block-size 16 --trust-remote-code --generation-config vllm \
  --attention-backend ROCM_AITER_UNIFIED_ATTN \
  --speculative-config '{"method":"mtp","num_speculative_tokens":2}' \
  --disable-prefix-caching \
  --tool-call-parser mimo --enable-auto-tool-choice --reasoning-parser mimo

# serve bench, one cell (input/output lengths and prompt counts per the tables)
vllm bench serve --backend openai-chat --model XiaomiMiMo/MiMo-V2.6-Flash-MOPD \
  --base-url http://127.0.0.1:8000 --endpoint /v1/chat/completions \
  --dataset-name random --random-input-len 131072 --random-output-len 32 \
  --num-prompts 2 --max-concurrency 1 --num-warmups 1 --ignore-eos --temperature 0 --seed 0

# gsm8k 25-shot against the running server
lm_eval --model local-chat-completions \
  --model_args model=XiaomiMiMo/MiMo-V2.6-Flash-MOPD,base_url=http://127.0.0.1:8000/v1/chat/completions,num_concurrent=32,max_retries=10,max_length=32000,timeout=60000 \
  --gen_kwargs '{"max_tokens":16000,"until":[]}' --batch_size auto \
  --tasks gsm8k --num_fewshot 25 --apply_chat_template

# tests
pytest tests/kernels/attention/test_cache.py -k asymmetric
pytest tests/models/test_mimo_v2_attention_backend.py

Bench serve

Mimo-V2.6 Flash at TP4:

workload	DiffKV (default)	AITER unified	delta
128k in / 32 out, 2 prompts, c1: TTFT p50	7633.8 ms	5317.7 ms	-30% (1.44x)
128k in / 32 out: output throughput	4.10 tok/s	5.96 tok/s	+45%
32k in / 128 out, 4 prompts, c1: TTFT p50	890.9 ms	738.2 ms	-17% (1.21x)
32k in / 128 out: TPOT p50	3.45 ms	2.94 ms	-15%
32k in / 128 out: output throughput	96.1 tok/s	116.0 tok/s	+21%
8k in / 256 out, 8 prompts, c8: ITL p50	8.42 ms	8.48 ms	parity
8k in / 256 out: output throughput	929 tok/s	947 tok/s	+2%

On MiMo-V2.6-Pro at TP8:

workload	DiffKV (default)	AITER unified	delta
128k in / 32 out, 2 prompts, c1: TTFT p50	9665.0 ms	8576.9 ms	-11% (1.13x)
128k in / 32 out: ITL p50	19.52 ms	11.77 ms	-40% (1.66x)
128k in / 32 out: output throughput	3.23 tok/s	3.68 tok/s	+14%

Acceptance rate:

	DiffKV (default)	AITER unified
MTP acceptance, 128k (acc % / length)	53.33 / 2.07	56.67 / 2.13
MTP acceptance, 32k	61.40 / 2.23	55.60 / 2.11
MTP acceptance, 8k c8	57.4-60.5	58.2-60.3

lm_eval

checkpoint	unified	DiffKV	note
Flash, gsm8k n=1319	0.9750 / 0.9742	0.9719 / 0.9712	strict / flexible

Needle-In-A-Haystack

Unified attention only; --max-model-len 270336, thinking off, greedy.

probe	config	result
niah	265,673 tokens, depths 0 / 50 / 100 %	3/3

Tests

# kernel fix: main's cache_kernels.cu vs this one
pytest tests/kernels/attention/test_cache.py -k asymmetric        # 33 failed -> 33 passed
# model change: main's backend selection vs this one (either kernel)
pytest tests/models/test_mimo_v2_attention_backend.py             # 2 failed, 5 passed -> 7 passed

The asymmetric kernel test runs MiMo's geometry (K 192 / V 128, 4 KV heads) at both block sizes the model uses (16 and 64) over a parametrized token count, in both layouts unified attention hands the kernel, for bf16 and fp8 caches and per-tensor and per-head scales. The whole file stays green (694 passed, 602 skipped).

The model test gained two checks: the selected backend's impl has to be unified attention's, not just its enum; and the cache the layer writes has to match a cache built in torch, per slot, with K in the first head_dim columns of a row and V in the remaining v_head_dim. Both fail on main. The kernel test's buffers carry slack past the cache, so a regressed kernel fails that comparison instead of writing out of bounds.

Details

reshape_and_cache_flash_kernel derived one head size, one page stride, one head stride and one element count from key/key_cache and reused all of them for the value side:

cache_t* v_dst_h = value_dst + head * head_stride;             // K's stride
vectorize_with_alignment(value_src, value_dst, n_elems, ...);  // K's width

Every caller until now was symmetric, so this was invisible. Unified attention is the first caller that passes v_head_size != head_size (MiMo is 192 / 128, packed as one (head_size + head_size_v) content dim per KV head), and for it V was read at K's per-head offset, copied at K's width - 64 elements past the end of each head - and written past the end of the cache on the last slot. The asymmetric case now gets its own kernel that addresses V with the value side's own head stride and length; the existing kernel and its NHD fast path are untouched, so symmetric callers run the same code as before.

Model side: the attention backend is chosen per layer, and the gate now consults a backend capability (supports_head_size_v) rather than listing backend names. DiffKV declares that it can serve any K/V width relation; unified attention declares v <= k, because it splits the page with split(head_size). A backend that cannot serve the shape still gets the previous behaviour: a warning and the DiffKV auto-pick.

Known limitations:

• Requires aiter >= 0.1.25 for the V-dim-masked unified-attention kernels; main pins 0.1.26.
• Unified attention serves V narrower than K, not wider (v_head_dim > head_dim cannot be expressed by its page split); such a request falls back to the DiffKV auto-pick with a warning.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[tjtanaa]

@vllmellm, please run NIAH dataset.

[tjtanaa]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97895 for commit 08070088f1b7.

[vllmellm]

@vllmellm, please run NIAH dataset.

added the NIAH section.

[vllmellm]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97904 for commit 0a90d841aeac.

[tjtanaa]

LGTM

[vllmellm]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98144 for commit 7e9ec1c7d10d.

[khluu]

stamp for CI change