[ROCm][MLA][Perf] Fuse decode QK-RoPE + Q-concat + KV-concat + KV-cache write for sparse MLA

[xaguilar-amd]

Summary

On ROCm, the sparse-MLA decode path currently runs the per-layer Q/KV preparation
as several back-to-back ops: main RoPE on q_pe/k_pe, the query
concat (ql_nope + RoPE'd q_pe), the KV concat, and the fp8 KV-cache write.
This PR collapses all of them into the single AITER kernel
fused_qk_rope_concat_and_cache_mla.

The fusion is implemented at the model level (eager op dispatch from the MLA
layer), not as an FX graph-rewrite pass since vLLM is moving away from the
graph-pass fusion mechanism, and several models already adopt the model-level
eager approach. Implementing it this way keeps this change aligned with the
current direction, and makes it easily portable to that new model-level mechanism.

The fusion is gated and falls back cleanly to the existing split-kernel path when
it is not applicable (non-ROCm, non-gfx9, no AITER, non-fp8 KV cache, the
fp8_ds_mla layout).

Before:

After:

Three coordinated decision sites, all keyed off one impl flag
(ROCMAiterMLASparseImpl.use_fused_qk_rope_cache):

1. Wrapper (mla.py) — when fusing, skip the in-place main RoPE and pass raw
   (pre-RoPE) q_pe/k_pe down, threading positions as a real argument.
2. KV write (ROCMAiterMLASparseImpl.do_kv_cache_update) — no-op when fusing
   (the write is folded into the fused kernel; the sparse path routes all tokens
   through forward_mqa, so num_mqa_tokens == q.size(0) and the fused call
   writes the full batch's KV, prefill included).
3. Decode (MLAAttention.forward_impl) — after the W_UK bmm, call the fused
   kernel to produce the pre-fused q and write the cache.

Why positions is a real argument (not a forward-context stash)

positions is threaded through MLAAttention.forward → the opaque custom op
(unified_mla_attention_with_output) / direct forward_impl → the fused kernel,
as an optional trailing tensor argument. An earlier stash into
forward_context.additional_kwargs inside the wrapper's forward proved not
to be torch.compile / CUDA-graph safe.

Compatibility / no-effect cases

• deepseek-v4 is unaffected: it uses its own standalone stack
  (DeepseekV4ROCMAiterMLAAttention / DeepseekV4ROCMAiterMLASparseBackend,
  registry tag ROCM_FLASHMLA_SPARSE_DSV4) with its own forward_mqa and cache
  write — it never touches MLAAttention, ROCMAiterMLASparseImpl,
  do_kv_cache_update, or the fused flag.
• deepseek-v3.2 / GLM sparse on ROCm run through the shared, modified path and
  are the intended beneficiaries.
• DeepseekV32Attention (nvidia) fully overrides forward; base MLA path
  untouched.
• Dense MLA / CUDA / non-fp8 / fp8_ds_mla all miss the gate → byte-for-byte
  unchanged.
• Spec-decode / MTP drafters reuse the same MLAAttention layers, so the
  positions argument covers them too.
• positions added as an optional trailing param (Tensor?) on
  MLAAttention.forward, forward_impl, and unified_mla_attention_with_output
  — backward compatible for every existing caller.
• Off-ROCm: _aiter_ops.py registration is @if_aiter_supported (no-op); the
  new op body lazy-imports AITER only when dispatched.

Test plan

Unit test (ROCm, gfx942)

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

Two tests cover the fused kernel:

1. Raw-kernel parity vs the split kernels (RoPE + concat_and_cache_mla) on a
   contiguous ql_nope, across is_neox (True/False), seq lengths, head counts
   (16/128), and block sizes (1/64):
   • bf16 q_out == concat(ql_nope, RoPE(q_pe))
   • fp8 q_out == scaled_fp8_quant(concat(...), q_scale) — verifies the kernel's
     folded q-side quant matches vLLM's downstream convention (tested with a
     non-unit q_scale so a divide-vs-multiply mistake can't alias).
   • kv_cache == concat_and_cache_mla(kv_c, RoPE(k_pe)) (fp8).
2. Production-helper test on a non-contiguous (transposed, production-layout)
   ql_nope: drives ROCMAiterMLASparseImpl.fused_qk_rope_concat_and_cache and
   checks fp8 q_out + KV parity. This guards the .contiguous() fix — the raw
   kernel silently produces wrong q_out for a strided ql_nope.

All configs pass.

Accuracy GLM-5.2-FP8 (gfx942, and TP=8)

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	30	exact_match	↑	0.9401	±	0.0065
		strict-match	30	exact_match	↑	0.9409	±	0.0065

Performance uplift GLM-5.2-FP8 (gfx942, TP=8, 3 repeats per experiment)

ISL	OSL	conc	TPOT_p50 baseline	TPOT_p50 fused	TPOT speedup	ITL_p50 baseline	ITL_p50 fused	ITL speedup	out_tok/s baseline	out_tok/s fused	tput ratio
1024	1024	1	15.03	14.49	1.037x	15.01	14.48	1.036x	65.93	68.38	1.037x
1024	1024	8	18.34	17.79	1.031x	18.34	17.72	1.035x	423.00	436.86	1.033x
1024	1024	32	24.84	24.17	1.028x	24.20	23.60	1.026x	1222.98	1253.28	1.025x
1024	1024	64	32.78	32.00	1.024x	31.39	30.66	1.024x	1824.20	1866.93	1.023x
8192	1024	1	15.85	15.32	1.035x	15.85	15.32	1.035x	60.72	62.77	1.034x
8192	1024	8	21.68	21.16	1.025x	19.34	18.77	1.031x	330.41	338.36	1.024x
8192	1024	32	39.66	37.94	1.045x	23.97	23.43	1.023x	714.48	743.42	1.040x
8192	1024	64	63.12	61.98	1.018x	29.19	28.49	1.024x	931.70	940.19	1.009x

Serving recipe:

export VLLM_ROCM_USE_AITER=1
export VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1
export VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4

vllm serve zai-org/GLM-5.2-FP8 \
  --enable_auto_tool_choice \
  --gpu_memory_utilization 0.83 \
  --kv_cache_dtype fp8 \
  --max-model-len 16384 \
  --max_num_batched_tokens 16384 \
  --no_enable_prefix_caching \
  --port 8004 \
  --reasoning_parser glm45 \
  --tensor_parallel_size 8 \
  --tool_call_parser glm47 \
  --trust_remote_code

Accuracy Deepseek-V3.2 (gfx942, TP=8)

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	30	exact_match	↑	0.9530	±	0.0058
		strict-match	30	exact_match	↑	0.9538	±	0.0058

Performance uplift DeepSeek-V3.2 (gfx942, TP=8)

ISL	OSL	conc	TPOT_p50 baseline	TPOT_p50 fused	TPOT speedup	ITL_p50 baseline	ITL_p50 fused	ITL speedup	out_tok/s baseline	out_tok/s fused	tput ratio
1024	1024	1	11.60	11.29	1.027x	11.59	11.21	1.034x	85.12	87.74	1.031x
1024	1024	8	15.82	15.38	1.029x	15.71	15.21	1.033x	488.21	507.87	1.040x
1024	1024	32	23.13	23.20	0.997x	22.86	22.26	1.027x	1328.22	1299.51	0.978x
1024	1024	64	29.76	30.67	0.970x	28.81	28.34	1.017x	2036.75	1970.10	0.967x

the small perf drop on mc32 and mc64 is most likely run variance noise and not a real perf. regression, since TPOT and ITL disagree. I can re-run the experiments a few times and provide averages if needed.

Serving recipe:

export VLLM_ROCM_MLA_SPARSE_MAX_SPLIT=1
export VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1
export VLLM_ROCM_USE_AITER=1
export VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4
export VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16=1

vllm serve deepseek-ai/DeepSeek-V3.2 \
  --tensor-parallel-size 8 \
  --gpu-memory-utilization 0.85 \
  --kv-cache-dtype fp8_e4m3 \
  --no_enable_prefix_caching \
  --max_num_batched_tokens 16384 \
  --max-num-seqs 156 \
  --block-size 64 \
  --enable-expert-parallel \
  --max_model_len 132096 \
  --hf_overrides '{"use_index_cache": true, "index_topk_freq": 4}'

Related work

• Complementary to the dense-path fusion efforts ([ROCm][Compile] Fuse RMSNorm + MXFP4 quant via AITER Triton kernels (DeepSeek-R1) #44437 F3, [ROCm][MLA] Fuse RoPE + MLA KV-cache write via AITER Triton kernel #45798).

This PR was created with the help of AI.

[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, @xaguilar-amd.

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

[Rohan138]

@xaguilar-amd can you use the existing AITER triton kernel fused_qk_rope_cat_and_cache?

[mergify]

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

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

[maeehart]

pre-commit failed on mypy, same three errors on 3.10 through 3.13. tests/kernels/core/test_fused_qk_rope_concat_and_cache_mla.py:290 passes a SimpleNamespace into fused_qk_rope_concat_and_cache, which expects ROCMAiterMLASparseImpl. The other two are vllm/models/deepseek_v32/amd/rocm.py:235 and :262 reading cos_sin_cache after rotary_emb became nn.Module | None on MLAAttention. Can you type the test helper, add the same rotary_emb assert you already have in attention.py to the AMD executor, and push?