ggml-cuda: enable MMA FlashAttention for head dim 256 on AMD RDNA - #26419
ggml-cuda: enable MMA FlashAttention for head dim 256 on AMD RDNA#26419srgtuszy wants to merge 3 commits into
Conversation
Assisted-by: Gemini
Assisted-by: Gemini
github-actions
Bot
added
ggml
|
Tested on 2 x R9700 (gfx1201), ROCm 7.2.1, kernel 6.17, Qwen3.5-122B-A10B IQ3_XXS (head dim 256, GQA 32/2), Speedup reproduces. pp512 tok/s,
pp4096 tok/s,
Against b10154 that is +2.1/+17.1/+17.2% at pp512 and +1.8/+16.8/+17.2% at pp4096. tg128 is flat (48.75/42.79/29.71 vs 48.69/42.58/29.73), expected since batch-1 decode does not clear the new
Gap to rocWMMA is wider than 12-20% hereAgainst b10182: +0.7/-10.5/-29.2% at pp512, -0.7/-9.5/-31.3% at pp4096. As a fraction of the rocWMMA deficit recovered, 55% at 16k and 26% at 65k for pp512, 58% and 24% for pp4096. I assumed pp512 was penalising the MMA kernel, but pp4096 gives the same answer, so it is not batch size. What is left against your setup is q8_0 KV vs f16, and IQ3_XXS 122B vs Q4_K_M 35B. If the quantised KV path is the cause that is worth knowing, since kv-quant at long context is common on 32GB cards. Not an argument against merging. +17% at depth with no tg regression is a clear win over the status quo. One correctness failure at head dim 192
Head dim 192, not 256. Widening the host gate to The reproduction evidence is thin and I would rather say so than dress it up: 1 failure in 3 full two-device runs on a build carrying a local allocator patch of mine, and 0 failures in 4 runs on a clean build of this PR. At a 1-in-3 rate, 4 clean runs happen about 20% of the time, so that A/B settles nothing in either direction. What made me look is source-level and does not depend on my build. On the
Nothing sits between 4 and 7. Both barriers that used to cover that window (698 and 1063) are now behind Happy to test patches for either. Both builds are deployed side by side here. |
When DKQ > 128 on AMD WMMA, K/V loads bypass LDS so the end-of-loop __syncthreads() barriers were skipped entirely. However, tile_mask still transits LDS and is read during softmax. Without a barrier at the end of the VKQ section, the next iteration's load_mask can overwrite tile_mask while a slow warp is still reading it from the current iteration. Fix: keep the __syncthreads() at the end of the V loop when the mask is active (ncols2 > 1 || mask_h), even when DKQ > 128. Reported-by: williamtwomey (PR ggml-org#26419 review) Tested: 5 consecutive runs of test-backend-ops -o FLASH_ATTN_EXT on gfx1201, 2920/2920 passed each time (14600/14600 total, 0 failures).
|
Thanks for testing @williamtwomey! You're right, when DKQ > 128, the LDS bypass skips the end-of-loop |
When DKQ > 128 on AMD WMMA, K/V loads bypass LDS so the end-of-loop __syncthreads() barriers were skipped entirely. However, tile_mask still transits LDS and is read during softmax. Without a barrier at the end of the VKQ section, the next iteration's load_mask can overwrite tile_mask while a slow warp is still reading it from the current iteration. Fix: keep the __syncthreads() at the end of the V loop when the mask is active (ncols2 > 1 || mask_h), even when DKQ > 128. Reported-by: williamtwomey (PR ggml-org#26419 review) Tested: 5 consecutive runs of test-backend-ops -o FLASH_ATTN_EXT on gfx1201, 2920/2920 passed each time (14600/14600 total, 0 failures).
|
Another data point, maybe it is useful: Validation on gfx1201 (Radeon AI PRO R9700): +62% prefill at 100k depth. Setup
Results
Decomposing per-token prefill cost as
It reads as a pure depth-term change: Extrapolated to this model's full 262,144-token context: ~23 min → ~11 min which is a big win for me. For now, I measured throughput, not correctness. I have not run perplexity or diffed outputs against the unpatched build, so please don't read this as validation of the kernel's numerics. Quick question about DKQ > 256: The guard here lands at |
This comment was marked as outdated.
This comment was marked as outdated.
|
with 2x R9700, powercapped at 230W each, with the docker config and model config below, actual long context tests with 500 tokens minimum, i get the following figures. I did not re-measure PP between different MTP configs since there is no effect of n-max during PP. RCCL + Graphs enabled. The Master baseline is eab8ee4. For the master baseline WMMA_FATTN was enabled, and removed for the PR test Compose and Model Config
Details
Prompt processing (PP) -- tokens/s, independent of parallel levelmtp=none
mtp=1
mtp=2
mtp=3
mtp=4
Decode (TG) -- combined tokens/s across concurrently-decoding slotsmtp=none parallel=1
mtp=none parallel=2
mtp=none parallel=3
mtp=1 parallel=1
mtp=1 parallel=2
mtp=1 parallel=3
mtp=2 parallel=1
mtp=2 parallel=2
mtp=2 parallel=3
mtp=3 parallel=1
mtp=3 parallel=2
mtp=3 parallel=3
mtp=4 parallel=1
mtp=4 parallel=2
mtp=4 parallel=3
Images
TL;DR, not sure if these results are conclusive, but for me it seems like the the current path with WMMA_FATTN is better in TG, but the PR has ~ +3.5% in PP. |
|
I've been keeping an eye out for PRs/commits that might impact RDNA 3.5 and specifically gfx1151 (Strix Halo). Benched this PR vs master and it looks like this regresses pp @ depth. Might be worth limiting the scope of this to RDNA4 on the DKQ <= 256 arm. tg/s was essentially flat, no impact. Data below:
llama-bench -m Qwen3.6-35B-A3B-UD-Q4_K_XL.gguf \
-ctk f16 -ctv f16 -fa 1 \
-b 4096 -ub 2048 -ngl 999 -lm none \
-p 512 -n 128 -r 3 \
-d 0,2048,8192,16384,32768,65536pp512 (t/s)
Each arm run 3+ times, alternating master / PR between runs; figures are means. tg128 was flat at every depth, within 0.2% (39.33 t/s on both arms at d=65,536). |
…fx1151 Assisted-by: Claude Opus 5
7 participants
Overview
Fixes #26220
After the removal of rocWMMA Flash Attention kernel path, the replacement had slower prompt processing at deep context on RDNA 4. The regression was caused by WMMA tensor cores not being used at all for head dim 256.
This PR widens
flash_attn_ext_f16device guard to unblock MMA at head dim 256. Without it, the tile kernel is used instead which works, but slows down with more context. The config table also needed a precision fix for DKQ=256.nbatch_combinewas too small, causing multi-pass FP16 accumulation errors in test-backend-ops.Additionally, for DKQ>128 on WMMA K/V is loaded directly from VRAM instead of going through LDS. Without it, MMA is actually slower than the tile kernel.
Additional information
Below are benchmark numbers to confirm no regressions and the speedup. It's worth to mention that, while prompt processing is faster than current master(a7a6d0d), it's still about 12-20% slower the old rocWMMA kernel depending on context size. There's definitely room for further improvement.
GPU: AMD Radeon AI PRO R9700 (gfx1201, RDNA4), ROCm 7.2.4
Baseline: master a7a6d0d, 3 runs per data point
test-backend-ops -o FLASH_ATTN_EXT: 2920/2920 passed
Qwen3.6-35B-A3B Q4_K_M (head dim 256)
Why LDS bypass is needed (Qwen3.6-35B-A3B, MMA enabled but without LDS bypass)
Llama-3.1-8B-Instruct Q4_K_M (head dim 128, regression check)
Requirements