diff --git a/python/sglang/srt/layers/attention/triton_ops/extend_attention.py b/python/sglang/srt/layers/attention/triton_ops/extend_attention.py index 870ca88cd..9945fec73 100644 --- a/python/sglang/srt/layers/attention/triton_ops/extend_attention.py +++ b/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) diff --git a/test/registered/attention/test_triton_attention_kernels.py b/test/registered/attention/test_triton_attention_kernels.py index dfc86cbf3..320bf5d93 100644 --- a/test/registered/attention/test_triton_attention_kernels.py +++ b/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 ):