root/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggerganov/llama.cpp.git synced 2024-12-26 11:24:35 +00:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

metal : fix F32 accumulation in FA vec kernel (#10232)

Browse Source
This commit is contained in:
Georgi Gerganov 2024-11-09 11:52:45 +02:00 committed by GitHub
parent f018acba22
commit bb38cdd8ba
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
ggml/src/ggml-metal.metal
View File

@ -3450,7 +3450,7 @@ kernel void kernel_flash_attn_ext_vec(
{ {
// each simdgroup processes 1 query and 4 keys // each simdgroup processes 1 query and 4 keys
for (short cc = 0; cc < C/4; ++cc) { for (short cc = 0; cc < C/4; ++cc) {
qk_t mqk = 0.0; qk_t mqka[4] = { 0.0, 0.0, 0.0, 0.0 };
device const kd4x4_t * pk = (device const kd4x4_t *) ((device const char *) k + ((ic + 4*cc + ty)*nb_12_1 + ikv2*nb_12_2 + ikv3*nb_12_3)); device const kd4x4_t * pk = (device const kd4x4_t *) ((device const char *) k + ((ic + 4*cc + ty)*nb_12_1 + ikv2*nb_12_2 + ikv3*nb_12_3));
@ -3461,13 +3461,14 @@ kernel void kernel_flash_attn_ext_vec(
k4x4_t mk; k4x4_t mk;
deq_k(pk + i/nl_k, i%nl_k, mk); deq_k(pk + i/nl_k, i%nl_k, mk);
mqk += mqka[0] += dot(mq[ii/NL][0], mk[0]);
dot(mq[ii/NL][0], mk[0]) + mqka[1] += dot(mq[ii/NL][1], mk[1]);
dot(mq[ii/NL][1], mk[1]) + mqka[2] += dot(mq[ii/NL][2], mk[2]);
dot(mq[ii/NL][2], mk[2]) + mqka[3] += dot(mq[ii/NL][3], mk[3]);
dot(mq[ii/NL][3], mk[3]);
} }
qk_t mqk = mqka[0] + mqka[1] + mqka[2] + mqka[3];
// simdgroup reduce // simdgroup reduce
// [ 0 .. 7] -> [ 0] // [ 0 .. 7] -> [ 0]
// [ 8 .. 15] -> [ 8] // [ 8 .. 15] -> [ 8]