[roberteg16]

Summary

Stacked on #1330. The base is #1330's branch, so the diff shows only this commit; it will be retargeted to gfx11 once #1330 lands. It needs ROCM_ATTN's gfx1151 impl from #1330, whose hooks call this kernel.

rdna35_rope_cache: fused decode q/k RMSNorm + RoPE + KV cache write.

• What it fuses: one launch for what a decode step runs between the qkv projection and attention:
  • the per-head q/k RMSNorm (RMSNorm or GemmaRMSNorm weights);
  • NeoX RoPE, full or partial;
  • Gemma 4's weightless V norm;
  • the K/V write into the packed HND cache.
• Integration: wired into vLLM's fuse_rope_kvcache and fuse_qk_norm_rope_kvcache passes through the hooks of ROCM_ATTN's gfx1151 impl. The passes needed several fixes:
  • the layer name as a LayerName pattern input;
  • one qk-norm pattern sized per match (Gemma 4 has 256- and 512-dim layers);
  • the fp8-query variant only when AITER's quant op exists;
  • impl-declared head sizes and the V norm;
  • a RoPE variant writing a contiguous query.
• Activation:
  • On by default from O2 when the model's attention resolves to ROCM_ATTN on gfx1151 and its activations and KV cache are fp16 or bf16.
  • Settings the user gives are kept (fuse_*=False, -rotary_embedding, splitting_ops), and a model served by another backend keeps its graph.
  • The hooks fall back to the unfused ops, warning once with the reason, for GPT-J RoPE, mixed dtypes or norm weights of another type.
• Guards:
  • the RoPE pass skips layers whose V head size differs from K's;
  • the [3, T] positions contract (rows equal, only row 0 read) is the caller's, and the passes never route mrope there;
  • every other contract is a TORCH_CHECK in the op.

Performance

rdna35_rope_cache

Measured with triton.testing.do_bench_cudagraph at T = nseq x M <= 32, contexts 128..1k (tools/rope_cache.py, tools/qknorm_rope_cache.py). It runs 2.0-2.7x faster than the kernels it replaces (Inductor's norm/RoPE and the Triton cache writer), and 1.26-1.41x faster than AITER's fused kernels in every cell.

End to end the gain is within noise, so this commit is not about end-to-end speed. It saves 2-4 us a layer on steps of 9-35 ms.

rocm-scripts regression_config.yaml (7f9f3aa), the out128 entries that fit a 32 GB board, fusion off vs on alternated per entry, median TPOT in ms:

entry	off	on	change
Gemma-2B-AWQ	9.468	9.462	-0.06 %
Gemma-3-4B-IT-w4a16	14.096	14.020	-0.54 %
Gemma-3-12B-IT-w4a16	35.036	34.938	-0.28 %
Gemma-4-E2B-IT	21.158	21.150	-0.04 %
Gemma-4-E4B-IT-AWQ-INT4	18.288	18.233	-0.30 %
Llama-2-7B-AWQ	17.209	17.192	-0.10 %
Llama-3.1-8B-AWQ-INT4	21.137	21.128	-0.04 %
Llama-3.1-8B-AWQ-INT4 EAGLE3	15.427	15.492	+0.42 %

Geomean +0.12 % decode speed. The gain matters more where steps are short or T is large: small models, batches, speculative decode at M = 8.

Test plan

• tests/kernels/attention/test_rdna35_rope_cache.py: 124 passed.
  • One case per model family of the target list (Qwen3/3.5, Gemma 3/4 incl. KV-shared layers, Llama 3.x, Phi long RoPE, Nemotron partial 78/128, mrope in decode, Granite NoPE, Ministral llama_4_scaling).
  • Up to 256 tokens (the passes' range), within 2 ulps of vLLM's native ops, which round after each RoPE multiply and add.
  • Padded slots, hybrid page sizes, a negative control, opcheck.
• tests/compile/passes/test_rdna35_rope_kvcache_fusion.py: 10 passed.
  • Both passes on ROCM_ATTN (gfx1151): LayerName, Gemma 4's two head sizes in one graph, the V norm, contiguous q/k.
  • The activation defaults: on for ROCM_ATTN, auto or selected; off for TRITON_ATTN, fp32, -rotary_embedding or fuse_*=False.
• gemma-4-E2B, mixed batches with chunked prefill: the fusion is applied, and the greedy text and the attention kernel/fallback/split counts are those of [ROCm][gfx1151] ROCM_ATTN: HIP decode attention on gfx1151 #1330 (700 / 280 / 35; one more warm-up compile range).
• Qwen3-0.6B with no flags: ROCM_ATTN with the fusion applied, correct text. With --attention-backend TRITON_ATTN it is not enabled.
• Model output, prompt logprobs, fusion off vs on, max |diff|: Qwen3-0.6B 0.16 (compiled vs eager: 0.15), Llama-3.2-1B identical, gemma-4-E2B 1.18 (compiled vs eager: 1.42).
• pre-commit (ruff, mypy, clang-format) clean.

Not duplicated

No other open PR adds a fused RoPE/KV-write kernel for gfx1151. AITER's fused_qk_norm_rope_cache_pts_quant_shuffle writes only a block-contiguous NHD cache and has no D=512, so it cannot serve ROCM_ATTN's packed HND cache on gfx1151.

AI assistance (Claude) was used for this change.