Repository navigation
[AMD][Perf] Tune extend attention block sizes for gfx950 (head_dim > 128) - #27793
Conversation
A larger query tile (BLOCK_M 64->128, num_warps 4->8, BLOCK_N unchanged) halves KV bytes streamed per kernel call on gfx950 for head_dim>128. MI350X head_dim 256: -36 to -38% kernel time (28->44% MFU), -29% _fwd_kernel in-server (median-of-3), GSM8K 0.970==0.970. Other AMD archs and head_dim<=128 keep the existing default via gfx95 gate.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces hardware-specific block size optimizations for AMD CDNA4 (gfx950) GPUs in the extend attention Triton kernel. Specifically, for head dimensions greater than 128 on gfx95, it increases the query tile size (BLOCK_M) to 128 and the number of warps to 8 to improve performance, while adding corresponding unit tests. Feedback suggests restricting this optimization to head dimensions between 128 and 256 (inclusive) to prevent register pressure or compilation failures on extremely large head dimensions (e.g., 576), and updating the unit tests to verify this fallback behavior.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if _is_hip: | ||
| BLOCK_M, BLOCK_N = (64, 64) | ||
| num_warps = 4 | ||
| if _is_gfx95 and Lq > 128: |
There was a problem hiding this comment.
For extremely large head dimensions (such as Lq = 576), using BLOCK_M = 128 on gfx95 can lead to excessive register pressure, potential register spilling, or even Triton compilation failures due to shared memory/register limits (similar to how CUDA limits BLOCK_M to 32 or 16 for Lq > 256).
To ensure safety and robustness across all supported head dimensions, consider restricting this optimization to 128 < Lq <= 256 so that larger head dimensions safely fall back to the default (64, 64, 4) configuration.
| if _is_gfx95 and Lq > 128: | |
| if _is_gfx95 and 128 < Lq <= 256: |
There was a problem hiding this comment.
Only head_dim 256 was validated, so the gate is now restricted to 128 < Lq <= 256; larger head dims fall back to the default. Added the 576 fallback assertion. Tests re-run on MI350X. Fixed in 31c5bff.
| # head_dim > 128: 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 | ||
| ) |
There was a problem hiding this comment.
To ensure that head dimensions larger than 256 (such as 576) correctly fall back to the default block sizes on gfx95, we should add an assertion for Lq = 576 in the block size selection test.
| # head_dim > 128: 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 > 128: tuned tile on gfx95 (up to 256), 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: should fall back to default (64, 64, 4) on all HIP archs | |
| self.assertEqual( | |
| ea._get_block_sizes_for_extend_attention(576, 576)[3:], (64, 64, 4) | |
| ) |
|
Hi @HaiShaw @merrymercy @Fridge003 @ispobock @Qiaolin-Yu @hebiao064 — when you have a moment, could someone add the run-ci / run-ci-extra labels and take a look? Small AMD-only change (one block-size branch + tests), happy to address any feedback. Thanks! |
|
/tag-and-rerun-ci |
|
The failing stage-a-test-1-gpu-small-amd lane is test_basic_sanity's temp=0 determinism check — a known AMD flake, unrelated to this change. The new block config only applies for head_dim > 128; this test runs Llama-3.1-8B (head_dim 128), so it takes the unchanged path. (And on the >128 path, only BLOCK_M/num_warps change — BLOCK_N/reduction order is unchanged, so it's numerically identical.) |
|
@amd-bot ci-status |
CI Status for PR #27793Merge verdict: Do not treat this as verified-green. All 8 executed/gate failures are unrelated to this PR (a determinism flake on AMD + two NPU flakes, plus their fast-fail cascade), so nothing is broken by the PR — but the PR's actual change is not exercised by any PR CI, so green proves nothing about it. The tuned block-size path is gated on Caution This PR's changed code path is not verified by PR CI. The new config Caution PR CI is incomplete (fast-fail cascade). AMD Changed files: Executed CI failure attribution: AMD: 1 executed failure (0 related) · Others (NPU): 3 executed failures (0 related). Aggregation/gate jobs ( AMD Executed Failures
Other (NPU) Executed Failures
Details / what to do before merge
Generated by amd-bot using Claude Code CLI |
3 participants
Motivation
While profiling long-context prefill on MI350X (Qwen3.5-style GQA models, head_dim 256), the Triton extend kernel turned out to be the dominant cost: about 62% of GPU time at roughly 28% MFU, far below what the memory system allows.
The reason is that
_get_block_sizes_for_extend_attentionuses one hardcoded config for every HIP arch (BLOCK_M=64, BLOCK_N=64, num_warps=4), while the CUDA side tunes per arch and head_dim. Each workgroup streams the whole prefix KV, so KV traffic scales with1/BLOCK_M. At BLOCK_M=64 with a 32k chunk we move ~70 GB of KV per call. Doubling BLOCK_M halves that, and gfx950 has enough CUs and registers to hold the larger tile.Modifications
One branch in
extend_attention.py: on gfx95 with head_dim > 128 use(BLOCK_M=128, BLOCK_N=64, num_warps=8). Everything else keeps the current default. CUDA untouched, no env vars, no buffers, no graph-shape changes.Also extends the existing tests: head_dim 256 in
test_extend_attention, plus a small HIP block-size selection test asserting the gate is on for gfx95 and off elsewhere.We tried BLOCK_N=128 too — it regresses (+19%), so it stays at 64. BLOCK_M=256 spills registers and is 2x slower; 128 is the sweet spot.
Note on #26283: it tunes the few-query verify shape (
max_len_extend <= 16) in this function; conditions don't overlap, and I'm happy to rebase if it lands first.Accuracy Tests
GSM8K 5-shot, 200 ex on 35B with NEXTN: 0.970 with and without, identical std and accept length (3.36). Kernel matches stock under allclose at head_dim 256. num_warps changes warp reduction order, so equivalent, not bit-identical.
Speed Tests (MI350X, BF16, M=32768, prefix 67k, median of 20)
In-server, 100K prompts, both arms concurrent, median of 3 traces:
_fwd_kernel14.4 → 10.3 s, full prefill 12.45 → 10.06 s. Gain shows in prefill compute; queue-bound serving medians don't move.Checklist
test/registered/attention/test_triton_attention_kernels.py: head_dim 256 parity + HIP block-size selection).extend_attention.py).CI States
Latest PR Test (Base): ✅ Run #27266246062
Latest PR Test (Extra): ❌ Run #27266245139