[AMD][Perf] Tune extend attention block sizes for gfx950 (head_dim > 128) (#27793)

python/sglang/srt/layers/attention/triton_ops/extend_attention.py

@@ -23,13 +23,14 @@ import triton.language as tl
 from sglang.srt.layers.attention.triton_ops.prefill_attention import (
     context_attention_fwd,
 )
-from sglang.srt.utils import is_cuda, is_hip
+from sglang.srt.utils import is_cuda, is_gfx95_supported, is_hip
 
 _is_cuda = is_cuda()
 if _is_cuda:
     CUDA_CAPABILITY = torch.cuda.get_device_capability()
 
 _is_hip = is_hip()
+_is_gfx95 = _is_hip and is_gfx95_supported()
 
 
 def _get_block_sizes_for_extend_attention(Lq: int, Lv: int):
@@ -61,8 +62,17 @@ def _get_block_sizes_for_extend_attention(Lq: int, Lv: int):
 
     # Determine BLOCK_M, BLOCK_N, and num_warps based on hardware
     if _is_hip:
-        BLOCK_M, BLOCK_N = (64, 64)
-        num_warps = 4
+        if _is_gfx95 and 128 < Lq <= 256:
+            # gfx950 (CDNA4), 128 < head_dim <= 256: a larger query tile halves KV bytes
+            # streamed per call (each workgroup reads the whole prefix); 8 warps
+            # hide the loads. Measured on MI350X head_dim 256: -36% kernel time,
+            # 28% -> 44% MFU, numerically equivalent (BLOCK_N reduction order
+            # unchanged). Other AMD archs / head dims keep the default below.
+            BLOCK_M, BLOCK_N = (128, 64)
+            num_warps = 8
+        else:
+            BLOCK_M, BLOCK_N = (64, 64)
+            num_warps = 4
     else:
         if _is_cuda and CUDA_CAPABILITY[0] == 12:
             # sm120 workstation Blackwell architecture (RTX Pro 6000) has a much smaller shared memory size (100K)

test/registered/attention/test_triton_attention_kernels.py

@@ -312,12 +312,32 @@ class TestTritonAttention(CustomTestCase):
     def test_extend_attention(self):
 
         # Define the varying parameter values
-        attention_values = [128, 96, 80, 13]
+        # 256 covers the head_dim > 128 block-size branch (tuned on gfx95)
+        attention_values = [256, 128, 96, 80, 13]
 
         # Loop through the values and call the method
         for value in attention_values:
             self._test_extend_attention_once(19, 12331, 12, 4, value)
 
+    def test_extend_attention_block_sizes(self):
+        from sglang.srt.layers.attention.triton_ops import extend_attention as ea
+
+        if not ea._is_hip:
+            self.skipTest("HIP-only block-size selection")
+        # head_dim <= 128 keeps the default config on all HIP archs
+        self.assertEqual(
+            ea._get_block_sizes_for_extend_attention(128, 128)[3:], (64, 64, 4)
+        )
+        # 128 < head_dim <= 256: tuned tile on gfx95, default elsewhere
+        expected = (128, 64, 8) if ea._is_gfx95 else (64, 64, 4)
+        self.assertEqual(
+            ea._get_block_sizes_for_extend_attention(256, 256)[3:], expected
+        )
+        # head_dim > 256: falls back to the default on all HIP archs
+        self.assertEqual(
+            ea._get_block_sizes_for_extend_attention(576, 576)[3:], (64, 64, 4)
+        )
+
     def _test_extend_attention_sliding_window_once(
         self, B, N_CTX, H_Q, H_KV, D, WINDOW_SIZE
     ):