Skip to content

[Perf] [AMD] attention: Restore the gfx950 extend tile for head_dim <= 128 #45099

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
HaiShaw merged 2 commits into from
Oct 21, 2026
Merged

[Perf] [AMD] attention: Restore the gfx950 extend tile for head_dim <= 128 #45099

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0f4a31a
[Perf] [AMD] attention: Restore the gfx950 extend tile for head_dim <…
aarnetalman Oct 19, 2026
d8e24af
Merge branch 'main' into rocm-extend-attention-gfx950-tile
HaiShaw Oct 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 5 additions & 6 deletions python/sglang/kernels/ops/attention/extend_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,12 +81,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:
Expand Down
Loading