[roberteg16]

Summary

vllm/v1/attention/ops/triton_prefill_attention.py is the encoder attention kernel — ViT towers via the mm-encoder backend, the encoder branches of TRITON_ATTN and ROCM_ATTN, and mimo_v2_omni's visual attention. (The paged LLM prefill path imports a same-named context_attention_fwd from prefix_prefill.py; that is a different module and is untouched here.)

It pads the head-dim tile to next_power_of_2(head_dim), which wastes most of the matmul on the non-power-of-2 head dims that vision towers use, and the KV tile vllm-project#58225 picks for RDNA is still wider than those shapes want. On gfx1151 this series is worth 2.19x on the head_dim=72 kernel and takes Gemma3-4B's TTFT down 21%.

Everything is gated on gfx115x, the only architecture measured; every other architecture keeps today's behaviour exactly. The split is mechanically vendor-neutral and should carry to CDNA and NVIDIA, but I have no hardware to confirm that, and an unmeasured regression on someone else's GPU is worse than a missed win. Widening the gate later is a branch in _get_head_dim_blocks / _get_encoder_launch.

Commits

commit	what it does
[ROCm] Add a gfx115x architecture predicate	_ON_GFX115X / on_gfx115x(). The other two commits gate on it, so it lands first.
[ROCm] Cover non-power-of-2 head dims with a split-D block on gfx115x	head_dim=72 runs the qk/pv dots over 128 lanes when only 5 K-passes of 16 are needed. Covers the head dim as a power-of-2 main block plus a power-of-2 tail block (72 → 64+16 = 80) behind a new constexpr BLOCK_DMODEL_TAIL; a tail of 0 is dead code and reproduces the single-block kernel exactly. _split_head_dim only takes the tail when it strictly narrows the extent — 112 would split into 64+64, covering the same 128 lanes for 12% more time.
[ROCm] Narrow the KV tile further for the split-D ViT head dims on gfx115x	The 64+16 split-D head dims want a 16-wide KV tile with 4 warps rather than the 32 that vllm-project#58225 picks. That tile iterates the KV loop more, so per-iteration load efficiency starts to matter, and at head_dim=72 the head stride is a multiple of 8 but not 16, which stops Triton attaching tt.divisibility=8 and leaves the D-contiguous loads scalar; a tl.multiple_of hint restores them. The two are one commit because neither pays on its own — see below.

Performance

Radeon 8060S (gfx1151), against upstream/main which already includes vllm-project#58225. Each cell carries that commit's own contribution in brackets.

Kernel — context_attention_fwd, B=1, S=3520, H=16, bf16, non-causal, ms/call. Each pass checks out every point in turn and the table is the median of 3 such passes; the spread across passes is ±1% or better except where noted:

head_dim	upstream	+split-D	+narrow tile & hint	total
64	1.727	1.745 (0.99x)	1.741 (1.00x)	0.99x
72	4.842	2.524 (1.92x)	2.213 (1.14x)	2.19x
80	3.511 ±8%	2.938 (1.20x)	2.347 (1.25x)	1.50x
96	3.578	3.222 (1.11x)	3.230 (1.00x)	1.11x
128	3.892	3.898 (1.00x)	3.902 (1.00x)	1.00x
256	13.746	13.744 (1.00x)	13.748 (1.00x)	1.00x

split-D moves only the non-power-of-2 head dims, as it must — elsewhere the tail is dead code. The narrow tile moves only the 64+16 split dims, 72 and 80.

Why the tile and the alignment hint are one commit. Measured separately at S=3520, median of 4 interleaved passes, they cancel out or regress on their own and only pay together:

	head_dim 72	head_dim 80
split-D alone	2.529	2.930
+ narrow tile, no hint	2.762 (0.92x)	2.338 (1.25x)
+ both	2.214 (1.14x)	2.337 (1.25x)

At 72 the narrow tile alone costs 9% and the hint alone costs 10%; together they are worth 1.14x. At 80 the head stride is already 16-byte aligned, so the hint is inert and the tile carries the win. Splitting them into two commits would put a measured regression in the middle of the series.

End to end — TTFT, single image at 1920x1080, max_tokens=1, prefix caching and the multimodal processor cache disabled and a different image per repetition so every call really re-runs the vision tower. Median of 5, ms:

model	ViT head_dim	upstream	+split-D	+narrow tile & hint	total
gemma-3-4b-it-quantized.w4a16 (SigLIP)	72	442.2	358.2 (−19.0%)	349.9 (−2.3%)	−20.9%
Qwen2.5-VL-7B-Instruct	80	1976.7	1948.5 (−1.4%)	1934.1 (−0.7%)	−2.2%

Gemma3 pads to a fixed 896x896 and pairs a 4B w4a16 language model with a 27-layer tower, so the vision encoder dominates its TTFT; Qwen2.5-VL-7B spends proportionally more time in the language model, which bounds what any encoder change can win there. Both runs were checked to really reach this kernel — 27 calls per request at shape (16, 72) and 32 at (16, 80), matching each tower's depth.

Why these values

A sweep of BLOCK_M x BLOCK_N x num_warps x waves_per_eu puts the optimum at BLOCK_N=16 with 4 warps for the 64+16 split dims only. Against vllm-project#58225's tile, median of 5 interleaved passes at S=3520:

head_dim	split	vllm-project#58225 tile	BLOCK_N=16, 4 warps
64	(64, 0)	1.818	1.943	0.94x
72	(64, 16)	2.909	2.309	1.26x
80	(64, 16)	2.934	2.341	1.25x
88	(64, 32)	3.202	3.350	0.96x
96	(64, 32)	3.225	3.385	0.95x
112	(128, 0)	3.672	8.009	0.46x
128	(128, 0)	3.896	9.236	0.42x

Hence the gate is exactly _split_head_dim(head_dim) == (64, 16); everything else keeps vllm-project#58225's tile. No waves_per_eu override helped anywhere, so none is set, and BLOCK_M keeps the generic selection.

benchmarks/kernels/benchmark_vit_prefill_attn.py is added here so this is reproducible; it launches _fwd_kernel directly so --bm/--bn/--nw/--ns/--we are honoured.

python benchmarks/kernels/benchmark_vit_prefill_attn.py --seq 3520 --heads 16 --head-dim 72 \
    --backend TRITON_ATTN --bm 128 --bn 16 --nw 4

Testing

pytest tests/kernels/attention/test_triton_prefill_attention.py   128 passed
pre-commit run --files <changed files>                            clean (ruff, ruff-format, mypy, SPDX)

The existing correctness tests only covered head_dim 128; they now also run 64, 72, 80 and 96, which between them exercise both tail widths the kernel can emit (72 and 80 split as 64+16, 96 as 64+32). Because the split is gated, a further test forces it on so the path keeps its coverage on CI machines that are not gfx115x, and _split_head_dim is checked exhaustively over head dims 1 to 512 in pure Python. Other new tests cover a view whose head stride is not 8-element aligned, and the tile selection on both sides of the split boundary. vllm-project#58225's own tests pass unchanged: this gate is chained behind _prefer_narrow_kv_tile, so mocking the platform mocks both.

Each commit compiles and passes its own tests in isolation (25, 116 and 128 passed), so the series stays bisectable. The split leaves every power-of-2 head dim bit-identical; the tile change then alters the order the softmax accumulates in, so bf16 output moves by up to one ULP.

Not a duplicate

This builds directly on vllm-project#58225, which caps BLOCK_N at 32 for all of RDNA3/RDNA4. That change covers head dims 96 and up on gfx115x, which is why this series claims nothing there. What is left is the 64+16 split dims, where the tile wants to be narrower still and the head dim needs the split in the first place.

vllm-project#56232, vllm-project#56253 and vllm-project#56276 tune triton_unified_attention.py (decoder prefill/decode) — the same argument applied to the other attention kernel, no file overlap. Encoder attention reaches this kernel under both ROCM_ATTN and TRITON_ATTN, so vllm-project#56232's backend reordering does not change whether this tune applies.

One known overlap: vllm-project#56307 adds the same _ON_GFX115X / on_gfx115x() predicate to vllm/platforms/rocm.py. Whichever lands first, the other drops that commit and rebases.

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#56450

[roberteg16]

Superseded by #1322. This PR was rebased onto current upstream, which now includes vllm-project#58225; GitHub would not reattach the force-pushed branch to this PR, so it was reopened as #1322 with the same branch and base.