cuda : take the short-query MMA FA shortcut only for the f16 K/V kernel

[sb32445]

Overview

Follow-up to #189. Since #189, test-backend-ops -o FLASH_ATTN_EXT aborts with an illegal memory access on prism (RTX 4070, sm_89), and on the MTP decode path of Ternary-Bonsai-2-27B the output changes and the draft acceptance drops from 71.6 % to 54.1 %.

The shortcut in ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1 (quantized V, 3 or 4 queries) calls the 64 column kernel without type_KV, so it always builds the f16 variant. For head size 128 or 256 with q4_0/q8_0 K/V, ggml_cuda_fattn_mma_kv_native_supported picks the in-place quantized kernel, and ggml_cuda_flash_attn_ext_get_f16_extra_data reserves no f16 scratch behind dst for it. With 3 or 4 queries the shortcut then converts K and V into memory behind the reserved range.

The change adds type_KV == GGML_TYPE_F16 to the condition. The in-place path takes the normal tile choice again, as before #189. The shortcut stays for the f16 kernel, where V is quantized but already converted.

Additional information

How it was found: compute-sanitizer --tool memcheck --report-api-errors no over the full FLASH_ATTN_EXT run reports Invalid __global__ write of size 2 bytes in dequantize_block<32,2,&dequantize_q4_0,__half>, 137 KB behind the compute buffer. A debug print at the entry of ggml_cuda_flash_attn_ext (not in this PR) with CUDA_LAUNCH_BLOCKING=1 pinned it to the case hsk=256 nq=3 kv=1024 nh=24 K=q4_0 V=q4_0 (the Bonsai shape: GQA 6, MTP with 3 queries). The abort in the test log shows up later, at a different case, because the write corrupts the neighbouring memory silently. See my comment on #189 for the first findings.

Checks (RTX 4070, CUDA arch 89, Release):

• test-backend-ops -o FLASH_ATTN_EXT: 4007/4007 on the prism head with this change. On the cuda: avoid short-query MMA FA crash with quantized V #189 commit with the change: 2994/2994, and compute-sanitizer reports 0 errors over the full run. The parent of cuda: avoid short-query MMA FA crash with quantized V #189 also passes 2994/2994.
• MTP decode (Ternary-Bonsai-2-27B PTQ1_0 + MTP Q8_0, q4_0 K/V, n-max 2, 114688 context, greedy, 4 prompts, 8 interleaved A/B pairs): prism with cuda: avoid short-query MMA FA crash with quantized V #189 94.14 tok/s, with this change 111.16 tok/s, +18.1 % (95 % CI [+18.0, +18.2] %), outputs differ. Against my build from before cuda: avoid short-query MMA FA crash with quantized V #189 (6 pairs): 111.22 against 111.10 tok/s (-0.09 %), outputs identical, acceptance 71.6 % in both.
• Not checked: other GPUs, and the crash that cuda: avoid short-query MMA FA crash with quantized V #189 was written for (I could not reproduce it), so I cannot say whether it is still covered.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES. The fix was written and checked with Claude Code (AI): finding the cause with compute-sanitizer, the one line change, the tests and the measurements. I review the change and can explain it: the f16 shortcut must not be taken while the kernel reads quantized K/V in place, because no f16 copy is reserved there.

[ZuluTwoThree]

Independent confirmation on two GPUs that the PR description lists as not checked: RTX 5080 (sm_120, CUDA 13.2) and RTX 3090 (sm_86, CUDA 12.0). This PR fixes a real-world llama-server failure I hit with Ternary-Bonsai-2-27B.

Symptom without this PR (RTX 5080, prism 8b0c19c8c)

llama-server -m Ternary-Bonsai-2-27B-PQ2_0.gguf -ngl 99 -fa on -ctk q4_0 -ctv q4_0 -c 262144 -b 2048 --parallel 1, temperature 0, a fresh server for every run, one request each:

-ub	prompt tokens	result
1024	6014	correct
1024	7015 / 8003 / 8601	output is a single repeated /, also in reasoning_content (8003 reproduced 3/3)
1024	19928 / 29860	correct
512	9160 / 19928	correct
512	30000 / 60000 / 100000	server aborts (SIGABRT); captured for 30000: CUDA error: an illegal memory access was encountered
256	9160	single repeated /

The abort is reported in rms_norm_mul_f32_cuda (from ggml_cuda_op_rms_norm_fused), but only after prompt processing had reached 29856 of 29860 tokens. So the last step was a 4-token tail, which matches the 3-to-4-query shortcut described here. The RMS norm kernel is only where the earlier out-of-bounds write is noticed. This also explains the confusing pattern in the table: whether the server's final tail ends up with 3 or 4 tokens depends on the prompt length and -ub. llama-cli with identical settings never showed the problem, because it doesn't split off that tail.

What did not help: GGML_CUDA_DISABLE_FUSION=1, GGML_CUDA_FWHT_FUSION=0, GGML_CUDA_GDN_GATHER_FUSION=0 (each tested against both failing cases). -fa off was clean at 9160 tokens (-ub 1024).

With this PR (prism 2a42998c5 + this diff)

Same flags. One continuously running server, prompts of increasing length, each with two facts hidden at 10 % and 90 % depth and a question about both at the end:

prompt tokens	RTX 5080 -ub 1024	RTX 3090 -ub 1024
2039 to 9150 (incl. 7018 / 8006 / 8604)	2/2	2/2
29716	2/2	2/2
99370	2/2	2/2
196701	2/2	2/2
238405	2/2 (with -ub 512)	2/2

All 13 lengths are correct on both cards, and the server is still alive afterwards. On the 5080 I also reran with -ub 512, including the earlier crashing 30k case: correct.

Caveat: I reproduced the failure only on the 5080. The 3090 was built with this PR from the start, so it shows the fix works there, not that the bug occurs there without it.

Builds: -DGGML_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=120a (5080) and =86 (3090), default GGML_CUDA_FA_QUANTS.

AI usage disclosure: the investigation and this comment were prepared with Claude Code (AI). All measurements were run on my own machines.