Repository navigation
CUDA/HIP: Flash Attention tuning (gfx1201) - #28102
Conversation
|
This looks really good! Such a simple change providing extra performance, good job on finding it! |
|
I'm not done yet!
|
6b24ee9 to
d68f876
Compare
There was a problem hiding this comment.
nice, looks good
Details
| GPU | Model | Microbatch size | Test | t/s master | t/s fattn-wmma-rdna4-256 | Speedup |
|---|---|---|---|---|---|---|
| AI PRO R9700 | gemma4 26B.A4B Q6_K | 8 | pp2048@d32768 | 327.19 | 308.26 | 0.94 |
| AI PRO R9700 | gemma4 26B.A4B Q6_K | 64 | pp2048@d32768 | 662.22 | 668.99 | 1.01 |
| AI PRO R9700 | gemma4 26B.A4B Q6_K | 512 | pp2048@d32768 | 1263.12 | 1256.57 | 0.99 |
| AI PRO R9700 | gemma4 26B.A4B Q6_K | 1024 | pp2048@d32768 | 1366.90 | 1495.99 | 1.09 |
| AI PRO R9700 | gpt-oss 20B MXFP4 MoE | 8 | pp2048@d32768 | 523.46 | 525.46 | 1.00 |
| AI PRO R9700 | gpt-oss 20B MXFP4 MoE | 64 | pp2048@d32768 | 1408.82 | 1364.30 | 0.97 |
| AI PRO R9700 | gpt-oss 20B MXFP4 MoE | 512 | pp2048@d32768 | 2588.42 | 2838.27 | 1.10 |
| AI PRO R9700 | gpt-oss 20B MXFP4 MoE | 1024 | pp2048@d32768 | 2869.95 | 3281.25 | 1.14 |
| AI PRO R9700 | lfm2moe 8B.A1B Q8_0 | 8 | pp2048@d32768 | 817.76 | 817.97 | 1.00 |
| AI PRO R9700 | lfm2moe 8B.A1B Q8_0 | 64 | pp2048@d32768 | 2947.91 | 2947.64 | 1.00 |
| AI PRO R9700 | lfm2moe 8B.A1B Q8_0 | 512 | pp2048@d32768 | 7976.07 | 7986.46 | 1.00 |
| AI PRO R9700 | lfm2moe 8B.A1B Q8_0 | 1024 | pp2048@d32768 | 9433.99 | 8928.84 | 0.95 |
| AI PRO R9700 | llama 8B Q8_0 | 8 | pp2048@d32768 | 312.91 | 312.85 | 1.00 |
| AI PRO R9700 | llama 8B Q8_0 | 64 | pp2048@d32768 | 1209.11 | 1181.70 | 0.98 |
| AI PRO R9700 | llama 8B Q8_0 | 512 | pp2048@d32768 | 1145.38 | 1705.73 | 1.49 |
| AI PRO R9700 | llama 8B Q8_0 | 1024 | pp2048@d32768 | 1064.61 | 1587.85 | 1.49 |
| AI PRO R9700 | qwen35 27B Q5_K_M | 8 | pp2048@d32768 | 75.64 | 72.25 | 0.96 |
| AI PRO R9700 | qwen35 27B Q5_K_M | 64 | pp2048@d32768 | 372.91 | 479.10 | 1.28 |
| AI PRO R9700 | qwen35 27B Q5_K_M | 512 | pp2048@d32768 | 468.45 | 707.39 | 1.51 |
| AI PRO R9700 | qwen35 27B Q5_K_M | 1024 | pp2048@d32768 | 475.40 | 708.34 | 1.49 |
<\details>
| bool use_stream_k = cc >= GGML_CUDA_CC_ADA_LOVELACE || amd_wmma_available(cc) || tiles_efficiency_percent < 75; | ||
| if (amd_wmma_available(cc) && ntiles_dst >= 2*max_blocks && tiles_efficiency_percent >= 75) { | ||
| use_stream_k = false; | ||
| } |
There was a problem hiding this comment.
we might want to refactor this into a helper function with table at this point.
|
can confirm the performance improvements on my R9700. Hope this PR will be merged soon. |
|
@IMbackK added the helper. |
JohannesGaessler
left a comment
There was a problem hiding this comment.
Please avoid piling on many unrelated changes into a single PR like this.
|
Yeah, I just noticed the GDA changes got swept in as well, I'll move them to a separate branch. |
|
@JohannesGaessler aight, cleaned it up to just the FATTN changes. |
It is a lot easier for me to check the performance with separate PRs. I'm available to work on this this weekend so I think we can get it done if you are as well. |
|
Aight, will try that. |
|
Actually, looking at the order the commits were done in it's not clear to me that the config change and the change to stream-k can be separated. But Claude is definitely wrong about the supposed race condition it fixed. |
|
imo the state of this pr was fine at 2d55b5f |
|
Yeah, shouldn't have trusted Claude on that one. Reverting that. |
|
mind doing a |
| // swizzle the tile stride for K and V based on the batch size. | ||
| constexpr int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2); | ||
| #if defined(AMD_WMMA_AVAILABLE) | ||
| constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : nbatch_V2 + 6; | ||
| #else | ||
| constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2); | ||
| #endif // defined(AMD_WMMA_AVAILABLE) |
There was a problem hiding this comment.
Again, please submit changes to the actual device code beyond changes to the config and the host-side orchestration separately. I would be surprised if this is actually the correct padding to minimize LDS bank conflicts.
f6264d8 to
a3b28f9
Compare
|
Rebased, squashed, hopefully it's good now. |
| if (GGML_CUDA_CC_IS_RDNA4(cc)) { | ||
| if (use_gqa_opt && gqa_ratio % 8 == 0) { | ||
| ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 8>(ctx, dst); | ||
| return; | ||
| } | ||
|
|
||
| if (use_gqa_opt && gqa_ratio % 4 == 0) { | ||
| ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 4>(ctx, dst); | ||
| return; | ||
| } | ||
|
|
||
| if (use_gqa_opt && gqa_ratio % 2 == 0) { | ||
| ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 2>(ctx, dst); | ||
| return; | ||
| } | ||
| } | ||
|
|
There was a problem hiding this comment.
I don't think these extra code branches are needed.
There was a problem hiding this comment.
Looking at the exact parameters for Qwen 3.5, what these seem to be doing is launch different kernels for GQA ratios that are not a power of 2, Qwen 3.5 has 6.
|
I observed some performance regressions on my RX 9060 XT with this PR. I pushed revised logic for whether or not to use stream-k, can you check how that performs on your system? This is the data I got: PW vs. master
JG vs. master
JG vs. PW
|
|
Sorry, the columns were swapped in the "PW vs. master" table, this should be the correct data. |
|
@JohannesGaessler sorry, took a while because benchmarking on my trusty old machine is flaky and I had to rerun a few suspicious results, but here's the full report: TL;DR: checks out fine. |
|
Regarding the host-side logic for how to handle a GQA ratio of 6: on my NVIDIA GPUs I found it to be faster to run the kernels with effectively a GQA ratio of 8 even though that results in some wasted work due to padding, on my RX 9060 XT using a GQA ratio of 2 (as is changed in this PR) seems to be ~5% faster. In principle there would be a breakeven point at a batch size of ~160 since a higher value in I'll check the performance on my RDNA 3.5 machine as well and after that this PR would be good to merge from my end. |
Assisted-by: Claude Assisted-by: Codex
Assisted-by: Claude Assisted-by: Codex
1873e6e to
461f7c1
Compare
RDNA 3.5 Performance
I checked RDNA 3.5 performance as well and adjusted the kernel selection logic accordingly. |
|
@IMbackK can you give a reapproval? |
* HIP: enable mma FA for head size 256 on RDNA4, tune configs Assisted-by: Claude Assisted-by: Codex * HIP: prefer whole-tile FA grids over stream-k on AMD WMMA Assisted-by: Claude Assisted-by: Codex * revise stream_k logic * revise kernel selection logic --------- Co-authored-by: Johannes Gäßler <johannesg@5d6.de> (cherry picked from commit 16378d9)
This reverts commit 39fa672 (upstream 16378d9) on gfx1151. With it in the build (b11016), Qwen3.8-Flash-Next answers prompts that interleave several images with text as pure garbage ("////...") and the MTP draft is rejected every round (0 of 16379). Text-only prompts up to 42k tokens and a single-image prompt are fine on the same build, and the identical prompt class was fine on b11008, which lacked this commit. The commit moves 256-wide prefill attention from the tile kernel onto the AMD WMMA kernel; the four other restored upstream commits stay in. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Undoes the b11018 revert of upstream 16378d9. The garbage output and 0 draft acceptance that got it reverted on 2026-09-11 turned out to be the RPC worker's full disk serving truncated cache entries (fixed in b11034); the FA tuning was innocent. With it back, the flash-attention sources differ from upstream only by the tile-kernel row compaction. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
* HIP: enable mma FA for head size 256 on RDNA4, tune configs Assisted-by: Claude Assisted-by: Codex * HIP: prefer whole-tile FA grids over stream-k on AMD WMMA Assisted-by: Claude Assisted-by: Codex * revise stream_k logic * revise kernel selection logic --------- Co-authored-by: Johannes Gäßler <johannesg@5d6.de>
* HIP: enable mma FA for head size 256 on RDNA4, tune configs Assisted-by: Claude Assisted-by: Codex * HIP: prefer whole-tile FA grids over stream-k on AMD WMMA Assisted-by: Claude Assisted-by: Codex * revise stream_k logic * revise kernel selection logic --------- Co-authored-by: Johannes Gäßler <johannesg@5d6.de>
The whole-tile FA path (introduced in ggml-org#28102) triggers an asynchronous ROCm illegal memory access on gfx1100 (RX 7900 XTX) with multi-GPU tensor-split and large ubatch sizes. Force stream-K on RDNA3 to restore the previously stable behaviour while keeping the RDNA4 optimisation.
* HIP: enable mma FA for head size 256 on RDNA4, tune configs Assisted-by: Claude Assisted-by: Codex * HIP: prefer whole-tile FA grids over stream-k on AMD WMMA Assisted-by: Claude Assisted-by: Codex * revise stream_k logic * revise kernel selection logic --------- Co-authored-by: Johannes Gäßler <johannesg@5d6.de>
* HIP: enable mma FA for head size 256 on RDNA4, tune configs Assisted-by: Claude Assisted-by: Codex * HIP: prefer whole-tile FA grids over stream-k on AMD WMMA Assisted-by: Claude Assisted-by: Codex * revise stream_k logic * revise kernel selection logic --------- Co-authored-by: Johannes Gäßler <johannesg@5d6.de>
ggml-org#28102 replay dropped fork's vec_dot_fattn_vec_KQ_{tq3_0,turbo3_0,turbo4_0}, dequantize_V_{tq3_0,turbo3_0,turbo4_0}, turbo4_decode_element + getter arms (static_assert 'bad type' on GB10 builds). Re-landed from main onto the post-ggml-org#28102/ggml-org#29152 upstream shape (185 lines, no other deltas).
…NA3.5 Upstream ggml-org#28552/ggml-org#28935 land the routed-MoE ncols picker on RDNA3.5, which Finding ggml-org#10 measured flat against our static J=48 cap, so the cap goes. Upstream ggml-org#28102 routes WMMA to MMA_F16 at D=256 and bypasses our tile config at prefill; guard it off on RDNA3.5 pending the G/N A/B. Assisted-by: Claude Opus 5.5
Hypothesis: every prefill FA launch for head_dim 256 was the scalar flash_attn_tile<256,256>, because fattn.cu gated the AMD WMMA (MMA_F16) path on D <= 128. FA was 59% of prefill at 32K depth and ~85% at 128K. Change (RDNA4 under HIP only; every other arch keeps the old expressions): - Port of upstream ggml-org/llama.cpp #28102 (gfx1201 FA tuning): WMMA for D=256 once Q rows * gqa_ratio_eff > 16; ncols2 chosen to divide the GQA ratio at D=256 (GQA 6 -> 2); RDNA D=256 config rows; stream-k only when tile efficiency < 75% (D=64 always). - Port of upstream #27870: divergent __syncthreads in the np>1 combine. - Ours: D=256 ncols=64 config stages Q in LDS with 32-half2 K/V chunks (Q_in_reg=false): VGPR spills 424 -> 48, FA kernel 5.44 -> 4.33 ms at 16K KV, 32.5 -> 27.7 ms at 64K. - test-backend-ops: model-shape cases (D=256, 4 KV heads, GQA 6; KV 512/4096/8192; nb 1..512 and 1024/1030 for -ub 1024; q8_0 KV). Benchmark (RX 9070 XT, Bonsai 2 27B PQ2_0, llama-bench -r 2): Before (deddbab): pp2048 @ 32K depth 518 t/s; pp512 1168; tg128 49.8 After: pp2048 @ 32K depth 906 t/s (+75%); pp512, tg128 same Earlier run (upstream config): @64k q8_0 340 -> 704, @128k 199 -> 498. Agent edit workload (ngram-mod + MTP draft 3, ~33-97-token verify batches, which now take the WMMA kernel): 307.2 -> 316.8 t/s, same acceptance. VRAM +32 MiB. Correctness: test-backend-ops FLASH_ATTN_EXT 2952/2952 (15/15 at the model shape). PPL (wikitext-2, 20 chunks) 7.8773 -> 7.8695 (-0.10%), identical at -ub 512 and -ub 1024; KLD vs previous 1.45e-4, same as the model's own ub1-vs-ub512 noise (1.57e-4): fp reduction order in attention. Decode (1 token) and MTP verify (<= 4 tokens) stay on the old kernels. Decision: KEEP Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* HIP: enable mma FA for head size 256 on RDNA4, tune configs Assisted-by: Claude Assisted-by: Codex * HIP: prefer whole-tile FA grids over stream-k on AMD WMMA Assisted-by: Claude Assisted-by: Codex * revise stream_k logic * revise kernel selection logic --------- Co-authored-by: Johannes Gäßler <johannesg@5d6.de>

Overview
So, got my new R9700 PRO. It's nice and shiny and has 32GB VRAM, so I decided to try out Qwen3.8 27B. That was a mistake. The prefill performance at longer contexts was abysmal, so I decided to do something about it. Also managed to fix a HS=256 bug in the general CUDA FA code which was preventing the selection of 256 kernels before.
Additional information
Before:
After:
Requirements