[roberteg16]

Summary

Two commits on gfx11, each reviewable on its own. The fused RoPE/KV-cache write that used to be part of this PR is now a separate draft PR, stacked on this one: #1342.

1. attention_benchmarks: honour the dtype, sliding window, KV out of cache. [Benchmarks] attention: keep the KV working set out of cache #1311 (--min-working-set-mb, so a small model at a short context does not stay in the 32 MiB MALL and read 35 % fast) plus two fixes the kernel work needed:
   • the benchmark ran the stand-in model's bf16 whatever --dtype asked;
   • BenchmarkConfig.sliding_window, with a block table holding only the blocks the window reaches.
2. ROCM_ATTN: HIP decode attention on gfx1151. No new backend: on gfx1151, ROCM_ATTN itself switches to this mechanism. Everywhere else it is unchanged.

   • Kernel: a HIP kernel for decode and speculative decode (M <= 8 query tokens per sequence) on gfx1151 (Strix Halo).

     • One workgroup per (kv head, row group, KV segment), packing every query row of a kv head so its KV is read once.
     • Both products on v_wmma_f32_16x16x16_{f16,bf16}, or a per-q-head dot decomposition where that wins.
     • Split-KV with an in-kernel merge, sliding windows, batched decode.
     • Head sizes 64, 128, 256 and 512.

   • Built into _rocm_C, no JIT at run time. csrc/rocm/rdna35_attn/variants.def lists 1280 variants:

     • every tuned configuration at M = 1..8, fp16 and bf16, one sequence and batched;
     • block size 16 plus the page sizes of gemma-4 (32) and the Qwen3.5/3.6 hybrids (528, 544, 784, 1056).

     vllm/v1/attention/ops/rdna35_hip_decode.py writes the list from its tuning tables, and cmake/rdna35_attn.cmake compiles one unit per (configuration, dtype).

   • What RocmAttentionBackend changes on gfx1151, and only there:

     • Triton's packed KV cache shape and stride order, with HND required;
     • head size 512 added;
     • RocmAttentionRdna35MetadataBuilder: Triton's builder, with the batch ordered decodes first;
     • RocmAttentionRdna35Impl, a TritonAttentionImpl subclass. TRITON_ATTN is not modified.

     The branches check the class, so ROCM_AITER_UNIFIED_ATTN (a subclass of ROCM_ATTN) keeps its own layout. registry.py and triton_attn.py are untouched.

   • Dispatch: the impl's forward launches the kernel when the call is a uniform decode it can serve. Anything else goes to TRITON_ATTN's forward on the same cache:

     • prefill and M > 8;
     • head sizes or configurations not built;
     • alibi, softcap, non-causal;
     • quantized KV or output.

     In a mixed batch the decodes go to the kernel and the prefills to Triton on a slice of the metadata, where the split measured faster.

   • Selection:

     • On gfx1151 ROCM_ATTN is ahead of TRITON_ATTN.
     • What ROCM_ATTN does not declare (sinks, KV connectors, encoder-decoder) goes to TRITON_ATTN through the normal selection.
     • Gemma 4, which forces one backend for its two head sizes, gets ROCM_ATTN on gfx1151.

   • Tuning tools in benchmarks/kernels/gfx1151_decode_attn/tools/ produce the tables and the measurements below.

Performance

Attention kernel against TRITON_ATTN

Same harness (attention_benchmarks, HND, batch 1, --min-working-set-mb 96), geomean over the contexts 128..32768; sliding-window configurations are included under their head size.

D	M	configs	fp16 vs Triton	fp16 median %roof	bf16 vs Triton	bf16 median %roof
64	1	4	1.29x	80.4 %	1.29x	80.0 %
64	2	4	1.28x	76.7 %	1.28x	76.5 %
64	3	4	1.25x	73.2 %	1.25x	72.7 %
64	4	4	1.34x	79.9 %	1.34x	79.3 %
128	1	10	1.21x	88.4 %	1.22x	88.4 %
128	2	10	1.23x	86.0 %	1.23x	85.9 %
128	3	10	1.23x	86.1 %	1.23x	86.0 %
128	4	10	1.29x	87.1 %	1.30x	86.9 %
256	1	13	1.63x	85.3 %	1.58x	81.4 %
256	2	13	1.66x	80.1 %	1.64x	79.0 %
256	3	13	1.64x	79.3 %	1.64x	78.7 %
256	4	13	1.76x	80.8 %	1.77x	80.4 %
512	1	5	2.96x	85.0 %	2.89x	81.9 %
512	2	5	3.38x	82.0 %	3.29x	79.4 %
512	3	5	3.68x	79.6 %	3.67x	79.5 %
512	4	5	3.95x	78.4 %	3.91x	78.6 %

At the page sizes of the hybrid models Triton is 2.0-6.5x slower than the kernel (e.g. 16/2/256 at block 1056: 3.4x at M=1, 5.8x at M=4, 16k context).

Test plan

• tests/kernels/attention/test_rdna35_hip_decode.py: 168 passed.
  • Every one of the 1280 variants built into _rocm_C matches the reference through the registry, for one sequence past a partial tile/page and for a batch with a padded sequence.
  • Kernel tests: layouts, split-KV merges, windows, graph replay, a negative control.
  • variants.def in sync with the tables.
  • ROCM_ATTN ahead of TRITON_ATTN only on gfx1151.
  • ROCM_ATTN's layout, impl and builder switch only on gfx1151, and ROCM_AITER_UNIFIED_ATTN keeps its own.
• Mixed-batch dispatch, gemma-4-E2B with chunked prefill, this PR alone in a clean worktree. It is identical to the earlier standalone RDNA35 backend and to the implementation before it, which hooked into TRITON_ATTN: the same greedy text on 12 prompts, and the same kernel/fallback/split call counts (700 / 245 / 35). The log reads ROCM_ATTN with the HND layout.
• Per layer, gemma-4-E2B eager, 64 tokens generated (measured with the follow-up PR applied; the attention dispatch does not depend on it). All 35 layers run the kernel on every call (Triton 0): sliding D=256 window 512, full D=512, KV-shared included.
• Qwen3-0.6B with no flags: ROCM_ATTN with HND, kernel active, correct text. With --attention-backend TRITON_ATTN the kernel does not run.
• Model output, attention kernel vs TRITON_ATTN, greedy: Llama-3.1-8B-AWQ + EAGLE3 identical on 3/3 prompts. Qwen2.5-0.5B-AWQ identical on 2/3; the third diverges where the kernel did before this PR.
• tests/v1/attention/test_rocm_attention_backends_selection.py, tests/kernels/attention/test_rocm_attention_selector.py, tests/kernels/attention/test_attention_selector.py on gfx1151: 41 passed, 3 failed, identical before and after. The 3 expect ROCM_ATTN with AITER off; with AITER installed in the test venv they get ROCM_AITER_UNIFIED_ATTN.
• pre-commit (ruff, mypy, clang-format, markdownlint) clean.
• End-to-end A/B of the attention backend against the gfx11 wheel on the regression boards.

Known, not changed here

• attention.py asserts get_name() == "TRITON_ATTN" for layers with chunk_lookback (rnj1). On ROCM_ATTN that fails on any ROCm device, gfx1151 included.
• Outside gfx1151 nothing changes. The legacy ROCM_ATTN path stores logits_soft_cap but chunked_prefill_paged_decode takes no softcap, so it is silently not applied. On gfx1151 softcap goes to Triton, which applies it.

Not duplicated

No other open PR adds an RDNA3.5 decode-attention kernel for gfx1151; it continues the gfx1151 attention work of #1295 and #1308. Commit 1 extends #1311. Whichever lands first, the other rebases onto it. The default HND layout for TRITON_ATTN on gfx1151 (#1315) is no longer part of this PR: ROCM_ATTN requires HND through its own get_required_kv_cache_layout.

AI assistance (Claude) was used for this change.

[roberteg16]

Supperseded by #1346