metal : add sparse FA (#28098)
* 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:
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user