[ROCm][Perf] Kimi-K3 AMD MLA: fuse the q-a and kv-a RMSNorms

[mpashkovskii]

Purpose

On the Kimi-K3 AMD MLA front-end, every token runs two separate RMSNorm
launches — q_a_layernorm(q_c) then kv_a_layernorm(kv_c) — once per MLA layer.

This PR collapses them into a single fused_q_kv_rmsnorm call
(models/common/ops/fused_qk_rmsnorm.py), a portable Triton kernel already in the
tree — no CUDA-only gate, and its one PDL branch is tl.constexpr-guarded on
current_platform.is_arch_support_pdl(), which returns False on ROCm — so it runs
on ROCm as-is with no new kernel.

The front-end norm fires on every token across all MLA layers, so collapsing two
launches into one removes a per-token kernel launch and its intermediate write. This
is a small, low-risk micro-optimization: verified at the kernel level (see Test
Result), throughput-neutral end-to-end on the workloads measured (the norms are a
small, non-critical-path share of GPU time, and the saving is launch overhead that
cudagraph capture already elides on the served decode path). It is offered as a clean
reduction in launch count and intermediate writes with no accuracy risk, not as a
headline speedup.

How it works

• MLAModules gains an optional q_kv_norm callback (layers/mla.py). When
  None (every other model/backend), the wrapper keeps its two separate norms —
  that path is byte-identical to before.
• The MLA wrapper's forward, on the q-LoRA path (q_lora_rank is not None),
  defers q_a_layernorm and, when q_kv_norm is set, normalizes q_c and kv_c
  together at the kv-split point. It rebinds q_c to the normed tensor so both the
  q-projection and any later indexer see the normalized value — matching the
  separate-norm semantics exactly.
• Kimi-K3 AMD KimiMLAAttention installs the callback, wrapping
  fused_q_kv_rmsnorm with the q_a/kv_a layernorm weights and config.rms_norm_eps
  (both norms share the same eps). Only wired on the q-LoRA path; the uncompressed
  path has no q norm and is unchanged.

This mirrors the pattern established by #50664 (optional MLAModules callback for the
K3 output gate; generic default, AMD installs the fused variant). The uncompressed
path (no q norm) and every other model that leaves q_kv_norm unset are untouched.

Relates to #50682 (ROCm/AMD Kimi-K3 gap and roadmap tracking → Performance
Optimization → Mixed Attention / MLA).

Not a duplicate

No AMD adapter or open PR wires fused_q_kv_rmsnorm into the AMD K3 path (checked
against #50682's PR list and open rocm+k3 PRs). The nearby fusions are different:
#50664 (MLA output gate), #50637 (AttnRes + RMSNorm), #50634 (KDA decode gate), #41095
(AITER dual-RMSNorm for DeepSeek MLA, not K3). #51772 (chunked-context K/V packing)
does not touch common/ops or this norm.

Files changed

• vllm/model_executor/layers/mla.py — optional q_kv_norm field on MLAModules;
  wrapper forward uses it on the q-LoRA path. No behavior change when unset.
• vllm/models/kimi_k3/amd/linear.py — KimiMLAAttention builds and passes the
  fused callback.
• tests/models/kimi_k3/test_amd_mla_qkv_norm.py — parity + empty-batch test.

Test Plan

Unit — the fused op must equal two independent RMSNorm modules (ROCm GPU):

.venv/bin/python -m pytest -q tests/models/kimi_k3/test_amd_mla_qkv_norm.py

Covers token counts {1, 7, 64, 512, 4096} (decode- through prefill-sized) plus an
empty batch.

Accuracy — this changes an activation path, so a full eval is required:

# serve K3 on the AMD box (gfx942), then:
make lm-eval-gsm8k     # gsm8k parity vs main (must match within tolerance)

End-to-end prefill perf (long input) — baseline main vs this PR, same seeds:

make bench-<name>      # 8k/1k profile; compare TTFT (prefill) before/after

Test Result

Environment: 8× MI325X (gfx942), moonshotai/Kimi-K3, int4_per_group_32, TP=8,
EP off. This is a small but fair micro-optimization: the kernel-level effect is
real and verified, and end-to-end it is throughput-neutral (within run-to-run noise)
on this prefill-heavy workload — it neither helps nor harms serving numbers, and it
carries no accuracy risk.

Kernel evidence (torch profiler traces) — the fused kernel does exactly what it
should. Each fused_q_kv_rmsnorm launch replaces the two separate q_a / kv_a
RMSNorm launches, and nothing else in the trace changes:

	candidate (fused)	baseline (two norms)
fused_q_kv_rmsnorm launches	384	0
plain q-a + kv-a RMSNorm launches	0	768
front-end norm GPU time	1.61 ms	3.16 ms

The plain-RMSNorm count drops by exactly 768 = 2 × 384 — every fused launch removes
the two norms it replaces — for ~1.55 ms less GPU time and 384 fewer kernel launches
across the profiling window. All other kernels are unchanged.

End-to-end serving (vllm bench serve, 8024/1024, conc 16). Baseline =
main (two norms); Candidate = this PR (fused):

Metric	Baseline	Candidate	Δ
Mean TTFT (ms)	4615.0	4617.2	+0.0%
Median TTFT (ms)	2559.4	2559.3	~0%
Mean TPOT (ms)	52.70	52.49	−0.4%
Output tok/s	279.35	280.34	+0.4%
Total tok/s	2468.3	2477.1	+0.4%

All deltas are within run-to-run noise (≤0.4%). Expected: the front-end norms are
~4.4% of GPU time and are not on the critical path (the workload is bound by the
projection GEMMs and the SiTU activation), and the saving here is kernel-launch
overhead — which full cudagraph capture already elides on the served decode path. So
the fusion is a clean win in eager mode and at the kernel level, but does not move
aggregate serving throughput at this operating point. It is included as a correct,
low-risk reduction in launch count and intermediate writes, not as a headline speedup.

Correctness

• Unit (test_amd_mla_qkv_norm.py): TODO PASS/FAIL (N passed — the fused op equals
  two independent RMSNorm modules across token counts {1, 7, 64, 512, 4096}).
• gsm8k parity vs main (1319 samples): TODO — baseline acc vs candidate acc; the
  fused op is fp32-accurate, so parity is expected.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

AI assistance was used for this change.

[mergify]

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

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