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 : fix sparse flash attention row addressing

Browse Source 
- 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
This commit is contained in:
Georgi Gerganov
2026-08-30 18:26:26 +03:00
parent 39a0be844a
commit c31b682bd1
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-metal/kernels/fa.metal
+4 -4
View File
@@ -1089,7 +1089,7 @@ kernel void kernel_flash_attn_ext_vec_idx(
const int i2 = tgpig[1];
const int i3 = tgpig[2];
device const half * pm = (device const half *) (mask + i1*args.nb31 + i2*args.nb32 + i3*args.nb33);
device const half * pm = (device const half *) ((device const char *) mask + i1*args.nb31 + i2*args.nb32 + i3*args.nb33);
device int * pidx = idx + ((i3*args.ne32 + i2)*args.ne31 + i1)*args.n_kv_max_padded;
const int n = args.ne30;
@@ -1294,14 +1294,14 @@ kernel void kernel_flash_attn_ext_vec(
const short ty = tiisg/NL;
// pointer to the mask
device const half * pm_base = (device const half *) (mask + (FC_flash_attn_ext_vec_has_sparse ? (iq1%args.ne31) : (iq1*Q))*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33);
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 + ((iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*args.ne31*args.n_kv_max_padded
+ (iq1%args.ne31)*args.n_kv_max_padded);
pidx = (device const int *) (idx + sizeof(int)*(((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;