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

cont : use sparse vec FA for prefill

Browse Source
This commit is contained in:
Georgi Gerganov
2026-08-31 21:32:14 +03:00
parent e7b7c42c91
commit 75d13dfd5c
2 changed files with 12 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-metal/ggml-metal-ops.cpp
+11 -16
View File
@@ -2859,11 +2859,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;
@@ -3281,7 +3276,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
@@ -3563,16 +3558,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;
tests/test-backend-ops.cpp
+1 -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);