From 75d13dfd5c75659ce3874df4edfb781be8520b7d Mon Sep 17 00:00:00 2001
From: Georgi Gerganov <ggerganov@gmail.com>
Date: Sun, 30 Aug 2026 18:58:36 +0300
Subject: [PATCH] cont : use sparse vec FA for prefill

---
 ggml/src/ggml-metal/ggml-metal-ops.cpp | 27 +++++++++++---------------
 tests/test-backend-ops.cpp             |  2 +-
 2 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp
index cb86862e7a..bdec057334 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp
@@ -2859,11 +2859,6 @@ static int ggml_metal_op_flash_attn_ext_n_kv_max_sparse(const ggml_tensor * op)
         return 0;
     }
 
-    // the sparse path is implemented for the vec kernels only
-    if (!ggml_metal_op_flash_attn_ext_use_vec(op)) {
-        return 0;
-    }
-
     // bound the size of the index lists
     if (n_kv_max > 4096) {
         return 0;
@@ -3281,7 +3276,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
         }
     }
 
-    if (!ggml_metal_op_flash_attn_ext_use_vec(op)) {
+    if (!use_sparse && !ggml_metal_op_flash_attn_ext_use_vec(op)) {
         // half8x8 kernel
         const int nqptg = OP_FLASH_ATTN_EXT_NQPSG; // queries per threadgroup
         const int ncpsg = OP_FLASH_ATTN_EXT_NCPSG; // cache values per simdgroup
@@ -3563,16 +3558,16 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
         // workgroups
         // each workgroup handles nsg*nkpsg cache values
         int32_t nwg = 1;
-        if (false) {
-            // for small KV caches, we could launch a single workgroup and write the results directly to dst/
-            // however, this does not lead to significant improvement, so disabled
-            nwg = 1;
-            nsg = 4;
-        } else if (use_sparse) {
-            // the reduce kernel sums over one workgroup per lane (32), so nwg must be 32;
-            // workgroups beyond the number of chunks emit empty partials which the reduce ignores
-            nsg = 1;
-            nwg = 32;
+        if (use_sparse) {
+            if (ne01 > 32) {
+                // large sparse batch
+                nwg = 1;
+                nsg = 4;
+            } else {
+                // small sparse batch
+                nsg = 1;
+                nwg = 32;
+            }
         } else {
             nwg = 32;
             nsg = 1;
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 6916913f93..06b7394867 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -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);