Repository navigation
vulkan : multi-row coopmat1 flash attention for prefill, with a V^T scratch - #159
Merged
Merged
Conversation
…cratch flash_attn_cm1_mr.comp: one workgroup of eight wave32 subgroups owns 64 rows (H heads x 64/H tokens that share a KV head). Q stays in registers for the whole KV loop and every 16-key K/V tile serves all 64 rows, where cm1 pins Br at 16 and streams the whole K/V range per 16 rows. K is read as per-head f16 rows, V from a per-head V^T scratch written by dequant_f16_transpose_vt.comp ([HSV][KV + 128], 64x64 LDS tiles). Keys masked in every output row have their V zeroed per tile, so a WMMA -0.0 product cannot move the output and NaN V in masked cells cannot leak. GGML_VK_FA_MR: default on for head size 256 on RADV RDNA3, =1 for every eligible shape (head size 128 or 256, f16 K after the contiguize pass, f32 accumulation, KV % 64 == 0, mask shared across heads), =0 off. GGML_VK_FA_MR_LAZY / _NOEB / _VPOST (default on) are schedule knobs. GGML_VK_FA_VT: V^T for the cm1 kernel too (default on RADV RDNA3 when NAN_SAFE_V and shmem staging are off), =0 also disables MR. New flag bits 64..1024; NAN_SAFE_V keeps 32. Measured in the fork (v0.7.8) on gfx1151, Qwen3.8-27B UD-Q4_K_XL: FA per pp2048 ubatch at d32768 2067 -> 954 ms, pp2048 +21.5% at d32768, +5.7% at d8192, KLD 0.000197 against the cm1 kernel. Assisted-by: Claude Opus 5.5
dzannotti
approved these changes
Oct 8, 2026
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.
This adds a coopmat1 prefill flash attention kernel that processes several query rows per subgroup, plus a transposed V scratch so the P x V product reads V in the layout the matrix cores want. It is the default for head size 256 on RADV RDNA3. The plain cm1 kernel also uses the transposed V on RADV RDNA3 when the NaN-safe path is off. Ported from the nathanw fork (v0.7.8), with the masked-V zeroing from the -0.0 fix built into the new kernel.
On a Strix Halo (gfx1151, RADV, Mesa 26.2.4), pp2048
-ub 2048, same binary withGGML_VK_FA_MR=0/GGML_VK_FA_VT=0as the off arm:Against current master the 27B goes from 284 to 346 t/s at 32k, and the 35B-A3B from about 1068 to 1461. Depth 0 is unchanged.
GGML_VK_FA_MR=0turns it off andGGML_VK_FA_MR=1enables it for head size 128 too.GGML_VK_FA_VT=0turns off the transposed V.GGML_VK_FA_MR_LOG=1logs each dispatch.Written with Claude, tested by me.
Assisted-by: Claude Opus 5.5