Skip to content

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

Merged
bri-prism merged 1 commit into
prismfrom
fix/cuda-fa-shortq-quant-kv
Oct 13, 2026
Merged

bri-prism merged 1 commit into
prismfrom
fix/cuda-fa-shortq-quant-kv

Conversation

@bri-prism

Copy link
Copy Markdown
Collaborator

What

One-line fix in ggml/src/ggml-cuda/fattn.cu: the short-query special case added in #189 (quantized V, 3-4 query rows) now passes type_KV through to ggml_cuda_flash_attn_ext_mma_f16_case. Also adds short-query FLASH_ATTN_EXT cases to test-backend-ops (nb 1-8, GQA 2 and 4, f16/q8_0/q4_0 K/V, head size 128 and 256).

Why

With q8_0 for both K and V, 3-4 row batches produced garbage on CUDA. Speculative decoding hits this on every verify step (target token plus 2 drafts), so MTP with -ctk q8_0 -ctv q8_0 gave 0% draft acceptance and corrupted output. Plain decoding, prompt processing and q8_0 on K only were fine. Metal was fine.

The special case called the F16 instantiation of the MMA kernel. That variant converts K/V to f16 into scratch right after dst. For K/V types that have the in-place quantized kernel (q8_0/q4_0 at head size 128/256), ggml_cuda_flash_attn_ext_get_alloc_size reserves no scratch, so the conversion wrote over neighbouring tensors.

Passing type_KV makes these types use the in-place kernel at the same 64/ncols2 tile, which is the instantiation that already exists for them. Other quantized V types still go through the F16 path, and for those the scratch is reserved.

Testing (H200, sm90)

  • New FA cases, kv 256/1024 and nb 1-5 subset: 120/120 pass with the fix. Before the fix the same run hung.
  • llama-server with an MTP draft head, --spec-draft-n-max 2, q8_0 K and V: before the fix, 0 drafts accepted and garbage text. After the fix, acceptance matches the f16 KV run (e.g. 15/16 and 14/14 on short code and JSON prompts), and the output matches plain f16 decoding wherever f16 MTP itself matches.

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).
@bri-prism

Copy link
Copy Markdown
Collaborator Author

Full new-case run at 1346e6d on H200: 216/216 FLASH_ATTN_EXT short-query cases pass (kv 256/1024/4096, nb 1-5 and 8, GQA 2/4, f16/q8_0/q4_0, hs 128/256).

@bri-prism
bri-prism merged commit 9f24523 into prism Oct 13, 2026
8 checks passed
sb32445 added a commit to sb32445/llama.cpp that referenced this pull request Oct 15, 2026
fattn.cu: take prism's PrismML-Eng#355 (pass type_KV so the in-place kernel is used);
it supersedes the ktrain hotfix "only for the f16 K/V kernel".
ggml-cuda.cu: keep both includes (l2-hint.cuh and moe-reduce.cuh).

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants