cuda: compile only the reachable native FlashAttention routes (follow-up to #55) - #72
Draft
Piggidragon wants to merge 8 commits into
Conversation
The route table is now three rows: D=256 32x2 for Q4_0/Q8_0, D=256 8x8 for every compiled type, and D=512 8x8 for Q4_0/Q8_0. The gate returns the tile shape it selected, so the compiled kernel set is exactly the selectable set. Everything the gate cannot reach is gone: mixed K/V pairs and the runtime-V kernel that served them, D=64 and D=128, D=512 for the extra-tier types, and the 14 column shapes the ncols switches never pick on this path. D=256 with GQA 8 is declined until the memory-safety question recorded in the PR is closed, and the device gate stops at Hopper, which is untested. 98 kernels become 6 in a default build and 485 become 9 with GGML_CUDA_FA_ALL_QUANTS. Assisted-by: Claude Opus 5
Both backends glob template-instances/fattn-mma*.cu, so they compiled the CUDA-only quant instances. HIP rejects the raw half2 initializers in the Q4 loader and fails the build; MUSA compiled them for a route it cannot take. Exclude the files from both globs, and gate the type manifest on the same condition so fattn.cu names no kernel those builds do not have. Assisted-by: Claude Opus 5
The equivalence cases inferred the route from the output allocation size, which is an allocator property rather than a dispatch contract: a workspace change could make them report the wrong route without changing dispatch. Read the CUDA backend's native-launch counter through get_proc_address instead, and add the D=256 GQA 8 case that must stay on the standard path. Assisted-by: Claude Opus 5
The focused job filtered on native_quants=1, which no test identity carries any more, so it selected nothing and the harness exits 0 on zero tests. Filter on the current identity and assert the exact native and fallback counts. The job also built and ran pull-request payloads on a persistent self-hosted runner; restrict it to push and workflow_dispatch. The inventory checker kept geometries in a set and only asked for "some", so it could not prove the documented inventory. Compare the built library against the exact set of cases the generated instances declare, and fail on a missing, unexpected or duplicated kernel, or on any mixed K/V or logit-softcap kernel. Assisted-by: Claude Opus 5
The documentation still advertised --flash-attn-native-quants, -fanq, the graph opt-in, the verbose route log and the fallback warning, all of which the final revision removed, and its measurements were taken with that switch. Describe the route table the code selects, and replace the tables with measurements of it: library size, correctness, per-row throughput on Ada and Ampere, and the memory result, which did not reproduce on this base. Assisted-by: Claude Opus 5
Piggidragon
marked this pull request as draft
September 5, 2026 21:28
The route was gated on ampere_mma_available(), so Turing fell back to the F16-casting path. The loaders need nothing Ampere has: they write the shared memory tile themselves and already force nstages = 0, and the MMA body is the same one the F16 path runs there. What differs is the tile. switch_ncols1() caps ncols1*ncols2 at 32 on Turing, so each row of the route table takes half the columns: 32x2 becomes 16x2 and 8x8 becomes 4x8. Those are the shapes the F16 path already uses for these geometries on Turing. Gate on turing_mma_available() and pick the width from ampere_mma_available(), which is the same question the escape in switch_ncols1() asks. Every row gains one shape, so a default build holds 12 native kernels and an all-quants build 18. The KV-length thresholds are unchanged. Turing is not measured yet, and keeping the Ada thresholds is what lets the equivalence cases assert the same route on every supported card. Assisted-by: Claude Opus 5
Turing takes the narrow tile widths and a different MMA config table, so its kernel bodies are a separate compile from the sm_89 arms. Add one 75-real arm to the compile matrix and make the architecture a matrix axis. Document the second tile column, the kernel counts it brings to 12 and 18, and that Turing is compiled but not measured. The build-size table keeps its numbers and says which widths they were taken with. Assisted-by: Claude Opus 5
A Quadro RTX 8000 has now run the route. The equivalence cases all pass, with the launch counter asserting the route on each one, and Qwen3.8-27B with a q4_0 cache takes it and reads correctly. The measurements are absolute, with no route-off build beside them, so the section says the route works on Turing and stops short of calling it faster. The cp.async argument is kept as reasoning, labelled as such. Assisted-by: Claude Opus 5
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.
Makes #55 ready for the 1.0 release by compiling only the route the dispatcher
can actually select, and closes the blocking items from the review on that PR.
The hand-tuned Q4, Q5 and Q8 loaders are untouched. What goes away is the
inventory around them that no dispatch reaches.
The route table
ggml_cuda_fattn_native_supported()now returns the tile shape it selected, sothe compiled kernel set is exactly the selectable set. Three rows, each at two
tile widths:
q4_0,q8_0q4_0,q4_1,q5_0,q5_1,q8_0q4_0,q8_0Each shape is what
switch_ncols1/switch_ncols2would pick inside thosebounds.
switch_ncols1()capsncols1 * ncols2at 32 on Turing, so every rowloses half its columns there; above Turing the AMD escape cannot fire and the
16 column shapes collapse to 2. The gate is
turing_mma_available()and thewidth comes from
ampere_mma_available(), which is the same question the escapeasks, so the compiled set stays exactly the selectable set on both.
Removed because nothing can select them: every mixed K/V pair and the runtime-V
kernel that served them, D=64 and D=128, D=512 for
q4_1/q5_0/q5_1, and theremaining unreachable column shapes.
fattn-mma-quant-decl.cuhdeclares thesame rows and nothing else, so a disagreement with the gate is a link error.
Sizes measured with the Ampere and Ada widths only, before Turing was added:
libggml-cuda.so01b141fc, default01b141fc, all-quantsThe Turing widths add one shape per row, so the current branch compiles 12
native kernels in a default build and 18 with all-quants. The library has
not been re-measured with them; the inventory checker asserts the counts.
Review items
FATTN_MMA_QUANT_AVAILABLEkeepsfattn.cufrom naming the kernels therepushandworkflow_dispatchggml_cuda_fattn_native_profitable(), with a case asserting the fallbackThe inline comment on the route detection in
test-backend-opsis addressedtoo: the cases no longer infer the route from the output allocation size. The
CUDA backend counts its native launches and exposes that through
get_proc_address, and each case asserts the delta.Validation
RTX 4070 (
sm_89) and RTX 3060 (sm_86), CUDA 13.3,sm_86;sm_89, Release.Turing results are in their own section below.
test-backend-ops -o FLASH_ATTN_EXT: 2936/2936 default and 3949/3949all-quants, on both GPUs.
-o NATIVE_QUANT_EQUIVALENCE: 9/9 and 15/15, withthe native and fallback counts asserted. Inventory: 12 and 18 cases, exact.
Regenerating the instance files reproduces the committed ones.
Turing
Quadro RTX 8000 (TU102,
sm_75), default build.-o NATIVE_QUANT_EQUIVALENCEpasses all nine cases: the six that must take thenative route and the three that must stay on the F16 path, each asserted against
the launch counter. Qwen3.8-27B (D=256, GQA 6) with a
q4_0cache takes theroute and generates coherent text, at 491-588 t/s for
pp512and about 28 t/sfor
tg64.Those are absolute numbers with no route-off build beside them, so they say the
route is correct on Turing, not that it is faster there. Turing keeps the Ada
thresholds in
ggml_cuda_fattn_native_profitable()for that reason. The D=256regression on Ampere comes from the native loaders forcing
nstages = 0, whichcosts a two-stage cp.async pipeline; Turing has no cp.async, so its F16 path
already runs at
nstages = 0and gives up nothing. That predicts no regressionrather than a gain, and it is reasoning, not measurement.
A
75-realarm in the CI matrix compiles the Turing kernel bodies, which take adifferent MMA config table from the sm_89 arms.
Throughput
Kernel-level, native against the F16-casting path with the cast kernel included
in both timings. Rows that stay on the F16 path in both builds move by at most
0.4% on the 4070 and 1.7% on the 3060, which is the noise floor.
q4_0, n_kv 16384q4_0, n_kv 16384q4_0, n_kv 1024q4_0, n_kv 1024q8_0, n_kv 512q8_0, n_kv 512q4_0, n_kv 1024q4_0, n_kv 1024q8_0, n_kv 1024q8_0, n_kv 1024q4_0, n_kv 4096q8_0, n_kv 4096Every row is faster on Ada. On Ampere the D=512 rows are the largest win of any
row on either card and the D=256 rows are slower.
The cause is the loading pipeline, not the loaders: every D=256 entry in the MMA
config table sets
nstages_target = 2and every D=512 entry sets1. Thenative loaders write the tile themselves, so they force
nstages = 0. At D=256that gives up a real two-stage cp.async pipeline, at D=512 there is none to give
up. Ada absorbs the loss and Ampere does not. The D=256 rows are kept on Ampere
anyway; staging the quantized tiles through cp.async is the follow-up that would
remove the tradeoff.
End to end, Qwen3.8-27B-UD-IQ2_M (D=256, 24 heads, 4 KV heads) with a
q4_0cache on one GPU, t/s, route asserted by the launch counter:
pp512pp2048 @ d16384tg64 @ d16384"off" is this branch with
ggml_cuda_fattn_native_profitable()returning false,so only the route differs.
Memory
The transient F16 copy this route removes did not change any measured
allocation on this base, which corrects the claim #50 and #55 lead with.
Reserve compute buffer for the model above at 16K context on one 4070 is
505.28 MiB with a
q4_0cache, 505.28 MiB withq8_0, and 505.02 MiB withf16, which has no copy to remove at all. Peak device memory sampled duringpp2048 @ d16384is 10419 MiB with the route on and with it off. Another nodesets the high-water mark and the copy never reaches it.
The route is worth taking for the throughput, not for the memory. A model with
more KV heads, or a tree where attention dominates the compute buffer, may still
show the saving; nothing here measures that.
Not changed
fattn_quant_load_width<GGML_TYPE_Q4_0>is 8 fornthreads == 128, which isevery route row except D=512 on Ampere and Ada. Width 16 measured faster on the
4070 in all six affected D=256 cases (-1.1% to -3.8%) and was neutral on the
3060. It is a hand-tuned constant, so it is left alone and reported here rather
than changed. The Turing rows land on width 8 as well and were not swept.
AI assistance: the pruning, the Turing tile widths, the tests and the
measurement runs 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:.