cuda: opt-in quantized-native MMA FlashAttention - #85
Open
Piggidragon wants to merge 15 commits into
Open
Piggidragon wants to merge 15 commits into
Piggidragon wants to merge 15 commits into
Conversation
The MMA FlashAttention kernel consumes a half2 shared-memory tile. These loaders write that tile straight from a quantized K/V row, so the visible window no longer has to be cast to F16 first. Each dequant() has to reproduce that type's F16 cast path bit for bit, not merely to within a few ulp, because the cast path is the reference this replaces. Which helper achieves that differs per type and is documented at each specialization: Q8_0 delegates to the vector path's half2 multiply, while Q4_0 and Q5_0 need their own, because the generic cast rounds once in float and the vector helper biases the quant as an integer, and the two disagree in the last bit. Q4_0, Q5_0 and Q8_0 are hand tuned, down to the per-thread run length in fattn_quant_load_width. Q4_1 and Q5_1 share the generic nibble loader in fattn-mma-quant-packed.cuh. Only symmetric K/V pairs are served. A mixed pair would need a kernel that switches on V's type at run time, so none is compiled. Co-authored-by: Anbeeld <anbeeld@gmail.com> Assisted-by: Claude Opus 5
…kernel flash_attn_ext_f16 gains type_K and type_V template parameters that default to GGML_TYPE_F16, so every existing instantiation is unchanged. Where they name a native type the tile load goes to the loader instead of the F16 copy, and launch_fattn is told it needs no F16 K/V, which is what removes the transient copy. Multi-stage cp.async loading is off for native tiles: the loader writes the tile itself, so there is no pipeline to stage. The tile is XOR swizzled with the same map the F16 loads use, so the MMA body reads it unchanged. ggml_cuda_fattn_native_supported() answers whether a kernel exists for a geometry and returns the tile shape it selected. The compiled set is exactly the selectable set: three rows at two tile widths each, 12 kernels in a default build and 18 with GGML_CUDA_FA_ALL_QUANTS. fattn-mma-quant-decl.cuh declares the same rows, so a gate that drifts from the table is a link error rather than a silent fallback. The route is CUDA only. HIP and MUSA exclude the generated instances from their source globs, and FATTN_MMA_QUANT_AVAILABLE keeps fattn.cu from naming kernels those builds do not have. The backend also counts the nodes that took the route and exposes that through get_proc_address, because route selection is otherwise only observable through allocation size or throughput, which are incidental. Assisted-by: Claude Opus 5
--flash-attn-native-quants turns the route on; without it every FlashAttention node keeps the established path. The graph opts in per node through a new flash_attn_ext op param, which the CUDA backend reads before it considers the geometry at all. Opt-in rather than automatic because the route is not uniformly faster. On Ampere the D=256 rows measured slower than the F16 path and the D=512 rows faster, and the split depends on the cache type. A caller who wants it asks for it, and a caller it does not suit turns it off, rather than the backend carrying a table of measured thresholds it has to keep true on hardware nobody here can test. The flag costs nothing when off: one integer read per node on the host side, before any launch. Measured at -0.4% to +0.1% against a build without it, which is inside the run-to-run scatter. Assisted-by: Claude Opus 5
Each NATIVE_QUANT_EQUIVALENCE case runs the attention twice in one graph, once over the quantized cache and once over an F16 copy of the same data, and compares them. That covers the numbers. It does not cover which path ran, and the route is otherwise only observable through allocation size or throughput, which are allocator and timing properties rather than a dispatch contract. So each case opts the graph in the way llm_graph_context does, reads the CUDA backend's native-launch counter through get_proc_address, and asserts the delta. The reference arm is pinned off, so the two halves of the graph differ in route and nothing else. With the graph opted in, expected_native is purely whether a kernel is compiled for the geometry: 10 native and 2 fallback cases in a default build, 13 and 8 with GGML_CUDA_FA_ALL_QUANTS, the difference being the extra-tier types that a default build reports as unsupported and skips. Assisted-by: Claude Opus 5
The failure this guards against is a build failure, not a wrong result. An extern-macro mistake can instantiate hundreds of full attention kernels in one translation unit while every runtime test still passes, and GGML_CUDA_FA_ALL_QUANTS changes which cache types get kernels at all. Nothing else in CI crosses that axis against this route. fattn-native-inventory.py reads the built library back and compares it against the exact set of cases the generated instance files declare. Missing, unexpected and duplicate kernels all fail, as does any mixed K/V, logit-softcap or sparse specialization, none of which the route can select. The compile matrix builds both tiers, both tile widths and the FA=OFF build that must compile the route out cleanly. Turing is its own arm because it takes the narrow tile shapes from a different MMA config table, and the tile loader asserts and the per-type load widths are functions of that table. The focused run needs a GPU, so it is restricted to push and workflow_dispatch on a self-hosted runner and never runs a pull-request payload. It asserts the case counts, because the harness exits 0 when it selects no tests at all. fattn-turing-model-test.sh is the real-model check that test-backend-ops cannot be: it reads the model geometry and refuses to measure a model the route cannot serve. Assisted-by: Claude Opus 5
Covers the flag, the route table, which cache types are in which build tier, and what widening any of it would owe in evidence. The build-size, correctness, route and inventory numbers are measured on this base. The throughput, end-to-end and memory numbers are not: they were taken before the rebase onto the current llama/dev and before the shared-memory tile gained the XOR swizzle, and the section they sit in says so rather than presenting them as current. Assisted-by: Claude Opus 5
github-actions
Bot
added
documentation
fattn_quant_load_width<GGML_TYPE_Q4_0> narrowed the per-thread run to 8 for every 128-thread config, which is every D=256 row. Nothing else narrows it, and the 16-wide default is faster on both architectures measured, at every D=256 row: 3060: -2.2% to -3.4% at GQA 6, -14.0% and -15.6% at GQA 2 4070: -3.1% to -7.2% at GQA 6, -3.1% at GQA 2 The D=512 row already ran at 16 and moves by noise. So the specialization goes and Q4_0 takes the same default as every other type. At D=256, GQA 6 and n_kv 16384 this brings the native path level with the F16-casting path measured at equal nstages, which leaves the cp.async pipeline the route gives up as the whole of the remaining Ampere regression there. Assisted-by: Claude Opus 5
The document blamed the loading pipeline alone. Measuring the F16-casting path with nstages forced to 0 removes the pipeline as a variable and shows that is wrong: with it held equal the Q8_0 loader costs between -4.6% and +8.6%, while Q4_0 at GQA 2 costs +31.9% to +41.1%. Record the three-way measurement, the load width that was part of it, and the nibble unpack that is the rest. Also note that forcing nstages to 0 made the F16 path itself faster at D=512 on Ampere, which is a pessimization in code this route does not touch. Assisted-by: Claude Opus 5
The document named it as the follow-up that would remove the Ampere tradeoff. It was prototyped and measured, and as written it is not a win: between -17.8% and +5.8% against the same path without it, over all three route rows and all five types on both cards. The split is by cache type, not by row. Q4_1 and Q5_1, whose blocks are 20 and 24 bytes, gain everywhere. Q4_0, Q5_0 and Q8_0, whose blocks are 18, 22 and 34, lose. A block that is not a multiple of 4 bytes is read with a 2-byte aligned access, which is cheap out of global memory and several sub-word loads out of shared. Record that, and that fixing it needs a reformatted staging layout rather than a copy of the raw bytes, so nobody spends the same week twice. Assisted-by: Claude Opus 5
The table it carried predated the rebase onto current llama/dev, the tile swizzle and the widened Q4_0 load run, and it claimed every row was faster on Ada, which is not true on this base. Measure instead: every route row against every cache type it admits, both cards, otherwise idle machine. Ada gains up to 40% at long context and loses 8-23% at short; Ampere loses on every D=256 row and gains on D=512 for Q4_0. The two generic loaders are the slowest of the five, by about 4 points for Q4_1 against Q4_0 and 6 for Q5_1 against Q5_0, which is what the hand tuning buys. The Turing section keeps its pre-rebase numbers and now says so; nobody here has that hardware. Assisted-by: Claude Opus 5
The section claimed the copy never reaches the high-water mark and that the route is not worth taking for memory. That was measured at one depth, 16K, which is the one depth where it is true. The copy is sized by the visible window, so measure across it. With the cache on the host and the GPU holding only the model, peak device memory with the route on is flat from 16K to 262K, 9985 to 10097 MiB, while with it off it climbs to 11141. At full context that is 1044 MiB saved, against 1024 MiB predicted from 2 * n_kv_heads * head_dim * n_kv * sizeof(F16). Prefill is also 0.5% to 5% faster there, which is the opposite sign from the device-resident matrix above. Say that the answer depends on where the cache lives rather than leaving the two results looking contradictory. Assisted-by: Claude Opus 5
The memory section measured one model at one depth with a pageable host cache. Replace it with the configuration the route is actually for: Qwen3.8-27B dense at its full 262144 context, cache on the host with pinned buffers and recurrent state offloaded, on one 4070, one 3060, and both, at two cache types. Device memory saved in every configuration, 898 to 2072 MiB at full context, matching 2 * n_kv_heads * head_dim * n_kv * sizeof(F16). Route-on memory barely moves with context while route-off climbs with it. Prefill is the part that is not unconditional: up to 7.9% faster on Ada and 7.1% slower on Ampere, with two-GPU rows between them. Decode is unaffected everywhere, within 0.4%, because a single-token query uses the vector kernel. 144 measurements, no failures. Assisted-by: Claude Opus 5
Assisted-by: Claude Opus 5
…CUDA_FA_QUANTS pair Assisted-by: Claude Opus 5
This was referenced Sep 17, 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
1 participant
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.
Reads a quantized K/V cache in place in the CUDA MMA FlashAttention kernel
instead of casting the visible attention window to F16 first. Output is
unchanged; what goes away is the transient F16 copy.
Supersedes #50 and #55, and replaces #72. Same work, scoped down: the hand
tuned symmetric loaders, behind a flag, with the automatic route table and its
measured thresholds removed.
What it is
--flash-attn-native-quantsturns it on. Without it every FlashAttention nodekeeps the established path, so nothing changes for anyone who does not ask.
With it on, the route is taken wherever a kernel is compiled for the geometry,
and nothing else is consulted. There is no profitability table, no KV-length
window and no host/device rule: if it does not suit a machine, turn the flag
off.
q4_0,q8_0q4_0,q8_0;q4_1,q5_0,q5_1when selectedq4_0,q8_0K and V must be the same type. A mixed pair would need a kernel that switches
on V's type at run time, so none is compiled.
q4_0,q5_0andq8_0havehand tuned loaders;
q4_1andq5_1share a generic nibble loader.That is 12 kernels in a default build, two more for each selected extra pair,
and 18 with
GGML_CUDA_FA_QUANTS=all.fattn-mma-quant-decl.cuhdeclares thesame rows, so a gate that drifts from the table is a link error rather than a
silent fallback.
Build selection
q4_0andq8_0kernels are compiled in every CUDA FlashAttention build.q4_1,q5_0andq5_1follow upstream'sGGML_CUDA_FA_QUANTS: a type getsits native kernels when its own
type-typepair is in the list, or when thelist is
all. That is the same selection that compiles its vector kernel.GGML_CUDA_FA_QUANTSq4_0-q4_0;q8_0-q8_0;f16-f16;bf16-bf16)q4_0,q8_0q5_0-q5_0q4_0,q8_0,q5_0all, or the deprecatedGGML_CUDA_FA_ALL_QUANTS=ONA cache of an unselected type still runs, on the F16-casting path.
fattn-mma-quant-types.hmaps each extra type to itsGGML_CUDA_FA_<T>_<T>definition, CMake leaves out the instance files of a type whose pair is not
selected, and
scripts/fattn-native-inventory.py --fa-quants <value>checks thebuilt library against the same value.
This branch is merged with current
llama/dev, which replacedGGML_CUDA_FA_ALL_QUANTSwithGGML_CUDA_FA_QUANTS(ggml-org#28079). The merge keepsupstream's vector kernel selection and adds the native route on top of it.
Correctness
RTX 4070 (
sm_89) and RTX 3060 (sm_86), CUDA 13.3, Release.After the
llama/devmerge and theGGML_CUDA_FA_QUANTSchange:alltest-backend-ops -o FLASH_ATTN_EXT-o NATIVE_QUANT_EQUIVALENCEEach equivalence case runs the attention twice in one graph, once over the
quantized cache and once over an F16 copy of the same data, and asserts both the
numbers and which path the dispatcher took, by reading the backend's
native-launch counter. Both builds run 13 cases declared native and 8 declared
fallback. In a default build the nine extra-tier cases run on the F16 path,
because upstream now accepts those cache types there, and a case declared native
expects the F16 path when the type's pair is not selected. Before the merge a
default build skipped them. Regenerating the instance files reproduces the
committed ones.
CMake picks the expected native instance files for the default list, the default
list plus
q5_0-q5_0,all, andGGML_CUDA_FA_ALL_QUANTS=ON.Cost of the flag
One integer read per node on the host side, before any launch. Against a build
without the flag at all, on the same work:
pp512 @ d16384pp2048 @ d16384tg64 @ d16384Every difference is inside one standard deviation. Selecting the route by flag
rather than automatically also costs nothing at the kernel: with the route
active either way, the two agree to within 0.15% and issue the same number of
native launches.
Build cost
CMAKE_CUDA_ARCHITECTURES=86-real;89-real, CUDA 13.3, Release. Measured beforethe
llama/devmerge; "all-quants" is whatGGML_CUDA_FA_QUANTS=allbuilds now.libggml-cuda.soThroughput
test-backend-ops perf -o FLASH_ATTN_EXT, us/run, lower is better. Every routerow against every cache type it admits, on an otherwise idle machine, native
against the F16-casting path with the cast kernel included in both timings.
This and the host-cache measurements below were taken before the
llama/devmerge and have not been repeated on it.
q4_0q4_0q4_0q4_0q4_1q4_1q4_1q4_1q5_0q5_0q5_0q5_0q5_1q5_1q5_1q5_1q8_0q8_0q8_0q8_0q4_0q4_0q4_0q4_0q8_0q8_0q8_0q8_0q4_0q4_0q8_0q8_0This is not uniformly faster, which is why it is opt-in.
Long context on Ada is where it pays: at D=256 GQA 2 and
n_kv 16384it is 38%to 40% faster, because it reads 4.5 or 8.5 bit weights where the other path
reads 16. Short context is where it does not: at
n_kv 1024every type costs 8%to 23%, because the dequant sits on the critical path with too little memory
traffic to hide behind.
Ampere loses on every D=256 row, by 11% to 42%. D=512 is the mirror image,
q4_0gaining 9% there and
q8_0losing 8% to 9% on both cards.The two types on the generic loader are the slowest of the five. Against their
hand tuned counterparts at the same geometry on Ampere,
q4_1costs about 4points more than
q4_0andq5_1about 6 more thanq5_0.Memory and speed with the cache on the host
This is what the route is for: the KV cache in host memory, the GPU holding only
the model and the compute buffers. The transient F16 copy is sized by the
visible attention window, so removing it saves device memory in proportion to
the context in use.
Qwen3.8-27B, dense,
-nkvo --kv-cpu-pinned --recurrent-state-offload, full262144 token context, RTX 4070 and RTX 3060. Peak device memory sampled at
250 ms; the route asserted on every row by the backend's native-launch counter.
144 measurements, no failures.
Full-context prefill
Ingesting all 262144 tokens,
-p 262144. One run each.q8_0q8_0q8_0q8_0q4_0q4_0q4_0q4_0Memory is saved in every configuration, 898 to 2072 MiB at full context. The
route-on figure barely moves with context while route-off climbs with it: on one
4070 with a
q4_0cache, 10104 MiB at 262144 against 11002. Two GPUs save abouttwice as much as one, because each card stages the window for the layers it
owns. The size matches the copy removed,
2 * n_kv_heads * head_dim * n_kv * sizeof(F16), 4 KiB per cached token for thismodel.
Speed is not. Prefill gains up to 7.9% on Ada and loses up to 7.1% on
Ampere, with the two-GPU rows in between because half the layers sit on each
card. Decode is untouched everywhere, within 0.4%, since a single-token query
uses the vector kernel and never reaches this route.
docs/quantized-native-flash-attention.mdcarries the per-depth ladders behindthese numbers, at 16384, 65536, 131072 and 262144.
Where the Ampere cost comes from
Three causes, separated by measuring the F16 path with
nstagesforced to 0,which removes the pipeline as a variable:
tile themselves, so
nstagesis 0 for them, while every D=256 config targets2. This is the largest share at D=256 GQA 6.
fattn_quant_load_width<GGML_TYPE_Q4_0>narrowed the per-thread run to 8 forevery 128-thread config, which is every D=256 row, where every other type
uses 16. The 16-wide default is faster on both cards at every D=256 row
measured (3060: -2.2% to -15.6%; 4070: -3.1% to -7.2%), so the specialization
is gone.
nstagesequal on both sides theq8_0loader costs between -4.6% and +8.6%, while
q4_0at GQA 2 costs +31.9% to+41.1%. That residue is integer work, and it is the part still open.
Staging the quantized tiles through cp.async, which the documentation used to
name as the fix for cause 1, was prototyped and rejected: it lands between
-17.8% and +5.8%, gaining only for the two cache types whose block size is a
multiple of 4 bytes.
docs/quantized-native-flash-attention.mdrecords why.Incidentally, at D=512 on Ampere, forcing
nstagesto 0 made the F16 pathitself faster, 3675 to 2939 us/run. That is
nstages_target = 1being apessimization in code this route does not touch.
AI assistance: the rebase onto
llama/dev, the scope reduction, the flagrestoration, the later
llama/devmerge, theGGML_CUDA_FA_QUANTSpairselection and the measurements on Ada and Ampere were carried out with Claude
Opus 5. The Turing hardware run was done separately by a third party on a
Quadro RTX 8000. Each commit carries
Assisted-by:.🤖 Generated with Claude Code
https://claude.ai/code/session_01Jsr2MJJ4wvLBkaQRo6zMAq