From 81b2b29f7c8064899df6fe5b15a246261a7220a7 Mon Sep 17 00:00:00 2001 From: JS van Dijk <267467744+hogeheer499-commits@users.noreply.github.com> Date: Sat, 25 Jul 2026 21:57:15 +0200 Subject: [PATCH] [Bugfix][Triton] Avoid RDNA4 unified attention LDS overflow Signed-off-by: JS van Dijk <267467744+hogeheer499-commits@users.noreply.github.com> --- .../ops/triton/attention/unified_attention.py | 11 +++++ .../attention/test_unified_attention.py | 48 +++++++++++++++++-- 2 files changed, 54 insertions(+), 5 deletions(-) diff --git a/aiter/ops/triton/attention/unified_attention.py b/aiter/ops/triton/attention/unified_attention.py index cfc18ca745e..2aca82f047b 100644 --- a/aiter/ops/triton/attention/unified_attention.py +++ b/aiter/ops/triton/attention/unified_attention.py @@ -230,6 +230,17 @@ def select_3d_config( ), "TILE_SIZE needs to be divisible by block_size" NUM_BLOCKS_GATHER_PER_TILE = TILE_SIZE // block_size + # gfx120x has a 64 KiB per-workgroup LDS limit. With the two-stage + # pipeline, K and V tiles at head_size=256 and TILE_SIZE=64 consume the + # full 64 KiB before Triton's scratch allocation, so compilation fails + # with a 65,792-byte shared-memory request. Keep two stages where the + # tiles leave headroom and fall back to one stage only at the limit. + rdna_lds_bytes = ( + 2 * TILE_SIZE * triton.next_power_of_2(head_size) * kv_cache_dtype.itemsize + ) + if DEVICE_ARCH in ("gfx1200", "gfx1201") and rdna_lds_bytes >= 64 * 1024: + attn_stages = 1 + # gfx1151 (RDNA3.5) decode is memory-latency-bound at bs=1: the default 2 # warps/workgroup leave unified_attention at only ~31% of the LPDDR5X # bandwidth roofline. 8 warps/workgroup reach ~59% (1.5-1.9x on bf16 decode) diff --git a/op_tests/triton_tests/attention/test_unified_attention.py b/op_tests/triton_tests/attention/test_unified_attention.py index 703c6042b80..97cdea587d4 100644 --- a/op_tests/triton_tests/attention/test_unified_attention.py +++ b/op_tests/triton_tests/attention/test_unified_attention.py @@ -6,22 +6,60 @@ import pytest import torch +import aiter.ops.triton.attention.unified_attention as unified_attention_module from aiter.ops.triton.attention.unified_attention import ( - unified_attention, is_2d_gluon_available, + select_3d_config, + unified_attention, ) -from aiter.ops.triton.utils.shuffle import shuffle_weight, shuffle_scale_batched +from aiter.ops.triton.utils._triton import arch_info +from aiter.ops.triton.utils.shuffle import shuffle_scale_batched, shuffle_weight +from aiter.ops.triton.utils.types import e4m3_dtype +from aiter.test_common import checkAllclose from op_tests.triton_tests.quant.test_quant_mxfp4 import ( torch_dynamic_mxfp4_quant, ) -from aiter.ops.triton.utils.types import e4m3_dtype -import aiter.ops.triton.utils._triton.arch_info as arch_info -from aiter.test_common import checkAllclose DEVICE_ARCH = arch_info.get_arch() IS_DEVICE_ARCH_GFX12 = DEVICE_ARCH in ("gfx1250",) +@pytest.mark.parametrize( + "device_arch, head_size, block_size, q_dtype, kv_cache_dtype, expected_num_stages", + [ + ("gfx1200", 256, 64, torch.bfloat16, torch.bfloat16, 1), + ("gfx1201", 256, 64, torch.bfloat16, torch.bfloat16, 1), + ("gfx1201", 256, 64, e4m3_dtype, e4m3_dtype, 2), + ("gfx1201", 256, 64, torch.bfloat16, e4m3_dtype, 2), + ("gfx1201", 256, 32, torch.bfloat16, torch.bfloat16, 2), + ("gfx1201", 128, 64, torch.bfloat16, torch.bfloat16, 2), + ("gfx1151", 256, 64, torch.bfloat16, torch.bfloat16, 2), + ], +) +def test_select_3d_config_respects_rdna_lds_limit( + monkeypatch, + device_arch: str, + head_size: int, + block_size: int, + q_dtype: torch.dtype, + kv_cache_dtype: torch.dtype, + expected_num_stages: int, +) -> None: + monkeypatch.setattr(unified_attention_module, "DEVICE_ARCH", device_arch) + + attn_config, _ = select_3d_config( + head_size=head_size, + block_size=block_size, + max_seqlen_k=8192, + target_num_prgms=64, + num_2d_prgms=8, + q_dtype=q_dtype, + kv_cache_dtype=kv_cache_dtype, + ) + + assert attn_config["num_stages"] == expected_num_stages + + def shuffle_kv_cache( key_cache: torch.Tensor, value_cache: torch.Tensor,