cuda: compile only the reachable native FlashAttention routes (follow-up to #55) - #72
Closed
Piggidragon wants to merge 13 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
… route test-backend-ops proves the kernels and the dispatch. It does not say whether a real model still reads well with the route active, or whether the route is worth anything on Turing, which is the open question there. The script reads the model geometry before it measures. The route only covers head dim 256 and 512, so a head-dim-128 model would time the unchanged F16 path and report a number that means nothing; it names the row the model lands on, or refuses to measure and says why. Then it generates once with a quantized cache and once with f16 for a side by side read, and benches prefill and generation. Generation is a batch of one and can never take the route, so it doubles as a control. --ab builds a second copy with the route disabled and reports the delta. That is the only honest speedup number, so it is opt-in rather than implied. Assisted-by: Claude Opus 5
The route is only taken for a range of KV lengths, and the range depends on the cache type. At D=256 with a GQA ratio above 4, q4_0 covers n_kv <= 1024 and n_kv >= 16384, while q8_0 only reaches 512. The script defaulted to q8_0 with a context of 1024, so half the run measured the F16 path it was supposed to be compared against. Default to q4_0, derive the context and the benchmark depths from the row the model lands on, and say which window is being measured. --deep moves to the long-context window, which is where the route showed its largest gain on Ada. Assisted-by: Claude Opus 5
The two dispatch helpers picked the remaining branch whenever the first one did not match: "not the wide ncols1" meant the narrow one, and "Q4_0/Q8_0 and not D=512" meant the D=256 GQA-2 row. A gate that drifted from the table would then launch a wrong-geometry kernel instead of hitting the GGML_ABORT below, so a supported() that grew a D=128 row would read past the K/V rows with the <256,256,32,2> kernel. Assert the shape instead. Also drop the Q5_0 load width. Its 32-wide branch keyed on the D=256 GQA-2 row, which is Q4_0/Q8_0 only, so Q5_0 can only reach ncols2 == 8 and the branch was dead in every build. The remaining value is the default. Assisted-by: Claude Opus 5
The comment and the docs read as if D=256 GQA 8 were a geometry the route cannot serve. It is not: switch_ncols2() maps every gqa_ratio > 4 to ncols2 == 8, so GQA 6 and 16 select the same <256,256,8,8> instance and keep the route, and the tests assert that for GQA 6. What PR 55 records is one measured Q8_0 case at that ratio with an open correctness and memory-safety question. The guard follows that undiagnosed case, not a property of the kernel, and covers every type until it is closed. Say so. The route table also listed the exclusion under supported(), which does not apply it. It lives in profitable(), so move it there. Assisted-by: Claude Opus 5
Two gaps in the native FlashAttention workflow. The Turing arm was pinned to ALL_QUANTS=OFF, so the Turing tile shape <256,256,4,8> was never compiled for Q4_1/Q5_0/Q5_1. Turing is the arm's reason to exist: it keys a different MMA config table, and the tile loader asserts and the per-type load widths are functions of that table. Add the arm that crosses it with the extra tier. The focused run then piped test-backend-ops into tee. GitHub runs the block under bash -e, not -o pipefail, so the step took its status from tee: a crash or a GGML_ABORT after the counted lines printed still passed. Set pipefail. Assisted-by: Claude Opus 5
Piggidragon
marked this pull request as ready for review
September 8, 2026 15:13
Author
|
Superseded by #85, which rebases this work onto current |
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:.