[AMD] Fix Triton 3.7 gfx950 extend-attention spills (#34741)

python/sglang/kernels/ops/attention/extend_attention.py

@@ -35,6 +35,14 @@ if _is_cuda:
 _is_hip = is_hip()
 _is_gfx95 = _is_hip and is_gfx95_supported()
 
+try:
+    _triton_version_parts = tuple(
+        int(part) for part in triton.__version__.split(".")[:2]
+    )
+except (AttributeError, ValueError):
+    _triton_version_parts = (0, 0)
+_is_triton_ge_37 = _triton_version_parts >= (3, 7)
+
 
 def _get_block_sizes_for_extend_attention(Lq: int, Lv: int):
     """
@@ -65,7 +73,14 @@ 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:
-        if _is_gfx95 and 128 < Lq <= 256:
+        if _is_gfx95 and _is_triton_ge_37 and Lq == 576 and Lv == 512:
+            # Triton 3.7's N64 codegen reaches 512 VGPRs and spills 472 bytes
+            # of scratch on gfx950. N32 keeps BLOCK_M/launch work unchanged,
+            # uses <=433 VGPRs without scratch, and restores the isolated
+            # late-prefill kernel from ~12.57 ms to ~5.24 ms.
+            BLOCK_M, BLOCK_N = (64, 32)
+            num_warps = 4
+        elif _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,

test/registered/attention/test_triton_attention_kernels.py

@@ -335,10 +335,70 @@ class TestTritonAttention(CustomTestCase):
         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)
-        )
+        # head_dim > 256 falls back to the default unless the automatic
+        # Triton-3.7 gfx950 Lq=576/Lv=512 spill workaround applies.
+        with unittest.mock.patch.object(ea, "_is_triton_ge_37", True):
+            expected = (64, 32, 4) if ea._is_gfx95 else (64, 64, 4)
+            self.assertEqual(
+                ea._get_block_sizes_for_extend_attention(576, 512)[3:],
+                expected,
+            )
+        with unittest.mock.patch.object(ea, "_is_triton_ge_37", False):
+            self.assertEqual(
+                ea._get_block_sizes_for_extend_attention(576, 512)[3:],
+                (64, 64, 4),
+            )
+
+    def test_extend_attention_triton37_lq576_n32(self):
+        from sglang.kernels.ops.attention import extend_attention as ea
+
+        if not (ea._is_gfx95 and ea._is_triton_ge_37):
+            self.skipTest("Triton >=3.7 gfx950-only spill workaround")
+
+        device = get_device()
+        dtype = torch.bfloat16
+        extend_lens = [32, 23]
+        prefix_lens = [64, 96]
+        h_q, h_kv, l_q, l_v = 12, 1, 576, 512
+        n_ext, n_prefix = sum(extend_lens), sum(prefix_lens)
+
+        q = torch.randn(n_ext, h_q, l_q, dtype=dtype, device=device)
+        k = torch.randn(n_ext, h_kv, l_q, dtype=dtype, device=device)
+        v = torch.randn(n_ext, h_kv, l_v, dtype=dtype, device=device)
+        k_buffer = torch.randn(n_prefix, h_kv, l_q, dtype=dtype, device=device)
+        v_buffer = torch.randn(n_prefix, h_kv, l_v, dtype=dtype, device=device)
+        qo_indptr = torch.tensor([0, 32, 55], dtype=torch.int32, device=device)
+        kv_indptr = torch.tensor([0, 64, 160], dtype=torch.int32, device=device)
+        kv_indices = torch.arange(n_prefix, dtype=torch.int64, device=device)
+        reference = torch.empty(n_ext, h_q, l_v, dtype=dtype, device=device)
+        candidate = torch.empty_like(reference)
+
+        def run(output):
+            extend_attention_fwd(
+                q,
+                k,
+                v,
+                output,
+                k_buffer,
+                v_buffer,
+                qo_indptr,
+                kv_indptr,
+                kv_indices,
+                None,
+                True,
+                None,
+                max(extend_lens),
+                1.0,
+                1.0,
+                sm_scale=1.0 / (l_q**0.5),
+                extend_seq_lens_cpu=extend_lens,
+            )
+
+        with unittest.mock.patch.object(ea, "_is_triton_ge_37", False):
+            run(reference)
+        with unittest.mock.patch.object(ea, "_is_triton_ge_37", True):
+            run(candidate)
+        torch.testing.assert_close(candidate, reference, atol=2e-2, rtol=1e-2)
 
     def test_compact_extend_attention_tile_count(self):
         self.assertEqual(