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

metal : reduce branches

Browse Source
This commit is contained in:
Georgi Gerganov 2024-01-21 11:38:17 +02:00
parent 528da7515e
commit 52ae085750
No known key found for this signature in database
GPG Key ID: 449E073F9DC10735
1 changed files with 8 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
ggml-metal.metal
View File

@ -2056,40 +2056,26 @@ kernel void kernel_flash_attn_ext_f16(
continue; continue;
} }
half4 s4 = 0.0f;
device const half4 * pk4 = (device const half4 *) ((device char *) k + (ic*nb11 + ik2*nb12 + ik3*nb13)); device const half4 * pk4 = (device const half4 *) ((device char *) k + (ic*nb11 + ik2*nb12 + ik3*nb13));
device const half4 * pv4 = (device const half4 *) ((device char *) v + (ic*nb21 + iv2*nb22 + iv3*nb23));
half4 s4 = 0.0h;
for (int64_t d = 0; d < D4; ++d) { for (int64_t d = 0; d < D4; ++d) {
s4 += pk4[d] * pq4[d]; s4 += pk4[d] * pq4[d];
} }
half s = s4.x + s4.y + s4.z + s4.w; half s = (s4.x + s4.y + s4.z + s4.w)*scale + mv;
s = s*scale + mv;
const half Mold = M; const half Mold = M;
half ms = 1.0f; M = max(M, s);
half vs = 1.0f;
if (s > M) { const half ms = exp(Mold - M);
M = s; const half vs = exp(s - M);
ms = exp(Mold - M);
// V = V*exp(Mold - M)
for (int64_t d = 0; d < D4; ++d) { for (int64_t d = 0; d < D4; ++d) {
V16[d] *= ms; V16[d] = V16[d]*ms + pv4[d]*vs;
}
} else {
vs = exp(s - M);
}
device const half4 * pv4 = (device const half4 *) ((device char *) v + (ic*nb21 + iv2*nb22 + iv3*nb23));
// V += v*exp(s - M)
for (int64_t d = 0; d < D4; ++d) {
V16[d] += pv4[d] * vs;
} }
S = S*ms + vs; S = S*ms + vs;