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

metal : optimize sparse FA + clean-up (#29377)

Browse Source 
* 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
This commit is contained in:
Georgi Gerganov
2026-09-24 22:44:36 +03:00
committed by GitHub
parent bced4595b8
commit cdc06426e7
3 changed files with 51 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.pi/gg/SYSTEM.md
+1
View File
@@ -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
ggml/src/ggml-metal/ggml-metal-ops.cpp
+23 -36
View File
@@ -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))
// 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);
//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;
//}
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;
ggml/src/ggml-metal/kernels/fa.metal
+27 -10
View File
@@ -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);