Repository navigation
Conversation
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
Assisted-by: Gemini
Assisted-by: Gemini
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).
d76c004 to
975671d
Compare
|
Thanks for looking into the 3.5 @justinappler! I've rebased onto Here's the latest benchmark results on a single Radeon AI Pro R9700: Qwen3.6-35B-A3B Q4_K_M (head dim 256, GQA):
|
|
Hi! My setup: Windows 11, 64 Gb RAM, Ryzen 5700G (all monitors on iGPU), RX 9070 XT fully available for inference. I independently arrived at the same change (enabling MMA for DKQ=256) and ran into a correctness problem that the PR does not cover. Look: Enabled MMA for 256, got the speedup and garbage output: Removed two guards (
But the output turned to garbage. On ~2k prompts garbage appeared in about half of the runs, on 50k in 100% of runs. Example response head: In } else {
constexpr bool oob_check = false;
for (; kb0 < kb0_stop-1; ++kb0) {
constexpr bool last_iter = false;
constexpr int k_VKQ_sup = nbatch_fa;
flash_attn_ext_f16_iter
<..., last_iter, oob_check, ...>(..., kb0, k_VKQ_sup);
}
constexpr bool last_iter = true;
constexpr int k_VKQ_sup = nbatch_fa; // <- full batch, not the remainder
flash_attn_ext_f16_iter
<..., last_iter, oob_check, ...>(..., kb0, k_VKQ_sup);
}The causal chain:
Fix: constexpr bool last_iter = true;
- constexpr int k_VKQ_sup = nbatch_fa;
+ // multi-stage cp.async loads cannot zero-fill, oob_check is not supported there
+ constexpr bool oob_check_last = nstages == 0;
+ const int k_VKQ_sup = ne11 - kb0*nbatch_fa;
flash_attn_ext_f16_iter
- <..., last_iter, oob_check, ...>
+ <..., last_iter, oob_check_last, ...>
I've tested it:
Caveat: I also checked reachability on clean master: took Llama-3.2-3B (head_dim 128, GQA 24/8, exactly the path MMA serves today), ran a cold 50k prompt, a short request after the long one, and a delta turn. Both builds (master and master+fix) give identical correct answers. So on current master at DKQ<=128 it does not reproduce, and the exact case this PR enables (DKQ=256) is the one I hit. On top of the fix I swept the RDNA (256,256,64) config: - GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 128, 2, 64, 128, 128, 64, 1, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 128, 2, 32, 64, 128, 64, 1, true);The fa curve: 128 -> crash (LDS exceeds 64KB), 96 -> -11.6%, 64 -> baseline, 32 -> +7.9% (pp32768, llama-bench), 32 is the floor (assert on multiple of 32). Mechanism: KQ_C splits into smaller tiles (4 -> 2), lower register pressure and finer stream-k granularity, the effect shows at deep context. So my final numbers (A/B interleaved 3+3 sessions)
Gen did not move at any step (batch=1 goes through the vec kernel). What this means for the PR
This is not AI written comment but I used it to translate/format some parts into English =) |
JohannesGaessler
left a comment
JohannesGaessler
left a comment
There was a problem hiding this comment.
This PR should definitely not be merged like this. It simply adds too much technical debt for what it does. I am also quite skeptical that an implementation loading directly from VRAM is actually the optimal solution. For RDNA in particular I realized when reviewing another PR that the 16 byte padding is conflicting with the way I implemented FP16 accumulation. In terms of scheduling, first I need to work out the details of #28536 , then I need to think about how to integrate AMD.
Assisted-by: Claude Opus 5
…fx1151 Assisted-by: Claude Opus 5
This branch has not been deployed
9 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