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:
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
+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));