diff --git a/.pi/gg/SYSTEM.md b/.pi/gg/SYSTEM.md index bd308ea963..8b8b8a1586 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 708703e511..54d5b98083 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)) + // 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; diff --git a/ggml/src/ggml-metal/kernels/fa.metal b/ggml/src/ggml-metal/kernels/fa.metal index f26d493d54..cf23ef04fa 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::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::value) { device const v4_t * pv4 = (device const v4_t *) (v + i11*args.nb21);