From e48c865d147ff42963ca7b7e82aa07478a2363d9 Mon Sep 17 00:00:00 2001
From: Cary Palmer <24235924+professorpalmer@users.noreply.github.com>
Date: Sun, 11 Oct 2026 07:49:34 -0500
Subject: [PATCH] cuda: reserve the f16 K/V scratch for the short-query MMA
 route (#346)

With a quantized V and 3-4 queries, ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1
takes the 64-column f16 tiles (#189), and launch_fattn converts K and V to f16 in
the op's extra buffer. ggml_cuda_flash_attn_ext_get_alloc_size reserved nothing
beyond dst when the in-place quantized K/V kernel exists for the head size, so
that conversion wrote past the op's allocation. Server prompt checkpoints
evaluate the last 4 prompt tokens as their own batch, which reaches this route.

- ggml_cuda_fattn_mma_short_query_f16(): the route's condition, used by the
  dispatch and by get_alloc_size.
- get_alloc_size reserves the f16 K/V for that route also on the in-place path.
- The MMA dispatch asserts that the buffer reserved the f16 K/V the route
  writes, so a mismatch aborts with a message instead of corrupting memory.
---
 ggml/src/ggml-cuda/fattn.cu | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu
index 629a2e868117..5cfaafd55d9a 100644
--- a/ggml/src/ggml-cuda/fattn.cu
+++ b/ggml/src/ggml-cuda/fattn.cu
@@ -5,13 +5,22 @@
 #include "fattn-vec.cuh"
 #include "fattn.cuh"
 
+// 3-4 queries with a quantized V take the 64-column f16 tiles (#189): launch_fattn converts K and V to f16 in the
+// op's extra buffer, also when an in-place quantized K/V kernel exists for the head size. ggml_cuda_flash_attn_ext_get_alloc_size
+// must reserve that buffer for these ops.
+static bool ggml_cuda_fattn_mma_short_query_f16(const int cc, const ggml_tensor * dst) {
+    const ggml_tensor * Q = dst->src[0];
+    const ggml_tensor * V = dst->src[2];
+    return turing_mma_available(cc) && ggml_is_quantized(V->type) && Q->ne[1] > 2 && Q->ne[1] <= 4;
+}
+
 template <int DKQ, int DV, int ncols2, ggml_type type_KV = GGML_TYPE_F16>
 static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
     const ggml_tensor * Q = dst->src[0];
     const ggml_tensor * V = dst->src[2];
 
-    if (turing_mma_available(cc) && ggml_is_quantized(V->type) && Q->ne[1] > 2 && Q->ne[1] <= 4) {
+    if (ggml_cuda_fattn_mma_short_query_f16(cc, dst)) {
         ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 64/ncols2, ncols2>(ctx, dst);
         return;
     }
@@ -579,7 +588,8 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d
 
     switch (kernel) {
         case BEST_FATTN_KERNEL_MMA_F16:
-            if (ggml_cuda_fattn_mma_kv_native_supported(dst)) {
+            if (ggml_cuda_fattn_mma_kv_native_supported(dst) &&
+                    !ggml_cuda_fattn_mma_short_query_f16(ggml_cuda_info().devices[device].cc, dst)) {
                 // In-place quantized K/V kernel: nothing to reserve beyond dst.
                 break;
             }
@@ -617,6 +627,12 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst
             ggml_cuda_flash_attn_ext_vec(ctx, dst);
             break;
         case BEST_FATTN_KERNEL_MMA_F16:
+            if (dst->buffer && ggml_cuda_fattn_mma_short_query_f16(ggml_cuda_info().devices[ggml_cuda_get_device()].cc, dst)) {
+                // the f16 K/V this route writes after dst must fit in what the buffer reserved for dst
+                const ggml_cuda_flash_attn_ext_f16_extra_data need = ggml_cuda_flash_attn_ext_get_f16_extra_data(dst, true, true);
+                GGML_ASSERT(need.end - (uintptr_t) dst->data <= ggml_backend_buffer_get_alloc_size(dst->buffer, dst) &&
+                    "flash attention: f16 K/V scratch not reserved for the short-query route");
+            }
             ggml_cuda_flash_attn_ext_mma_f16(ctx, dst);
             break;
     }