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.
Assisted-by: Claude Assisted-by: Codex
Assisted-by: Claude Assisted-by: Codex
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. |
|
Just leaving this here.
Device 0: AMD Radeon RX 9070 XT, gfx1201 (0x1201), VMM: no, Wave Size: 32, VRAM: 16304 MiB
Device 0: AMD Radeon RX 9070 XT, gfx1201 (0x1201), VMM: no, Wave Size: 32, VRAM: 16304 MiB
But I'm noticing a strange slowdown on pp2048 @ d16384. |
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