Skip to content

cuda: compile only the reachable native FlashAttention routes (follow-up to #55) - #72

Draft
Piggidragon wants to merge 8 commits into
GenerelSchwerz:codex/pr50-native-quant-integrationfrom
Piggidragon:llama/fattn-native-route-table
Draft

cuda: compile only the reachable native FlashAttention routes (follow-up to #55)#72
Piggidragon wants to merge 8 commits into
GenerelSchwerz:codex/pr50-native-quant-integrationfrom
Piggidragon:llama/fattn-native-route-table

Conversation

@Piggidragon

@Piggidragon Piggidragon commented Sep 5, 2026

Copy link
Copy Markdown

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, so
the compiled kernel set is exactly the selectable set. Three rows, each at two
tile widths:

Head dim GQA ratio Query batch Cache types Tile (sm_80+) Tile (Turing)
256 2 > 16 q4_0, q8_0 32x2 16x2
256 > 4, not 8 > 4 q4_0, q4_1, q5_0, q5_1, q8_0 8x8 4x8
512 > 4 > 4 q4_0, q8_0 8x8 4x8

Each shape is what switch_ncols1/switch_ncols2 would pick inside those
bounds. switch_ncols1() caps ncols1 * ncols2 at 32 on Turing, so every row
loses 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 the
width comes from ampere_mma_available(), which is the same question the escape
asks, 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 the
remaining unreachable column shapes. fattn-mma-quant-decl.cuh declares the
same 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:

Build libggml-cuda.so vs base Native kernels
base 01b141fc, default 122,779,344 B 0
#55, default 158,431,640 B +29.04% 98
this, default 125,252,112 B +2.01% 6
base 01b141fc, all-quants 182,828,488 B 0
#55, all-quants 420,927,440 B +130.23% 485
this, all-quants 186,695,944 B +2.12% 9

The 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

# Item Change
1 HIP and MUSA compile the CUDA-only instances Excluded from both globs; FATTN_MMA_QUANT_AVAILABLE keeps fattn.cu from naming the kernels there
2 Self-hosted runner executes PR payloads Focused job restricted to push and workflow_dispatch
3 Focused test selects zero cases and passes Filters on the current identity, asserts exact counts (6 native / 3 fallback default, 9 / 6 all-quants)
4 D=256 GQA 8 enabled despite an open question Declined in ggml_cuda_fattn_native_profitable(), with a case asserting the fallback
5 Documentation describes a removed interface Rewritten for the automatic route
6 Generated matrix much larger than the route table Above
7 Inventory checker cannot prove exactness Exact comparison against the declared cases; fails on missing, unexpected, duplicate, mixed K/V or logit-softcap kernels
8 Performance evidence does not validate the final route Re-measured with an explicit route assertion, below
9 Hardware gate broader than the validation Stops below Hopper; Ampere and Ada measured for throughput, Turing verified correct on hardware

The inline comment on the route detection in test-backend-ops is addressed
too: 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/3949
all-quants, on both GPUs. -o NATIVE_QUANT_EQUIVALENCE: 9/9 and 15/15, with
the 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_EQUIVALENCE passes all nine cases: the six that must take the
native 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_0 cache takes the
route and generates coherent text, at 491-588 t/s for pp512 and about 28 t/s
for 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=256
regression on Ampere comes from the native loaders forcing nstages = 0, which
costs a two-stage cp.async pipeline; Turing has no cp.async, so its F16 path
already runs at nstages = 0 and gives up nothing. That predicts no regression
rather than a gain, and it is reasoning, not measurement.

A 75-real arm in the CI matrix compiles the Turing kernel bodies, which take a
different 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.

Route row n_q 4070 (Ada) 3060 (Ampere)
D=256, GQA 6, q4_0, n_kv 16384 512 -21.5% +0.3%
D=256, GQA 6, q4_0, n_kv 16384 2048 -15.8% +2.4%
D=256, GQA 6, q4_0, n_kv 1024 512 -6.8% +6.7%
D=256, GQA 6, q4_0, n_kv 1024 2048 -3.7% +10.8%
D=256, GQA 6, q8_0, n_kv 512 512 -3.5% +9.3%
D=256, GQA 6, q8_0, n_kv 512 2048 +0.9% +13.9%
D=256, GQA 2, q4_0, n_kv 1024 512 -17.2% +1.9%
D=256, GQA 2, q4_0, n_kv 1024 2048 -9.5% +7.5%
D=256, GQA 2, q8_0, n_kv 1024 512 -17.0% +1.0%
D=256, GQA 2, q8_0, n_kv 1024 2048 -7.0% +9.2%
D=512, GQA 16, q4_0, n_kv 4096 512 -12.2% -15.6%
D=512, GQA 16, q8_0, n_kv 4096 512 -6.4% -5.1%

Every 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 = 2 and every D=512 entry sets 1. The
native loaders write the tile themselves, so they force nstages = 0. At D=256
that 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_0
cache on one GPU, t/s, route asserted by the launch counter:

Test 4070 off 4070 on 3060 off 3060 on
pp512 1179.46 1178.58 534.05 533.54
pp2048 @ d16384 982.61 1022.77 453.35 450.35
tg64 @ d16384 33.27 33.34 17.53 17.62

"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_0 cache, 505.28 MiB with q8_0, and 505.02 MiB with
f16, which has no copy to remove at all. Peak device memory sampled during
pp2048 @ d16384 is 10419 MiB with the route on and with it off. Another node
sets 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 for nthreads == 128, which is
every 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:.

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
@github-actions github-actions Bot added documentation Improvements or additions to documentation testing CUDA ggml devops labels Sep 5, 2026
@Piggidragon
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CUDA devops documentation Improvements or additions to documentation ggml testing

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant