[ROCm][Kimi-K3][Perf] Fuse MLA decode KV-cache write and Q-prep via AITER

[rbrugaro-amd]

Purpose

Kimi-K3 MLA decode on ROCm runs three kernels per layer per step:

1. vllm::concat_and_cache_mla_kernel — the fp8 KV-cache write
2. at::native::CatArrayBatchedCopy — the [ql_nope | q_pe] decode-query concat
3. aiter::scaled_quant_kernel — the static fp8 quant of that query

This collapses all three into the single AITER fused_qk_rope_concat_and_cache_mla
launch that #47757 already wires up for sparse MLA: 15.5 µs → 5.77 µs per MLA
layer per decode step, 2.62× on the fused region, measured on MI355X TP=8.

Before fusion trace:

After fusion trace:

K3 MLA is NoPE (mla_use_nope), while that kernel always applies RoPE and takes
the cos/sin caches as required arguments. It is fed an identity rotation
(cos=1, sin=0) so the rotation degenerates to a no-op — the same approach
ROCm/ATOM uses to run K3 through this kernel.

The identity cache is a single row. AITER's per-head and _opt decode kernels
clamp pos into [0, cos_cache.size(0)), so real positions pass through
untouched; at K3's max_position_embeddings of 1 048 576 a full-length constant
cache would otherwise cost ~268 MB to express a provable no-op. AITER's general
decode kernel does not clamp (cache_kernels.cu:2868), so the K3 wrapper
disables fusion for the head counts that would select it (kv_lora_rank * num_heads < 2048, i.e. ≤3 heads/rank, reachable at TP=32).

Dependencies

• Depends on [ROCm][MLA][Perf] Fuse decode QK-RoPE + Q-concat + KV-concat + KV-cache write for sparse MLA #47757. That PR adds the _aiter_ops.py binding, the
  MLAAttention.forward_impl hook, and the positions plumbing this builds on.
  When it merges I will rebase and its commits
  will drop out of this diff, leaving only the final commit.
• Its commits are included here unmodified, with original authorship and
  sign-offs, purely so CI can run. Please review only the last commit.

Why this is not a duplicate

No open or merged PR fuses the MLA KV-cache write with Q-prep for Kimi-K3.

• [ROCm][MLA][Perf] Fuse decode QK-RoPE + Q-concat + KV-concat + KV-cache write for sparse MLA #47757 — same kernel, but sparse MLA on deepseek_v32. This is the dense
  path, and the NoPE handling is new.
• [ROCm] Fuse decode rope + MLA KV-cache write + fp8 query assembly via AITER #55230 — same kernel family, DeepSeek-R1, edits shared layers/mla.py.
• [ROCm][MLA] Fuse RoPE + MLA KV-cache write via AITER Triton kernel #45798 / [ROCm] Add AITER RoPE + KV cache fusion for MLA prefill and decode #38313 / [ROCm] Add AITER RoPE+KV cache and dual RMSNorm fusion for MLA #41099 / [ROCm] Add unified AITER RoPE+KV cache and dual RMSNorm fusion for MLA #41101 — the same fusion as a torch.compile
  graph-rewrite pass. Kimi-K3 is not @support_torch_compile, so those passes
  never fire for it.
• [Perf][ROCm][Attention] Fuse Kimi-K3 MLA chunked-context gather on AITER #53166 — K3 ROCm MLA, but chunked-context gather, a different op.
• [AMD][kimik3][ROCm][Perf] Fuse MLA q/kv RMSNorm in AMD Kimi-K3 MLA wrapper #52494 (merged) — the sibling q/kv RMSNorm fusion in the same wrapper.

What this adds on top of #47757

• FusedQKRopeCacheMLAMixin in rocm_aiter_mla.py, extracted from
  ROCMAiterMLASparseImpl so the dense AITER MLA impl used by K3 shares one
  implementation with the sparse one instead of duplicating it.
• MLAAttention.can_fuse_qk_rope_cache() as the single batch-independent
  predicate, consumed by _fuses_qk_rope_cache, the base MLA wrapper, and the K3
  wrapper. Previously the wrapper kept a second hand-maintained copy of the
  condition; that is how RoPE ownership can drift between the site that skips the
  eager rotation and the site that performs the fused call.
• Dense MLA fuses only decode-only batches, and only when the rotation is a
  declared no-op (is_identity). A dense impl that fuses some batches but not
  others would otherwise leave the wrapper applying RoPE eagerly on all of them
  while the kernel applied it again on some.
• DCP, PCP and HiSparse excluded. The fused call covers neither prefill rows
  nor DCP padded tokens, and skipping the standalone write would also skip the
  PCP latent all-gather (a collective every rank must reach) and the HiSparse
  cursor advance.

The fuse/no-fuse decision is derived only from state that both update_kv_cache
and forward_impl already hold, so nothing is carried across the two custom-op
boundaries between them.

Test Plan

All numbers below are from the stock upstream nightly
vllm/vllm-openai-rocm:nightly-rocm100-0bfc7a15d095fe83ecc82b50561a93c177fece2d
(2026-09-16), 8× MI355X, Kimi-K3 fp4, TP=8.

• arm A (baseline) = that image, unmodified
• arm B (fused) = that image with this PR's .py files copied over
  site-packages — the only difference between the arms

Every run asserts, from inside the running container, that the fused marker is
present in arm B and absent in arm A, so neither arm can silently A/B against
itself.

Unit tests

pytest tests/kernels/core/test_fused_qk_rope_concat_and_cache_mla.py -v

E2E serving

export VLLM_ROCM_USE_AITER=1
vllm serve <Kimi-K3> --served-model-name moonshotai/Kimi-K3 \
  --trust-remote-code --load-format fastsafetensors --language-model-only \
  --tensor-parallel-size 8 --gpu-memory-utilization 0.9 \
  --kv-cache-dtype fp8 --max-model-len 32768 --no-enable-prefix-caching \
  --host 0.0.0.0 --port 8010

vllm bench serve --backend openai-chat --endpoint /v1/chat/completions \
  --model moonshotai/Kimi-K3 --tokenizer <Kimi-K3> --trust-remote-code \
  --host 127.0.0.1 --port 8010 \
  --dataset-name random --random-input-len 1024 --random-output-len 8192 \
  --random-range-ratio 0 --ignore-eos \
  --num-prompts <conc> --max-concurrency <conc> \
  --seed 42 --percentile-metrics ttft,tpot,itl,e2el --metric-percentiles 50,90,99

--block-size is left to vLLM, which derives 1536 to match the KDA
mamba page size.

Results

1. Unit tests — PASS

44 passed, including 4 new test_nope_identity_cos_sin_leaves_pe_unrotated
cases covering both is_neox settings and block_size 1/64, with positions up
to 2²⁰ to exercise the clamp. Mutation check: flipping the identity cache from
cos=1 to cos=0.5 fails all 4, so they are not vacuous.

2. Decode kernel census — 3 kernels → 1

Rank-0 torch profile, 4 engine steps, 24 MLA layers, 96 = 24 × 4.

kernel	baseline	fused
vllm::concat_and_cache_mla_kernel	96 × 4.85 µs = 465.5 µs	0
at::native::CatArrayBatchedCopy	96 × 6.43 µs = 617.1 µs	0
aiter::scaled_quant_kernel	96 × 4.65 µs = 446.4 µs	0
aiter::fuse_qk_rope_concat_and_cache_mla_kernel	0	96 × 5.77 µs = 553.8 µs

Per MLA layer per decode step: 15.5 µs → 5.77 µs. Launch count 2988 → 2940
per step (−48, −1.6%).

All 8 ranks: 12 941 µs → 4 933 µs = 2.62× (per-rank range 2.58–2.67×,
exactly 96 fused kernels on every rank).

The dispatch lands on AITER's _opt decode kernel, not per_head — per_head
requires kv_cache.size(1) == 1, and vLLM derives block size 1536 here. The
_opt gate is kv_lora_rank × num_heads ≥ 2048; at TP=8 that is 512 × 12 =
6144. _opt is also the position-clamping kernel the single-row identity cache
requires.

3. GSM8K — no accuracy change

Full 1319 questions, 5-shot, temperature=0, no --limit. Three A/B pairs, two
on an earlier nightly and one on the image above:

pair	baseline	fused	delta
1	0.9629	0.9719	+0.90 pp
2	0.9674	0.9704	+0.30 pp
3 (this image)	0.9704	0.9629	−0.75 pp
mean	0.9669	0.9684	+0.15 pp

The delta changes sign across pairs with identical code, swinging ±0.9 pp
against a binomial se_diff of ~0.7 pp. Read as "no measurable accuracy change,
±0.9 pp run-to-run", not as a benefit. All six runs sit far above the 0.90 gate.

4. E2E serving — unchanged within noise

ISL 1024 / OSL 8192 (decode-dominant, where this fusion has its largest share of
wall-clock), 3 repeats per arm with arm order alternating between repeats.
±rng is max−min across repeats, as a percentage.

conc	metric	baseline	fused	ratio	A ±rng	B ±rng
1	TPOT p50 (ms)	15.78	15.74	1.003×	0.27%	0.16%
1	ITL p50 (ms)	15.79	15.74	1.003×	0.31%	0.20%
1	ITL p90 (ms)	15.90	15.84	1.004×	0.21%	0.20%
1	out tok/s	63.28	63.47	1.003×	0.27%	0.16%
8	TPOT p50 (ms)	21.27	21.22	1.002×	0.40%	0.53%
8	ITL p50 (ms)	21.30	21.26	1.002×	0.32%	0.50%
8	ITL p90 (ms)	21.45	21.40	1.002×	0.10%	0.47%
8	out tok/s	374.76	375.51	1.002×	0.40%	0.53%
32	TPOT p50 (ms)	30.93	30.84	1.003×	1.07%	1.48%
32	ITL p50 (ms)	31.26	31.10	1.005×	0.78%	1.65%
32	ITL p90 (ms)	31.89	31.80	1.003×	0.40%	0.06%
32	out tok/s	1025.94	1029.03	1.003×	1.06%	1.46%

Every ratio is 1.002–1.005×, i.e. ≤0.5%, and at conc 8 and 32 smaller than the
arms' own run-to-run spread. At conc 1 — the quietest point, and the one with
least batching to amortise fixed per-layer cost — the ~0.3% delta does exceed
the spread and is consistent across all six paired repeat comparisons, including
the repeat where the baseline ran first. That is suggestive, but it is a 0.3%
effect measured with 3 repeats and no A-vs-A null control, so this PR does not
claim it.

This is the expected outcome by construction: the fusion removes 0.53% of a
decode step's device time, on 24 MLA layers, in a model that is roughly 1:3
MLA:KDA. For reference, cross_device_reduce_2stage alone is ~46% of rank-0
decode device time in the same trace.

Compatibility / no-effect cases

• Non-ROCm, non-AITER, non-fp8 KV cache — use_fused_qk_rope_cache is never
  set, can_fuse_qk_rope_cache() short-circuits on its first guard.
• Dense MLA with a real RoPE (DeepSeek-R1 and friends on ROCm AITER) —
  refused by the is_identity guard, so the wrapper keeps applying RoPE eagerly
  exactly as today. Without that guard the flag alone would make the wrapper skip
  the eager rotation on batches the kernel then declines to fuse.
• Sparse MLA ([ROCm][MLA][Perf] Fuse decode QK-RoPE + Q-concat + KV-concat + KV-cache write for sparse MLA #47757's path) — preserved. always_fuses_qk_rope_cache is
  True there, so can_fuse and _fuses reduce to [ROCm][MLA][Perf] Fuse decode QK-RoPE + Q-concat + KV-concat + KV-cache write for sparse MLA #47757's original condition on
  every batch. The only new refusals are PCP and HiSparse, which [ROCm][MLA][Perf] Fuse decode QK-RoPE + Q-concat + KV-concat + KV-cache write for sparse MLA #47757 did not
  exclude and where the fused call would have written un-gathered latents at
  un-remapped slots.
• DCP — excluded; unchanged from today.
• CUDA / CPU / XPU — untouched.

I enumerated impl × kv-dtype × rope-kind × pcp × hisparse × batch-mix × dcp
(768 combinations) and checked that RoPE is applied exactly once in every one,
and that positions is present whenever the fused branch runs.

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

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @rbrugaro-amd.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[rbrugaro-amd]

superseded by #57640 which I implemented only touching amd vendor files