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

metal : add sparse FA (#28098)

Browse Source 
* 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

* 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

* cont : use sparse vec FA for prefill

* 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

* 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

* qwen4 : enable sparse attention

* cont : adjust nsg

* cont : sync test-backend-ops

* cont : disable Qwen4 for now

* cont : clean-up + tests
This commit is contained in:
Georgi Gerganov
2026-09-03 13:51:13 +03:00
committed by GitHub
parent 0df017d6dd
commit 7bb0fc18f6
9 changed files with 426 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-metal/ggml-metal-device.cpp
+25 -2
View File
@@ -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);
ggml/src/ggml-metal/ggml-metal-device.h
+5
View File
@@ -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,
ggml/src/ggml-metal/ggml-metal-impl.h
+13
View File
@@ -458,8 +458,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;
ggml/src/ggml-metal/ggml-metal-ops.cpp
+157 -19
View File
@@ -2857,6 +2857,65 @@ 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;
}
// 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 +3086,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 +3181,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;
@@ -3206,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
@@ -3378,13 +3464,18 @@ 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);
int nqptg = cfg.Q; // queries per threadgroup
// 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 +3485,39 @@ 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_ASSERT(ne30 == ne11);
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);
@@ -3455,11 +3578,26 @@ 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;
if (use_sparse) {
if (ne01 > 32) {
// large sparse batch
nwg = 1;
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
nwg = 32;
nsg = 1;
while (2*nwg*nsg*ncpsg < n_kv_max_padded && nsg < 4) {
nsg *= 2;
}
}
} else {
nwg = 32;
nsg = 1;
@@ -3484,7 +3622,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 +3648,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 +3662,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);
@@ -3530,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);
ggml/src/ggml-metal/ggml-metal-ops.h
+1
View File
@@ -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);
ggml/src/ggml-metal/ggml-metal.cpp
+1
View File
@@ -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:
ggml/src/ggml-metal/kernels/fa.metal
+204 -11
View File
@@ -1071,6 +1071,112 @@ 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;
constexpr short NLOCAL = 32; // max finite positions kept in registers per thread
const int i1 = tgpig[0];
const int i2 = tgpig[1];
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 + (((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;
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 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) {
if (isfinite((float) pm[i])) {
if (nloc < NLOCAL) {
local[nloc] = i;
nloc++;
}
cnt++;
}
}
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 (; 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;
}
}
}
}
// 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 +1197,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 +1207,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]],
@@ -1137,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
@@ -1207,6 +1315,14 @@ kernel void kernel_flash_attn_ext_vec(
// 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);
// 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 +
((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;
// ALiBi
@@ -1265,11 +1381,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 {
@@ -1280,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) {
@@ -1294,9 +1422,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 +1439,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<kd4_t, k4_t>::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<kd4_t, k4_t>::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<kd4_t, k4_t>::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 +1582,40 @@ kernel void kernel_flash_attn_ext_vec(
}
}
if (is_same<vd4_t, v4_t>::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<vd4_t, v4_t>::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<vd4_t, v4_t>::value) {
device const v4_t * pv4 = (device const v4_t *) (v + ic*args.nb21);
pv4 += ty*NS20/4 + tx;
src/models/qwen4exp.cpp
+3
View File
@@ -744,6 +744,9 @@ 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);
// 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);
tests/test-backend-ops.cpp
+17 -1
View File
@@ -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);
@@ -10295,6 +10295,22 @@ static std::vector<std::unique_ptr<test_case>> 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));