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