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

metal : single-pass flash attention sparse index compaction

Browse Source 
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
This commit is contained in:
Georgi Gerganov
2026-08-30 21:38:46 +03:00
parent 75d13dfd5c
commit 5c755617f7
1 changed files with 29 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-metal/kernels/fa.metal
+29 -10
View File
@@ -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;
}
}
}
}