From c31b682bd11247ad645e2dc85d606d0a888d2c23 Mon Sep 17 00:00:00 2001
From: Georgi Gerganov <ggerganov@gmail.com>
Date: Sun, 30 Aug 2026 18:26:26 +0300
Subject: [PATCH] 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
---
 ggml/src/ggml-metal/kernels/fa.metal | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/ggml/src/ggml-metal/kernels/fa.metal b/ggml/src/ggml-metal/kernels/fa.metal
index aeacd5f904..89ee51de63 100644
--- a/ggml/src/ggml-metal/kernels/fa.metal
+++ b/ggml/src/ggml-metal/kernels/fa.metal
@@ -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;