ver4a/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

vulkan: enable coopmat2 FA gqa and split_k optimizations more often (#12931)

Browse source 
The grouped query attention optmization doesn't require a power of two ratio,
the only thing relying on it was the modulo operation written as bitwise &.

split_k need not depend on gqa_ratio - enable it any time there's only one
workgroup in the X dimension. The shader gets the split index from the x coord,
and multiple workgroups in the X dimension (pre-split) indicates a larger
FA operation that wouldn't need splitting.
This commit is contained in:
Jeff Bolz 2025-04-16 13:37:25 -05:00 • committed by GitHub
parent b43d89e311
commit 015022bb53
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 7 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

2
ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp
View file

@ -131,7 +131,7 @@ ACC_TYPE perElemOpStoreCol0(const in uint32_t r, const in uint32_t c, const in A
// Load the slope matrix, indexed by Q's dimension 2.
ACC_TYPE perElemOpComputeSlope(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem, const in uint32_t iq2)
{
const uint32_t h = iq2 + (r & (p.gqa_ratio - 1));
const uint32_t h = iq2 + (r % p.gqa_ratio);
const ACC_TYPE base = ACC_TYPE(h < p.n_head_log2 ? p.m0 : p.m1);
const int exph = int(h < p.n_head_log2 ? h + 1 : 2*(h - p.n_head_log2) + 1);