From 600330789cb1b3404838d3fc36c02208273be2e6 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Sun, 30 Aug 2026 16:47:15 +0300 Subject: [PATCH 01/10] metal : support n_kv_max sparse mask hint in flash attention vec kernel - add kernel_flash_attn_ext_vec_idx: compacts finite mask entries into a per-row index list (Hillis-Steele scan, one threadgroup per row) - extend vec FA kernel with optional sparse index gathering (FC slot 5) - add host-side gate: sparse path when n_kv_max > 0, mask present, supported head sizes / KV types, n_kv_max <= 4096 - new buffer region extra_idx for the index list - pipeline getter extended with has_sparse param - add test cases: head sizes, quant types, nb>1, nr23 variants, sinks, ALiBi, softcap, permute, v_view_of_k, no-mask fallback Note: multi-row (nb*nr23[1] > 1) cases still failing - rid mapping in the store phase needs revisiting for the sparse path. Assisted-by: pi:llama.cpp/Qwen3.8-27B --- ggml/src/ggml-metal/ggml-metal-device.cpp | 27 ++- ggml/src/ggml-metal/ggml-metal-device.h | 5 + ggml/src/ggml-metal/ggml-metal-impl.h | 14 ++ ggml/src/ggml-metal/ggml-metal-ops.cpp | 151 +++++++++++++++-- ggml/src/ggml-metal/ggml-metal-ops.h | 1 + ggml/src/ggml-metal/ggml-metal.cpp | 1 + ggml/src/ggml-metal/kernels/fa.metal | 193 ++++++++++++++++++++-- tests/test-backend-ops.cpp | 55 +++++- 8 files changed, 420 insertions(+), 27 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index b8d2ef9ce27..c296d17b157 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -1577,6 +1577,26 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext( return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_idx( + ggml_metal_library_t lib, + const ggml_tensor * op) { + assert(op->op == GGML_OP_FLASH_ATTN_EXT); + assert(op->src[3]); + + char name[256]; + + snprintf(name, 256, "kernel_flash_attn_ext_vec_idx"); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr); + } + + GGML_UNUSED(op); + + return res; +} + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec( ggml_metal_library_t lib, const ggml_tensor * op, @@ -1585,6 +1605,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v bool has_bias, bool has_scap, bool has_kvpad, + bool has_sparse, int32_t nqpsg, int32_t ne, int32_t nsg, @@ -1614,13 +1635,14 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v dv, qne_suffix); - snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d", + snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_sparse=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d", base, has_mask, has_sinks, has_bias, has_scap, has_kvpad, + has_sparse, ns10, ns20, nsg, nwg); @@ -1633,7 +1655,8 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v ggml_metal_cv_set_bool(cv, has_sinks, FC_FLASH_ATTN_EXT_VEC + 1); ggml_metal_cv_set_bool(cv, has_bias, FC_FLASH_ATTN_EXT_VEC + 2); ggml_metal_cv_set_bool(cv, has_scap, FC_FLASH_ATTN_EXT_VEC + 3); - ggml_metal_cv_set_bool(cv, has_kvpad, FC_FLASH_ATTN_EXT_VEC + 4); + ggml_metal_cv_set_bool(cv, has_kvpad, FC_FLASH_ATTN_EXT_VEC + 4); + ggml_metal_cv_set_bool(cv, has_sparse, FC_FLASH_ATTN_EXT_VEC + 5); ggml_metal_cv_set_int32(cv, ns10, FC_FLASH_ATTN_EXT_VEC + 20); ggml_metal_cv_set_int32(cv, ns20, FC_FLASH_ATTN_EXT_VEC + 21); diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index ae4871d3586..31fc07d44d4 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -201,6 +201,10 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att int32_t ns10, int32_t ns20); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_idx( + ggml_metal_library_t lib, + const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec( ggml_metal_library_t lib, const struct ggml_tensor * op, @@ -209,6 +213,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att bool has_bias, bool has_scap, bool has_kvpad, + bool has_sparse, int32_t nqpsg, int32_t ne, int32_t nsg, diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index bdcd9c9e3d8..f6d6867c6fd 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -107,6 +107,7 @@ #define FC_SUM_ROWS 1400 #define FC_UPSCALE 1500 #define FC_GATED_DELTA_NET 1600 +#define FC_FLASH_ATTN_EXT_VEC_IDX 1700 // op-specific constants #define OP_FLASH_ATTN_EXT_NQPSG 8 @@ -458,8 +459,21 @@ typedef struct { float m1; int32_t n_head_log2; float logit_softcap; + int32_t n_kv_max_padded; } ggml_metal_kargs_flash_attn_ext_vec; +typedef struct { + int32_t ne30; + int32_t ne31; + int32_t ne32; + int32_t ne33; + uint64_t nb31; + uint64_t nb32; + uint64_t nb33; + int32_t n_kv_max; + int32_t n_kv_max_padded; +} ggml_metal_kargs_flash_attn_ext_vec_idx; + typedef struct { int32_t nrows; } ggml_metal_kargs_flash_attn_ext_vec_reduce; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index bc8b3c8d485..bf29e6e720e 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2857,6 +2857,70 @@ static bool ggml_metal_op_flash_attn_ext_use_kv_f16(const ggml_tensor * op) { } } +// returns the n_kv_max hint if the sparse path is available for this op, or 0 otherwise +// the mask (src[3]) remains the single source of truth: finite entries are the valid KV positions, +// n_kv_max is only an upper bound on their number per mask row, used to size the index lists +static int ggml_metal_op_flash_attn_ext_n_kv_max_sparse(const ggml_tensor * op) { + assert(op->op == GGML_OP_FLASH_ATTN_EXT); + + int32_t n_kv_max = 0; + memcpy(&n_kv_max, ((const int32_t *) op->op_params) + 4, sizeof(n_kv_max)); + + if (n_kv_max <= 0) { + return 0; + } + + // the sparse indices are gathered from the mask + if (!op->src[3]) { + return 0; + } + + // the sparse path is implemented for the vec kernels only + if (!ggml_metal_op_flash_attn_ext_use_vec(op)) { + return 0; + } + + // bound the size of the index lists + if (n_kv_max > 4096) { + return 0; + } + + // vec kernel instantiations exist for these (type, dk, dv) combinations only + const int64_t dk = op->src[1]->ne[0]; + const int64_t dv = op->src[2]->ne[0]; + + const bool dk_dv_ok = (dk == 32 && dv == 32) || + (dk == 64 && dv == 64) || + (dk == 96 && dv == 96) || + (dk == 128 && dv == 128) || + (dk == 192 && dv == 128) || + (dk == 192 && dv == 192) || + (dk == 256 && dv == 256) || + (dk == 320 && dv == 256) || + (dk == 512 && dv == 512) || + (dk == 576 && dv == 512); + + if (!dk_dv_ok) { + return 0; + } + + switch (op->src[1]->type) { + case GGML_TYPE_F16: + case GGML_TYPE_BF16: + case GGML_TYPE_F32: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + break; + default: + return 0; + } + + return n_kv_max; +} + // in some models (e.g. MLA-based), V is a view of K (the first ne20 elements of each K row); // the dequantized V is then a view of the dequantized K and does not need its own dequant or scratch // - ref: https://github.com/ggml-org/llama.cpp/pull/13435 @@ -3027,6 +3091,24 @@ size_t ggml_metal_op_flash_attn_ext_extra_kv_f16(const ggml_tensor * op) { return k_size + v_size; } +// size of the sparse index lists: one list of KV indices per mask row, +// padded with -1 up to a multiple of OP_FLASH_ATTN_EXT_VEC_NCPSG +size_t ggml_metal_op_flash_attn_ext_extra_idx(const ggml_tensor * op) { + assert(op->op == GGML_OP_FLASH_ATTN_EXT); + + GGML_TENSOR_LOCALS( int32_t, ne3, op->src[3], ne); + + const int n_kv_max = ggml_metal_op_flash_attn_ext_n_kv_max_sparse(op); + + if (n_kv_max <= 0) { + return 0; + } + + const int n_kv_max_padded = GGML_PAD(n_kv_max, OP_FLASH_ATTN_EXT_VEC_NCPSG); + + return GGML_PAD(sizeof(int32_t)*(size_t) n_kv_max_padded*ne31*ne32*ne33, 16); +} + int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); @@ -3104,7 +3186,16 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { ggml_metal_buffer_id bid_kv_f16 = bid_tmp; bid_kv_f16.offs += ggml_metal_op_flash_attn_ext_extra_tmp(op); - const bool use_kv_f16 = ggml_metal_op_flash_attn_ext_use_kv_f16(op); + // sparse path: gather the finite mask entries into index lists and run the vec kernels over them + const int n_kv_max_sparse = ggml_metal_op_flash_attn_ext_n_kv_max_sparse(op); + const bool use_sparse = n_kv_max_sparse > 0; + const int n_kv_max_padded = use_sparse ? GGML_PAD(n_kv_max_sparse, OP_FLASH_ATTN_EXT_VEC_NCPSG) : 0; + + // the vec kernels dequantize the KV inline; no need for the F16 dequant pass in the sparse path + const bool use_kv_f16 = !use_sparse && ggml_metal_op_flash_attn_ext_use_kv_f16(op); + + ggml_metal_buffer_id bid_idx = bid_kv_f16; + bid_idx.offs += ggml_metal_op_flash_attn_ext_extra_kv_f16(op); ggml_metal_buffer_id bid_k = bid_src1; ggml_metal_buffer_id bid_v = bid_src2; @@ -3378,12 +3469,15 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { #undef FATTN_SMEM } else { // half4x4 kernel - auto cfg = ggml_metal_tuning::fa_vec_pick( - props_dev->device_id, - props_dev->gpu_family, - (int) op->src[1]->type, - (int) ne00, (int) ne20, // dk, dv (ne00 == dk for FA) - ne11, ne01); + // sparse: the index lists are per query row, so a threadgroup can share KV with Q == 1 only + auto cfg = use_sparse + ? ggml_metal_tuning::fa_vec_baseline_cfg((int) ne00, (int) ne20) + : ggml_metal_tuning::fa_vec_pick( + props_dev->device_id, + props_dev->gpu_family, + (int) op->src[1]->type, + (int) ne00, (int) ne20, // dk, dv (ne00 == dk for FA) + ne11, ne01); int nqptg = cfg.Q; // queries per threadgroup const int ncpsg = OP_FLASH_ATTN_EXT_VEC_NCPSG; // cache values per simdgroup !! sync with kernel template arguments !! const int nhptg = 1; // heads per threadgroup @@ -3394,7 +3488,37 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { bool need_sync = false; - const bool has_kvpad = ne11 % ncpsg != 0; + const bool has_kvpad = !use_sparse && ne11 % ncpsg != 0; + + if (use_sparse) { + assert(ggml_metal_op_flash_attn_ext_extra_idx(op) != 0); + + ggml_metal_kargs_flash_attn_ext_vec_idx args0 = { + /*.ne30 =*/ ne30, + /*.ne31 =*/ ne31, + /*.ne32 =*/ ne32, + /*.ne33 =*/ ne33, + /*.nb31 =*/ nb31, + /*.nb32 =*/ nb32, + /*.nb33 =*/ nb33, + /*.n_kv_max =*/ n_kv_max_sparse, + /*.n_kv_max_padded =*/ n_kv_max_padded, + }; + + auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_vec_idx(lib, op); + + ggml_metal_encoder_set_pipeline(enc, pipeline0); + ggml_metal_encoder_set_bytes (enc, &args0, sizeof(args0), 0); + ggml_metal_encoder_set_buffer (enc, bid_src3, 1); + ggml_metal_encoder_set_buffer (enc, bid_idx, 2); + + int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline0), 256); + nth = std::max(32, (nth/32)*32); + + ggml_metal_encoder_dispatch_threadgroups(enc, ne31, ne32, ne33, nth, 1, 1); + + need_sync = true; + } if (has_kvpad) { assert(ggml_metal_op_flash_attn_ext_extra_pad(op) != 0); @@ -3460,6 +3584,11 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { // however, this does not lead to significant improvement, so disabled nwg = 1; nsg = 4; + } else if (use_sparse) { + // the reduce kernel sums over one workgroup per lane (32), so nwg must be 32; + // workgroups beyond the number of chunks emit empty partials which the reduce ignores + nsg = 1; + nwg = 32; } else { nwg = 32; nsg = 1; @@ -3484,7 +3613,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.nb01 =*/ nb01, /*.nb02 =*/ nb02, /*.nb03 =*/ nb03, - /*.ne11 =*/ ne11, + /*.ne11 =*/ use_sparse ? n_kv_max_padded : ne11, /*.ne_12_2 =*/ ne12, /*.ne_12_3 =*/ ne13, /*.ns10 =*/ ns10, @@ -3510,9 +3639,10 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.m1 =*/ m1, /*.n_head_log2 =*/ n_head_log2, /*.logit_softcap =*/ logit_softcap, + /*.n_kv_max_padded =*/ n_kv_max_padded, }; - auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nqptg, cfg.NE, nsg, nwg, use_kv_f16, ns10, ns20); + auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, use_sparse, nqptg, cfg.NE, nsg, nwg, use_kv_f16, ns10, ns20); GGML_ASSERT(nsg*32 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); @@ -3523,6 +3653,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_buffer (enc, bid_v, 3); ggml_metal_encoder_set_buffer (enc, bid_src3, 4); 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); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index 159a628d04a..f8fe50b468e 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -43,6 +43,7 @@ size_t ggml_metal_op_flash_attn_ext_extra_pad(const struct ggml_tensor * op); size_t ggml_metal_op_flash_attn_ext_extra_blk(const struct ggml_tensor * op); size_t ggml_metal_op_flash_attn_ext_extra_tmp(const struct ggml_tensor * op); size_t ggml_metal_op_flash_attn_ext_extra_kv_f16(const struct ggml_tensor * op); +size_t ggml_metal_op_flash_attn_ext_extra_idx(const struct ggml_tensor * op); int ggml_metal_op_concat (ggml_metal_op_t ctx, int idx); int ggml_metal_op_repeat (ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index 4d58dc821cf..3bd6abd06fd 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -232,6 +232,7 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_ res += ggml_metal_op_flash_attn_ext_extra_blk(tensor); res += ggml_metal_op_flash_attn_ext_extra_tmp(tensor); res += ggml_metal_op_flash_attn_ext_extra_kv_f16(tensor); + res += ggml_metal_op_flash_attn_ext_extra_idx(tensor); } break; case GGML_OP_CUMSUM: case GGML_OP_ARGSORT: diff --git a/ggml/src/ggml-metal/kernels/fa.metal b/ggml/src/ggml-metal/kernels/fa.metal index e95dec258a3..aeacd5f9046 100644 --- a/ggml/src/ggml-metal/kernels/fa.metal +++ b/ggml/src/ggml-metal/kernels/fa.metal @@ -1071,6 +1071,93 @@ constant int32_t FC_flash_attn_ext_vec_ns10 [[function_constant(FC_FLASH_ATTN_EX constant int32_t FC_flash_attn_ext_vec_ns20 [[function_constant(FC_FLASH_ATTN_EXT_VEC + 21)]]; constant int32_t FC_flash_attn_ext_vec_nsg [[function_constant(FC_FLASH_ATTN_EXT_VEC + 22)]]; constant int32_t FC_flash_attn_ext_vec_nwg [[function_constant(FC_FLASH_ATTN_EXT_VEC + 23)]]; +constant bool FC_flash_attn_ext_vec_has_sparse [[function_constant(FC_FLASH_ATTN_EXT_VEC + 5)]]; + +// compress the finite entries of each KQ mask row into a list of KV indices (ascending order), +// padded with -1 up to n_kv_max_padded (a multiple of OP_FLASH_ATTN_EXT_VEC_NCPSG) +// one threadgroup per mask row; the mask remains the single source of truth for the values +kernel void kernel_flash_attn_ext_vec_idx( + constant ggml_metal_kargs_flash_attn_ext_vec_idx & args, + device const half * mask, + device int * idx, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + constexpr short NW = N_SIMDWIDTH; + + const int i1 = tgpig[0]; + const int i2 = tgpig[1]; + const int i3 = tgpig[2]; + + device const half * pm = (device const half *) (mask + i1*args.nb31 + i2*args.nb32 + i3*args.nb33); + device int * pidx = idx + ((i3*args.ne32 + i2)*args.ne31 + i1)*args.n_kv_max_padded; + + const int n = args.ne30; + const int q = n/ntg.x; + const int r = n%ntg.x; + + // each thread handles a contiguous slice of the mask row + const int r0 = q*tiitg + min((int) tiitg, r); + const int r1 = r0 + q + (tiitg < r ? 1 : 0); + + // count the finite entries in the slice + int cnt = 0; + for (int i = r0; i < r1; ++i) { + cnt += isfinite((float) pm[i]) ? 1 : 0; + } + + const short sgitg = tiitg/NW; + const short tiisg = tiitg%NW; + + threadgroup int tcount[8]; + + // simd_sum is a collective: all lanes must evaluate it + const int sg_sum = simd_sum(cnt); + if (tiisg == 0) { + tcount[sgitg] = sg_sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + int total = 0; + for (short s = 0; s < ntg.x/NW; ++s) { + total += tcount[s]; + } + + // base offset of this thread's slice in the output list (exclusive scan within the simdgroup) + int sg_base = 0; + for (short s = 0; s < sgitg; ++s) { + sg_base += tcount[s]; + } + + // exclusive prefix scan of the per-thread counts within the simdgroup + int incl = cnt; + for (int d = 1; d < NW; d <<= 1) { + const int v = simd_shuffle_up(incl, d); + if (tiisg >= d) { + incl += v; + } + } + const int base = sg_base + (incl - cnt); + + // write the finite positions in order; if the hint is violated, keep only the first n_kv_max entries + int j = 0; + for (int i = r0; i < r1; ++i) { + if (base + j >= args.n_kv_max) { + break; + } + if (isfinite((float) pm[i])) { + pidx[base + j] = i; + j++; + } + } + + // pad the tail of the list with -1 + const int count = min(total, args.n_kv_max); + for (int i = count + tiitg; i < args.n_kv_max_padded; i += ntg.x) { + pidx[i] = -1; + } +} template< typename q4_t, // query types in shared memory @@ -1091,6 +1178,7 @@ template< short NE = 4, // head elements per thread short Q = OP_FLASH_ATTN_EXT_VEC_NQPSG, // queries per threadgroup short C = OP_FLASH_ATTN_EXT_VEC_NCPSG> // cache items per threadgroup + kernel void kernel_flash_attn_ext_vec( constant ggml_metal_kargs_flash_attn_ext_vec & args, device const char * q, @@ -1100,6 +1188,7 @@ kernel void kernel_flash_attn_ext_vec( device const char * sinks, device const char * pad, device char * dst, + device const char * idx, threadgroup half * shmem_f16 [[threadgroup(0)]], uint3 tgpig[[threadgroup_position_in_grid]], ushort tiisg[[thread_index_in_simdgroup]], @@ -1205,7 +1294,15 @@ kernel void kernel_flash_attn_ext_vec( const short ty = tiisg/NL; // pointer to the mask - device const half * pm_base = (device const half *) (mask + iq1*Q*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33); + device const half * pm_base = (device const half *) (mask + (FC_flash_attn_ext_vec_has_sparse ? (iq1%args.ne31) : (iq1*Q))*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33); + + // sparse indices: the list of finite mask entries per query row + // the sparse path requires Q == 1 (enforced by the host) + device const int * pidx = nullptr; + if (FC_flash_attn_ext_vec_has_sparse) { + pidx = (device const int *) (idx + ((iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*args.ne31*args.n_kv_max_padded + + (iq1%args.ne31)*args.n_kv_max_padded); + } float slope = 1.0f; @@ -1265,11 +1362,22 @@ kernel void kernel_flash_attn_ext_vec( } if (FC_flash_attn_ext_vec_has_mask) { - FOR_UNROLL (short qq = 0; qq < Q; ++qq) { - if ((iq1*Q + qq) < args.ne01) { - sm[qq*C + tiisg] = pm[qq][ic + tiisg]; - } else { - sm[qq*C + tiisg] = -MAXHALF; + if (FC_flash_attn_ext_vec_has_sparse) { + FOR_UNROLL (short qq = 0; qq < Q; ++qq) { + const int i11 = pidx[ic + tiisg]; + if ((iq1*Q + qq) < args.ne01 && i11 >= 0) { + sm[qq*C + tiisg] = pm[qq][i11]; + } else { + sm[qq*C + tiisg] = -MAXHALF; + } + } + } else { + FOR_UNROLL (short qq = 0; qq < Q; ++qq) { + if ((iq1*Q + qq) < args.ne01) { + sm[qq*C + tiisg] = pm[qq][ic + tiisg]; + } else { + sm[qq*C + tiisg] = -MAXHALF; + } } } } else { @@ -1294,9 +1402,13 @@ kernel void kernel_flash_attn_ext_vec( // Q*K^T { - device const k4_t * pk4 = (device const k4_t *) (k + ic*args.nb11); + device const k4_t * pk4 = nullptr; - pk4 += ty*NS10/4 + tx; + if (!FC_flash_attn_ext_vec_has_sparse) { + pk4 = (device const k4_t *) (k + ic*args.nb11); + + pk4 += ty*NS10/4 + tx; + } qk_t mqk[Q][C/NE]; FOR_UNROLL (short qq = 0; qq < Q; ++qq) { @@ -1307,7 +1419,35 @@ kernel void kernel_flash_attn_ext_vec( // each simdgroup processes Q queries and NE (NW/NL) cache elements FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { - if (is_same::value) { + 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]; + if (i11 >= 0) { + if (is_same::value) { + device const k4_t * pk4s = (device const k4_t *) (k + i11*args.nb11) + tx; + FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) { + const k4_t k_elem = pk4s[ii*NL]; + FOR_UNROLL (short qq = 0; qq < Q; ++qq) { + mqk[qq][cc] += dot((float4) k_elem, (float4) sq4[qq*PK4 + ii*NL + tx]); + } + } + } else { + device const kd4_t * pk = (device const kd4_t *) (k + i11*args.nb11); + + k4_t mk; + + FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) { + const short i = ii*NL + tx; + + deq_k_t4(pk + i/nl_k, i%nl_k, mk); + + FOR_UNROLL (short qq = 0; qq < Q; ++qq) { + mqk[qq][cc] += dot((float4) mk, (float4) sq4[qq*PK4 + i]); + } + } + } + } + } else if (is_same::value) { FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) { const k4_t k_elem = pk4[cc*NE*NS10/4 + ii*NL]; FOR_UNROLL (short qq = 0; qq < Q; ++qq) { @@ -1422,7 +1562,40 @@ kernel void kernel_flash_attn_ext_vec( } } - if (is_same::value) { + 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]; + if (i11 >= 0) { + if (is_same::value) { + device const v4_t * pv4 = (device const v4_t *) (v + i11*args.nb21); + + pv4 += tx; + + FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { + const v4_t v_elem = pv4[ii*NL]; + FOR_UNROLL (short qq = 0; qq < Q; ++qq) { + lo[qq][ii] += o4_t(float4(v_elem)*float4(ss[qq*C + cc*NE + ty])); + } + } + } else { + device const vd4_t * pv4 = (device const vd4_t *) (v + i11*args.nb21); + + FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { + const short i = ii*NL + tx; + + v4_t mv; + + deq_v_t4(pv4 + i/nl_v, i%nl_v, mv); + + FOR_UNROLL (short qq = 0; qq < Q; ++qq) { + lo[qq][ii] += o4_t(float4(mv)*float4(ss[qq*C + cc*NE + ty])); + } + } + } + } + } + } else if (is_same::value) { device const v4_t * pv4 = (device const v4_t *) (v + ic*args.nb21); pv4 += ty*NS20/4 + tx; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 1b0eaca8f0b..6af51fcf7cf 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -1502,11 +1502,6 @@ struct test_case { double err = ud->tc->err(f1.data(), f2.data(), f1.size()); if (err > ud->tc->max_err(ud->backend1)) { printf("[%s] ERR = %.9f > %.9f ", ggml_op_desc(t1), err, ud->tc->max_err(ud->backend1)); - //for (int i = 0; i < (int) f1.size(); i++) { - // printf("%5d %9.6f %9.6f, diff = %9.6f\n", i, f1[i], f2[i], f1[i] - f2[i]); - //} - //printf("\n"); - //exit(1); ud->ok = false; } return true; @@ -10280,6 +10275,56 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 2}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 768)); test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 2304)); test_cases.emplace_back(new test_flash_attn_ext(256, 256, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512, false)); + + // Sparse mask hint: more head sizes + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(320, 256, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(192, 128, 4, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext( 64, 64, 4, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext( 96, 96, 4, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + + // Sparse mask hint: quantized KV + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_1, GGML_TYPE_Q4_1, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q5_0, GGML_TYPE_Q5_0, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q5_1, GGML_TYPE_Q5_1, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16, {0, 1, 2, 3}, true, false, 512)); + + // Sparse mask hint: no mask (dense fallback, hint ignored) + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + + // Sparse mask hint: with sinks + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, true, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + + // Sparse mask hint: with ALiBi bias + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 1.0f, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + + // Sparse mask hint: with logit softcap + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 50.0f, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + + // Sparse mask hint: larger batch + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 8, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + + // Sparse mask hint: V is a view of K + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 512)); + + // Sparse mask hint: permuted KV layout + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, true, false, 512)); + + // Sparse mask hint: batch > 1 + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, { 8, 1}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + + // Sparse mask hint: nr23 != [8,1] + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 4, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {16, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 2}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 4, 2}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 2, 1}, 4096, 3, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); // more V-is-sub-view-of-K cases: other head shapes, and full views with equal head sizes test_cases.emplace_back(new test_flash_attn_ext(320, 256, 1, {32, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); From 619e14f51b567cf2ca802e4fe655f9e1fbc2c2d5 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Sun, 30 Aug 2026 18:26:26 +0300 Subject: [PATCH 02/10] metal : fix sparse flash attention row addressing - kernel_flash_attn_ext_vec_idx: mask param is half* but nb31 is a byte stride, so the per-row mask offset was scaled by 2x; cast to char* before applying the byte strides - kernel_flash_attn_ext_vec: sparse pidx param is char* so the per-row element offset was under-scaled by sizeof(int); scale it by sizeof(int) to get the correct byte offset - fixes the multi-row (nb*nr23[1] > 1) sparse flash attention failures Assisted-by: pi:llama.cpp/DeepSeek-v4-0731 --- ggml/src/ggml-metal/kernels/fa.metal | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-metal/kernels/fa.metal b/ggml/src/ggml-metal/kernels/fa.metal index aeacd5f9046..89ee51de63b 100644 --- a/ggml/src/ggml-metal/kernels/fa.metal +++ b/ggml/src/ggml-metal/kernels/fa.metal @@ -1089,7 +1089,7 @@ kernel void kernel_flash_attn_ext_vec_idx( const int i2 = tgpig[1]; const int i3 = tgpig[2]; - device const half * pm = (device const half *) (mask + i1*args.nb31 + i2*args.nb32 + i3*args.nb33); + device const half * pm = (device const half *) ((device const char *) mask + i1*args.nb31 + i2*args.nb32 + i3*args.nb33); device int * pidx = idx + ((i3*args.ne32 + i2)*args.ne31 + i1)*args.n_kv_max_padded; const int n = args.ne30; @@ -1294,14 +1294,14 @@ kernel void kernel_flash_attn_ext_vec( const short ty = tiisg/NL; // pointer to the mask - device const half * pm_base = (device const half *) (mask + (FC_flash_attn_ext_vec_has_sparse ? (iq1%args.ne31) : (iq1*Q))*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33); + device const half * pm_base = (device const half *) (mask + iq1*Q*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33); // sparse indices: the list of finite mask entries per query row // the sparse path requires Q == 1 (enforced by the host) device const int * pidx = nullptr; if (FC_flash_attn_ext_vec_has_sparse) { - pidx = (device const int *) (idx + ((iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*args.ne31*args.n_kv_max_padded - + (iq1%args.ne31)*args.n_kv_max_padded); + pidx = (device const int *) (idx + sizeof(int)*(((iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*args.ne31*args.n_kv_max_padded + + (iq1%args.ne31)*args.n_kv_max_padded)); } float slope = 1.0f; From 17f6ff315485d8a67100ca28ff8ca0827ca052fe Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Sun, 30 Aug 2026 18:58:36 +0300 Subject: [PATCH 03/10] cont : use sparse vec FA for prefill --- ggml/src/ggml-metal/ggml-metal-ops.cpp | 27 +++++++++++--------------- tests/test-backend-ops.cpp | 2 +- 2 files changed, 12 insertions(+), 17 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index bf29e6e720e..2d7eb8d3bb4 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2875,11 +2875,6 @@ static int ggml_metal_op_flash_attn_ext_n_kv_max_sparse(const ggml_tensor * op) return 0; } - // the sparse path is implemented for the vec kernels only - if (!ggml_metal_op_flash_attn_ext_use_vec(op)) { - return 0; - } - // bound the size of the index lists if (n_kv_max > 4096) { return 0; @@ -3297,7 +3292,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { } } - if (!ggml_metal_op_flash_attn_ext_use_vec(op)) { + if (!use_sparse && !ggml_metal_op_flash_attn_ext_use_vec(op)) { // half8x8 kernel const int nqptg = OP_FLASH_ATTN_EXT_NQPSG; // queries per threadgroup const int ncpsg = OP_FLASH_ATTN_EXT_NCPSG; // cache values per simdgroup @@ -3579,16 +3574,16 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { // workgroups // each workgroup handles nsg*nkpsg cache values int32_t nwg = 1; - if (false) { - // for small KV caches, we could launch a single workgroup and write the results directly to dst/ - // however, this does not lead to significant improvement, so disabled - nwg = 1; - nsg = 4; - } else if (use_sparse) { - // the reduce kernel sums over one workgroup per lane (32), so nwg must be 32; - // workgroups beyond the number of chunks emit empty partials which the reduce ignores - nsg = 1; - nwg = 32; + if (use_sparse) { + if (ne01 > 32) { + // large sparse batch + nwg = 1; + nsg = 4; + } else { + // small sparse batch + nsg = 1; + nwg = 32; + } } else { nwg = 32; nsg = 1; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 6af51fcf7cf..1151fcb0618 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -191,7 +191,7 @@ static void init_tensor_kq_mask(ggml_tensor * tensor, float min = -1.0f, float m static void init_tensor_kq_mask_sparse(ggml_tensor * tensor, int64_t n_kv_max) { GGML_ASSERT(tensor->type == GGML_TYPE_F16); - GGML_ASSERT(n_kv_max > 1 && n_kv_max <= tensor->ne[0]); + GGML_ASSERT(n_kv_max > 0 && n_kv_max <= tensor->ne[0]); const int64_t ne0 = tensor->ne[0]; const int64_t nrows = ggml_nrows(tensor); From f4280161a4d752f516c9cf63a8d62b26041f2fbf Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Sun, 30 Aug 2026 21:38:46 +0300 Subject: [PATCH 04/10] metal : single-pass flash attention sparse index compaction The idx kernel previously read the mask row twice: once to count the finite entries (for the prefix scan) and again to recover their positions. Since the kernel is memory-bound, this doubled the mask traffic. Keep the finite positions in a per-thread register array during the count pass and write them out directly, avoiding the second mask read. A dense mask with more than NLOCAL finite entries in a slice falls back to re-reading the mask to write the remaining positions. Assisted-by: pi:llama.cpp/DeepSeek-v4-0731 --- ggml/src/ggml-metal/kernels/fa.metal | 39 +++++++++++++++++++++------- 1 file changed, 29 insertions(+), 10 deletions(-) diff --git a/ggml/src/ggml-metal/kernels/fa.metal b/ggml/src/ggml-metal/kernels/fa.metal index 89ee51de63b..2f7aed3204e 100644 --- a/ggml/src/ggml-metal/kernels/fa.metal +++ b/ggml/src/ggml-metal/kernels/fa.metal @@ -1084,6 +1084,7 @@ kernel void kernel_flash_attn_ext_vec_idx( ushort tiitg[[thread_index_in_threadgroup]], ushort3 ntg[[threads_per_threadgroup]]) { constexpr short NW = N_SIMDWIDTH; + constexpr short NLOCAL = 32; // max finite positions kept in registers per thread const int i1 = tgpig[0]; const int i2 = tgpig[1]; @@ -1100,10 +1101,18 @@ kernel void kernel_flash_attn_ext_vec_idx( const int r0 = q*tiitg + min((int) tiitg, r); const int r1 = r0 + q + (tiitg < r ? 1 : 0); - // count the finite entries in the slice - int cnt = 0; + // count the finite entries in the slice and keep their positions in registers (single mask read) + int cnt = 0; // total finite entries in the slice + int nloc = 0; // finite entries kept in registers + int local[NLOCAL]; for (int i = r0; i < r1; ++i) { - cnt += isfinite((float) pm[i]) ? 1 : 0; + if (isfinite((float) pm[i])) { + if (nloc < NLOCAL) { + local[nloc] = i; + nloc++; + } + cnt++; + } } const short sgitg = tiitg/NW; @@ -1142,13 +1151,23 @@ kernel void kernel_flash_attn_ext_vec_idx( // write the finite positions in order; if the hint is violated, keep only the first n_kv_max entries int j = 0; - for (int i = r0; i < r1; ++i) { - if (base + j >= args.n_kv_max) { - break; - } - if (isfinite((float) pm[i])) { - pidx[base + j] = i; - j++; + for (; j < nloc && base + j < args.n_kv_max; ++j) { + pidx[base + j] = local[j]; + } + + // a dense mask may have more than NLOCAL finite entries in a slice; re-read the mask to write the rest + if (cnt > nloc && base + nloc < args.n_kv_max) { + int j2 = 0; + for (int i = r0; i < r1; ++i) { + if (isfinite((float) pm[i])) { + if (j2 >= nloc) { + pidx[base + j2] = i; + } + j2++; + if (base + j2 >= args.n_kv_max) { + break; + } + } } } From 35eee1137d1c2682b75cdb10940e1fdc8be8d2c4 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Sun, 30 Aug 2026 21:39:06 +0300 Subject: [PATCH 05/10] tests : add perf cases for sparse flash attention prefill Measure the sparse vec FA kernel across KV sizes, n_kv_max hints and batch sizes. Run with: ./build/bin/test-backend-ops -b MTL0 -o FLASH_ATTN_EXT -p "n_kv_max=[1-9]" perf Assisted-by: pi:llama.cpp/DeepSeek-v4-0731 --- tests/test-backend-ops.cpp | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 1151fcb0618..e4b8a4fa22c 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10522,6 +10522,16 @@ static std::vector> make_test_cases_eval() { static std::vector> make_test_cases_perf() { std::vector> test_cases; + // Sparse flash attention prefill: long-KV with a sparse mask (n_kv_max hint). + // the vec FA kernel iterates only n_kv_max valid entries per row instead of the full KV. + for (int64_t kv : { 2048, 4096, 8192, 16384 }) { + for (int64_t n_kv_max : { 256, 512, 640, 1024 }) { + for (int64_t nb : { 1, 4, 16, 32 }) { + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 32, {8, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, n_kv_max)); + } + } + } + // SWIGLU at a 27B-class FFN width, fused [gate|up] vs split operands // note: same bytes either way, so a backend that indexes them differently shows it here for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) { From abad554ce2d4d7364d861843a6b0d9430917fc83 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Mon, 31 Aug 2026 18:32:20 +0300 Subject: [PATCH 06/10] qwen4 : enable sparse attention --- src/models/qwen4exp.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 8f0e47b1fef..9dc508b72ba 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -744,7 +744,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = mctx_cur->get_v(ctx0, il); - ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, 0, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, top_k->ne[0], kq_scale, il); cb(cur, "kqv_out", il); // the rotation is its own inverse, so undo it on the value side of the output From 8b2442b36686e0fd53f1150e5cc24abc0cfad7ad Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Mon, 31 Aug 2026 20:34:54 +0300 Subject: [PATCH 07/10] 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 2d7eb8d3bb4..4396bfceb31 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -3578,11 +3578,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 2f7aed3204e..c966c039720 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) { From 8e215df139f1837ef4e55f96f4136617aa6954c3 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 2 Sep 2026 19:48:26 +0300 Subject: [PATCH 08/10] cont : sync test-backend-ops --- tests/test-backend-ops.cpp | 65 +++----------------------------------- 1 file changed, 5 insertions(+), 60 deletions(-) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index e4b8a4fa22c..7165fd824da 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -1502,6 +1502,11 @@ struct test_case { double err = ud->tc->err(f1.data(), f2.data(), f1.size()); if (err > ud->tc->max_err(ud->backend1)) { printf("[%s] ERR = %.9f > %.9f ", ggml_op_desc(t1), err, ud->tc->max_err(ud->backend1)); + //for (int i = 0; i < (int) f1.size(); i++) { + // printf("%5d %9.6f %9.6f, diff = %9.6f\n", i, f1[i], f2[i], f1[i] - f2[i]); + //} + //printf("\n"); + //exit(1); ud->ok = false; } return true; @@ -10275,56 +10280,6 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 2}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 768)); test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 2304)); test_cases.emplace_back(new test_flash_attn_ext(256, 256, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512, false)); - - // Sparse mask hint: more head sizes - test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext(320, 256, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext(192, 128, 4, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext( 64, 64, 4, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext( 96, 96, 4, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - - // Sparse mask hint: quantized KV - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_1, GGML_TYPE_Q4_1, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q5_0, GGML_TYPE_Q5_0, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q5_1, GGML_TYPE_Q5_1, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16, {0, 1, 2, 3}, true, false, 512)); - - // Sparse mask hint: no mask (dense fallback, hint ignored) - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - - // Sparse mask hint: with sinks - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, true, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - - // Sparse mask hint: with ALiBi bias - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 1.0f, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - - // Sparse mask hint: with logit softcap - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 50.0f, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - - // Sparse mask hint: larger batch - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 8, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - - // Sparse mask hint: V is a view of K - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 512)); - - // Sparse mask hint: permuted KV layout - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, true, false, 512)); - - // Sparse mask hint: batch > 1 - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, { 8, 1}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - - // Sparse mask hint: nr23 != [8,1] - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 4, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {16, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 2}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 4, 2}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); - test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 2, 1}, 4096, 3, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); // more V-is-sub-view-of-K cases: other head shapes, and full views with equal head sizes test_cases.emplace_back(new test_flash_attn_ext(320, 256, 1, {32, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); @@ -10522,16 +10477,6 @@ static std::vector> make_test_cases_eval() { static std::vector> make_test_cases_perf() { std::vector> test_cases; - // Sparse flash attention prefill: long-KV with a sparse mask (n_kv_max hint). - // the vec FA kernel iterates only n_kv_max valid entries per row instead of the full KV. - for (int64_t kv : { 2048, 4096, 8192, 16384 }) { - for (int64_t n_kv_max : { 256, 512, 640, 1024 }) { - for (int64_t nb : { 1, 4, 16, 32 }) { - test_cases.emplace_back(new test_flash_attn_ext(128, 128, 32, {8, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, n_kv_max)); - } - } - } - // SWIGLU at a 27B-class FFN width, fused [gate|up] vs split operands // note: same bytes either way, so a backend that indexes them differently shows it here for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) { From 46437c9c511988c46b94bd0b71bd5b0c6ed638ca Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 2 Sep 2026 20:15:16 +0300 Subject: [PATCH 09/10] cont : disable Qwen4 for now --- src/models/qwen4exp.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 9dc508b72ba..31cce900d52 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -744,7 +744,10 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = mctx_cur->get_v(ctx0, il); - ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, top_k->ne[0], kq_scale, il); + // TODO: enable sparse attention when we are ready + // ref: https://github.com/ggml-org/llama.cpp/pull/27970 + //ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, top_k->ne[0], kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, 0, kq_scale, il); cb(cur, "kqv_out", il); // the rotation is its own inverse, so undo it on the value side of the output From 1d818e88154fabb1c2018453feb6b56beba8ea34 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 3 Sep 2026 10:16:34 +0300 Subject: [PATCH 10/10] cont : clean-up + tests --- ggml/src/ggml-metal/ggml-metal-impl.h | 1 - ggml/src/ggml-metal/ggml-metal-ops.cpp | 8 +++++--- ggml/src/ggml-metal/kernels/fa.metal | 6 +++--- tests/test-backend-ops.cpp | 16 ++++++++++++++++ 4 files changed, 24 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index f6d6867c6fd..30e40f527f9 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -107,7 +107,6 @@ #define FC_SUM_ROWS 1400 #define FC_UPSCALE 1500 #define FC_GATED_DELTA_NET 1600 -#define FC_FLASH_ATTN_EXT_VEC_IDX 1700 // op-specific constants #define OP_FLASH_ATTN_EXT_NQPSG 8 diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 4396bfceb31..b93bcfd2e9a 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -3473,7 +3473,9 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { (int) op->src[1]->type, (int) ne00, (int) ne20, // dk, dv (ne00 == dk for FA) ne11, ne01); - int nqptg = cfg.Q; // queries per threadgroup + + int nqptg = cfg.Q; // queries per threadgroup + const int ncpsg = OP_FLASH_ATTN_EXT_VEC_NCPSG; // cache values per simdgroup !! sync with kernel template arguments !! const int nhptg = 1; // heads per threadgroup @@ -3488,6 +3490,8 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { if (use_sparse) { assert(ggml_metal_op_flash_attn_ext_extra_idx(op) != 0); + GGML_ASSERT(ne30 == ne11); + ggml_metal_kargs_flash_attn_ext_vec_idx args0 = { /*.ne30 =*/ ne30, /*.ne31 =*/ ne31, @@ -3666,8 +3670,6 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { GGML_ASSERT(smem <= props_dev->max_theadgroup_memory_size); if (nwg == 1) { - assert(ggml_metal_op_flash_attn_ext_extra_tmp(op) == 0); - // using 1 workgroup -> write the result directly into dst ggml_metal_encoder_set_buffer(enc, bid_pad, 6); ggml_metal_encoder_set_buffer(enc, bid_dst, 7); diff --git a/ggml/src/ggml-metal/kernels/fa.metal b/ggml/src/ggml-metal/kernels/fa.metal index c966c039720..d0e928d732c 100644 --- a/ggml/src/ggml-metal/kernels/fa.metal +++ b/ggml/src/ggml-metal/kernels/fa.metal @@ -1091,7 +1091,7 @@ kernel void kernel_flash_attn_ext_vec_idx( const int i3 = tgpig[2]; device const half * pm = (device const half *) ((device const char *) mask + i1*args.nb31 + i2*args.nb32 + i3*args.nb33); - device int * pidx = idx + ((i3*args.ne32 + i2)*args.ne31 + i1)*args.n_kv_max_padded; + device int * pidx = idx + (((int64_t)i3*args.ne32 + i2)*args.ne31 + i1)*args.n_kv_max_padded; const int n = args.ne30; const int q = n/ntg.x; @@ -1319,8 +1319,8 @@ kernel void kernel_flash_attn_ext_vec( // the sparse path requires Q == 1 (enforced by the host) device const int * pidx = nullptr; if (FC_flash_attn_ext_vec_has_sparse) { - pidx = (device const int *) (idx + sizeof(int)*(((iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*args.ne31*args.n_kv_max_padded - + (iq1%args.ne31)*args.n_kv_max_padded)); + pidx = (device const int *) idx + + ((int64_t)(iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*args.ne31*args.n_kv_max_padded + (iq1%args.ne31)*args.n_kv_max_padded; } float slope = 1.0f; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 7165fd824da..4f2ef34eca7 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10280,6 +10280,22 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 2}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 768)); test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 2304)); test_cases.emplace_back(new test_flash_attn_ext(256, 256, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 1, { 8, 1}, 4096, 4, true, false, 8.0f, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + + // sparse mask with large batch size + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 2048)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 1}, 4096, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 512)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 1}, 4096, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 2048)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 1, { 8, 1}, 4096, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 1, { 8, 1}, 4096, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 1, { 8, 1}, 4096, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 2048)); + + // sparse mask + quantized cache + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, false, 512)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 1, { 8, 1}, 4096, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, false, 512)); // more V-is-sub-view-of-K cases: other head shapes, and full views with equal head sizes test_cases.emplace_back(new test_flash_attn_ext(320, 256, 1, {32, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true));