From cd3a4883201f26b4acbbd3aeb096bc38ca484779 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Mon, 31 Aug 2026 20:34:54 +0300 Subject: [PATCH] cont : adjust nsg --- ggml/src/ggml-metal/ggml-metal-ops.cpp | 14 ++++++++++++-- ggml/src/ggml-metal/kernels/fa.metal | 5 +++-- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index bdec057334..1fd984cbd9 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -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; diff --git a/ggml/src/ggml-metal/kernels/fa.metal b/ggml/src/ggml-metal/kernels/fa.metal index 2f7aed3204..c966c03972 100644 --- a/ggml/src/ggml-metal/kernels/fa.metal +++ b/ggml/src/ggml-metal/kernels/fa.metal @@ -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) {