cuda: reserve the f16 K/V scratch for the short-query MMA route (#346)

[professorpalmer]

Overview

Fixes #346: an out-of-bounds write in launch_fattn with quantized K/V when prompt checkpoints are on.

Cause. With a quantized V and 3-4 queries, ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1 takes the 64-column f16 tiles (#189). That route calls ggml_cuda_flash_attn_ext_mma_f16_case without type_KV, so launch_fattn converts K and V to f16 in the op's extra buffer, also when the in-place quantized K/V kernel exists for the head size (q4_0 / q8_0, D = 128 / 256, from #221). ggml_cuda_flash_attn_ext_get_alloc_size reserves nothing beyond dst on that in-place path, so the f16 K/V is written past the op's allocation. Server prompt checkpoints evaluate the last 4 prompt tokens as their own batch, which reaches this route; with --ctx-checkpoints 0 it is not reached. This matches the memcheck trace in #346 (dequantize_block_q4_0 from launch_fattn, ggml_cuda_flash_attn_ext_mma_f16_case<256, 256, 8, 8, F16>).

Change (ggml/src/ggml-cuda/fattn.cu only):

• ggml_cuda_fattn_mma_short_query_f16(): the route's condition, now used by the dispatch and by get_alloc_size.
• get_alloc_size reserves the f16 K/V for that route also on the in-place path.
• The MMA dispatch asserts that the buffer reserved the f16 K/V the route will write. A mismatch now aborts with a message instead of corrupting memory.

The #189 routing itself is unchanged.

Additional information

RTX 4070 (sm_89), CUDA 13, Windows, on prism-b10770-6684606 (the build in #346) with the same change:

build	FLASH_ATTN_EXT(hsk=256,hsv=256,nh=4,nr23=[6,1],kv=16384,nb=3,...,type_K=q4_0,type_V=q4_0)	test-backend-ops -o FLASH_ATTN_EXT -b CUDA0
assert only (no get_alloc_size change)	GGML_ASSERT ... f16 K/V scratch not reserved for the short-query route	stops at that case
this PR	OK	4007/4007

On Ada the overflow did not fault in my earlier repro attempts (the write lands in memory the process owns), so the assert is what makes it visible here. @Powerful-Plantain, a test of this branch on your 3070 setup (checkpoints on, same command) would confirm it on the card that crashed.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES - the analysis, the change and the tests above were done with Claude Code (Claude Opus 5.5).

[professorpalmer]

Closing: #355 fixes the same root cause in the base.

#355 passes type_KV into the short-query (3-4 rows, quantized V) case, so that route no longer takes the F16 instantiation that converted K/V into scratch ggml_cuda_flash_attn_ext_get_alloc_size had not reserved. That conversion was the write past the f16 buffer in #346. #355 also adds the short-query FLASH_ATTN_EXT cases.

Check on the current base (e3de715222), RTX 4070 (sm_89), CUDA 13, Windows: test-backend-ops -o FLASH_ATTN_EXT -b CUDA0: 4335/4335 passed, 0 FAIL. That includes 243 cases with nb 3-4 and quantized V, among them D 256 with q4_0 and q8_0.

One note: on MSVC, tests/test-backend-ops.cpp does not compile on this base (__builtin_popcountll in arclab_init_hadamard, from #356). I used a portable popcount for this check. I will open a separate small PR for that.

Thanks to @Powerful-Plantain for the reproduction and the memcheck runs on #346.