Repository navigation
Conversation
The fattn-mma-f16 and fattn-tile kernels ingest f16 KV data layer by layer. For non-f16 (quantized) KV types this requires a conversion buffer, pre-allocated per ggml-org#23907. This size of this buffer is equal to one layer of f16 KV and scales with ctx length. For quantized KV types, this negates large part of VRAM savings obtained by quantization in the first place. Implement serial split-KV - Cap the buffer to a fixed size (64 MiB, env override GGML_CUDA_FATTN_CONVERT_BYTES for testing only), process KV in chunks re-using the buffer. For each processed KV chunk fattn kernel emits an unnormalized partial and (max_row_val, row_sum), these get folded into running accumulators by flash_attn_combine_results kernel found in fattn-common. KV chunking is enabled based on two decisions: Firstly, in addition to picking the best fattn kernel, ggml_cuda_flash_attn_ext_get_alloc_size() also sets the kernel_supports_chunking flag to true, currently enabled for fattn-mma-f16 and fattn-tile kernels. Secondly, ggml_cuda_fattn_type_chunk_enabled() disables/enables chunking depending on other criteria such as KV type. Currently, it returns true for quantized KV types only. For launches with KV types f16 and f32, launches with fattn-vec kernel and launches with FA convert buffer size below the GGML_CUDA_FATTN_CONVERT_BYTES cap, the outputs are bit-identical to master. Before this commit, the scratch buffer costs one layer of f16 KV. For Qwen3.6-27B at ctx 131072, for example, this negates 13.3% of the VRAM savings at q8_0 and 8.7% at q4_0 KV quantization. After this commit the cost is bounded at 64 MiB max, regardless of context length. The f16 conversion of quantized KV chunk, similar to master, picks between to_fp16_cuda and to_fp16_nc_cuda conversion kernels depending on the KV chunk's allocation contiguity. Assisted-by: Claude Fable 5
4 tasks done
Contributor
|
Context: I am the maintainer for the files changed by this PR and this is one of my saved replies.
|
Author
|
Sorry for the wall of text, I guess I have been brooding on this feature for too long and my verbose prose is directed more at my level of understanding of the FA code. I still think it would be good to find a solution to the issue. The amount of VRAM wasted due to the f16 scratch buffer being ctx length dependent is too much: |
This was referenced Oct 2, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
The fattn-mma-f16 and fattn-tile kernels ingest f16 KV data layer by layer. For
non-f16 (quantized) KV types this requires a conversion buffer, pre-allocated
per #23907. This size of this buffer is equal to one layer of f16 KV and scales
with ctx length. For quantized KV types, this negates large part of VRAM savings
obtained by quantization in the first place.
Implement serial split-KV - Cap the buffer to a fixed size (64 MiB, env override
GGML_CUDA_FATTN_CONVERT_BYTES for testing only), process KV in chunks re-using
the buffer. For each processed KV chunk fattn kernel emits an unnormalized
partial and (max_row_val, row_sum), these get folded into running accumulators
by flash_attn_combine_results kernel found in fattn-common.
KV chunking is enabled based on two decisions: Firstly, in addition to picking
the best fattn kernel, ggml_cuda_flash_attn_ext_get_alloc_size() also sets the
kernel_supports_chunking flag to true, currently enabled for fattn-mma-f16 and
fattn-tile kernels. Secondly, ggml_cuda_fattn_type_chunk_enabled()
disables/enables chunking depending on other criteria such as KV type.
Currently, it returns true for quantized KV types only.
For launches with KV types f16 and f32, launches with fattn-vec kernel and
launches with FA convert buffer size below the GGML_CUDA_FATTN_CONVERT_BYTES
cap, the outputs are bit-identical to master.
For master, the scratch buffer costs one layer of f16 KV. For Qwen3.6-27B at
ctx 131072, for example, this negates 13.3% of the VRAM savings at q8_0 and 8.7%
at q4_0 KV quantization. For this PR the cost is bounded at 64 MiB max,
regardless of context length.
The f16 conversion of quantized KV chunk, similar to master, picks between
to_fp16_cuda and to_fp16_nc_cuda conversion kernels depending on the KV chunk's
allocation contiguity.
Additional information
when the f16 convert buffer does not fit
this PR
fusing the conversion into fattn-mma-f16 kernel. Covers q8_0 and q4_0 KV
types on fattn-mma-f16 only. Both PRs are complementary: fused KV pairs
need no buffer, and every other quantized type uses the buffer cap.
PR is not yet submitted due to 1 PR for new contributors rule,
see fork: throughput, two GPUs | VRAM | write-up
VRAM
tokens on qwen35 27B, 470 MiB at 131k ctx weight-independent and cannot be
reduced by lowering ub) to 64MiB, regardless of context size.
CUDA0 compute buffer, llama-server-ngl 99 -fa on -ub 512 --no-warmup,KV q8_0/q8_0, RTX 3090, qwen35 27B (weight-independent; other models and
KV pairs in the four-model ladder below), master
36b101543vs PR54b94d245, measured 2026-09-22/23:Click to expand - CUDA0 compute buffer (MiB), four models, context ladder, all KV pairs, default and 32 MiB caps (RTX 3090, llama-server -ub 512)
VRAM: CUDA0 compute buffer (MiB), KV q8_0/q8_0, -ub 512
VRAM: CUDA0 compute buffer (MiB), KV q4_0/q4_0, -ub 512
VRAM: CUDA0 compute buffer (MiB), KV q8_0/q4_0, -ub 512
VRAM: CUDA0 compute buffer (MiB), KV f16/f16, -ub 512
Performance
paths
bound, by reducing VRAM pressure if chunk fits into L2 cache. Tested on RTX
5070 Laptop, 32 MiB L2, cap of 32 MiB, yields 1.59x agains master for llama 8B
Speedup = PR t/s / master t/s at the default 64 MiB cap. llama-bench,
pp1024 at depth 32768, uniform Q4_0 weights,
-fa 1, PR54b94d245vsmaster
36b101543, builds from .devops/cuda.Dockerfile, 3 reps, measured2026-09-22/23 (RTX 3090) and 2026-09-22..25 (RTX 5070 Laptop, 115 W Dynamic Boost):
Click to expand - Full llama-bench matrix, both GPUs: ub ladders 1-512 for q8_0/q8_0, q4_0/q4_0, q8_0/q4_0; q8_0 at d65536 and d131072; pp512 and tg128 depth ladders; f16/f16 controls; columns PR 64MiB (default cap) and PR 32MiB (env override)
KV q8_0/q8_0, pp1024@d32768
KV q8_0/q8_0, pp1024@d65536
KV q8_0/q8_0, pp1024@d131072
KV q8_0/q8_0, pp512 depth ladder
KV q8_0/q8_0, tg128 depth ladder
KV q4_0/q4_0, pp1024@d32768
KV q8_0/q4_0, pp1024@d32768
KV f16/f16, pp1024@d16384
KV f16/f16, pp1024@d32768
KV f16/f16, tg128 depth ladder
Correctness
test-backend-ops test -o FLASH_ATTN_EXT -b CUDA0(GGML_CUDA_FA_ALL_QUANTSbuild) at tip
54b94d245vs master36b101543, RTX 3090: 3969/3969 at thedefault budget, 3969/3969 at GGML_CUDA_FATTN_CONVERT_BYTES=1048576 and
3969/3969 at =262144 (forces multi-chunk stream-k execution including
both fixup variants); master 3955/3955. RTX 5070 Laptop: 3969/3969 /
3969/3969 / 3969/3969 vs master 3955/3955. The delta of 14 is exactly
the cases this PR adds.
test-backend-ops test -b CUDA0(same build): RTX 3090 15644/15644 vs15630/15630; RTX 5070 Laptop 15644/15644 vs 15630/15630.
RTX 3090 ERROR SUMMARY: 0 errors (exit 0); RTX 5070 Laptop ERROR SUMMARY: 0 errors (exit 0).
benchmark.sh --testsin the .devops-derived toolchain image (both GPUs 2026-09-25):VALIDATE-RTX3090-sm86-24GB.txt,
VALIDATE-RTX5070LaptopGPU-sm120-8GB.txt.
cases + 8 chunk-shaped small-kv cases. The plain suite exercises chunking
on CUDA with default 64 MiB cap.
36b101543, uniformQ4_0 weights,
benchmark.sh --ppl --kld(renders:PPL-RTX3090-sm86-24GB.md).
Every statistic is bit-identical to master when not capped. When capped,
llama and qwen at n_ctx 32768 on the RTX 3090, |delta| <= 0.011: chunk maxima
differ, causing a rescale during accumulation.
Perplexity: wikitext-2 test, n_ctx=8192, 9 chunks (73728 tokens)
Perplexity: wikitext-2 test, n_ctx=16384, 9 chunks (147456 tokens)
Perplexity: wikitext-2 test, n_ctx=32768, 9 chunks (294912 tokens)
KL divergence vs the master f16-KV run: wikitext-2 test, n_ctx=4096, 4 chunks (16384 tokens)
Requirements