From 1346e6d9691518f94cbb4e3bb52ab475c1afc0db Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Tue, 13 Oct 2026 12:22:39 -0700 Subject: [PATCH] cuda: keep quantized K/V type in the short-query MMA FA path The quantized-V special case for 3-4 query rows (#189) called the F16 instantiation of the MMA kernel. For K/V types that have an in-place quantized kernel (q8_0/q4_0, head size 128/256), the F16 variant converts K/V to f16 into scratch past dst, but get_alloc_size reserves none for that path, so the conversion overwrote neighbouring tensors. With q8_0 K and V this corrupted 3-4 row batches such as speculative-decoding verify. Pass type_KV so the in-place kernel is used, and add short-query FA tests (nb 1-8, GQA 2/4, f16/q8_0/q4_0, head size 128/256). --- ggml/src/ggml-cuda/fattn.cu | 4 +++- tests/test-backend-ops.cpp | 14 ++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 629a2e868117..32f6440e667d 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/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(ctx, dst); + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); return; } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 81cb1c4127f6..6f6d7b4b5538 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10255,6 +10255,20 @@ static std::vector> 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 nr23 : { std::array{4, 1}, std::array{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;