From 666b08b4a5e908117e2fad8c5af20b8ac202349e Mon Sep 17 00:00:00 2001
From: Alex Nails <alex.nails@radixark.ai>
Date: Mon, 24 Aug 2026 00:29:39 -0700
Subject: [PATCH] [ROCm] Extend the gfx950 extend-attention tile to head_dim <=
 128: -43% kernel, -14% TTFT, bit-identical (#34461)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 .../kernels/ops/attention/extend_attention.py | 11 +++++-----
 .../test_triton_attention_kernels.py          | 22 ++++++++++---------
 2 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/python/sglang/kernels/ops/attention/extend_attention.py b/python/sglang/kernels/ops/attention/extend_attention.py
index ac558426c..7ffc2bde3 100644
--- a/python/sglang/kernels/ops/attention/extend_attention.py
+++ b/python/sglang/kernels/ops/attention/extend_attention.py
@@ -80,12 +80,11 @@ def _get_block_sizes_for_extend_attention(Lq: int, Lv: int):
             # 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,
-            # 28% -> 44% MFU, numerically equivalent (BLOCK_N reduction order
-            # unchanged). Other AMD archs / head dims keep the default below.
+        elif _is_gfx95 and Lq <= 256:
+            # gfx950 (CDNA4), head_dim <= 256: every workgroup streams the whole
+            # prefix, so a larger query tile halves the KV bytes read per call;
+            # BLOCK_M / num_warps = 16 rows per warp is exactly one MFMA tile at
+            # matrix_instr_nonkdim=16. Measured on MI350X at head_dim 64, 128, 256.
             BLOCK_M, BLOCK_N = (128, 64)
             num_warps = 8
         else:
diff --git a/test/registered/attention/test_triton_attention_kernels.py b/test/registered/attention/test_triton_attention_kernels.py
index 411a0154e..ff90b20a7 100644
--- a/test/registered/attention/test_triton_attention_kernels.py
+++ b/test/registered/attention/test_triton_attention_kernels.py
@@ -326,22 +326,24 @@ class TestTritonAttention(CustomTestCase):
 
         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
+        # head_dim <= 256: tuned tile on gfx95, default elsewhere. 64 is gpt-oss,
+        # 128 is the llama/qwen family, 256 is gemma -- all measured on MI350X.
         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
-        )
+        for head_dim in (64, 128, 256):
+            self.assertEqual(
+                ea._get_block_sizes_for_extend_attention(head_dim, head_dim)[3:],
+                expected,
+            )
         # head_dim > 256 falls back to the default unless the automatic
         # Triton-3.7 gfx950 Lq=576/Lv=512 spill workaround applies.
+        self.assertEqual(
+            ea._get_block_sizes_for_extend_attention(576, 576)[3:], (64, 64, 4)
+        )
         with unittest.mock.patch.object(ea, "_is_triton_ge_37", True):
-            expected = (64, 32, 4) if ea._is_gfx95 else (64, 64, 4)
+            expected_spill = (64, 32, 4) if ea._is_gfx95 else (64, 64, 4)
             self.assertEqual(
                 ea._get_block_sizes_for_extend_attention(576, 512)[3:],
-                expected,
+                expected_spill,
             )
         with unittest.mock.patch.object(ea, "_is_triton_ge_37", False):
             self.assertEqual(