[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 vision towers use, and the KV tile vllm-project#58225 picks for RDNA is still wider than those shapes want. On gfx1151 this is worth 2.19x on the head_dim=72 kernel and takes Gemma3-4B's TTFT down 21%.

Everything is inside one if on_gfx1151() in context_attention_fwd that only overrides launch values. Every other part keeps today's behaviour exactly: the launch arguments are computed as they are now, and only reassigned inside that branch.

    BLOCK_DMODEL = triton.next_power_of_2(Lk)
    BLOCK_DMODEL_TAIL = 0
    HEAD_STRIDE_ALIGNED_8 = False

    if on_gfx1151() and q.dtype in (torch.bfloat16, torch.float16):
        main = BLOCK_DMODEL // 2
        tail = max(16, triton.next_power_of_2(Lk - main))
        if Lk != BLOCK_DMODEL and main + tail < BLOCK_DMODEL:
            BLOCK_DMODEL, BLOCK_DMODEL_TAIL = main, tail

        HEAD_STRIDE_ALIGNED_8 = ...   # from the real strides of Q/K/V/O

        if HEAD_STRIDE_ALIGNED_8 and (BLOCK_DMODEL, BLOCK_DMODEL_TAIL) == (64, 16):
            BLOCK_N, num_warps = 16, 4

Exactly one existing line in the wrapper changes (BLOCK_DMODEL= in the launch). Four more change in _fwd_kernel, hoisting the head-axis offsets into locals so the alignment hint has something to annotate. Everything else is additive, and the new path is constexpr dead code when BLOCK_DMODEL_TAIL == 0.

Commits

commit	what it does
[ROCm] Split the head-dim block for non-power-of-2 head dims on gfx1151	Covers the head dim as a power-of-2 main block plus a power-of-2 tail (72 → 64+16 = 80) behind a new constexpr BLOCK_DMODEL_TAIL; a tail of 0 reproduces the single-block kernel exactly. Skipped where it would not narrow the extent — 112 would be 64+64, the same 128 lanes for 12% more time.
[ROCm] Narrow the KV tile for the 64 + 16 head dims on gfx1151	Two things, and only one is about the tile. The tl.multiple_of hint on the head-axis offsets keeps the D-contiguous loads vectorized at head_dim=72, where the head stride is 8 mod 16 and Triton will not infer the alignment; it is set from the real strides and is unconditional. The narrow tile (BLOCK_N=16, 4 warps) is then taken only where that hint is active.

Performance

Radeon 8060S (gfx1151), against upstream/main which already includes vllm-project#58225. Each pass checks out every point in turn and tables are the median of 3 such passes; spread is ±1% or better except where noted.

Kernel — context_attention_fwd, B=1, S=3520, H=16, bf16, non-causal, ms/call:

head_dim	upstream	+split	+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

The hint is unconditional; the tile is what depends on it. At every head dim whose stride is already 16-byte aligned the hint is a tautology the compiler had anyway, and it measures exactly that. Each head dim on the tile it is given, median of 4 interleaved passes, ms/call:

head_dim	no hint	with hint
64	1.814	1.813	1.00x
72	2.887	2.302	1.25x
80	2.322	2.325	1.00x
96	3.204	3.203	1.00x
128	3.871	3.885	1.00x
256	13.556	13.563	1.00x

Only head_dim=72, the one that is 8 mod 16, learns anything from it. The narrow tile is the part that needs it: with vectorized loads BLOCK_N=16 and 4 warps are worth 1.25x at both 72 (2.895 → 2.293) and 80 (2.924 → 2.331), but without them they are a regression — at head_dim=72 with scalar loads the narrow tile runs 2.676 ms against 2.449 for the tile vllm-project#58225 picks. Hence if HEAD_STRIDE_ALIGNED_8 and ...: a head_dim of 72 in a view whose head stride is padded keeps the wider tile.

Sweeping BLOCK_N x num_warps confirms vllm-project#58225's tile is right for everything else: 64 wants 32/4, and 96, 128 and 256 want 32/8, where BLOCK_N=16 is 22% to 45% slower even with the warp count it prefers.

End to end — TTFT, single 1920x1080 image, 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	+narrow tile & hint	total
gemma-3-4b-it-quantized.w4a16 (SigLIP)	72	442.8	362.2 (−18.2%)	350.8 (−3.1%)	−20.8%
Qwen2.5-VL-7B-Instruct	80	1946.7	1934.4 (−0.6%)	1907.1 (−1.4%)	−2.0%

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. Both runs were checked to really reach this kernel — 27 calls per request at (16, 72) and 32 at (16, 80), matching each tower's depth.

Why these values

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 narrower tile is taken exactly for the 64 + 16 dims. 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 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). The launch configuration is asserted through the capture harness vllm-project#58225 added, extended with a head_dim and an on_gfx1151 knob — so the gfx1151 branch is covered on CI agents that are not gfx1151, and vllm-project#58225's own tests stop depending on the host part. A further test runs the split kernel for real against the SDPA reference with the predicate mocked on.

Each commit compiles and passes its own tests in isolation (115 and 128 passed). 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.

Model evaluation

The change perturbs bf16 output by up to one ULP (the KV tile alters the order the softmax accumulates in), so it is worth showing what that does downstream.

Multimodal, which is what this kernel serves. Five synthetic images with a known answer through Qwen2.5-VL-7B-Instruct, greedy, 32 tokens, TRITON_ATTN mm-encoder backend — 32 encoder-kernel calls per request, confirmed by instrumenting context_attention_fwd. The generated text is identical between upstream/main and this branch on all five:

[PASS] red/circle: 'Red circle.'      [PASS] yellow/square: 'Yellow square.'
[PASS] green/circle: 'Green circle.'  [PASS] purple/circle: 'Purple circle.'
[PASS] blue/square: 'Blue square.'

GSM8K, as a regression check only. tests/evals/gsm8k/gsm8k_eval.py, 600 questions, 5-shot, temperature 0, same model. It is worth being explicit that this benchmark does not reach this kernel: a text-only prompt makes zero calls to context_attention_fwd, measured the same way, because there is no vision tower to run. Running upstream/main three times gives 0.1250 / 0.1350 / 0.1400, so the harness spreads 1.5 accuracy points on identical code; this branch scores 0.1417, inside that spread. (The absolute number is low because the harness's bare few-shot completion prompt does not make this instruct model stop — 230 of 256 tokens per question — not because of anything here.)

Not a duplicate

This builds on vllm-project#58225, which caps BLOCK_N at 32 for all of RDNA3/RDNA4 and already covers head dims 96 and up on gfx1151 — which is why nothing is claimed there. What is left is the 64 + 16 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 applies.

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.