cuda: avoid short-query MMA FA crash with quantized V

[Wyzix33]

Overview

Fix a CUDA illegal memory access in short-query MMA FlashAttention when using a quantized V cache.

For MMA-capable GPUs, Q lengths 3–4 with quantized V are now routed to the existing ncols=64 specialization instead of the crashing ncols=32 path.

Reproducer

Tested on:

RTX 4080 / SM89
CUDA 12.4
Windows / MSVC
Bonsai 2 27B PTQ1_0
FlashAttention enabled

Failing shape:

Q length: 4
K: f16
V: q8_0
24 Q heads / 4 KV heads
KV length: 1280

The ncols=32 MMA specialization produced a CUDA illegal memory access in launch_fattn. The same cache worked with the ncols=64 specialization.

Validation

After the change:

-fa on -ctk q8_0 -ctv q5_1
GSM8K: 1319/1319 completed
no CUDA crashes
no repeated-slash corruption
clean-build smoke test: 50/50 completed

This is a narrow dispatch fallback, not a low-level fix for the ncols=32 kernel.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure:

[Copilot]

🟡 Changes recommended

The new branch selects an unsupported 64-column specialization on Turing GPUs.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adjusts CUDA FlashAttention dispatch to avoid the failing 32-column MMA path for short queries with quantized V caches.

Changes:

• Detects quantized V caches with query lengths 3-4.
• Routes matching workloads to the 64-column specialization.

File summaries

File	Description
ggml/src/ggml-cuda/fattn.cu	Adds the short-query MMA dispatch fallback.

Review details

• Files reviewed: 1/1 changed files
• Comments generated: 1
• Review effort level: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[bri-prism]

Agent review: posted by the maintainer's coding agent at their request.

No findings in this source pass. The fallback is narrowly scoped to quantized V and query lengths 3-4, using the existing larger MMA specialization.

CUDA compilation and device execution were not performed here. Before merging, please retain boundary coverage at query lengths 2, 3, 4 and 5 with quantized and unquantized V, including the applicable grouped-query specializations. The reported model evaluation is useful evidence but was not rerun in this review.

Reviewed commit: 19836568c4f1b504d8ee579b7ab7fcf4edbd8f78.

[Wyzix33]

Validation was performed on build 1983656 using an RTX 4080 (SM89) with FlashAttention enabled.

Test configuration:

Model: Ternary-Bonsai-2-27B-PTQ1_0.gguf (Qwen3.5 27B)
Attention layout: 24 query heads / 4 KV heads (GQA 6:1)
Prefilled context: 1280 tokens
Test matrix: Q = {2, 3, 4, 5} × V = {f16, q8_0, q5_1}

All 12 configurations completed successfully with no CUDA errors.

This validates the quantized-V fallback path for Q = 3–4, as well as the boundary cases Q = 2 and Q = 5.

The Turing-specific issue remains unresolved, since testing was performed only on Ada / RTX 4080 hardware.

[bri-prism]

Agent benchmark follow-up, posted at the maintainer's request.

Pinned head 1983656 against merge-base 1a07bfa. Same public PTQ1 model on both arms, full GPU offload, flash attention, q4_0 K/V, batch/microbatch 512, 8 CPU threads. Three alternating baseline/candidate pairs, three repetitions per invocation, short/empty starting context.

GPU	pp512 before → after, tok/s	Paired change	tg128 before → after, tok/s	Paired change
RTX 3090	750.97 → 749.47	-0.2%	59.65 → 59.53	-0.2%
RTX 4090	1566.08 → 1568.74	+0.2%	90.18 → 90.41	+0.2%
H100 SXM	1225.54 → 1229.74	+0.3%	88.48 → 88.50	+0.0%
RTX 5090	1874.75 → 1877.34	+0.1%	119.58 → 119.61	+0.0%

Selected CPU-reference backend checks passed on the compared arms.

The short-query quantized-V attention checks and small-batch timings were included. No crash was observed in the selected completed cases; this does not exhaust all attention shapes.

The percentages describe these paired runs; small changes should not be interpreted as established improvements. No long-context, multi-slot serving, or end-to-end logit-parity claim is made.

PQ2 control measurements (paired pp512 / tg128 changes):

• RTX 3090: -0.7% / -0.1%.
• RTX 4090: -0.4% / +0.4%.
• H100 SXM: -0.2% / -0.0%.
• RTX 5090: -0.8% / +0.0%.

[sb32445]

I tried this on an RTX 4070 (Ada, cc 8.9) with Ternary-Bonsai-2-27B and found two problems on the MTP decode path. I could not reproduce the original crash, so this is only about the side effects.

Setup: the patch applied cleanly on current prism (2459f68; the PR is based on 1a07bfa). Default CUDA build (no GGML_CUDA_FA_ALL_QUANTS), CUDA 13.4, Release, -DCMAKE_CUDA_ARCHITECTURES=89. A = prism, B = prism + this patch, same other code.

Speed. Ternary-Bonsai-2-27B PTQ1_0 + MTP head, --spec-type draft-mtp --spec-draft-n-max 2, q4_0 K/V, context 16384, 4 greedy prompts, 8 interleaved A/B pairs:

• A 99.7 tok/s, B 85.5 tok/s (median), paired difference -14.2 % (95 % CI [-14.65, -13.83] %).
• Draft acceptance 71.6 % -> 54.9 %, and the outputs are no longer identical.
• A/A control (same binary on both sides, 8 pairs): +0.84 %, CI [-0.05, +2.53] %, outputs identical, acceptance 71.6 % on both sides.

The MTP verify step has 3 queries, so with q4_0 V it takes the new Q->ne[1] > 2 && <= 4 branch. The changed acceptance suggests the results differ from the old path, not only that it is slower.

test-backend-ops. test-backend-ops test -b CUDA0 -o FLASH_ATTN_EXT: A passes 2994/2994. B aborts with CUDA error: an illegal memory access was encountered at FLASH_ATTN_EXT(hsk=64,hsv=64,nh=4,nr23=[1,1],kv=512,nb=75,...,type_K=q8_0,type_V=q8_0,permute=[0,1,2,3]). That case passes when run alone with -p, so the failing call is probably an earlier one; I did not narrow it down. It aborts the same way with GGML_CUDA_DISABLE_GRAPHS=1 CUDA_LAUNCH_BLOCKING=1.

What I could not check. The crash this PR fixes (K f16, V quantized, Q 3-4) needs mixed K/V types. With the default build those fall back to the CPU (see #267), so I could not trigger it. With matching types (q8_0/q8_0, q4_0/q4_0), llama-batched-bench -npp 1280 -b N -ub N for N = 2..5 does not crash on prism. Only an RTX 4070 was tested.

The measurements and scripts were run with Claude Code; I can post the exact commands if useful.

[sb32445]

@bri-prism @Wyzix33 after this merged, test-backend-ops -o FLASH_ATTN_EXT aborts with an illegal memory access on an RTX 4070 (sm_89). The plain prism head reproduces it, none of our patches are involved.

test-backend-ops test -b CUDA0 -o FLASH_ATTN_EXT
FLASH_ATTN_EXT(hsk=64,hsv=64,nh=4,nr23=[1,1],kv=512,nb=75,mask=1,sinks=0,max_bias=0,logit_softcap=0,prec=f32,type_K=q8_0,type_V=q8_0,permute=[0,1,2,3],kv_view=0,v_is_view_of_k=0): CUDA error: an illegal memory access was encountered

I built three commits (CUDA arch 89, Release, same flags):

commit	FLASH_ATTN_EXT
c53526a09 (#252, the parent of this PR)	2994/2994 passed
0d1825501 (this PR, #189)	aborts at the case above
8b0c19c8c (prism head, with #317)	aborts at the same case

So it comes in with this commit, and #317 does not fix it. The cases before it in the list (same shape, permute=[0,2,1,3]) are OK. The failing case passes when run alone (-p), and so do the subsets I tried (all q8_0 cases with hsk=64: 229/229, all hsk=64: 1137/1137, plus hsk=40 and the CUDA-graph cases at the start of the list), also with CUDA_LAUNCH_BLOCKING=1 it is always this case. It only aborts in the full FLASH_ATTN_EXT run, three times in a row, so the trigger is an earlier case or the allocation state, not this shape itself. I have not found the cause yet and have not tested other GPUs.

The diff itself is six lines in ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1 (an early return for quantized V with Q->ne[1] of 3 or 4). It calls ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 64/ncols2, ncols2> without type_KV, while the other calls in the function pass it. I do not know whether that matters here, since the quantized K/V can already be converted to f16 at this point.

In the #317 thread the plan was to land #317 and close #189 in its favour, because #189 also hit an illegal memory access in testing. Since #189 was merged anyway, it may need a revert or a fix.

The checks were run with Claude Code (AI assistance).

[sb32445]

@bri-prism @Wyzix33 follow-up: the cause is in the shortcut added here, and I opened #337 with a one line fix.

The shortcut calls the 64 column MMA kernel without type_KV, so it always builds the f16 variant. For head size 128 or 256 with q4_0/q8_0 K/V the in-place quantized kernel is used and no f16 scratch is reserved behind dst; with 3 or 4 queries K and V are then converted into memory behind the reserved range. compute-sanitizer reports the invalid write in dequantize_block<32,2,&dequantize_q4_0,__half> at hsk=256 nq=3 K=q4_0 V=q4_0 nh=24 (GQA 6), which is the Bonsai MTP verify shape. The abort I first posted (the q8_0 case at nb=75) was only where the corrupted memory showed up.

The effect is not only the test abort: on Ternary-Bonsai-2-27B with MTP n-max 2 on an RTX 4070, #189 gives 94.14 tok/s and a draft acceptance of 54.1 %; with #337 111.16 tok/s (+18.1 %, 8 interleaved pairs) and 71.6 %, the same as before #189. Outputs differ with #189 and are identical to the pre-#189 build with #337.

The fix and the checks were made with Claude Code (AI assistance); details are in #337.

[sb32445]

For reference: the same one line change was found independently by Leeleit on an RTX 5090 (sm_120, CUDA 13.4): Leeleit@42000f6

Their report, with Ternary-Bonsai-2-27B-PQ2_0 and q8_0 K/V (llama-perplexity, -c 512 -b 512 --chunks 6 -fa on): at -ub 512 both builds give 5.1364. At -ub 4 the unpatched build gives 14.96 and then nan, at -ub 3 nan; with the condition the values are 5.1365 and 5.1363. With an MTP head (--spec-draft-n-max 2, q8_0 K/V) the generated text was loops and garbage at a draft acceptance of about 0.2 before and clean afterwards.

That adds a second GPU generation and a quality metric to my RTX 4070 numbers (MTP decode 94.14 -> 111.16 tok/s, acceptance 54.1 -> 71.6 %). I did not reproduce their numbers myself.