Make KVarN usable on ROCm RDNA3 (toward #122: D256 race, portable prefill, f32 WMMA prototype) - #149
Open
raufaser wants to merge 12 commits into
Open
Make KVarN usable on ROCm RDNA3 (toward #122: D256 race, portable prefill, f32 WMMA prototype)#149raufaser wants to merge 12 commits into
raufaser wants to merge 12 commits into
Conversation
flash_attn_ext_f16_process_tile reuses tile_Q as combine staging across k00 iterations; the end-of-iteration barrier only fired for np > 1. D256/ncols=64 (nbatch_combine=64, DV/2=128, np=1) corrupted output. Sync on all but the last iteration, plus sync between process_tile calls that reuse tile_Q. Enable D256 in the RDNA WMMA route policy and device guard.
RDNA WMMA VKQ accumulators are fp16 (~3e-4/call vs ~1e-5 portable), compounding through depth into a KLD collapse. Prefer portable-native direct-record attention for HIP prompt-prefill (nq > 16); decode stays on WMMA. Opt out with GGML_KVARN_AMD_PROMPT_PORTABLE=0.
Replace the per-token 7-stage shared-memory reduction tree with warp shuffles plus one cross-warp step; rescale the half accumulator in fp32; unroll the V-load loop.
GGML_KVARN_TEST_NKV_LADDER_ONLY runs GPU-native vs CPU-materialized attention error over D128/256/512 and growing n_kv.
Serve QB=8 queries per block from one shared K/V token stream (no-tail path; tail keeps QB=1): ~2.7x prefill (pp4096 107 -> 291 t/s) at ladder-identical accuracy. Also replace the per-token shared-memory reduction tree with warp shuffles, resolve each token once per block instead of per thread, and add an env-gated kernel attribute print (GGML_KVARN_PORTABLE_ATTRS).
RDNA3 WMMA VKQ accumulators are fp16 (~3e-4/call), compounding over depth into a KLD collapse (2.14 at 32k). Mirror the proven DV=80/112 fp32-PV tiles for DV=128/256: 16-wide A, one wmma_f32 per K step, persistent f32 accumulator. Also fix the RDNA3 VKQ_C entry count for 16-row A tiles (was sized for 32-row tiles); the count is unchanged for all pre-existing configs. gfx1100, Qwen3.6-27B-Q5_K_S, kvarn6: ladder RMSE ~1e-5, 32k KLD 0.022 vs portable 0.023 (was 2.14), same-top 97.0 percent, prefill pp4096 286 -> 497 t/s. Gated behind GGML_KVARN_AMD_PROMPT_PORTABLE=0; default serving path unchanged.
Owner
|
Please change the target branch to v0.4.7, and ensure changes are compatible with it. |
The three direct final-meta stores in flash_attn_ext_f16_process_tile were gated '!is_kvarn_kv', so KVarN whole-tile blocks never wrote body_meta (softmax max/rowsum). Split tiles get meta from the stream-k fixup, but flash_attn_stream_k_fixup_general skips tiles whose K range aligns exactly to tile boundaries; those rows kept zero meta (den=0) and the tail merge silently discarded their correct body values. Symptom: KVarN WMMA prompt path clean at ub<=64, garbage at ub>=96. Removing the gate makes whole-tile blocks publish their (max, rowsum) like dense FA. Validated: ub512 KLD on 4B/27B/35B-MoE all at portable parity; 32k KLD 2.118 -> 0.029.
# Conflicts: # ggml/src/ggml-cuda/fattn-kvarn-portable.cuh
…s standard) Make F32-WMMA direct-record attention the default for HIP KVarN prompt-prefill (was: portable-native by default, WMMA opt-in). The ub-geometry fix (whole-tile body_meta) plus the DV=128/256 fp32 accumulators make WMMA both the fast and the exact route (~1e-5 ladder RMSE, 32k KLD at portable parity), validated at the serving default (-ub 512): 4B KLD 0.006822 (bit-exact vs fp16 base), 27B 0.004814, 35B-MoE 0.003526. Portable-native remains for unsupported shapes and as an explicit opt-in (GGML_KVARN_AMD_PROMPT_PORTABLE=1). Decode (nq<=16) unchanged. HIP-gated; CUDA/Vulkan and all non-KVarN paths untouched.
The GGML_KVARN_PORTABLE_ATTRS debug print used hipFuncAttributes / hipFuncGetAttributes unconditionally, which do not exist on CUDA and broke the CUDA compile (even though the block only runs when the env var is set). Use the HIP API under GGML_USE_HIP and the CUDA API (cudaFuncAttributes / cudaFuncGetAttributes, same fields) otherwise. No behavior change on HIP; CUDA now compiles. Found while validating the F32-WMMA flip on thermis/4090.
Resolve portable.cuh by taking upstream's complete optimized D64 rewrite. Our interim D64 hand-merge (RECORD_DIM/v_original_domain port) and portable QB-batching/warp-shuffle are superseded (upstream now covers D64 natively; portable is fallback-only post-flip so peak fallback perf is acceptable; re-propose batching as a follow-up if fallback perf matters). Our CUDA func-attributes fix is not needed (upstream removed the HIP-only debug block). Preserved: F32-WMMA default flip (dispatch, auto-merged) and ub-geometry fix (fattn-mma-f16, auto-merged).
Follow-up to the v0.4.7 (complete optimized D64) merge, which superseded our interim portable QB-batching: the flip's fallback dispatch still called the now-removed ggml_cuda_flash_attn_ext_kvarn_portable_batched (undeclared on the merged tree). Route the fallback through the standard portable kernel (correct, unbatched). Portable is fallback-only post-flip.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
Toward #122 (KVarN on ROCm is unusable: wrong results plus prefill far behind). This branch makes KVarN work on RDNA3 step by step, validated on gfx1100 (RX 7900 XTX, ROCm 7.14) with Qwen3.6-27B-Q5_K_S and kvarn6:
fattn-mma-f16.cuh(tile_Q reused across k00 iterations, barrier only fired for np > 1), plus inter-tile sync. Enables the RDNA D256 route.GGML_KVARN_AMD_PROMPT_PORTABLE=0.GGML_KVARN_TEST_NKV_LADDER_ONLY): GPU-native vs CPU-materialized error over D128/256/512 and n_kv 256 to 8192.wmma_f32per K step, persistent f32 acc, mirroring DV=80/112), plus theVKQ_Centry-count fix it required. Still gated behindPORTABLE=0, default serving path unchanged.dst_final_metastores inflash_attn_ext_f16_process_tilewere gated!is_kvarn_kv, so KVarN whole-tile blocks never published their (max, rowsum). Split tiles got meta from the stream-k fixup, but that skips boundary-aligned tiles, so those rows kept zero meta and the tail merge silently discarded their correct body values. Symptom: KVarN WMMA prompt path clean at ub<=64, garbage at ub>=96. Removing the gate makes whole-tile blocks publish their meta like dense FA. Validated: ub512 KLD on 4B / 27B / 35B-MoE at portable parity (4B bit-exact); 32k KLD 2.118 -> 0.029.Results: ladder RMSE ~1e-5; 32k KLD 0.022 vs portable 0.023 (was 2.14 on the fp16 path), same-top 97.0%; prefill pp4096 286 (portable) vs 497 (f32-WMMA); decode ~33.5 both; reference q8_0 at 804.
Base: v0.4.7
Retargeted to
v0.4.7per maintainer request. The branch mergesv0.4.7(D=64 KV-head support, etc.) with this work; the only merge conflict wasfattn-kvarn-portable.cuh, resolved by porting v0.4.7's D=64 support (RECORD_DIM, thev_original_domainoriginal->rotated V Hadamard, D=64 dispatch) onto the portable QB-batching + warp-shuffle kernel.fattn-kvarn-dispatch.cu: D64 stripped from capabilities on non-CUDA backends "until each backend has independent runtime qualification"). So on this gfx1100 build D64 is never dispatched; the D=64 kernel here is compile-validated and mirrors v0.4.7's, but runtime qualification on HIP is out of scope. D128/256/512 are the validated ROCm paths.Additional information
ctestKVarN subset, ladder env above,llama-perplexityKLD flow andllama-bench -p 4096from AGENTS.md.PORTABLE=0); two flaky AVs under long KLD runs share one fault offset with dumps saved, zero occurrences on portable,PORTABLE=1is the fallback. (The ub>=96 serving divergence previously listed here is now fixed by the ub>64 commit above.)Requirements
Update 2026-09-09: F32-WMMA is now the default + CUDA build fix (2 new commits)
This supersedes the "f32-WMMA default flip (still behind
PORTABLE=0)" note above.bad6f86c1ÔÇö F32-WMMA is now the default for HIP KVarN prompt-prefill(was: portable-native by default, WMMA opt-in). The ub-geometry fix plus the
DV=128/256 fp32 accumulators make WMMA both fast and exact, validated at the
serving default (
-ub 512, gfx1100, kvarn6, ngl99): 4B KLD 0.006822 (bit-exactvs fp16 base), 27B 0.004814, 35B-MoE 0.004093; 32k KLD 0.029 (was 2.118);
ladder D128/256/512 ~1e-5. Portable-native remains for unsupported shapes and
as an explicit opt-in (
GGML_KVARN_AMD_PROMPT_PORTABLE=1). Decode (nq<=16)unchanged. HIP-gated; CUDA/Vulkan untouched.
edc1d2f64ÔÇö CUDA build fix (pre-existing, found by the thermis RTX 4090build): the
GGML_KVARN_PORTABLE_ATTRSdebug print used HIP-onlyfunc-attribute APIs, breaking nvcc. Now HIP API on HIP / CUDA API on CUDA
(debug-only, no behavior change). This PR as-merged did not build on CUDA;
this fix is required.
GGML_CUDA=ON): full build clean; KVarN/FA ctest(
kvarn|fattn-route|cuda-fattn) 14/14 pass, incl.test-kvarn. The flip iscompiled out on CUDA (HIP-gated) ÔÇö expected no-op, confirmed.
Community validation needed: gfx1151 (Strix Halo) / gfx1201 (RDNA4)
We only have gfx1100 (RX 7900 XTX) plus CUDA/4090. If you run gfx1151 or
gfx1201, please validate this PR on your hardware and report back:
GGML_HIP=ON,GPU_TARGETSset for your arch at initial configure).as-default is accurate at your serving
-ub(state the-ubvalue).GGML_KVARN_TEST_NKV_LADDER_ONLY=1).ctest -R "kvarn|fattn-route".are the expected ROCm paths. Thank you!
Update 2026-09-09 (pm): merged v0.4.7 "complete optimized D64", conflict resolved
origin/v0.4.7advanced (e1f6d6fe6, complete optimized D64) after this PR'searlier v0.4.7 merge, conflicting in
portable.cuh. Resolved by merging(
cc566c899) and taking upstream's complete D64 rewrite.(upstream covers D64 natively; portable is fallback-only post-flip, so peak
fallback perf is acceptable; re-propose batching as a follow-up if it matters).
Our CUDA func-attributes fix proved unnecessary (upstream removed the HIP-only
debug block).
(both auto-merged). One 3-line follow-up (
37ce9e337) drops the now-deadbatched-portable call (fallback uses standard portable — correct, unbatched).
pre-merge (4B 0.006822 bit-exact, 27B 0.004814, 35B 0.004093); key-target build
clean. PR is MERGEABLE/CLEAN.