HIP: fix turbo KV decode crash under graph capture; batch-aware VEC/TILE FA routing

[KaiFelixBennett]

With GGML_HIP_GRAPHS=ON, any turbo KV cache type crashes on the first decode step on RDNA4:

FLASH_ATTN_EXT failed: operation not permitted when stream is capturing

Two coupled causes:

1. The force-VEC condition for quantized KV applies to all batch sizes, so large prefill batches also run the sequential VEC kernel (~188 t/s pp2048 on gfx1201).
2. launch_fattn's f16 dequant temp buffers (K_f16/V_f16) use raw cudaMalloc/cudaFree/cudaStreamSynchronize, which are illegal during graph capture.

The raw alloc exists for a good reason on gfx1201 (no VMM): the legacy pool would retain the multi-GB dequant buffer permanently and negate the quantized-KV VRAM savings. So the fix has to be capture-aware rather than always-pool.

Changes

• fattn.cu: only force VEC for small batches (Q->ne[1] <= 8). Decode stays on the graph-safe VEC path (inline dequant, no temp buffer); large prefill batches fall through to the TILE/MMA kernel.
• fattn-common.cuh: check cudaStreamIsCapturing; allocate K_f16/V_f16 from the ggml pool while capturing (pool alloc/free make no CUDA calls), keep raw alloc + immediate free for large eager prefill. This is safe because ggml resets graph warmup whenever tensor sizes change, so the pool buffer is warmed at the right size before any capture.

Results

gfx1201 (Radeon AI PRO R9700, 32 GB), Windows 11, HIP SDK 7.1, Gemma-4-31B-it Q4_K_M:

Config	pp2048	tg128	decode crash
HIP_GRAPHS=OFF, turbo4	188 t/s	~22 t/s	no
HIP_GRAPHS=ON, unpatched, turbo4	—	—	yes (first step)
HIP_GRAPHS=ON + this patch, turbo4	735 t/s	22.9 t/s	no

Long context with -b 2048 -ub 512: turbo3/turbo3 tg128 @ d131072 = 9.38 ± 0.93 t/s; the full 256K context loads in ~22.9 GB. Quality checked with KLD vs an f16 baseline and needle-in-a-haystack (9/9 for q8_0/turbo4 and turbo3/turbo3). Full methodology, raw data and a one-command repro build for gfx1201: https://github.com/KaiFelixBennett/gemma4-turboquant-rdna4

Possibly related: #12 (turbo3 crash on an RX 9070 XT — same gfx1201 family, closed as stale).

Originally developed and benchmarked at 7d9715f; rebased onto 73eb521 (the fattn files were untouched in between) and built warning-free there. Can re-run any of the benchmarks on this hardware if useful.

[KaiFelixBennett]

Note for triage: #156 touches the same region (it removes hip_f16_alloc entirely by extending VEC D=512 coverage for quantized KV). The two approaches solve different problems and will conflict textually: removing the f16-dequant path forces quantized-KV prefill onto the sequential VEC kernel (the ~188 t/s path on gfx1201), while this PR keeps the fast TILE/MMA prefill (~735 t/s measured) and makes its temp allocation capture-aware instead. If #156 lands first I am happy to rebase; the capture-awareness here would still be needed for any remaining f16-dequant fallback.

[TheTom]

Thanks @KaiFelixBennett, the design reads right: the capture path adopts what the stock CUDA build already does unconditionally (pool alloc for K_f16/V_f16), so HIP-under-capture is upstream-proven behavior rather than something new, and the ne[1] <= 8 threshold covers decode plus MTP spec-decode while letting eager prefill keep the raw alloc/free that protects VRAM on gfx1201.

One question before merging: the warmup-safety argument assumes the first eval at any new size runs eagerly before capture begins. Did you verify a fresh-context, first-token decode with GGML_HIP_GRAPHS=ON (not just steady-state)? If ggml ever captures the first eval at a new size, a cold pool miss would cudaMalloc during capture and reproduce the original crash on a different path.

Minor, can be a follow-up: when decode cannot use VEC (head_dim 192 or a K stride mismatch), the captured TILE path pool-allocates the full dequant buffer and the legacy pool retains it permanently, which is the exact OOM the old raw-alloc comment defended against. Worth a sentence in the comment noting those configs trade memory for graph compatibility.

[KaiFelixBennett]

@TheTom Verified both ways — mechanism and a cold-start run.

ggml never captures the first eval at a new size. ggml_backend_cuda_graph_compute gates capture behind a two-eval warmup (graph->warmup_complete, per graph-key): the eval where node ne/nb first change runs eager (use_cuda_graph = false, no cudaStreamBeginCapture); capture only fires on the next eval with identical props, and any later size change resets it ("warmup reset") and runs eager again. So a captured eval's sizes always match the immediately preceding eager eval — which is what warms the pool. That's what the "ggml resets graph warmup" note refers to.

For VEC-eligible head dims (Gemma) the pool path isn't even in the captured graph — decode routes to the VEC kernel, which doesn't call launch_fattn. The capture-aware pool branch is only reachable under capture in the head_dim==192 / K-stride-mismatch fallthrough — your minor point.

Cold-start, instrumented (Gemma-4-31B, q8_0-K/turbo4-V, -fa on, GGML_HIP_GRAPHS=ON, --no-warmup, fresh ctx):

D ggml_backend_cuda_graph_compute: CUDA graph warmup complete
D CUDA Graph id 42 reused   <- captured decode + token output

grep -c over the full log: FLASH_ATTN_EXT failed = 0, operation not permitted when stream is capturing = 0, out of memory = 0. Warmup completes, then the captured decode graph is reused — exactly the ordering your question probes. Redacted excerpt: https://github.com/KaiFelixBennett/gemma4-turboquant-rdna4/blob/master/benchmarks/results/coldstart-hip-graph.md

Minor point: pushed a commit adding a sentence to the launch_fattn comment noting the head_dim==192 / K-stride-mismatch configs trade permanent pool retention for graph compatibility (VRAM, not a crash). Squashable on the final rebase.

[TheTom]

Both answered exactly as hoped. The warmup-gating trace is the key point: capture only fires on an eval whose sizes match the immediately preceding eager eval, so the pool is always warmed at the right size before any cudaMalloc-free capture, and a size change resets back to eager. That closes the cold-start window I was worried about, and your instrumented run (zero capture failures, captured decode reused after warmup) confirms it on real hardware. The added comment on the head_dim 192 / K-stride fallthrough is exactly right too: VRAM tradeoff, not a crash.

Thanks for the thorough verification and the clean writeup. Merging.