CS348Project/llama.cpp
5
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2025-11-07 09:57:00 +00:00
Code Issues Packages Projects Releases 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
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);