Repository navigation
cuda: GGML_CUDA_FA_PREFILL_F16=N, f16 prefill for quantized K/V from pool memory (decode stays in place) - #330
Conversation
… to N cells to f16 in pool memory for prefill-sized FA batches Decode keeps the in-place quantized read. The f16 copies are sized by the actual KV length and nothing is reserved at load. Off by default.
|
Tested this PR on an H200: the prefill gain transfers to Hopper, including with MTP depth 1. Same Release binary,
Gains above compare median prompt throughput. At 60k, q4 AR improved 1987.7 → 2314.1 tok/s, q4/MTP1 1752.7 → 1996.1, and q8 AR 1843.4 → 2330.9. Every individual quantized-cache pair showed a positive prefill gain; q4 AR's 60k paired range was +11.9% to +16.5%. Full request latency at 60k fell 13.2%, 11.6%, and 19.8%, respectively. Validation:
Memory caveat: load-time process memory was identical off/on, but the enabled pool retained an additional 46 MiB at 16k, 110 MiB at 32k, and 232 MiB at 60k. Logical scratch frees do not return the VMM pool's physical high-water allocation. The cap limits padded cells per operation, rather than aggregate bytes across contexts. These requests extended a synthetic repeated-code prompt; actual evaluated token counts were 4096 / 12800 / 16896 / 27744 after prefix reuse and hybrid checkpoint replay. Output parity is scoped to this workload. MTP1 verification has only two queries, so it stays below the new >8 threshold; wider speculative verification can also take the conversion path. Server generation timing includes the final MTP draft-prefill hook, which is why full HTTP latency was retained. This looks useful as an opt-in on H200, especially at depth. It does not establish Ampere/Ada performance or reproduce the Windows/Turing setup. |
|
@bri-prism, thank you for the H200 run. The paired AB/BA order, the token-ID parity and the Nsight split make the result easy to trust. Memory caveat: correct. The VMM pool keeps its physical high-water allocation, so the f16 copy stays reserved after the largest prefill batch. That is why the switch is opt-in, and why Ada and Ampere: I will add Ada numbers here. I will run |
|
Ada data, as promised (RTX 4070 12 GB, Windows, sm_89 build of this PR's code): 1. Kernel,
On Ada the in-kernel q8_0 reads cost almost nothing at depth, so the f16 kernel is no faster there. This does not include the bulk conversion. 2. Serving A/B, Bonsai 2 27B PTQ1_0, q8_0 K/V, all-VRAM cache (64k window), one 60k prompt, switch off against 3. One interaction to fix: speculative verify batches wider than 8 queries. With a lookup drafter (ngram-mod in front of MTP, 4. Memory, as in your caveat. On our 12 GB recipe, the tiered KV cache fills VRAM up to a 1000 MiB margin. With the switch on, the server stopped at 8k context when the pool grew: |
|
Ada result for the 64-query minimum (point 3 above). Same RTX 4070 and recipe: all-VRAM cache (64k window), one 60k prompt, then 7 greedy decode jobs at 30k (4 file rewrites with lookup drafting, 3 plain-text answers). One run of each arm, in this order:
The minimum keeps the prefill result and removes the decode loss. The change is one condition in static const int64_t min_q = [] {
const char * e = getenv("GGML_CUDA_FA_PREFILL_F16_MIN_Q");
return e ? std::max<int64_t>(1, atoll(e)) : (int64_t) 64;
}();
...
return max_kv > 0 && Q->ne[1] >= min_q && K->ne[1] <= max_kv && K->type == V->type && ggml_is_quantized(K->type);If you agree, I will push it to this PR as one commit. The Turing A/B follows when the RTX 2060 is free. |
|
@bri-prism, an update, in one comment for both cards. 1. A bug we found on our 8 GB preset is not in this PR. On an RTX 2060 SUPER, fresh prompts longer than about 30k sometimes got a 1-token answer. The 8 GB preset sets this PR's switch (
Every run with the shared pool off since then: 0 of 53 requests, fresh prompts from 4k to 126k. Cause: the two contexts run on two CUDA streams that got the same LIFO pool, so a block freed by one stream could be reused by the other before the first was done with it. We fixed it on our side (an event orders the two streams at graph boundaries: 0 of 18 against 18 of 18 on the 2060). Nothing changes in this PR. 2. Turing data for the 64-query minimum (RTX 2060 SUPER, lookup drafting on, the switch at 32768, current rule
Fresh prefill, 2 runs each: 30k 345.2 / 345.2 against 344.8 / 344.5 tok/s; 60k 255.0 against 254.4. Text identical between the two engines in each round. So on Turing the minimum is neutral (within 1.5%), and prefill keeps the gain this PR measured there. On Ada (my earlier comment) the minimum removes an 8% decode loss on copy-heavy work (195.4 -> 212.4 tok/s, prefill unchanged). A 30k file-rewrite row for Turing follows when a clean rerun finishes. The minimum is still one commit; I will push it here if you agree. |
llama-server log — prism-b10770, 30K prompt, CUDA_LAUNCH_BLOCKING=10.01.556.460 I cmn common_param: common_params_print_info: verbosity = 3 (adjust with the |
|
Thank you for the log. The error is reported in
Two values from your log may matter. The prompt is 30,180 tokens with a 65,536-token slot, so the K/V view is about 30k cells at that point. For Bonsai 2 27B (4 KV heads x 256), the f16 copy of K and V at 30k cells is about 120 MiB of pool memory per attention op. That is close to the default limit of the f16 path (32,768 cells), so a failed or short pool allocation there is one candidate. |
|
Thanks for the quick reply — and apologies, this should have been a separate issue rather than a comment on your PR. Full report moved to #346. Answers to your questions:
Same request crashes at 30K and 45K; 45K succeeds with |
3 participants
Opt-in, off by default:
GGML_CUDA_FA_PREFILL_F16=Nlets prefill-sized flash-attention batches (more than 8 queries) over a quantized K/V cache of at most N cells convert K/V to f16 in pool memory sized by the actual KV length and run the f16 MMA kernel. Decode keeps the in-place quantized read from #221. Single commit on currentprism(6bfcd79), applies without conflicts.Why
The in-place q4_0/q8_0 read (#221) is the right default: it saves the f16 copy that
ggml_cuda_flash_attn_ext_get_alloc_sizewould otherwise reserve behinddstfor the full context at load (512 MiB at a 131k window for Bonsai 2 27B), and it wins decode. On Turing it costs prefill at depth: the dequantization inside the kernel is slower there than one bulk conversion per chunk. Taking the copy from the pool for prefill only gets both: nothing reserved at load, the f16 kernel for prefill, and the cap keeps the transient copy within the free memory the deployment has.What it does
ggml_cuda_fattn_prefill_f16(dst): true when the variable is set,Q->ne[1] > 8,K->ne[1] <= N, and K/V are the same quantized type.ggml_cuda_fattn_mma_kv_native_supportedreturns false for those batches, so dispatch takes the f16 MMA path.ggml_cuda_flash_attn_ext_get_alloc_sizereserves nothing for them;launch_fattnallocates the f16 copies from the pool (ggml_cuda_pool_alloc<half>) instead of the space behinddst.Receipts
RTX 2060 SUPER 8 GB (sm_75, PCIe 3.0, stock clocks), Windows, CUDA 13.4, Bonsai 2 27B
PTQ1_0, q4_0 K/V, 131,072 window with 58,112 cells on the device, greedy,bench/quick_tps.py-style requests at 4k -> 16k -> 32k -> 60k (each extending the previous prompt), two runs:GGML_CUDA_FA_PREFILL_F16=65536test-backend-ops -o FLASH_ATTN_EXT -b CUDA0on sm_75 with the variable set to 1048576 (path forced): 2994/2994.Not measured on Ampere or Ada, where the in-place read may already win prefill; the variable keeps the current behaviour unless set.