Superminaren/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-04 02:37:27 +02:00
Code Issues Packages Projects Releases Wiki Activity

cont : adjust nsg

Browse Source
This commit is contained in:
Georgi Gerganov
2026-08-31 21:26:45 +03:00
parent 7916e96d01
commit cd3a488320
2 changed files with 15 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-metal/ggml-metal-ops.cpp
+12 -2
View File
@@ -3562,11 +3562,21 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
if (ne01 > 32) {
// large sparse batch
nwg = 1;
nsg = 4;
nsg = 1;
if (n_kv_max_padded == 640) {
nsg = 4; // 640 % (4*32) == 0
} else {
while (2*nwg*nsg*ncpsg < n_kv_max_padded && nsg < 4) {
nsg *= 2;
}
}
} else {
// small sparse batch
nsg = 1;
nwg = 32;
nsg = 1;
while (2*nwg*nsg*ncpsg < n_kv_max_padded && nsg < 4) {
nsg *= 2;
}
}
} else {
nwg = 32;
ggml/src/ggml-metal/kernels/fa.metal
+3 -2
View File
@@ -1245,8 +1245,8 @@ kernel void kernel_flash_attn_ext_vec(
//const short T = PK + NSG*SH; // shared memory size per query in (half)
//threadgroup q_t * sq = (threadgroup q_t *) (shmem_f16 + 0*PK); // holds the query data
threadgroup q4_t * sq4 = (threadgroup q4_t *) (shmem_f16 + 0*PK); // same as above but in q4_t
//threadgroup q_t * sq = (threadgroup q_t *) (shmem_f16 + 0*PK); // holds the query data
threadgroup q4_t * sq4 = (threadgroup q4_t *) (shmem_f16 + 0*PK); // same as above but in q4_t
threadgroup s_t * ss = (threadgroup s_t *) (shmem_f16 + sgitg*SH + Q*NSG*PK); // scratch buffer for attention
threadgroup s4_t * ss4 = (threadgroup s4_t *) (shmem_f16 + sgitg*SH + Q*NSG*PK); // same as above but in s4_t
threadgroup half * sm = (threadgroup half *) (shmem_f16 + sgitg*SH + 2*Q*C + Q*NSG*PK); // scratch buffer for mask
@@ -1407,6 +1407,7 @@ kernel void kernel_flash_attn_ext_vec(
}
}
// skip -INF mask
{
bool any_finite = false;
FOR_UNROLL (short qq = 0; qq < Q; ++qq) {