Commit cdc06426e for llama.cpp
commit cdc06426e70c23a1ac2bce40c85b252fec344702
Author: Georgi Gerganov <ggerganov@gmail.com>
Date: Thu Sep 24 22:44:36 2026 +0300
metal : optimize sparse FA + clean-up (#29377)
* metal : cache sparse FA indices in shared memory
Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp
* metal : simplify shared memory size calculation
Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp
* pi : update general
* metal : unroll sparse index load
Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp
diff --git a/.pi/gg/SYSTEM.md b/.pi/gg/SYSTEM.md
index bd308ea96..8b8b8a158 100644
--- a/.pi/gg/SYSTEM.md
+++ b/.pi/gg/SYSTEM.md
@@ -7,6 +7,7 @@ General:
- Don't try to build or run the code unless you are explicitly asked to do so
- Use the `gh` CLI tool when querying PRs, issues, or other GitHub resources
- When [MODEL] is needed, first try to get it from the `PI_MODEL_NAME` env var before asking the user
+- Never read the `AGENTS.md` file
Coding:
- When in doubt, always refer to the CONTRIBUTING.md file of the project
diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp
index 708703e51..54d5b9808 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp
@@ -3494,34 +3494,20 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
const int is_q = !use_kv_f16 && ggml_is_quantized(op->src[1]->type) ? 1 : 0;
- // 2*(2*ncpsg)
- // ncpsg soft_max values + ncpsg mask values
- //
- // 16*32*(nsg)
- // the shared memory needed for the simdgroups to load the KV cache
- // each thread loads (dequantizes) 16 head elements, there are 32 threads in th SG
- //
-#define FATTN_SMEM(nsg) (GGML_PAD((nqptg*(ne00 + 2*GGML_PAD(ne20, 64) + 2*(2*ncpsg)) + is_q*(16*32*(nsg)))*(sizeof(float)/2), 16))
-
- //int64_t nsgmax = 4;
- //
- //if (is_q) {
- // nsgmax = 2;
- // while (true) {
- // const size_t smem = FATTN_SMEM(nsgmax);
- // if (smem > props_dev->max_theadgroup_memory_size) {
- // break;
- // }
- // nsgmax *= 2;
- // }
- // nsgmax /= 2;
- //}
+ // shared memory layout (halfs unless noted):
+ // queries/attn/result: Q*(DK + 2*PAD2(DV,64) + 4*C)
+ // quantized KV scratch: 16*32*NSG (only when is_q)
+ const int64_t dv_pad = GGML_PAD(ne20, 64);
+
+ auto fa_smem = [&](int32_t nsg) -> size_t {
+ const size_t smem_half = nqptg*(ne00 + 2*dv_pad + 4*ncpsg) + is_q*(16*32*nsg);
+ return GGML_PAD(smem_half*sizeof(ggml_fp16_t), 16);
+ };
// simdgroups per threadgroup (a.k.a. warps)
- //nsg = ne01 <= nqptg ? MAX(4, MIN(nsgmax, MIN(ne11/ncpsg, (int64_t) pipeline.maxTotalThreadsPerThreadgroup/32))) : 4;
int32_t nsg = ne00 >= 512 ? 8 : 4;
- const size_t smem = FATTN_SMEM(nsg);
+ const size_t smem = fa_smem(nsg);
const int32_t ns10 = nb11_attn/nb10_attn;
const int32_t ns20 = nb21_attn/nb20_attn;
@@ -3577,7 +3563,6 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
ggml_metal_encoder_dispatch_threadgroups(enc, (ne01 + nqptg - 1)/nqptg, ne02, ne03, 32, nsg, 1);
-#undef FATTN_SMEM
} else {
// half4x4 kernel
// sparse: the index lists are per query row, so a threadgroup can share KV with Q == 1 only
@@ -3679,14 +3664,18 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
// note: for simplicity assume the K is larger or equal than V
GGML_ASSERT(ne10 >= ne20);
- // ne00 + 2*ncpsg*(nsg)
- // for each query, we load it as f16 in shared memory (ne00)
- // and store the soft_max values and the mask
- //
- // ne20*(nsg)
- // each simdgroup has a full f32 head vector in shared mem to accumulate results
- //
-#define FATTN_SMEM(nsg) (GGML_PAD(((GGML_PAD(ne00, 128) + 4*ncpsg + 2*GGML_PAD(ne20, 128))*(nsg)*nqptg)*(sizeof(float)/2), 16))
+ // shared memory layout (halfs unless noted):
+ // queries: Q*NSG*PAD2(ne00, 128)
+ // attn + mask: NSG*4*Q*C
+ // results: 2*NSG*Q*PAD2(ne20, 128)
+ // sparse idx: NSG*C ints (only when use_sparse)
+ const int64_t dk_pad = GGML_PAD(ne00, 128);
+ const int64_t dv_pad = GGML_PAD(ne20, 128);
+
+ auto fa_vec_smem = [&](int64_t nsg, int32_t nqptg) -> size_t {
+ const size_t smem_half = (size_t) (dk_pad + 4*ncpsg + 2*dv_pad)*nqptg*nsg;
+ return GGML_PAD(smem_half*sizeof(ggml_fp16_t) + (use_sparse ? (size_t) nsg*ncpsg*sizeof(int) : 0), 16);
+ };
int64_t nsg = 1;
@@ -3722,7 +3711,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
}
// fall back to baseline (Q=1) if the tuned config exceeds threadgroup memory
- if ((size_t) FATTN_SMEM(nsg) > props_dev->max_theadgroup_memory_size) {
+ if (fa_vec_smem(nsg, nqptg) > props_dev->max_theadgroup_memory_size) {
cfg = ggml_metal_tuning::fa_vec_baseline_cfg((int) ne00, (int) ne20);
nqptg = cfg.Q; // = 1
}
@@ -3779,9 +3768,8 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
ggml_metal_encoder_set_buffer (enc, bid_src4, 5);
ggml_metal_encoder_set_buffer (enc, use_sparse ? bid_idx : bid_src0, 8);
- const size_t smem = FATTN_SMEM(nsg);
+ const size_t smem = fa_vec_smem(nsg, nqptg);
- //printf("smem: %zu, max: %zu, nsg = %d, nsgmax = %d\n", smem, props_dev->max_theadgroup_memory_size, (int) nsg, (int) nsgmax);
GGML_ASSERT(smem <= props_dev->max_theadgroup_memory_size);
if (nwg == 1) {
@@ -3827,7 +3815,6 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, 32*nwg, 1, 1);
}
}
-#undef FATTN_SMEM
}
return 1;
diff --git a/ggml/src/ggml-metal/kernels/fa.metal b/ggml/src/ggml-metal/kernels/fa.metal
index f26d493d5..cf23ef04f 100644
--- a/ggml/src/ggml-metal/kernels/fa.metal
+++ b/ggml/src/ggml-metal/kernels/fa.metal
@@ -1248,17 +1248,24 @@ kernel void kernel_flash_attn_ext_vec(
constexpr short NL = NW/NE; // note: this can be adjusted to support different head sizes and simdgroup work loads
constexpr short SH = 4*Q*C; // shared memory per simdgroup
+ const int SMEM_Q = Q*NSG*PK;
+ const int SMEM_S = NSG*SH;
+ const int SMEM_O = 2*NSG*Q*PV;
+ const int SMEM = SMEM_Q + SMEM_S + SMEM_O;
+
static_assert(DK4 % NL == 0, "DK4 must be divisible by NL");
static_assert(DV4 % NL == 0, "DV4 must be divisible by NL");
- //const short T = PK + NSG*SH; // shared memory size per query in (half)
+ threadgroup q4_t * sq4 = (threadgroup q4_t *) shmem_f16; // holds the query data
+ threadgroup s_t * ss = (threadgroup s_t *) (shmem_f16 + SMEM_Q + sgitg*SH); // scratch buffer for attention
+ threadgroup s4_t * ss4 = (threadgroup s4_t *) (shmem_f16 + SMEM_Q + sgitg*SH); // same as above but in s4_t
+ threadgroup half * sm = (threadgroup half *) (shmem_f16 + SMEM_Q + sgitg*SH + 2*Q*C); // scratch buffer for mask
+ threadgroup o4_t * so4 = (threadgroup o4_t *) (shmem_f16 + SMEM_Q + SMEM_S + 2*sgitg*Q*PV); // scratch buffer for the results
- //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
- threadgroup o4_t * so4 = (threadgroup o4_t *) (shmem_f16 + 2*sgitg*Q*PV + Q*NSG*PK + NSG*SH); // scratch buffer for the results
+ // sparse indices for the current block
+ threadgroup int * spidx = FC_flash_attn_ext_vec_has_sparse
+ ? (threadgroup int *) (shmem_f16 + SMEM) + sgitg*C
+ : nullptr;
// store the result for all queries in shared memory (the O matrix from the paper)
so4 += tiisg;
@@ -1388,10 +1395,20 @@ kernel void kernel_flash_attn_ext_vec(
ic = 0;
}
+ // load the sparse KV indices for the current block into shared memory
+ if (FC_flash_attn_ext_vec_has_sparse) {
+ FOR_UNROLL (short ii = 0; ii < C/NW; ++ii) {
+ const short i = ii*NW + tiisg;
+
+ spidx[i] = pidx[ic + i];
+ }
+ simdgroup_barrier(mem_flags::mem_threadgroup);
+ }
+
if (FC_flash_attn_ext_vec_has_mask) {
if (FC_flash_attn_ext_vec_has_sparse) {
FOR_UNROLL (short qq = 0; qq < Q; ++qq) {
- const int i11 = pidx[ic + tiisg];
+ const int i11 = spidx[tiisg];
if ((iq1*Q + qq) < args.ne01 && i11 >= 0) {
sm[qq*C + tiisg] = pm[qq][i11];
} else {
@@ -1449,7 +1466,7 @@ kernel void kernel_flash_attn_ext_vec(
FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) {
if (FC_flash_attn_ext_vec_has_sparse) {
// the KV rows are gathered from the index list; -1 entries are padding
- const int i11 = pidx[ic + NE*cc + ty];
+ const int i11 = spidx[NE*cc + ty];
if (i11 >= 0) {
if (is_same<kd4_t, k4_t>::value) {
device const k4_t * pk4s = (device const k4_t *) (k + i11*args.nb11) + tx;
@@ -1593,7 +1610,7 @@ kernel void kernel_flash_attn_ext_vec(
if (FC_flash_attn_ext_vec_has_sparse) {
FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) {
// the KV rows are gathered from the index list; -1 entries are padding
- const int i11 = pidx[ic + NE*cc + ty];
+ const int i11 = spidx[NE*cc + ty];
if (i11 >= 0) {
if (is_same<vd4_t, v4_t>::value) {
device const v4_t * pv4 = (device const v4_t *) (v + i11*args.nb21);