[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)