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;