[Bugfix][Triton] Avoid RDNA4 unified attention LDS overflow

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)

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,