ikawrakow/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-01-26 17:20:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use fp32 for K*Q in Metal FA implementation (#62)

Browse Source 
Co-authored-by: Iwan Kawrakow <iwan.kawrakow@gmail.com>
This commit is contained in:
Kawrakow
2024-09-25 13:08:55 +03:00
committed by GitHub
parent 15b231cb74
commit 1cdb6993ee
1 changed files with 8 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -2904,11 +2904,11 @@ kernel void kernel_flash_attn_ext_vec_f16(
const short iv3 = iq3 / rv3;
// load the queries from shared memory into local memory
half4 mq[D4];
float4 mq[D4];
for (short ii = 0; ii < D4; ii += NW) {
short i = ii + tiisg;
mq[i] = sq4[i];
mq[i] = (float4)sq4[i];
}
// pointer to the mask
@@ -2934,11 +2934,11 @@ kernel void kernel_flash_attn_ext_vec_f16(
for (short ii = 0; ii < D4; ii += NW) {
const short i = ii + tiisg;
half4x4 mk;
mk[0] = pk4[i + 0*(nb11/8)];
mk[1] = pk4[i + 1*(nb11/8)];
mk[2] = pk4[i + 2*(nb11/8)];
mk[3] = pk4[i + 3*(nb11/8)];
float4x4 mk;
mk[0] = (float4)pk4[i + 0*(nb11/8)];
mk[1] = (float4)pk4[i + 1*(nb11/8)];
mk[2] = (float4)pk4[i + 2*(nb11/8)];
mk[3] = (float4)pk4[i + 3*(nb11/8)];
mqk += (float4) (mq[i] * mk);
}
@@ -2960,6 +2960,7 @@ kernel void kernel_flash_attn_ext_vec_f16(
ss4[cc] = mqk;
}
}
}