[roberteg16]

Summary

On gfx1151 the 2D unified-attention kernel schedules better when the KV head is the fastest varying grid axis rather than the query block. This launches the 2D grid as (kv_heads, q_blocks) there, worth up to 31% of attention kernel time on long prefill.

The kernel gains a USE_SWAPPED_GRID constexpr defaulting to False, so every other platform compiles and launches exactly what it does today. The whole change is 2 lines of existing code plus the new gating helper.

Base

This is stacked on rocm-2d-launch-config, so the diff here shows only the swap. It is textually independent of that series — they share no hunks — but the swap is only meaningful on top of it: without the 2D launch tuning gfx11 does not route to TRITON_ATTN at all, and the head_size >= 80 boundary is the one that series already established.

Output is unchanged, bit for bit

The swap only relabels which grid axis carries which index. Every program still computes the same (q_block, kv_head) pair and accumulates over KV in the same order, so there is no reassociation and no change in results — unlike the tile/block-size tuning in the base series, which does alter reduction blocking.

Verified rather than asserted: a 3968-token Qwen3-8B-shaped prefill run with the swap on and off produces outputs that are torch.equal over the full [3968, 32, 128] tensor (max abs diff 0.0, zero differing elements).

Results

These are kernel microbenchmarks. benchmarks/attention_benchmarks operates on synthetic tensors; the configs carry only head geometry (num_q_heads, num_kv_heads, head_dim, block_size) and load no weights. The model names are shorthand for that geometry — quantization labels in particular have no effect here.

geometry	head_dim	shape	base	swapped	run 1	run 2
Qwen3-1.7B	128	q8k	14.80 ms	11.28 ms	+31.1%**	+40.9%
Llama-2-7B	128	q8k	107.54 ms	93.59 ms	+14.9%**	+15.4%
Qwen2.5-7B	128	q3968s4096	5.59 ms	4.99 ms	+12.0%	+11.0%
Qwen3-1.7B	128	q3968s4k	3.15 ms	2.96 ms	+6.4%	+4.0%
Qwen3-8B	128	q3968s4096	6.19 ms	5.90 ms	+5.0%	+4.2%
Qwen3-Omni-30B-A3B	256	q8k	36.33 ms	34.71 ms	+4.7%	+5.7%
Llama-2-7B	128	q1920s2k	2.20 ms	2.25 ms	−2.3%	−4.7%
Qwen3-4B	80	q4000s4128	5.09 ms	5.16 ms	−1.4%	−1.6%

Testing

pytest tests/kernels/attention/test_triton_unified_attention.py    1874 passed (gfx1151)
pre-commit run --files <the two changed files>                     clean

AI assistance

This change was prepared with AI assistance (Claude). Every line was reviewed by the submitter, and the measurements above were run and checked by hand.

[roberteg16]

Upstream PR vllm-project#56253