Repository navigation
[ROCm][Perf] Swap the unified-attention 2D grid order on gfx1151 - #56253
Draft
roberteg16 wants to merge 5 commits into
Draft
roberteg16 wants to merge 5 commits into
roberteg16 wants to merge 5 commits into
Conversation
roberteg16
requested review from
AndreasKaratzas,
WoosukKwon,
dllehr-amd,
mgoin,
shen-shanshan,
tdoublep,
tjtanaa,
tlrmchlsmth,
yewentao256 and
zyongye
as code owners
September 10, 2026 10:25
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
This was referenced Sep 10, 2026
Closed
[ROCm] Encoder prefill attention on gfx115x: split-D head dims and a narrower KV tile
ROCm/vllm#1294
Closed
roberteg16
changed the title
2 tasks done
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
roberteg16
force-pushed
the
rocm-swapped-grid
branch
from
September 13, 2026 19:39
9bc3ddc to
9903ea6
Compare
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
Give gfx11 its own tile selection so the generic path is left untouched, and match TILE_SIZE to the KV cache block size there. A tile deeper than the block straddles two blocks, which are not contiguous in the paged cache, so every tile load becomes a scatter: on Qwen3-8B-AWQ at an 8k prefill a 32-deep tile over 16-token blocks costs 21% of end-to-end TTFT against a matched tile (6440 ms -> 5086 ms). Matching unconditionally would overflow the 64KB of LDS a gfx11 workgroup gets, since a block_size of 1056 rounds up to a 2048-deep tile whose K and V halves alone need megabytes and engine startup dies with "OutOfResources: shared memory". So the largest tile that fits is derived from head_size and element_size, and the tile is shrunk towards it. Gemma3 keeps its tuned decode tile of 32, which already fits. The gating is gfx11 rather than is_navi(), which also matches gfx10xx and gfx12xx, or is_rocm(), which also matches CDNA. Only gfx11 was measured. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
Triton has no launch configuration for the 2D unified-attention path on ROCm, so it falls back on defaults chosen for NVIDIA. Pick num_warps, waves_per_eu and a per-phase pipeline depth for gfx11, and bound that depth by the 64KB LDS budget: each stage holds an independent K and V tile, so a wide head overflows shared memory and Triton aborts with "OutOfResources". The estimate is derived from the actual tile shapes rather than tabulated, so it stays correct as the tile selection changes. Only the 2D path is touched; the 3D path and every non-gfx11 target keep Triton's defaults, and an explicit launch_num_warps / launch_num_stages from the caller still overrides everything. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
The generic query block (BLOCK_M=16 below 16 queries per KV head) leaves long prefill on gfx1151 badly under-tiled. Extract the choice into _select_query_block so architectures are opted in explicitly and everything else keeps the generic block. Only gfx1151 is opted in, because that is the only part benchmarked here. The helper's shape is modelled on vllm-project#52684, which tunes the same code path for gfx1100 and is still open. The two are complementary and conflict only textually. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
On gfx11 the ROCM_ATTN custom paged-attention kernel is only used for a narrow slice of configurations: head_size 128, block_size 16, a GQA ratio between 3 and 16, an unquantized KV cache and no sliding window, alibi or sinks. Everything outside that falls back to an internal Triton path. Rank TRITON_ATTN above ROCM_ATTN there so gfx11 lands on the unified attention kernel the preceding commits tune. Measured end to end on gfx1151 against upstream, mean TTFT: Qwen3-8B-AWQ -9.6%, Llama-3.2-1B -6.4%, Gemma-3-12B -34.4%, Qwen3.5-35B-A3B -15.0% at an 8k prefill, with decode throughput equal or better in every case. The KV-connector carve-out that keeps ROCM_ATTN out of the list entirely is preserved, and MLA, sparse and every non-gfx11 target are untouched. Gated on gfx11 alone; the ROCm fork applies the same reordering to gfx12, but RDNA4 has not been measured here. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
gfx1151 schedules the 2D kernel better when the KV head is the fastest varying grid axis rather than the query block. The kernel gains a USE_SWAPPED_GRID constexpr, defaulting to False, so every other platform compiles and launches exactly what it does today. head_size 64 regresses under the swapped order, so _use_swapped_grid draws the line at 80. 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 the result is unchanged bit for bit. Verified on a 3968 token Qwen3-8B-shaped prefill, torch.equal over the full output tensor. Measured on gfx1151 (Radeon 8060S) with benchmarks/attention_benchmarks, against the same tree with the swap disabled. This is a kernel-level harness on synthetic tensors: the configs carry only head geometry, no weights are loaded, and the model names below are shorthand for that geometry. Attention kernel time, 20 repeats after 10 warm-up iterations: Qwen3-1.7B q8k 14.80 ms -> 11.28 ms +31.1% Llama-2-7B q8k 107.54 ms -> 93.59 ms +14.9% Qwen2.5-7B-AWQ q3968s4096 5.59 ms -> 4.99 ms +12.0% Qwen3-1.7B q3968s4k 3.15 ms -> 2.96 ms +6.4% Qwen3-8B q3968s4096 6.19 ms -> 5.90 ms +5.0% Qwen3-Omni-30B q8k 36.33 ms -> 34.71 ms +4.7% A second independent A/B pass reproduced these within about a percentage point, except Qwen3-1.7B q8k, which came out at +40.9% with a much wider spread; the +31.1% above is the conservative reading. The Qwen3-8B geometry is present in the matrix twice under two names, so it is really four measurements spanning +3.8% to +5.0%. Two shapes lose. Llama-2-7B at q1920s2k costs 2.3% (4.7% on the second pass), while the same model at q8k gains 15%: the win needs a long enough sequence to amortise the reordered scheduling. Qwen3-4B-AWQ at q4000s4128 is 1.4% down, inside the noise band but consistent across both passes, and it sits exactly on the head_size 80 boundary. Everything else is flat, decode included: the swap is gated on the 2D path. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Robert Esclapez Garcia <robert.garcia@amd.com>
roberteg16
force-pushed
the
rocm-swapped-grid
branch
from
September 23, 2026 09:26
9903ea6 to
8e70bbe
Compare
roberteg16
marked this pull request as draft
September 24, 2026 08:18
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Important Reviewer Notice
Depends on #56232, but cannot stack them since branches are coming from outside the repository. #56232 needs merging first and then rebase this one over main.
Only last commit is meaningful, the rest are coming from #56232
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_GRIDconstexpr defaulting toFalse, 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 toTRITON_ATTNat all, and thehead_size >= 80boundary 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.equalover the full[3968, 32, 128]tensor (max abs diff 0.0, zero differing elements).Results
These are kernel microbenchmarks.
benchmarks/attention_benchmarksoperates 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.q8kq8kq3968s4096q3968s4kq3968s4096q8kq1920s2kq4000s4128Testing
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.
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.