cuda: keep quantized K/V type in the short-query MMA FA path

ggml/src/ggml-cuda/fattn.cu

@@ -11,8 +11,10 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con
     const ggml_tensor * Q = dst->src[0];
     const ggml_tensor * V = dst->src[2];
 
+    // keep type_KV: the F16 instantiation converts K/V to f16 into scratch that
+    // ggml_cuda_flash_attn_ext_get_alloc_size does not reserve for the in-place quantized path
     if (turing_mma_available(cc) && ggml_is_quantized(V->type) && Q->ne[1] > 2 && Q->ne[1] <= 4) {
-        ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 64/ncols2, ncols2>(ctx, dst);
+        ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 64/ncols2, ncols2, type_KV>(ctx, dst);
         return;
     }
 

tests/test-backend-ops.cpp

@@ -10255,6 +10255,20 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         }
     }
 
+    // short-query cases (nb 2..8): speculative-decoding verify batches (target token + a few drafts)
+    // with quantized K/V and GQA, e.g. Qwen3.5 hsk=256 with 4 query heads per KV head
+    for (int kv : { 256, 1024, 4096, }) {
+        for (int hs : { 128, 256, }) {
+            for (int nb : { 1, 2, 3, 4, 5, 8, }) {
+                for (std::array<int64_t, 2> nr23 : { std::array<int64_t, 2>{4, 1}, std::array<int64_t, 2>{2, 1}, }) {
+                    for (ggml_type type_KV : { GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, }) {
+                        test_cases.emplace_back(new test_flash_attn_ext(hs, hs, 4, nr23, kv, nb, true, false, 0, 0, GGML_PREC_F32, type_KV, type_KV));
+                    }
+                }
+            }
+        }
+    }
+
     for (int hsk : { 40, 64, 72, 80, 96, 128, 192, 256, 320, 512, 576 }) {
         for (int hsv : { 40, 64, 72, 80, 96, 128, 192, 256, 512 }) {
             if (hsk != 192 && hsk != 320 && hsk != 576 && hsk != hsv) continue;