From 054f67a90b1fda8497891ffc8d199e595a1bfa7e Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Wed, 9 Sep 2026 14:20:55 -0700 Subject: [PATCH 1/5] [ROCm] Match the unified-attention KV tile to the block size on gfx11 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 Signed-off-by: Robert Esclapez Garcia --- .../test_triton_unified_attention.py | 74 +++++++++++++++++++ .../attention/ops/triton_unified_attention.py | 71 +++++++++++++++++- 2 files changed, 143 insertions(+), 2 deletions(-) diff --git a/tests/kernels/attention/test_triton_unified_attention.py b/tests/kernels/attention/test_triton_unified_attention.py index e30d9225eac8..2c8ef6399acd 100644 --- a/tests/kernels/attention/test_triton_unified_attention.py +++ b/tests/kernels/attention/test_triton_unified_attention.py @@ -9,6 +9,7 @@ from vllm.triton_utils import tl, triton from vllm.utils.math_utils import next_power_of_2 from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.ops import triton_unified_attention as triton_ua from vllm.v1.attention.ops.triton_attention_helpers import ( apply_softcap, compute_tile_loop_bounds, @@ -937,3 +938,76 @@ def test_softcap_does_not_overflow_on_large_scores() -> None: ref = soft_cap * torch.tanh(scores / soft_cap) assert torch.isfinite(out).all(), out torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3) + + +def _patch_arch( + monkeypatch: pytest.MonkeyPatch, + gfx11: bool = False, + gfx1151: bool = False, +) -> None: + """Pretend to be a given AMD architecture, so these tests need no GPU.""" + monkeypatch.setattr(triton_ua, "_ON_GFX11", gfx11) + monkeypatch.setattr(triton_ua, "_ON_GFX1151", gfx1151) + + +@pytest.mark.parametrize("is_prefill", [True, False]) +@pytest.mark.parametrize("head_size", [64, 80, 128, 256, 512]) +@pytest.mark.parametrize("block_size", [16, 32, 1056]) +def test_get_tile_size_off_gfx11_is_unchanged( + monkeypatch: pytest.MonkeyPatch, + is_prefill: bool, + head_size: int, + block_size: int, +) -> None: + """Off gfx11 the tile selection must match the pre-existing behaviour.""" + _patch_arch(monkeypatch) + tile = triton_ua._get_tile_size( + head_size, -1, element_size=2, is_prefill=is_prefill, block_size=block_size + ) + assert tile == (32 if is_prefill else 16) + + +@pytest.mark.parametrize("head_size", [128, 256, 512]) +@pytest.mark.parametrize("block_size", [16, 32, 1056]) +def test_get_tile_size_gfx11_fits_lds( + monkeypatch: pytest.MonkeyPatch, head_size: int, block_size: int +) -> None: + """One stage of K+V tiles must fit the 64KB gfx11 LDS budget. + + Regression guard for ``OutOfResources: shared memory``: a KV cache + ``block_size`` of 1056 rounds up to a 2048-deep tile, which overflows by + a wide margin at every head size. + """ + _patch_arch(monkeypatch, gfx11=True) + tile = triton_ua._get_tile_size( + head_size, -1, element_size=2, is_prefill=False, block_size=block_size + ) + assert tile & (tile - 1) == 0, "AMD Triton requires a power-of-2 TILE_SIZE" + assert 2 * tile * head_size * 2 <= triton_ua._GFX11_LDS_BUDGET + + +def test_get_tile_size_gfx11_preserves_gemma3(monkeypatch: pytest.MonkeyPatch) -> None: + """Gemma3's decode tile must survive the gfx11 power-of-2 rewrite.""" + _patch_arch(monkeypatch, gfx11=True, gfx1151=True) + for head_size in (128, 256): + tile = triton_ua._get_tile_size( + head_size, 1024, element_size=2, is_prefill=False, block_size=16 + ) + assert tile == 32 + + +def test_get_tile_size_gfx11_matches_block_size( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The tile must match the KV block, or every load straddles two blocks. + + A deeper tile than ``block_size`` spans blocks that are not contiguous in + the paged cache; on Qwen3-8B at an 8k prefill that cost 21% of TTFT. + """ + _patch_arch(monkeypatch, gfx11=True, gfx1151=True) + for head_size in (80, 128, 256): + for is_prefill in (True, False): + tile = triton_ua._get_tile_size( + head_size, -1, element_size=2, is_prefill=is_prefill, block_size=16 + ) + assert tile == 16 diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index b0231dd0c09a..c4fdb42c3151 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -30,6 +30,12 @@ ) from vllm.v1.kv_cache_interface import KVQuantMode +if current_platform.is_rocm(): + from vllm.platforms.rocm import _ON_GFX11, _ON_GFX1151 +else: + _ON_GFX11 = False + _ON_GFX1151 = False + logger = init_logger(__name__) is_batch_invariant = envs.VLLM_BATCH_INVARIANT float8_info = torch.finfo(current_platform.fp8_dtype()) @@ -786,13 +792,66 @@ def _is_gemma3_attention(head_size: int, sliding_window: int) -> bool: return sliding_window == 1024 and head_size in (128, 256) +# gfx11 exposes 64KB of LDS per workgroup. +_GFX11_LDS_BUDGET = 65536 + + +def _gfx11_tile_size( + head_size: int, + sliding_window: int, + element_size: int, + block_size: int, +) -> int: + """Select the KV tile size on gfx11. + + Match the tile to the KV cache block size. A tile deeper than the block + straddles two blocks, which are not contiguous in the paged cache, so + every tile load turns into a scatter; on Qwen3-8B at an 8k prefill a + 32-deep tile over 16-token blocks costs 21% of end-to-end TTFT against + a matched tile. Gemma3 keeps its tuned decode tile of 32. + + One software-pipeline stage holds a K and a V tile of + ``tile_size * head_size * element_size`` bytes each, so the tile is then + shrunk until that pair fits the LDS budget. Without it Triton aborts + engine startup with ``OutOfResources: shared memory`` — a ``block_size`` + of 1056 rounds up to a 2048-deep tile, which overflows at every head size. + + Args: + head_size: Attention head dimension. + sliding_window: Window size, used to recognise Gemma3. + element_size: Bytes per KV element. + block_size: KV cache block size. + + Returns: + A power-of-2 tile size that fits the budget. + + """ + # Note: tile size must be at least 32 for fp8 (element_size == 1). + min_tile = 16 if element_size >= 2 else 32 + + if _is_gemma3_attention(head_size, sliding_window): + tile_size = 32 + else: + tile_size = triton.next_power_of_2(block_size) + + max_tile = _GFX11_LDS_BUDGET // (2 * head_size * element_size) + while tile_size > min_tile and tile_size > max_tile: + tile_size //= 2 + + return tile_size + + def _get_tile_size( head_size: int, sliding_window: int, element_size: int, is_prefill: bool, + block_size: int, ) -> int: """Select tile size with Gemma3-specific optimization.""" + if _ON_GFX11: + return _gfx11_tile_size(head_size, sliding_window, element_size, block_size) + if _is_gemma3_attention(head_size, sliding_window): # Gemma3: use 32 for decode (default is 16) return 32 @@ -981,10 +1040,18 @@ def unified_attention( chunk_lookback = -1 TILE_SIZE_PREFILL = _get_tile_size( - head_size, sliding_window_val, q.element_size(), is_prefill=True + head_size, + sliding_window_val, + q.element_size(), + is_prefill=True, + block_size=block_size, ) TILE_SIZE_DECODE = _get_tile_size( - head_size, sliding_window_val, q.element_size(), is_prefill=False + head_size, + sliding_window_val, + q.element_size(), + is_prefill=False, + block_size=block_size, ) # Wider KV tile for the tuned large-head path (see above). Only the 2D From da904b904026e74d3f51401b8dd4ec8553a256eb Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Wed, 9 Sep 2026 14:21:45 -0700 Subject: [PATCH 2/5] [ROCm] Give the unified-attention 2D path a gfx11 launch config 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 Signed-off-by: Robert Esclapez Garcia --- .../test_triton_unified_attention.py | 48 +++++++++++ .../attention/ops/triton_unified_attention.py | 79 +++++++++++++++++++ 2 files changed, 127 insertions(+) diff --git a/tests/kernels/attention/test_triton_unified_attention.py b/tests/kernels/attention/test_triton_unified_attention.py index 2c8ef6399acd..79c8fabf8459 100644 --- a/tests/kernels/attention/test_triton_unified_attention.py +++ b/tests/kernels/attention/test_triton_unified_attention.py @@ -1011,3 +1011,51 @@ def test_get_tile_size_gfx11_matches_block_size( head_size, -1, element_size=2, is_prefill=is_prefill, block_size=16 ) assert tile == 16 + + +@pytest.mark.parametrize("head_size", [80, 128, 256, 512]) +@pytest.mark.parametrize("block_m", [16, 64, 128]) +def test_cap_num_stages_for_gfx11_lds(head_size: int, block_m: int) -> None: + """Pipeline depth must never push the estimated LDS past the budget.""" + tile_size, element_size = 32, 2 + capped = triton_ua._cap_num_stages_for_gfx11_lds( + 3, block_m, tile_size, head_size, element_size + ) + assert 1 <= capped <= 3 + q_bytes = block_m * head_size * element_size + s_bytes = block_m * tile_size * 4 + kv_bytes = capped * 2 * tile_size * head_size * element_size + # A single stage is the floor, so the tightest shapes still exceed the + # budget on paper; the tile shrink in _get_tile_size is what keeps them + # launchable. + assert capped == 1 or q_bytes + s_bytes + kv_bytes <= triton_ua._GFX11_LDS_BUDGET + + +@pytest.mark.parametrize("max_seqlen_q", [1, 4096]) +@pytest.mark.parametrize("head_size", [64, 128, 256]) +@pytest.mark.parametrize("block_m", [16, 64, 128]) +def test_gfx11_launch_config( + monkeypatch: pytest.MonkeyPatch, + max_seqlen_q: int, + head_size: int, + block_m: int, +) -> None: + """Every launch config must be 4 warps and a depth that fits the budget.""" + _patch_arch(monkeypatch, gfx11=True) + tile_size, element_size = 32, 2 + config = triton_ua._gfx11_launch_config( + max_seqlen_q, head_size, element_size, block_m, tile_size + ) + assert config["num_warps"] == 4 + assert config["waves_per_eu"] == 2 + assert config["num_stages"] == triton_ua._cap_num_stages_for_gfx11_lds( + config["num_stages"], block_m, tile_size, head_size, element_size + ) + + +def test_gfx11_launch_config_decode_depth(monkeypatch: pytest.MonkeyPatch) -> None: + """Only a narrow-head decode is allowed a deep pipeline.""" + _patch_arch(monkeypatch, gfx11=True) + assert triton_ua._gfx11_launch_config(1, 64, 2, 16, 16)["num_stages"] == 3 + assert triton_ua._gfx11_launch_config(1, 128, 2, 16, 16)["num_stages"] == 1 + assert triton_ua._gfx11_launch_config(4096, 64, 2, 16, 32)["num_stages"] == 1 diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index c4fdb42c3151..e1a096b995cd 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -863,6 +863,79 @@ def _get_tile_size( return 16 if element_size >= 2 else 32 +def _cap_num_stages_for_gfx11_lds( + num_stages: int, + block_m: int, + tile_size: int, + head_size: int, + element_size: int, +) -> int: + """Cap the K/V software-pipeline depth to fit the gfx11 LDS budget. + + Each stage holds independent K and V tiles, so LDS grows with + ``num_stages``. Wide heads overflow the budget at the depths chosen + below and Triton aborts with ``OutOfResources: shared memory``. + Estimated layout:: + + Q tile : block_m * head_size * element_size + S accum: block_m * tile_size * 4 (fp32) + K + V : num_stages * 2 * tile_size * head_size * element_size + + Args: + num_stages: Requested pipeline depth. + block_m: Query rows per block. + tile_size: KV tile depth. + head_size: Attention head dimension. + element_size: Bytes per KV element. + + Returns: + The capped depth. Never below 1: a single stage that still does not + fit cannot be shortened any further. + + """ + q_bytes = block_m * head_size * element_size + s_bytes = block_m * tile_size * 4 + kv_bytes_per_stage = 2 * tile_size * head_size * element_size + remaining = _GFX11_LDS_BUDGET - q_bytes - s_bytes + return max(1, min(num_stages, remaining // kv_bytes_per_stage)) + + +def _gfx11_launch_config( + max_seqlen_q: int, + head_size: int, + element_size: int, + block_m: int, + tile_size: int, +) -> dict[str, int]: + """Launch parameters for the 2D unified-attention kernel on gfx11. + + Triton's defaults are tuned for NVIDIA and leave the 2D path an order of + magnitude below the bf16 peak on RDNA3. The 3D path is not covered. + + Args: + max_seqlen_q: Longest query in the batch; ``1`` means decode. + head_size: Attention head dimension. + element_size: Bytes per query element. + block_m: Query rows per block. + tile_size: KV tile depth the kernel will launch with. + + Returns: + Keyword arguments for the kernel launch. + + """ + # A wide head lets the K/V tiles dominate LDS, so only a narrow-head + # decode can afford a deep pipeline. + num_stages = (1 if head_size >= 80 else 3) if max_seqlen_q == 1 else 1 + + return { + "num_warps": 4, + "num_stages": _cap_num_stages_for_gfx11_lds( + num_stages, block_m, tile_size, head_size, element_size + ), + "waves_per_eu": 2, + } + + def unified_attention( q, k, @@ -1153,7 +1226,13 @@ def unified_attention( grid = (total_num_q_blocks, num_kv_heads, num_par_softmax_segments) tile_size = TILE_SIZE_DECODE + # Lowest precedence first: platform tuning, then the caller's explicit + # ``launch_*`` overrides. An empty dict leaves the Triton defaults. launch_kwargs: dict[str, int] = {} + if not use_3d and _ON_GFX11: + launch_kwargs = _gfx11_launch_config( + max_seqlen_q, head_size, q.element_size(), BLOCK_M, tile_size + ) if launch_num_warps is not None: launch_kwargs["num_warps"] = launch_num_warps if launch_num_stages is not None: From 21516e2e5170fcf6c7b61357e2040a97329abbda Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Wed, 9 Sep 2026 14:21:50 -0700 Subject: [PATCH 3/5] [ROCm] Widen the unified-attention query block for long gfx11 prefill 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/vllm#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 Signed-off-by: Robert Esclapez Garcia --- .../test_triton_unified_attention.py | 88 ++++++++++++++++--- .../attention/ops/triton_unified_attention.py | 68 ++++++++++++-- 2 files changed, 135 insertions(+), 21 deletions(-) diff --git a/tests/kernels/attention/test_triton_unified_attention.py b/tests/kernels/attention/test_triton_unified_attention.py index 79c8fabf8459..b840aa825d01 100644 --- a/tests/kernels/attention/test_triton_unified_attention.py +++ b/tests/kernels/attention/test_triton_unified_attention.py @@ -1031,31 +1031,93 @@ def test_cap_num_stages_for_gfx11_lds(head_size: int, block_m: int) -> None: assert capped == 1 or q_bytes + s_bytes + kv_bytes <= triton_ua._GFX11_LDS_BUDGET -@pytest.mark.parametrize("max_seqlen_q", [1, 4096]) -@pytest.mark.parametrize("head_size", [64, 128, 256]) -@pytest.mark.parametrize("block_m", [16, 64, 128]) +@pytest.mark.parametrize( + ("max_seqlen_q", "nq_per_kv", "head_size", "expected"), + [ + # gfx1151 long prefill: wide head takes 64, narrow head takes 128. + (256, 4, 128, (64, 16, True)), + (4096, 8, 80, (64, 8, True)), + (4096, 4, 64, (128, 32, True)), + # Below the long-prefill threshold, and decode, stay generic. + (255, 4, 128, (16, 4, False)), + (1, 4, 128, (16, 4, False)), + # Extreme GQA still respects num_queries_per_kv. + (4096, 32, 128, (64, 2, True)), + ], +) +def test_select_query_block_gfx1151( + monkeypatch: pytest.MonkeyPatch, + max_seqlen_q: int, + nq_per_kv: int, + head_size: int, + expected: tuple[int, int, bool], +) -> None: + _patch_arch(monkeypatch, gfx11=True, gfx1151=True) + assert triton_ua._select_query_block(max_seqlen_q, nq_per_kv, head_size) == expected + + +def test_select_query_block_untuned_arch_is_generic( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An architecture nobody tuned must keep the generic query block.""" + _patch_arch(monkeypatch, gfx11=True) + assert triton_ua._select_query_block(8192, 4, 128) == (16, 4, False) + + +@pytest.mark.parametrize( + ("gfx1151", "max_seqlen_q", "num_kv_heads", "head_size", "block_m", "tuned"), + [ + # Decode: only a narrow head can afford a deep pipeline. + (True, 1, 8, 64, 16, False), + (True, 1, 8, 128, 16, False), + # Prefill on the generic query block stays shallow. + (True, 4096, 8, 128, 16, False), + # Tuned gfx1151 prefill: wide head, then narrow head. + (True, 4096, 8, 128, 64, True), + (True, 4096, 8, 64, 128, True), + # Gemma-2B style MQA opts back out of the deep pipeline. + (True, 4096, 1, 256, 64, True), + # A tuned query block on an unmeasured gfx11 part keeps the defaults. + (False, 4096, 8, 128, 64, True), + ], +) def test_gfx11_launch_config( monkeypatch: pytest.MonkeyPatch, + gfx1151: bool, max_seqlen_q: int, + num_kv_heads: int, head_size: int, block_m: int, + tuned: bool, ) -> None: - """Every launch config must be 4 warps and a depth that fits the budget.""" - _patch_arch(monkeypatch, gfx11=True) + """Every launch config must be 4 warps and fit the gfx11 LDS budget.""" + _patch_arch(monkeypatch, gfx11=True, gfx1151=gfx1151) tile_size, element_size = 32, 2 config = triton_ua._gfx11_launch_config( - max_seqlen_q, head_size, element_size, block_m, tile_size + max_seqlen_q, + num_kv_heads, + head_size, + element_size, + block_m, + tile_size, + tuned, ) assert config["num_warps"] == 4 - assert config["waves_per_eu"] == 2 + assert config["num_stages"] >= 1 + assert config["waves_per_eu"] in (2, 4, 6) assert config["num_stages"] == triton_ua._cap_num_stages_for_gfx11_lds( config["num_stages"], block_m, tile_size, head_size, element_size ) -def test_gfx11_launch_config_decode_depth(monkeypatch: pytest.MonkeyPatch) -> None: - """Only a narrow-head decode is allowed a deep pipeline.""" - _patch_arch(monkeypatch, gfx11=True) - assert triton_ua._gfx11_launch_config(1, 64, 2, 16, 16)["num_stages"] == 3 - assert triton_ua._gfx11_launch_config(1, 128, 2, 16, 16)["num_stages"] == 1 - assert triton_ua._gfx11_launch_config(4096, 64, 2, 16, 32)["num_stages"] == 1 +def test_gfx11_launch_config_gfx1151_prefill_is_tuned( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The tuned gfx1151 prefill path is where the speedup comes from.""" + _patch_arch(monkeypatch, gfx11=True, gfx1151=True) + config = triton_ua._gfx11_launch_config(4096, 8, 64, 2, 128, 32, True) + assert config == {"num_warps": 4, "num_stages": 3, "waves_per_eu": 4} + + # head_size 256 with a single KV head is Gemma-2B: shallow on purpose. + mqa = triton_ua._gfx11_launch_config(4096, 1, 256, 2, 64, 32, True) + assert mqa == {"num_warps": 4, "num_stages": 1, "waves_per_eu": 4} diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index e1a096b995cd..3bd74650af6f 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -863,6 +863,33 @@ def _get_tile_size( return 16 if element_size >= 2 else 32 +def _select_query_block( + max_seqlen_q: int, num_queries_per_kv: int, head_size: int +) -> tuple[int, int, bool]: + """Pick ``(BLOCK_M, BLOCK_Q, tuned)`` for the 2D query blocking. + + Long prefill on gfx1151 is served far better by a wider query block than + the generic selection gives it. Architectures are opted in explicitly; + everything else keeps the generic block. + + Returns: + The query block size, the derived per-KV-head block, and whether the + tuned path was taken (the launch config keys off the flag). + + """ + if max_seqlen_q >= 256 and _ON_GFX1151: + block_m = max( + 64 if head_size >= 80 else 128, + triton.next_power_of_2(num_queries_per_kv), + ) + return block_m, block_m // num_queries_per_kv, True + + block_m = ( + 16 if num_queries_per_kv <= 16 else triton.next_power_of_2(num_queries_per_kv) + ) + return block_m, block_m // num_queries_per_kv, False + + def _cap_num_stages_for_gfx11_lds( num_stages: int, block_m: int, @@ -902,10 +929,12 @@ def _cap_num_stages_for_gfx11_lds( def _gfx11_launch_config( max_seqlen_q: int, + num_kv_heads: int, head_size: int, element_size: int, block_m: int, tile_size: int, + tuned_query_block: bool, ) -> dict[str, int]: """Launch parameters for the 2D unified-attention kernel on gfx11. @@ -914,25 +943,43 @@ def _gfx11_launch_config( Args: max_seqlen_q: Longest query in the batch; ``1`` means decode. + num_kv_heads: KV head count, used to spot MQA. head_size: Attention head dimension. element_size: Bytes per query element. block_m: Query rows per block. tile_size: KV tile depth the kernel will launch with. + tuned_query_block: Whether ``_select_query_block`` took its tuned + gfx1151 path. Returns: Keyword arguments for the kernel launch. """ - # A wide head lets the K/V tiles dominate LDS, so only a narrow-head - # decode can afford a deep pipeline. - num_stages = (1 if head_size >= 80 else 3) if max_seqlen_q == 1 else 1 + waves_per_eu = 2 + if max_seqlen_q == 1: + # A wide head lets the K/V tiles dominate LDS, so decode cannot + # afford a deep pipeline. + num_stages = 1 if head_size >= 80 else 3 + else: + num_stages = 1 + if tuned_query_block: + # A wider query block pays for a deeper pipeline and more + # occupancy per EU. Only reachable on gfx1151, the sole + # architecture measured here. + num_stages = 3 + waves_per_eu = 6 if head_size >= 80 else 4 + if head_size == 256 and num_kv_heads == 1: + # Gemma-2B style MQA: the single KV head starves the deeper + # pipeline, so keep it shallow. + num_stages = 1 + waves_per_eu = 4 return { "num_warps": 4, "num_stages": _cap_num_stages_for_gfx11_lds( num_stages, block_m, tile_size, head_size, element_size ), - "waves_per_eu": 2, + "waves_per_eu": waves_per_eu, } @@ -1066,10 +1113,9 @@ def unified_attention( num_queries_per_kv = num_query_heads // num_kv_heads head_size = q.shape[2] - BLOCK_M = ( - 16 if num_queries_per_kv <= 16 else triton.next_power_of_2(num_queries_per_kv) + BLOCK_M, BLOCK_Q, tuned_query_block = _select_query_block( + max_seqlen_q, num_queries_per_kv, head_size ) - BLOCK_Q = BLOCK_M // num_queries_per_kv # Tuned launch parameters; ``None`` lets Triton pick its defaults. launch_num_warps: int | None = None @@ -1231,7 +1277,13 @@ def unified_attention( launch_kwargs: dict[str, int] = {} if not use_3d and _ON_GFX11: launch_kwargs = _gfx11_launch_config( - max_seqlen_q, head_size, q.element_size(), BLOCK_M, tile_size + max_seqlen_q, + num_kv_heads, + head_size, + q.element_size(), + BLOCK_M, + tile_size, + tuned_query_block, ) if launch_num_warps is not None: launch_kwargs["num_warps"] = launch_num_warps From fa7ce9981319b0f6c3956f77b719446428b4a42b Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Wed, 9 Sep 2026 14:21:55 -0700 Subject: [PATCH 4/5] [ROCm] Prefer TRITON_ATTN over ROCM_ATTN on gfx11 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 Signed-off-by: Robert Esclapez Garcia --- .../attention/test_attention_selector.py | 26 +++++++++++++++++++ vllm/platforms/rocm.py | 10 ++++++- 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/tests/kernels/attention/test_attention_selector.py b/tests/kernels/attention/test_attention_selector.py index 53867a5edf93..81e71f60e1e7 100644 --- a/tests/kernels/attention/test_attention_selector.py +++ b/tests/kernels/attention/test_attention_selector.py @@ -947,3 +947,29 @@ def test_rswa_selection_does_not_reuse_causal_result(blackwell_selection): get_attn_backend(256, torch.bfloat16, None) config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN assert get_attn_backend(256, torch.bfloat16, None).get_name() == "TRITON_ATTN" + + +@pytest.mark.skipif(RocmPlatform is None, reason="requires ROCm") +@pytest.mark.parametrize("gfx11", [True, False]) +def test_rocm_attn_priority_gfx11(monkeypatch: pytest.MonkeyPatch, gfx11: bool): + """gfx11 ranks TRITON_ATTN above ROCM_ATTN; other ROCm parts do not.""" + import vllm.platforms.rocm as rocm + + monkeypatch.setattr(rocm, "on_gfx11", lambda: gfx11) + order = rocm._get_backend_priorities(use_mla=False, use_sparse=False) + triton_idx = order.index(AttentionBackendEnum.TRITON_ATTN) + rocm_idx = order.index(AttentionBackendEnum.ROCM_ATTN) + assert (triton_idx < rocm_idx) is gfx11 + + +@pytest.mark.skipif(RocmPlatform is None, reason="requires ROCm") +def test_rocm_attn_still_dropped_for_kv_connector(monkeypatch: pytest.MonkeyPatch): + """The KV-connector carve-out must survive the gfx11 reordering.""" + import vllm.platforms.rocm as rocm + + monkeypatch.setattr(rocm, "on_gfx11", lambda: True) + order = rocm._get_backend_priorities( + use_mla=False, use_sparse=False, use_kv_connector=True + ) + assert AttentionBackendEnum.ROCM_ATTN not in order + assert AttentionBackendEnum.TRITON_ATTN in order diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 9089d025e0e5..f4aa9636be21 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -471,7 +471,13 @@ def _get_backend_priorities( backends = [] # Keep ROCM_ATTN disabled for KV connectors until connector transfer # semantics are validated for its asymmetric native K/V cache views. - if not use_kv_connector: + allow_rocm_attn = not use_kv_connector + # On gfx11 the ROCM_ATTN custom paged-attention kernel only covers + # head_size 128 with block_size 16, a GQA ratio of 3-16 and an unquantized + # KV cache (see use_rocm_custom_paged_attention); anything else falls back + # to its internal Triton path. Rank TRITON_ATTN above it there. + rocm_attn_first = allow_rocm_attn and not on_gfx11() + if rocm_attn_first: backends.append(AttentionBackendEnum.ROCM_ATTN) if rocm_aiter_ops.is_mha_enabled(): backends.append(AttentionBackendEnum.ROCM_AITER_FA) @@ -480,6 +486,8 @@ def _get_backend_priorities( elif rocm_aiter_ops.is_rdna_aiter_enabled(): backends.insert(0, AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN) backends.append(AttentionBackendEnum.TRITON_ATTN) + if allow_rocm_attn and not rocm_attn_first: + backends.append(AttentionBackendEnum.ROCM_ATTN) backends.append(AttentionBackendEnum.TURBOQUANT) return backends From 8e70bbeb3205c05c351a132f781114071c333eb7 Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Mon, 7 Sep 2026 10:00:01 -0700 Subject: [PATCH 5/5] [ROCm] Swap the unified-attention 2D grid order on gfx1151 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 Signed-off-by: Robert Esclapez Garcia --- .../test_triton_unified_attention.py | 21 +++++++++++++ .../attention/ops/triton_unified_attention.py | 30 +++++++++++++++++-- 2 files changed, 48 insertions(+), 3 deletions(-) diff --git a/tests/kernels/attention/test_triton_unified_attention.py b/tests/kernels/attention/test_triton_unified_attention.py index b840aa825d01..3366511f2a6a 100644 --- a/tests/kernels/attention/test_triton_unified_attention.py +++ b/tests/kernels/attention/test_triton_unified_attention.py @@ -1121,3 +1121,24 @@ def test_gfx11_launch_config_gfx1151_prefill_is_tuned( # head_size 256 with a single KV head is Gemma-2B: shallow on purpose. mqa = triton_ua._gfx11_launch_config(4096, 1, 256, 2, 64, 32, True) assert mqa == {"num_warps": 4, "num_stages": 1, "waves_per_eu": 4} + + +@pytest.mark.parametrize("head_size", [64, 80, 128, 256, 512]) +@pytest.mark.parametrize("gfx11", [False, True]) +def test_use_swapped_grid_off_gfx1151_is_default_order( + monkeypatch: pytest.MonkeyPatch, head_size: int, gfx11: bool +) -> None: + """Only gfx1151 was measured, so nothing else may reorder the grid.""" + _patch_arch(monkeypatch, gfx11=gfx11) + assert not triton_ua._use_swapped_grid(head_size) + + +@pytest.mark.parametrize( + "head_size,swapped", [(64, False), (80, True), (128, True), (256, True)] +) +def test_use_swapped_grid_gfx1151_excludes_head_size_64( + monkeypatch: pytest.MonkeyPatch, head_size: int, swapped: bool +) -> None: + """head_size 64 regresses under the swapped order and must be excluded.""" + _patch_arch(monkeypatch, gfx11=True, gfx1151=True) + assert triton_ua._use_swapped_grid(head_size) is swapped diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index 3bd74650af6f..15727044c6d8 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -295,6 +295,9 @@ def kernel_unified_attention( # instead of letting them override it. Default False preserves the # original (causal AND SW) OR mm_prefix behavior for all other models. MM_PREFIX_CLAMP_SW: tl.constexpr = False, + # Read the 2D grid as (kv_heads, q_blocks) rather than the default + # (q_blocks, kv_heads). See ``_use_swapped_grid`` for the gating. + USE_SWAPPED_GRID: tl.constexpr = False, ): # Per-(token, head) scale caches: used iff KV_QUANT_MODE in {2, 3}. USE_PER_TOKEN_HEAD_SCALES: tl.constexpr = (KV_QUANT_MODE >= 2) and ( @@ -308,8 +311,12 @@ def kernel_unified_attention( "USE_TD requires BLOCK_SIZE to be a multiple of TILE_SIZE", ) - q_block_global_idx = tl.program_id(0) - kv_head_idx = tl.program_id(1) + if USE_SWAPPED_GRID: + kv_head_idx = tl.program_id(0) + q_block_global_idx = tl.program_id(1) + else: + q_block_global_idx = tl.program_id(0) + kv_head_idx = tl.program_id(1) segm_idx = tl.program_id(2) if IS_3D else 0 ( @@ -890,6 +897,16 @@ def _select_query_block( return block_m, block_m // num_queries_per_kv, False +def _use_swapped_grid(head_size: int) -> bool: + """Whether to launch the 2D grid as ``(kv_heads, q_blocks)``. + + gfx1151 schedules the 2D kernel better when the KV head is the fastest + varying axis, but head_size 64 regresses under that order, so it keeps the + default. Every other architecture keeps the default too. + """ + return _ON_GFX1151 and head_size >= 80 + + def _cap_num_stages_for_gfx11_lds( num_stages: int, block_m: int, @@ -1264,9 +1281,15 @@ def unified_attention( segm_expsum_ptr = softmax_segm_expsum if use_3d else None num_segments = num_par_softmax_segments if use_3d else 1 + use_swapped_grid = not use_3d and _use_swapped_grid(head_size) + grid: tuple[Any, ...] if not use_3d: - grid = (total_num_q_blocks, num_kv_heads) + grid = ( + (num_kv_heads, total_num_q_blocks) + if use_swapped_grid + else (total_num_q_blocks, num_kv_heads) + ) tile_size = TILE_SIZE_PREFILL else: grid = (total_num_q_blocks, num_kv_heads, num_par_softmax_segments) @@ -1366,6 +1389,7 @@ def unified_attention( USE_TD=use_td, USE_TD_QO=use_td_qo, MM_PREFIX_CLAMP_SW=mm_prefix_clamp_sliding_window, + USE_SWAPPED_GRID=use_swapped_grid, **launch_kwargs, )