FA4 consumer Blackwell (sm_120) integration

thad0ctor · 2026-06-08T23:43:59Z

Adds FlashAttention-4 (CuTeDSL) support for consumer/workstation Blackwell — RTX 50-series / RTX PRO 6000 / DGX Spark (compute capability 12.x). Dispatch and tile selection are auto-tuned for the arch; no environment variables are required for normal use.

Background

Consumer/workstation Blackwell (sm_120) ships 5th-gen tensor cores with new fp4/fp8 data types but — unlike datacenter Blackwell (sm_100) — keeps the Ampere-style warp-level mma.sync programming model: no WGMMA, no tcgen05, no Tensor Memory (confirmed against the CUDA 13.x toolkit and ptxas SASS — the tensor-memory subsystem is absent in silicon, not just unexposed). So FA4's sm_100 path does not apply; sm_120 reuses the SM80-base kernels (cp.async, plus a TMA + warp-specialized forward) and needs its own dispatch, tiling, and correctness work.

This builds on the in-flight upstream FA4 sm_120 kernels (Dao-AILab Dao-AILab#2553 / Dao-AILab#2349 / Dao-AILab#2389), which landed initial kernels but left public-API routing, correctness, and coverage gaps:

Area	Upstream (Dao-AILab#2553 / Dao-AILab#2349 / Dao-AILab#2389)	This PR
Dense forward	kernels exist; public routing hits dispatcher/runtime gaps	works across fixed + varlen `sm_120` paths via the public API, auto-tuned dispatch
Backward	not usable end-to-end for the coverage target	supported head dims work through the public API, incl. D256 dense/causal
Packed GQA	crashes / disabled on several `sm_120` GQA shapes	works where valid; invalid cases auto-downgrade; pack_gqa + SplitKV decode now works
Paged KV	SM80-base path takes the args but doesn't implement `sm_120` correctly	real paged-KV forward through `head_dim ≤ 256`, validated
`head_dim > head_dim_v`	can route into the `sm_120` TMA hang	routed to the non-TMA path; hang-guard tests included
fp8 KV-cache decode	—	new memory-efficient fp8 e4m3/e5m2 decode (≈1.6–1.9× at GQA ratio ≤ 4, half the KV bandwidth)
Correctness / hardening	—	8 real bugs found and fixed (below); strict arch-gating (non-`sm_120` byte-identical to `main`); CodeRabbit follow-ups

What's supported

Forward & backward: dense, causal, local/sliding-window; MHA / GQA / MQA; variable-length; paged-KV; block sparsity; score_mod / mask_mod; learnable sink. Head dims 64/96/128/192/256.
fp8 KV-cache decode (e4m3/e5m2): for seqlen_q == 1 with a quantized K/V cache + bf16/fp16 query (pass fp8 k/v + per-(batch, kv_head) k_descale/v_descale). Auto-routes to a memory-efficient GEMV decode kernel; ~1.6–1.9× faster than bf16 at GQA ratio ≤ 4 while halving KV bandwidth, within ~2e-3 of an fp8-quantized reference.
torch.compile compatible: the FA4 (CuTe-DSL) public entry points are made opaque to TorchDynamo (graph-break → run eagerly) and tensor max_seqlen is coerced to a host int in the varlen backward, so full-model torch.compile works for forward, backward, and variable-length with output bit-identical to eager — no symbolic/fake tensors leak into the DSL const() path. Non-FA4 / non-sm120 kernels are behaviourally unchanged.

Performance vs FA2 (and SDPA)

Two sm_120 parts: RTX 5090 (consumer, 170 SM, 32 GB) and RTX PRO 6000 Blackwell (workstation, 188 SM, 96 GB). bf16. In-process round-wise interleaved A/B — each round times the impls back-to-back under the same clock state, so the ratios are clock-robust; medians reported.

Legend. D = head dim · B = batch size · r = GQA ratio = query-heads per KV-head (r1/MHA = one KV head per query head; r8 = 8 query heads share 1 KV head; MQA = all heads share 1) · dense = full attention, causal = causal mask · prefill = seqlen_q == seqlen_k, decode = seqlen_q == 1 (one new query against a length-Sk KV cache) · ratio = FA4 / FA2 throughput, >1 = FA4 faster (so 1.15 = FA4 15% faster, 0.40 = FA4 60% slower). Config labels are (Hq/Hkv). Short seqlens have higher run-to-run variance — read the trend.

Prefill forward — per seqlen (B=1)

The headline is D256; D128 is the other common case. (Full matrix incl. D64/D96/MHA and B=2 in the collapsibles below.)

D256 GQA (16/2), FA4/FA2:

card · mode	S512	S1024	S2048	S4096	S8192	S16384	S32768
5090 dense	1.34	1.34	1.07	0.95	1.01	1.04	1.06
6000 dense	1.66	1.32	1.03	1.16	1.13	1.15	1.15
5090 causal	1.18	1.31	1.06	0.95	1.02	1.05	1.05
6000 causal	1.46	1.26	1.05	1.06	1.12	1.13	1.14

D128 GQA (32/4), FA4/FA2:

card · mode	S512	S1024	S2048	S4096	S8192	S16384	S32768
5090 dense	1.07	1.04	1.03	1.02	1.03	1.02	1.03
6000 dense	1.46	1.08	1.01	1.01	1.00	1.00	1.00
5090 causal	1.07	1.20	1.07	1.06	1.04	1.02	1.02
6000 causal	1.34	1.17	1.07	1.04	1.02	1.00	0.99

Geomean rollup (over seqlen 512–32k, FA4/FA2)

	5090 B1	5090 B2	6000 B1	6000 B2
D128 GQA fwd dense	1.03	1.03	1.07	1.02
D128 GQA fwd causal	1.06	1.06	1.08	1.04
D256 GQA fwd dense	1.11	1.03	1.21	1.17
D256 GQA fwd causal	1.08	1.05	1.16	1.14
D128 GQA bwd dense / causal	1.01 / 1.02	1.01 / 1.02	1.01 / 1.01	1.01 / 1.02
D256 GQA bwd dense / causal	1.06 / 1.01	1.07 / 1.01	1.00 / 1.00	1.01 / 1.01
D256 MHA bwd causal (floor)	0.94	0.96	0.93	0.94

Full per-seqlen prefill matrix — RTX 5090 (all configs, B=1 & B=2, fwd+bwd, dense+causal)

--- PREFILL FWD dense   B=1  (fa4/fa2, median[min..max], >1=FA4 faster) ---
shape                           S512            S1024            S2048            S4096            S8192           S16384           S32768
D64  MHA  (32/32)   1.072[1.03-1.09] 1.016[1.01-1.02] 0.983[0.98-0.99] 0.980[0.97-0.99] 0.984[0.95-1.00] 0.980[0.97-0.99] 0.923[0.88-0.98]  geo=0.990
D64  GQA8 (32/4)    1.088[1.03-1.22] 1.052[1.04-1.06] 1.019[1.01-1.02] 1.004[1.00-1.01] 0.996[0.94-1.05] 0.997[0.99-1.00] 0.995[0.98-1.04]  geo=1.021
D96  MHA  (32/32)   1.056[0.95-1.07] 0.969[0.95-1.16] 0.931[0.93-0.94] 0.920[0.89-0.94] 0.928[0.92-0.93] 0.926[0.92-0.93] 0.927[0.91-0.95]  geo=0.950
D96  GQA8 (32/4)    1.101[1.09-1.13] 1.052[1.04-1.06] 1.021[1.01-1.02] 1.005[1.00-1.01] 1.002[0.99-1.01] 0.997[0.99-1.00] 0.997[0.96-1.04]  geo=1.024
D128 MHA  (32/32)   1.066[1.04-1.07] 1.031[1.02-1.04] 1.007[1.00-1.01] 0.997[0.99-1.01] 0.995[0.99-1.01] 0.997[0.98-1.00] 0.994[0.98-1.00]  geo=1.012
D128 GQA4 (32/8)    1.070[1.06-1.08] 1.042[1.03-1.05] 1.030[1.02-1.04] 1.019[1.01-1.02] 1.026[1.02-1.05] 1.026[1.02-1.03] 1.014[0.99-1.02]  geo=1.032
D128 GQA8 (32/4)    1.065[1.06-1.08] 1.039[1.00-1.11] 1.033[1.03-1.04] 1.022[1.01-1.04] 1.027[1.02-1.03] 1.024[1.02-1.03] 1.026[1.03-1.03]  geo=1.034
D256 MHA  (16/16)   1.293[1.28-1.30] 1.377[1.37-1.39] 1.058[1.06-1.06] 0.940[0.93-0.94] 1.004[1.00-1.01] 1.037[1.03-1.04] 1.050[1.04-1.06]  geo=1.099
D256 GQA4 (16/4)    1.328[1.31-1.39] 1.350[1.33-1.36] 1.070[1.07-1.07] 0.946[0.94-0.95] 1.007[0.99-1.02] 1.044[1.03-1.06] 1.059[1.06-1.06]  geo=1.106
D256 GQA8 (16/2)    1.336[1.33-1.36] 1.343[1.33-1.37] 1.071[1.06-1.08] 0.948[0.95-0.95] 1.012[1.01-1.03] 1.042[1.04-1.05] 1.056[1.05-1.06]  geo=1.106

--- PREFILL FWD causal  B=1  (fa4/fa2, median[min..max], >1=FA4 faster) ---
shape                           S512            S1024            S2048            S4096            S8192           S16384           S32768
D64  MHA  (32/32)   1.075[1.03-1.09] 1.354[1.34-1.37] 1.309[1.26-1.32] 1.128[1.11-1.15] 0.974[0.94-0.99] 0.986[0.98-0.99] 0.880[0.88-0.88]  geo=1.089
D64  GQA8 (32/4)    1.037[1.01-1.06] 1.086[1.08-1.12] 1.007[0.99-1.02] 1.029[1.01-1.04] 1.012[1.00-1.02] 1.029[1.02-1.03] 1.016[1.01-1.03]  geo=1.031
D96  MHA  (32/32)   0.968[0.84-1.06] 1.340[1.33-1.36] 1.205[1.19-1.23] 1.115[1.10-1.14] 0.990[0.96-1.01] 0.918[0.91-0.92] 0.887[0.88-0.89]  geo=1.050
D96  GQA8 (32/4)    1.031[1.00-1.07] 1.100[1.06-1.11] 1.040[1.03-1.05] 1.010[1.00-1.02] 1.008[1.00-1.02] 1.017[1.01-1.04] 1.006[1.00-1.07]  geo=1.030
D128 MHA  (32/32)   1.043[1.02-1.14] 1.264[1.26-1.28] 1.170[1.14-1.18] 1.111[1.09-1.12] 1.021[1.01-1.05] 0.969[0.96-0.98] 0.950[0.94-0.97]  geo=1.071
D128 GQA4 (32/8)    1.078[1.05-1.08] 1.207[1.20-1.30] 1.057[1.05-1.07] 1.028[0.99-1.03] 1.007[1.00-1.01] 1.020[1.01-1.03] 1.015[0.97-1.04]  geo=1.057
D128 GQA8 (32/4)    1.066[1.05-1.08] 1.199[1.18-1.20] 1.071[1.06-1.07] 1.056[1.05-1.06] 1.035[1.03-1.07] 1.024[0.99-1.05] 1.021[1.02-1.03]  geo=1.066
D256 MHA  (16/16)   1.149[1.13-1.27] 1.456[1.42-1.56] 1.013[1.00-1.16] 0.905[0.89-0.91] 0.971[0.96-1.00] 1.027[1.00-1.03] 1.049[1.01-1.05]  geo=1.070
D256 GQA4 (16/4)    1.190[1.17-1.20] 1.332[1.32-1.34] 1.037[1.02-1.11] 0.937[0.93-0.94] 1.007[0.97-1.04] 1.044[1.04-1.05] 1.058[1.05-1.07]  geo=1.080
D256 GQA8 (16/2)    1.176[1.17-1.19] 1.310[1.30-1.32] 1.062[1.06-1.07] 0.948[0.95-0.95] 1.017[0.97-1.07] 1.045[1.03-1.05] 1.053[1.05-1.08]  geo=1.082

--- PREFILL FWD dense   B=2  (fa4/fa2, median[min..max], >1=FA4 faster) ---
shape                           S512            S1024            S2048            S4096            S8192           S16384           S32768
D64  MHA  (32/32)   1.011[0.96-1.02] 1.083[1.07-1.09] 0.971[0.96-0.98] 0.883[0.88-0.90] 0.971[0.95-0.98] 0.977[0.97-0.98] 0.919[0.90-0.92]  geo=0.972
D64  GQA8 (32/4)    1.066[1.06-1.09] 1.034[1.02-1.05] 1.010[1.01-1.02] 1.002[1.00-1.01] 1.007[0.97-1.06] 0.995[0.96-1.00] 0.994[0.96-1.00]  geo=1.015
D96  MHA  (32/32)   0.994[0.99-1.00] 0.938[0.92-0.94] 0.917[0.91-0.93] 0.900[0.88-0.92] 0.924[0.91-0.94] 0.924[0.91-0.93] 0.925[0.91-0.93]  geo=0.931
D96  GQA8 (32/4)    1.065[1.05-1.08] 1.031[1.03-1.04] 1.012[1.00-1.02] 0.998[0.99-1.00] 1.001[0.99-1.01] 0.995[0.97-1.01] 0.996[1.00-1.00]  geo=1.014
D128 MHA  (32/32)   1.047[1.04-1.07] 1.016[1.01-1.03] 0.998[0.99-1.00] 0.994[0.99-1.01] 0.995[0.99-1.00] 0.995[0.99-1.00] 0.996[1.00-1.00]  geo=1.006
D128 GQA4 (32/8)    1.091[1.07-1.12] 1.040[1.03-1.04] 1.022[0.99-1.02] 1.012[0.99-1.02] 1.024[1.02-1.04] 1.024[0.98-1.03] 1.015[1.00-1.02]  geo=1.032
D128 GQA8 (32/4)    1.096[1.08-1.11] 1.042[1.04-1.04] 1.025[1.02-1.03] 1.014[1.00-1.04] 1.027[1.02-1.03] 1.015[1.01-1.02] 1.016[1.02-1.02]  geo=1.033
D256 MHA  (16/16)   1.018[1.01-1.03] 1.014[1.01-1.02] 0.935[0.93-0.94] 1.009[0.99-1.02] 1.040[1.04-1.04] 1.055[1.05-1.06] 1.061[1.06-1.07]  geo=1.018
D256 GQA4 (16/4)    1.062[1.05-1.07] 1.040[1.04-1.04] 0.946[0.94-0.95] 0.999[0.99-1.01] 1.048[1.04-1.05] 1.062[1.06-1.07] 1.071[1.07-1.07]  geo=1.032
D256 GQA8 (16/2)    1.058[1.05-1.08] 1.035[1.02-1.05] 0.946[0.94-0.95] 1.003[1.00-1.01] 1.049[1.05-1.05] 1.062[1.06-1.07] 1.072[1.07-1.07]  geo=1.031

--- PREFILL FWD causal  B=2  (fa4/fa2, median[min..max], >1=FA4 faster) ---
shape                           S512            S1024            S2048            S4096            S8192           S16384           S32768
D64  MHA  (32/32)   1.219[1.20-1.23] 1.231[1.21-1.24] 1.135[1.12-1.15] 0.964[0.96-0.97] 0.874[0.86-0.90] 0.959[0.95-0.96] 0.870[0.87-0.87]  geo=1.026
D64  GQA8 (32/4)    1.037[1.02-1.05] 1.047[1.04-1.06] 1.008[1.00-1.02] 1.008[1.00-1.01] 1.014[1.01-1.01] 1.016[1.01-1.02] 1.011[1.01-1.01]  geo=1.020
D96  MHA  (32/32)   1.202[1.19-1.21] 1.150[1.09-1.17] 1.070[1.07-1.07] 0.982[0.97-0.99] 0.916[0.91-0.93] 0.891[0.89-0.89] 0.875[0.87-0.88]  geo=1.005
D96  GQA8 (32/4)    1.028[1.01-1.03] 1.094[1.09-1.10] 1.017[1.01-1.03] 1.008[1.00-1.01] 1.005[0.97-1.02] 1.005[1.00-1.01] 0.998[1.00-1.00]  geo=1.022
D128 MHA  (32/32)   1.190[1.17-1.20] 1.181[1.17-1.19] 1.081[1.08-1.08] 1.017[1.01-1.03] 0.972[0.97-0.98] 0.952[0.95-0.95] 0.941[0.94-0.94]  geo=1.043
D128 GQA4 (32/8)    1.159[1.15-1.17] 1.153[1.14-1.17] 1.054[1.05-1.06] 1.027[1.02-1.03] 1.024[1.02-1.02] 1.008[1.00-1.03] 1.014[0.99-1.02]  geo=1.061
D128 GQA8 (32/4)    1.076[1.05-1.09] 1.146[1.12-1.16] 1.062[1.06-1.07] 1.047[1.04-1.06] 1.032[1.00-1.07] 1.022[1.01-1.03] 1.007[0.98-1.02]  geo=1.055
D256 MHA  (16/16)   1.032[1.00-1.04] 1.015[0.97-1.04] 1.003[1.00-1.01] 0.957[0.95-0.97] 1.014[0.99-1.02] 1.022[0.95-1.08] 1.050[1.03-1.07]  geo=1.013
D256 GQA4 (16/4)    1.064[1.03-1.08] 1.042[1.04-1.05] 1.028[1.03-1.04] 0.974[0.93-1.01] 1.044[1.02-1.07] 1.046[1.02-1.08] 1.063[1.05-1.07]  geo=1.037
D256 GQA8 (16/2)    1.068[1.05-1.07] 1.058[1.05-1.06] 1.042[1.04-1.04] 1.006[1.00-1.01] 1.048[1.03-1.06] 1.062[1.06-1.07] 1.068[1.04-1.08]  geo=1.050

--- PREFILL BWD dense   B=1  (fa4/fa2, median[min..max], >1=FA4 faster) ---
shape                           S512            S1024            S2048            S4096            S8192           S16384           S32768
D64  MHA  (32/32)   1.078[1.07-1.35] 1.057[0.96-1.16] 1.101[1.06-1.15] 0.998[0.97-1.01] 0.900[0.89-0.91] 0.878[0.87-0.88] 0.869[0.86-0.88]  geo=0.979
D64  GQA8 (32/4)    0.993[0.97-1.04] 0.994[0.89-1.11] 1.083[1.08-1.09] 1.029[1.02-1.07] 0.909[0.89-0.92] 0.881[0.88-0.88] 0.869[0.87-0.87]  geo=0.963
D96  MHA  (32/32)   1.048[0.89-1.20] 1.027[0.90-1.04] 1.083[1.07-1.09] 1.011[1.00-1.02] 0.956[0.94-0.97] 0.939[0.94-0.94] 0.933[0.93-0.93]  geo=0.998
D96  GQA8 (32/4)    0.972[0.78-1.08] 0.983[0.95-0.99] 1.072[1.06-1.09] 1.011[0.98-1.02] 0.963[0.96-0.97] 0.940[0.94-0.94] 0.929[0.93-0.96]  geo=0.980
D128 MHA  (32/32)   1.105[1.05-1.39] 1.058[0.98-1.20] 1.029[0.99-1.04] 1.009[1.00-1.02] 1.013[1.01-1.02] 1.008[1.01-1.01] 1.007[1.01-1.01]  geo=1.032
D128 GQA4 (32/8)    1.009[0.85-1.03] 1.026[0.94-1.05] 1.015[0.97-1.02] 1.003[0.99-1.01] 1.009[1.01-1.01] 1.007[1.01-1.01] 1.005[1.00-1.03]  geo=1.011
D128 GQA8 (32/4)    1.015[1.01-1.02] 1.014[0.95-1.06] 1.016[1.00-1.05] 1.012[0.99-1.04] 1.013[1.01-1.02] 1.009[1.01-1.01] 1.006[1.00-1.01]  geo=1.012
D256 MHA  (16/16)   1.089[1.02-1.28] 1.015[0.96-1.03] 0.959[0.91-0.96] 1.029[1.02-1.04] 1.066[1.06-1.07] 1.086[1.07-1.12] 1.088[1.08-1.10]  geo=1.046
D256 GQA4 (16/4)    1.009[0.86-1.02] 1.042[0.98-1.08] 0.965[0.95-1.00] 1.031[1.02-1.04] 1.089[1.09-1.09] 1.079[1.08-1.08] 1.082[1.08-1.08]  geo=1.042
D256 GQA8 (16/2)    1.056[0.91-1.17] 1.024[1.02-1.03] 0.984[0.93-1.01] 1.054[1.05-1.06] 1.092[1.09-1.10] 1.108[1.11-1.11] 1.109[1.11-1.11]  geo=1.060

--- PREFILL BWD causal  B=1  (fa4/fa2, median[min..max], >1=FA4 faster) ---
shape                           S512            S1024            S2048            S4096            S8192           S16384           S32768
D64  MHA  (32/32)   1.075[1.05-1.19] 1.137[1.09-1.24] 1.092[1.08-1.11] 1.014[1.00-1.02] 0.939[0.93-0.95] 0.902[0.89-0.91] 0.886[0.86-0.90]  geo=1.002
D64  GQA8 (32/4)    0.988[0.98-1.02] 1.075[0.93-1.11] 1.049[1.02-1.11] 1.038[1.02-1.04] 0.942[0.93-0.96] 0.914[0.91-0.92] 0.902[0.86-0.91]  geo=0.985
D96  MHA  (32/32)   1.132[1.03-1.32] 1.157[1.13-1.18] 1.111[1.10-1.12] 1.035[1.01-1.05] 0.965[0.95-0.98] 0.937[0.93-0.94] 0.927[0.92-0.93]  geo=1.034
D96  GQA8 (32/4)    1.049[0.96-1.23] 1.084[0.96-1.24] 1.089[1.04-1.17] 1.019[1.01-1.04] 0.979[0.96-0.99] 0.940[0.94-0.95] 0.920[0.91-0.95]  geo=1.010
D128 MHA  (32/32)   1.096[1.07-1.11] 1.070[1.07-1.13] 1.029[0.99-1.06] 1.013[1.01-1.03] 1.007[1.00-1.01] 1.000[1.00-1.00] 1.004[1.00-1.02]  geo=1.031
D128 GQA4 (32/8)    1.030[0.97-1.06] 1.025[0.93-1.05] 1.017[0.98-1.03] 1.012[0.98-1.03] 0.999[0.99-1.02] 1.000[1.00-1.01] 0.997[0.98-1.02]  geo=1.011
D128 GQA8 (32/4)    1.057[0.96-1.11] 1.016[0.96-1.14] 1.026[0.89-1.10] 1.012[0.98-1.04] 1.014[1.00-1.02] 1.002[1.00-1.05] 0.998[1.00-1.01]  geo=1.018
D256 MHA  (16/16)   0.947[0.75-1.12] 0.862[0.85-0.88] 0.887[0.86-0.90] 0.924[0.91-0.95] 0.965[0.96-0.98] 0.997[0.94-1.00] 1.002[0.98-1.03]  geo=0.939
D256 GQA4 (16/4)    0.896[0.85-1.05] 0.970[0.88-0.99] 1.017[1.00-1.05] 0.974[0.94-1.00] 1.023[1.01-1.03] 1.058[1.05-1.06] 1.064[1.03-1.08]  geo=0.999
D256 GQA8 (16/2)    0.923[0.82-1.09] 1.026[0.97-1.03] 1.027[1.00-1.03] 0.983[0.94-1.01] 1.032[1.02-1.04] 1.058[1.06-1.06] 1.059[1.06-1.06]  geo=1.014

--- PREFILL BWD dense   B=2  (fa4/fa2, median[min..max], >1=FA4 faster) ---
shape                           S512            S1024            S2048            S4096            S8192           S16384           S32768
D64  MHA  (32/32)   1.080[1.07-1.09] 1.115[1.11-1.12] 1.026[1.01-1.05] 0.944[0.93-0.96] 0.872[0.87-0.87] 0.863[0.86-0.87] 0.860[0.86-0.87]  geo=0.961
D64  GQA8 (32/4)    1.002[0.99-1.02] 1.078[1.07-1.08] 1.038[1.00-1.05] 0.969[0.96-0.98] 0.884[0.88-0.89] 0.869[0.86-0.87] 0.861[0.86-0.86]  geo=0.954
D96  MHA  (32/32)   1.045[1.04-1.06] 1.087[1.02-1.10] 1.004[0.98-1.02] 0.939[0.93-0.96] 0.934[0.93-0.94] 0.928[0.92-0.93] 0.925[0.92-0.93]  geo=0.978
D96  GQA8 (32/4)    0.974[0.94-1.01] 1.080[1.05-1.19] 1.009[0.98-1.02] 0.958[0.95-0.96] 0.938[0.93-0.94] 0.924[0.92-0.93] 0.919[0.91-0.93]  geo=0.970
D128 MHA  (32/32)   1.093[1.08-1.12] 1.043[1.00-1.05] 1.012[0.97-1.03] 1.008[1.00-1.02] 1.009[1.01-1.01] 1.007[1.01-1.01] 1.011[1.00-1.01]  geo=1.026
D128 GQA4 (32/8)    1.010[0.90-1.19] 1.030[0.94-1.08] 1.013[1.00-1.03] 1.006[1.00-1.01] 1.016[0.99-1.03] 1.006[1.00-1.02] 1.010[1.00-1.04]  geo=1.013
D128 GQA8 (32/4)    1.013[0.90-1.02] 1.027[0.96-1.05] 1.023[1.02-1.03] 1.011[0.86-1.07] 1.013[1.01-1.01] 1.010[1.00-1.01] 1.006[0.99-1.01]  geo=1.014
D256 MHA  (16/16)   1.011[0.99-1.03] 0.982[0.98-0.98] 1.023[1.01-1.03] 1.063[1.05-1.11] 1.080[1.03-1.09] 1.090[1.04-1.14] 1.084[1.02-1.14]  geo=1.047
D256 GQA4 (16/4)    1.019[0.99-1.03] 0.981[0.97-0.98] 1.033[1.01-1.11] 1.068[1.06-1.08] 1.106[1.10-1.11] 1.089[1.09-1.09] 1.082[1.06-1.10]  geo=1.053
D256 GQA8 (16/2)    1.012[0.91-1.12] 0.980[0.95-1.00] 1.052[1.02-1.07] 1.085[1.06-1.09] 1.113[1.11-1.13] 1.107[1.09-1.13] 1.110[1.09-1.13]  geo=1.065

--- PREFILL BWD causal  B=2  (fa4/fa2, median[min..max], >1=FA4 faster) ---
shape                           S512            S1024            S2048            S4096            S8192           S16384           S32768
D64  MHA  (32/32)   1.104[1.09-1.14] 1.135[1.13-1.15] 1.057[0.99-1.10] 0.976[0.97-0.98] 0.896[0.89-0.91] 0.883[0.88-0.89] 0.872[0.87-0.90]  geo=0.984
D64  GQA8 (32/4)    1.028[1.00-1.08] 1.083[0.94-1.21] 1.063[1.05-1.10] 1.005[0.98-1.07] 0.913[0.90-0.92] 0.901[0.90-0.93] 0.896[0.89-0.90]  geo=0.981
D96  MHA  (32/32)   1.124[1.09-1.13] 1.141[1.06-1.26] 1.059[1.01-1.08] 0.974[0.94-1.01] 0.935[0.93-0.95] 0.922[0.88-0.97] 0.919[0.91-0.93]  geo=1.007
D96  GQA8 (32/4)    1.044[0.87-1.12] 1.109[1.05-1.18] 1.050[1.01-1.10] 0.952[0.93-0.98] 0.945[0.93-0.96] 0.930[0.89-0.94] 0.918[0.90-0.93]  geo=0.990
D128 MHA  (32/32)   1.087[0.91-1.25] 1.041[0.97-1.04] 1.022[1.02-1.02] 1.007[1.00-1.01] 1.002[1.00-1.00] 1.001[1.00-1.00] 1.000[1.00-1.00]  geo=1.022
D128 GQA4 (32/8)    1.022[1.01-1.04] 1.024[1.02-1.05] 1.022[0.97-1.04] 1.009[1.01-1.02] 1.004[1.00-1.01] 0.999[0.99-1.01] 0.996[1.00-1.02]  geo=1.011
D128 GQA8 (32/4)    1.029[0.98-1.03] 1.043[1.02-1.08] 1.035[1.03-1.05] 1.021[1.00-1.03] 1.011[1.01-1.01] 1.005[1.00-1.03] 0.997[0.99-1.03]  geo=1.020
D256 MHA  (16/16)   0.934[0.93-0.97] 0.911[0.91-0.92] 0.921[0.90-0.93] 0.955[0.93-0.96] 0.982[0.97-0.99] 0.997[0.99-1.00] 1.000[1.00-1.00]  geo=0.956
D256 GQA4 (16/4)    0.913[0.86-0.93] 1.055[1.03-1.07] 0.999[0.98-1.01] 1.003[1.00-1.01] 1.031[1.03-1.03] 1.049[1.05-1.05] 1.057[1.06-1.07]  geo=1.014
D256 GQA8 (16/2)    0.923[0.80-0.94] 1.015[0.95-1.03] 0.973[0.95-0.98] 1.014[0.99-1.03] 1.039[1.03-1.04] 1.053[1.03-1.06] 1.060[1.05-1.07]  geo=1.010

Full per-seqlen prefill matrix — RTX PRO 6000 (all configs, B=1 & B=2, fwd+bwd, dense+causal)

--- PREFILL FWD dense   B=1  (fa4/fa2, median[min..max], >1=FA4 faster) ---
shape                           S512            S1024            S2048            S4096            S8192           S16384           S32768
D64  MHA  (32/32)   1.096[1.08-1.11] 1.045[1.03-1.07] 0.928[0.92-0.93] 0.832[0.83-0.83] 0.948[0.95-0.95] 0.938[0.94-0.94] 0.954[0.95-0.95]  geo=0.960
D64  GQA8 (32/4)    1.125[1.11-1.14] 1.076[1.07-1.09] 1.016[1.01-1.02] 0.994[0.99-1.00] 0.989[0.98-0.99] 0.990[0.99-0.99] 0.989[0.99-0.99]  geo=1.024
D96  MHA  (32/32)   1.400[1.38-1.41] 0.971[0.96-1.08] 0.960[0.95-0.96] 0.936[0.93-0.94] 0.936[0.94-0.94] 0.933[0.93-0.93] 0.932[0.93-0.93]  geo=0.999
D96  GQA8 (32/4)    1.439[1.43-1.49] 1.075[1.07-1.08] 1.023[1.02-1.03] 1.004[1.00-1.01] 0.999[1.00-1.00] 0.999[1.00-1.00] 0.997[1.00-1.00]  geo=1.068
D128 MHA  (32/32)   1.440[1.43-1.45] 1.041[1.03-1.05] 1.018[1.01-1.02] 1.015[1.01-1.02] 1.007[1.01-1.01] 1.007[1.01-1.01] 1.007[1.01-1.01]  geo=1.068
D128 GQA4 (32/8)    1.471[1.46-1.48] 1.073[1.07-1.08] 1.012[1.01-1.01] 1.005[1.00-1.01] 1.002[1.00-1.00] 1.003[1.00-1.00] 1.000[1.00-1.00]  geo=1.071
D128 GQA8 (32/4)    1.456[1.45-1.48] 1.076[1.06-1.09] 1.014[1.01-1.02] 1.005[1.00-1.01] 1.003[1.00-1.00] 1.003[1.00-1.00] 1.003[1.00-1.00]  geo=1.071
D256 MHA  (16/16)   1.642[1.63-1.65] 1.318[1.30-1.34] 1.018[1.01-1.02] 1.131[1.13-1.13] 1.108[1.10-1.11] 1.135[1.13-1.14] 1.138[1.14-1.14]  geo=1.199
D256 GQA4 (16/4)    1.671[1.63-1.70] 1.303[1.29-1.31] 1.030[1.03-1.03] 1.156[1.15-1.16] 1.125[1.12-1.13] 1.151[1.15-1.15] 1.153[1.15-1.16]  geo=1.213
D256 GQA8 (16/2)    1.659[1.64-1.67] 1.315[1.30-1.32] 1.030[1.03-1.03] 1.155[1.15-1.16] 1.129[1.13-1.13] 1.152[1.15-1.15] 1.154[1.15-1.16]  geo=1.215

--- PREFILL FWD causal  B=1  (fa4/fa2, median[min..max], >1=FA4 faster) ---
shape                           S512            S1024            S2048            S4096            S8192           S16384           S32768
D64  MHA  (32/32)   1.128[1.11-1.14] 1.288[1.28-1.30] 1.122[1.11-1.13] 0.904[0.90-0.91] 0.861[0.83-0.89] 0.912[0.91-0.91] 0.893[0.89-0.89]  geo=1.005
D64  GQA8 (32/4)    1.107[1.08-1.12] 1.115[1.10-1.13] 1.056[1.05-1.07] 1.010[1.01-1.02] 1.010[1.01-1.01] 1.021[1.02-1.02] 1.012[1.01-1.01]  geo=1.046
D96  MHA  (32/32)   1.260[1.23-1.27] 1.261[1.25-1.29] 1.174[1.16-1.19] 0.997[0.99-1.00] 0.950[0.94-0.95] 0.892[0.89-0.89] 0.874[0.87-0.87]  geo=1.047
D96  GQA8 (32/4)    1.324[1.30-1.34] 1.123[1.12-1.13] 1.059[1.05-1.07] 1.000[0.99-1.01] 0.998[0.99-1.00] 1.006[1.00-1.01] 0.995[0.99-1.00]  geo=1.067
D128 MHA  (32/32)   1.288[1.27-1.30] 1.192[1.18-1.20] 1.146[1.14-1.15] 1.056[1.05-1.06] 0.995[0.99-1.00] 0.962[0.96-0.96] 0.951[0.95-0.95]  geo=1.078
D128 GQA4 (32/8)    1.370[1.36-1.38] 1.170[1.17-1.18] 1.066[1.06-1.07] 1.033[1.03-1.04] 0.897[0.89-0.90] 0.995[0.99-0.99] 0.990[0.99-0.99]  geo=1.065
D128 GQA8 (32/4)    1.344[1.32-1.36] 1.170[1.15-1.18] 1.067[1.06-1.07] 1.036[1.03-1.04] 1.015[1.01-1.02] 0.997[1.00-1.00] 0.988[0.99-0.99]  geo=1.082
D256 MHA  (16/16)   1.445[1.41-1.46] 1.207[1.19-1.23] 1.005[0.99-1.01] 0.986[0.98-0.99] 1.069[1.07-1.07] 1.105[1.10-1.11] 1.122[1.12-1.12]  geo=1.126
D256 GQA4 (16/4)    1.456[1.45-1.48] 1.251[1.24-1.25] 1.047[1.04-1.06] 1.047[1.04-1.05] 1.104[1.10-1.11] 1.126[1.12-1.13] 1.136[1.13-1.14]  geo=1.160
D256 GQA8 (16/2)    1.463[1.46-1.48] 1.256[1.24-1.27] 1.054[1.05-1.07] 1.060[1.06-1.06] 1.117[1.11-1.12] 1.129[1.13-1.13] 1.138[1.14-1.14]  geo=1.167

--- PREFILL FWD dense   B=2  (fa4/fa2, median[min..max], >1=FA4 faster) ---
shape                           S512            S1024            S2048            S4096            S8192           S16384           S32768
D64  MHA  (32/32)   1.028[1.02-1.04] 0.900[0.89-0.91] 0.918[0.92-0.92] 0.839[0.82-0.85] 0.942[0.94-0.94] 0.939[0.94-0.94] 0.953[0.95-0.95]  geo=0.930
D64  GQA8 (32/4)    1.092[1.09-1.11] 1.046[1.04-1.07] 1.001[1.00-1.00] 0.995[0.99-1.00] 0.994[0.99-1.00] 0.992[0.99-0.99] 0.991[0.99-0.99]  geo=1.015
D96  MHA  (32/32)   1.004[1.00-1.01] 0.967[0.96-0.98] 0.938[0.93-0.94] 0.938[0.93-0.94] 0.933[0.93-0.93] 0.930[0.93-0.93] 0.931[0.93-0.93]  geo=0.948
D96  GQA8 (32/4)    1.075[1.06-1.10] 1.038[0.95-1.05] 1.006[0.99-1.01] 1.003[1.00-1.01] 1.001[1.00-1.00] 1.000[1.00-1.00] 0.999[1.00-1.00]  geo=1.017
D128 MHA  (32/32)   1.057[1.05-1.06] 1.030[1.03-1.04] 1.053[1.05-1.06] 1.007[1.01-1.01] 1.006[1.01-1.01] 1.008[1.01-1.01] 1.006[1.00-1.01]  geo=1.024
D128 GQA4 (32/8)    1.110[1.10-1.12] 1.024[1.02-1.03] 1.008[1.01-1.01] 1.000[1.00-1.00] 1.005[1.00-1.01] 1.003[1.00-1.01] 1.000[1.00-1.00]  geo=1.021
D128 GQA8 (32/4)    1.107[1.09-1.12] 1.026[1.02-1.03] 1.009[1.01-1.01] 1.004[1.00-1.01] 1.005[1.00-1.01] 1.002[1.00-1.00] 1.000[1.00-1.00]  geo=1.021
D256 MHA  (16/16)   1.429[1.42-1.44] 1.006[1.00-1.01] 1.125[1.12-1.13] 1.104[1.10-1.11] 1.131[1.13-1.13] 1.135[1.13-1.14] 1.137[1.14-1.14]  geo=1.146
D256 GQA4 (16/4)    1.429[1.41-1.45] 1.038[1.03-1.04] 1.161[1.16-1.16] 1.122[1.11-1.13] 1.153[1.15-1.15] 1.153[1.15-1.15] 1.153[1.15-1.15]  geo=1.168
D256 GQA8 (16/2)    1.425[1.40-1.44] 1.032[1.03-1.04] 1.159[1.15-1.16] 1.132[1.13-1.14] 1.155[1.15-1.16] 1.154[1.15-1.16] 1.154[1.15-1.15]  geo=1.168

--- PREFILL FWD causal  B=2  (fa4/fa2, median[min..max], >1=FA4 faster) ---
shape                           S512            S1024            S2048            S4096            S8192           S16384           S32768
D64  MHA  (32/32)   1.217[1.21-1.24] 1.156[1.13-1.17] 0.968[0.97-0.98] 0.870[0.86-0.89] 0.807[0.80-0.81] 0.899[0.90-0.90] 0.889[0.89-0.89]  geo=0.962
D64  GQA8 (32/4)    1.096[1.09-1.10] 1.084[1.07-1.10] 1.033[0.91-1.04] 1.008[0.98-1.01] 1.013[1.01-1.02] 1.015[1.01-1.02] 1.009[1.01-1.01]  geo=1.036
D96  MHA  (32/32)   1.161[1.15-1.17] 1.104[1.05-1.12] 1.007[1.00-1.01] 0.953[0.94-0.97] 0.903[0.90-0.91] 0.879[0.88-0.88] 0.868[0.87-0.87]  geo=0.976
D96  GQA8 (32/4)    1.095[1.08-1.11] 1.189[1.18-1.20] 1.030[1.02-1.04] 1.000[0.99-1.00] 0.999[1.00-1.00] 0.999[1.00-1.00] 0.991[0.99-0.99]  geo=1.041
D128 MHA  (32/32)   1.163[1.15-1.17] 1.140[1.12-1.15] 1.080[1.08-1.08] 1.003[1.00-1.00] 0.968[0.97-0.97] 0.954[0.95-0.95] 0.947[0.95-0.95]  geo=1.033
D128 GQA4 (32/8)    1.154[1.06-1.17] 1.098[1.09-1.12] 1.050[1.05-1.05] 1.018[1.01-1.02] 1.001[1.00-1.00] 0.996[1.00-1.00] 0.988[0.99-0.99]  geo=1.042
D128 GQA8 (32/4)    1.084[1.07-1.09] 1.104[1.06-1.22] 1.056[1.05-1.07] 1.026[1.02-1.05] 1.005[1.00-1.01] 0.996[0.99-1.00] 0.984[0.98-0.98]  geo=1.035
D256 MHA  (16/16)   1.373[1.36-1.40] 1.024[1.00-1.05] 0.976[0.97-0.98] 1.048[1.04-1.05] 1.093[1.09-1.10] 1.113[1.11-1.11] 1.126[1.12-1.13]  geo=1.102
D256 GQA4 (16/4)    1.392[1.38-1.41] 1.076[1.06-1.09] 1.020[1.02-1.02] 1.088[1.08-1.09] 1.121[1.12-1.12] 1.133[1.13-1.13] 1.142[1.14-1.14]  geo=1.134
D256 GQA8 (16/2)    1.388[1.38-1.40] 1.072[1.06-1.08] 1.025[1.02-1.03] 1.105[1.10-1.11] 1.127[1.13-1.13] 1.136[1.14-1.14] 1.143[1.14-1.14]  geo=1.138

--- PREFILL BWD dense   B=1  (fa4/fa2, median[min..max], >1=FA4 faster) ---
shape                           S512            S1024            S2048            S4096            S8192           S16384           S32768
D64  MHA  (32/32)   1.113[1.08-1.45] 1.076[1.07-1.33] 0.912[0.91-0.92] 0.873[0.85-0.89] 0.796[0.79-0.80] 0.789[0.79-0.79] 0.785[0.78-0.79]  geo=0.898
D64  GQA8 (32/4)    1.002[0.99-1.01] 1.032[1.01-1.17] 0.917[0.88-0.98] 0.906[0.89-0.93] 0.805[0.80-0.81] 0.796[0.79-0.80] 0.790[0.79-0.79]  geo=0.888
D96  MHA  (32/32)   1.068[1.04-1.08] 1.071[1.06-1.09] 0.859[0.83-0.88] 0.864[0.85-0.87] 0.855[0.85-0.86] 0.852[0.85-0.85] 0.852[0.85-0.85]  geo=0.912
D96  GQA8 (32/4)    0.985[0.85-1.19] 1.015[1.01-1.17] 0.878[0.87-0.91] 0.866[0.86-0.88] 0.863[0.86-0.86] 0.858[0.86-0.86] 0.854[0.85-0.85]  geo=0.901
D128 MHA  (32/32)   1.097[1.09-1.10] 1.052[0.99-1.06] 1.006[0.98-1.04] 0.989[0.98-0.99] 0.998[1.00-1.00] 0.994[0.99-0.99] 0.993[0.99-0.99]  geo=1.018
D128 GQA4 (32/8)    1.015[1.01-1.02] 1.014[0.99-1.09] 1.024[1.01-1.06] 0.992[0.98-1.00] 1.000[1.00-1.00] 0.998[1.00-1.00] 0.992[0.99-0.99]  geo=1.005
D128 GQA8 (32/4)    1.014[1.01-1.03] 1.014[0.99-1.10] 1.022[1.02-1.03] 1.000[0.98-1.02] 1.007[1.01-1.01] 1.001[1.00-1.00] 0.994[0.99-0.99]  geo=1.007
D256 MHA  (16/16)   1.119[0.92-1.15] 0.925[0.84-0.93] 1.036[1.03-1.06] 0.998[0.96-1.00] 1.022[1.02-1.02] 1.015[1.01-1.02] 1.010[1.01-1.01]  geo=1.016
D256 GQA4 (16/4)    1.012[0.99-1.04] 1.022[0.93-1.03] 1.049[1.05-1.05] 1.005[0.99-1.01] 1.015[1.01-1.02] 1.018[1.02-1.02] 1.014[1.01-1.02]  geo=1.019
D256 GQA8 (16/2)    1.039[0.96-1.15] 0.907[0.86-0.98] 1.044[1.01-1.08] 1.004[1.00-1.01] 1.018[1.02-1.02] 1.005[1.00-1.01] 1.005[1.00-1.01]  geo=1.002

--- PREFILL BWD causal  B=1  (fa4/fa2, median[min..max], >1=FA4 faster) ---
shape                           S512            S1024            S2048            S4096            S8192           S16384           S32768
D64  MHA  (32/32)   1.112[0.89-1.42] 1.078[1.05-1.13] 0.989[0.87-1.00] 0.949[0.94-0.95] 0.846[0.84-0.85] 0.827[0.83-0.83] 0.812[0.81-0.81]  geo=0.938
D64  GQA8 (32/4)    1.026[0.99-1.07] 1.018[1.00-1.03] 0.953[0.89-1.03] 0.958[0.95-0.98] 0.870[0.87-0.87] 0.845[0.84-0.85] 0.828[0.83-0.83]  geo=0.925
D96  MHA  (32/32)   1.204[1.04-1.46] 1.112[1.01-1.32] 1.009[0.91-1.02] 0.915[0.91-0.94] 0.884[0.88-0.89] 0.864[0.86-0.87] 0.848[0.85-0.85]  geo=0.969
D96  GQA8 (32/4)    1.066[1.04-1.25] 1.007[0.90-1.05] 1.003[0.93-1.04] 0.932[0.93-0.95] 0.896[0.89-0.90] 0.873[0.87-0.88] 0.853[0.85-0.85]  geo=0.944
D128 MHA  (32/32)   1.113[1.10-1.13] 1.048[1.04-1.06] 1.026[0.96-1.03] 1.005[0.99-1.01] 1.002[0.99-1.01] 0.996[0.99-1.00] 0.994[0.99-0.99]  geo=1.026
D128 GQA4 (32/8)    1.025[1.01-1.04] 0.993[0.87-1.05] 1.043[1.02-1.10] 1.009[0.99-1.02] 1.002[1.00-1.00] 1.000[1.00-1.00] 0.994[0.99-1.00]  geo=1.009
D128 GQA8 (32/4)    1.027[1.01-1.07] 0.994[0.99-1.00] 1.043[1.03-1.05] 1.021[1.01-1.04] 1.012[1.01-1.01] 1.005[1.00-1.01] 0.997[1.00-1.00]  geo=1.014
D256 MHA  (16/16)   0.951[0.90-0.97] 0.907[0.89-0.99] 0.893[0.88-0.92] 0.919[0.91-0.93] 0.937[0.94-0.94] 0.945[0.95-0.95] 0.944[0.94-0.94]  geo=0.928
D256 GQA4 (16/4)    0.927[0.90-0.94] 0.966[0.92-0.99] 0.995[0.99-1.01] 0.984[0.98-0.99] 1.010[1.01-1.01] 1.018[1.02-1.02] 1.015[1.01-1.02]  geo=0.987
D256 GQA8 (16/2)    0.936[0.93-0.94] 1.020[0.84-1.06] 1.006[1.00-1.01] 1.000[0.97-1.01] 1.017[1.01-1.02] 1.021[1.02-1.02] 1.017[1.02-1.02]  geo=1.002

--- PREFILL BWD dense   B=2  (fa4/fa2, median[min..max], >1=FA4 faster) ---
shape                           S512            S1024            S2048            S4096            S8192           S16384           S32768
D64  MHA  (32/32)   1.105[1.07-1.15] 0.969[0.97-0.99] 0.923[0.92-0.94] 0.854[0.84-0.87] 0.795[0.79-0.80] 0.789[0.79-0.79] 0.782[0.78-0.78]  geo=0.882
D64  GQA8 (32/4)    1.030[0.99-1.07] 0.937[0.84-0.94] 0.956[0.95-0.96] 0.875[0.87-0.88] 0.802[0.80-0.80] 0.795[0.79-0.79] 0.788[0.79-0.79]  geo=0.879
D96  MHA  (32/32)   1.096[1.09-1.10] 0.893[0.87-0.90] 0.883[0.88-0.91] 0.844[0.84-0.84] 0.856[0.86-0.86] 0.854[0.85-0.85] 0.853[0.85-0.85]  geo=0.894
D96  GQA8 (32/4)    1.013[1.00-1.03] 0.899[0.89-0.91] 0.915[0.91-0.92] 0.854[0.85-0.86] 0.865[0.86-0.87] 0.860[0.86-0.86] 0.857[0.86-0.86]  geo=0.893
D128 MHA  (32/32)   1.073[1.06-1.08] 1.018[1.00-1.11] 0.993[0.97-1.01] 0.989[0.99-0.99] 0.998[1.00-1.00] 0.996[0.99-1.00] 0.993[0.99-0.99]  geo=1.008
D128 GQA4 (32/8)    1.011[0.86-1.03] 1.038[1.03-1.09] 1.000[0.99-1.02] 0.989[0.98-0.99] 1.001[1.00-1.00] 0.996[1.00-1.00] 0.992[0.99-0.99]  geo=1.004
D128 GQA8 (32/4)    1.004[0.99-1.05] 1.033[0.98-1.08] 1.012[1.00-1.02] 0.999[1.00-1.00] 1.004[1.00-1.01] 0.999[1.00-1.00] 0.992[0.99-0.99]  geo=1.006
D256 MHA  (16/16)   0.968[0.95-0.97] 1.054[1.02-1.12] 1.010[1.00-1.01] 1.028[1.03-1.03] 1.022[1.02-1.02] 1.016[1.02-1.02] 1.023[1.02-1.02]  geo=1.017
D256 GQA4 (16/4)    1.011[1.00-1.02] 1.061[1.01-1.07] 1.010[0.99-1.02] 1.035[1.03-1.04] 1.018[1.02-1.02] 1.020[1.02-1.02] 1.018[1.01-1.02]  geo=1.025
D256 GQA8 (16/2)    0.950[0.83-0.95] 1.057[1.05-1.07] 1.021[1.02-1.03] 1.030[1.03-1.03] 1.020[1.02-1.02] 1.011[1.01-1.01] 1.006[1.00-1.01]  geo=1.013

--- PREFILL BWD causal  B=2  (fa4/fa2, median[min..max], >1=FA4 faster) ---
shape                           S512            S1024            S2048            S4096            S8192           S16384           S32768
D64  MHA  (32/32)   1.108[0.87-1.15] 1.024[0.95-1.05] 0.983[0.95-1.02] 0.914[0.90-0.93] 0.836[0.83-0.84] 0.824[0.82-0.83] 0.810[0.81-0.81]  geo=0.922
D64  GQA8 (32/4)    1.033[1.02-1.28] 0.980[0.96-0.99] 1.004[1.00-1.02] 0.940[0.93-0.95] 0.858[0.86-0.86] 0.838[0.84-0.84] 0.826[0.82-0.83]  geo=0.922
D96  MHA  (32/32)   1.142[1.14-1.15] 1.015[0.99-1.04] 0.936[0.90-0.98] 0.883[0.88-0.89] 0.873[0.87-0.88] 0.858[0.86-0.86] 0.848[0.85-0.85]  geo=0.932
D96  GQA8 (32/4)    1.059[1.06-1.10] 0.995[0.98-1.01] 0.974[0.97-0.98] 0.905[0.90-0.91] 0.885[0.88-0.89] 0.866[0.87-0.87] 0.855[0.85-0.86]  geo=0.931
D128 MHA  (32/32)   1.058[0.96-1.23] 1.028[0.97-1.04] 1.015[1.01-1.03] 0.995[0.99-1.00] 0.999[1.00-1.00] 0.997[1.00-1.00] 0.993[0.99-1.00]  geo=1.012
D128 GQA4 (32/8)    1.027[0.99-1.14] 1.043[0.99-1.08] 1.012[0.99-1.02] 0.994[0.99-1.00] 1.001[1.00-1.00] 1.000[1.00-1.00] 0.997[1.00-1.00]  geo=1.011
D128 GQA8 (32/4)    1.016[0.96-1.04] 1.042[1.01-1.06] 1.037[1.03-1.08] 1.011[1.01-1.01] 1.011[1.01-1.01] 1.006[1.01-1.01] 1.000[1.00-1.00]  geo=1.017
D256 MHA  (16/16)   0.930[0.80-0.97] 0.922[0.85-0.93] 0.928[0.92-0.93] 0.941[0.94-0.94] 0.943[0.94-0.94] 0.948[0.95-0.95] 0.946[0.95-0.95]  geo=0.937
D256 GQA4 (16/4)    0.927[0.92-0.95] 1.047[1.04-1.06] 0.971[0.93-0.98] 1.008[1.01-1.01] 1.016[1.01-1.02] 1.021[1.02-1.02] 1.016[1.02-1.02]  geo=1.000
D256 GQA8 (16/2)    0.943[0.94-0.95] 1.005[1.00-1.01] 1.007[0.98-1.03] 1.020[1.02-1.02] 1.023[1.02-1.02] 1.024[1.02-1.03] 1.018[1.02-1.02]  geo=1.005

Prefill vs PyTorch SDPA (best torch backend)

FA4 matches or beats the best PyTorch SDPA on prefill (FA4/SDPA, >1 = FA4 faster); on the 6000 D256 it leads by 14–18%:

cell	5090 FA4/SDPA	6000 FA4/SDPA
D128 GQA8 causal S8192 / S16384	1.04 / 1.05	1.01 / 1.00
D256 GQA8 dense S8192 / S16384	1.01 / 1.07	1.14 / 1.18

Decode (`seqlen_q == 1`)

Decode is memory-bandwidth-bound, so absolute latency (not a FLOP ratio) is the meaningful metric. fp8 K/V also halves cache memory (1 byte/elem vs 2). Latency at Sk=16384 (µs, lower = better):

cell	card	FA4-bf16	FA4-fp8	FA2-bf16	fp8 vs FA2
D128 r2 B64	5090	5371	3089	6871	2.22×
D128 r2 B64	6000	5315	2967	5467	1.84×
D128 r4 B64	5090	3242	1912	3773	1.97×
D128 r4 B64	6000	2939	1855	2735	1.47×
D256 r2 B16	5090	1450	955	1377	1.44×
D256 r4 B64	6000	6200	3607	5299	1.47×
D128 r8 B64	6000	2429	3350	1395	0.42× (loss)

Across the full decode sweep (batch 16/64/128 × Sk 4k–64k): fp8 decode wins at GQA r2/r4 with batch ≥ 64 — geomean ~1.4–2.0× vs FA2, up to ~2.2× — and scales up with context length, while halving KV memory. It loses at GQA r8 (~0.3–0.5×: the packed GQA-decode GEMV becomes compute-bound) and at tiny-batch + short-context. bf16 decode is at parity for MHA/r2 and trails at r4/r8.

PyTorch SDPA is not a viable GQA-decode baseline — it has no GQA-decode path, so it materializes the expanded K/V (17–34 GB here): where it fits it is 23–45× slower than FA4-fp8 (e.g. 101 ms vs 3 ms), and otherwise OOMs. FA4 is the only efficient GQA-decode path of the three.

Achieved throughput (prefill forward, TFLOP/s)

Algorithmic TFLOP/s = standard attention FLOPs ÷ measured time; the FLOP count is ncu-verified exact (the hardware tensor-instruction count matches the analytical formula — dense/backward exact, causal +0.4%).

headline cell	5090 FA4	5090 FA2	6000 FA4	6000 FA2
D256 GQA8 dense S8192	179	178	258	228
D128 GQA8 causal S16384	191	191	253	254
D256 GQA8 dense S16384 (bwd)	137	123	162	160

FA4 forward sustains ~180–191 TF/s on the 5090 (~86–91% of its ~209 TF/s bf16 dense peak) and ~245–261 TF/s on the 6000.

Honest sub-parity regimes (so they aren't a surprise to a reviewer)

Decode GQA r8 (~0.3–0.5× FA2, both cards, bf16 and fp8): packed GQA-decode GEMV is compute-bound — a known limitation / future-work item.
Long-causal MHA / D64 / D96 at S ≥ 8192 (~0.85–0.95×): the 1-CTA/SM occupancy floor (255 regs/thread). GQA — the common case — stays at parity.
D256-MHA causal backward (~0.93–0.96×): the acc_dK/dV register wall (intrinsic; sm_120 lacks the TMEM datacenter Blackwell uses to hold accumulators off-register).
bf16 decode at r4, small batch trails FA2; fp8 + batch ≥ 64 recovers it.

Correctness

Validated against PyTorch SDPA across the full tests/cute matrix, which is green on sm_120 (with documented skips for the unsupported configs below). Every sm_120 change is arch-gated: SM80 / SM90 / SM100 / SM110 are byte-identical to main (statically audited across the diff; their dedicated kernels are untouched).

Real bugs found and properly fixed during development (not merely guarded):

learnable_sink + SplitKV double-counted the sink across splits → wrong forward output. Fixed: the sink is folded into LSE only in split 0, so it's counted exactly once.
Negative-offset sliding-window backward (window_size=(None,-X)/(-X,None)) produced wrong dK/dV (empty/inverted m-block prune). Fixed: the prune is skipped per-side for negative bounds (forward was already correct).
PDL on sm_120 (flash_fwd_combine.py): arch>=90 wrongly enabled griddepcontrol on the sm_80-compiled target. Fixed to exclude sm_120.
Varlen forward TypeError (head_dim ≤ 128 without max_seqlen_q): a decode-tile guard compared a None seqlen before its short-circuit, breaking the entire general varlen forward (a regression from a decode-tile optimization). Fixed (guard reorder).
Paged-KV SplitKV wrote every split into partial-slot 0 (missing split_idx) → NaN/garbage with cache_seqlens + num_splits > 1. Fixed (thread split_idx + empty-split guard).
pack_gqa + SplitKV (varlen and non-varlen): the SplitKV partial epilogue used the unpacked layout while the buffers were packed → rows scattered to wrong partial slots. Fixed with a pack-aware partial epilogue (which also restores the pack_gqa decode speedup, ≈1.24× geomean / up to 1.8× at GQA r8).
Block-sparse + mask_mod forward had a write-after-read shared-memory race (no inter-block barrier on the single-stage K/V buffer) → nondeterministic wrong output at multi-tile seqlens. Fixed (barrier).
Block-sparse dispatch omitted q_subtile_factor on sm_120 → out-of-bounds read (wrong output + illegal memory access) when the sparse block exceeded the kernel tile. Fixed.

Independently double-checked: every fix re-verified cross-device (RTX 5090 + RTX PRO 6000), compute-sanitizer clean on the formerly-IMA paths, and the arch-gating re-audited end-to-end.

Known limitations / floors on sm_120 (documented in `flash_attn/cute/README.md`)

Causal MHA (qhead_per_kvhead==1) at seqlen ≥ 8192 is ~0.95× FA2 — an intrinsic register/occupancy wall (255 regs/thread → 1 CTA/SM). GQA (the common case) is at parity or faster.
fp8 KV decode regresses below bf16 at GQA ratio ≥ 8 (GEMV becomes compute-bound); still halves KV memory and is the only fp8-cache path.
learnable_sink runs single-split (no decode SplitKV speedup); negative-offset windows are forward-only; deterministic=True backward is unsupported. (Proper fixes for the last items are follow-ups.)
fp8 is forward/decode-only (no fp8 backward); fp8 prefill MMA is not viable on sm_120 (QK/PV m16n8k32 thread-layout mismatch).

Testing

New sm_120 test coverage: test_fp8_decode_sm120.py, test_flash_attn_sm120_local.py, test_paged_kv_sm120.py, test_flash_attn_sm120_dgtdv.py, test_flash_attn_bwd_sm120_pack_gqa.py, test_flash_attn_bwd_sm120_postprocess.py. The test_flash_attn.py suite is green on sm_120 modulo documented skips for the unsupported configs above.

Summary by CodeRabbit

New Features
- Added SM120 (Consumer Blackwell) fp8 KV-cache decode kernel for seqlen_q=1.
- Added SM120 paged-KV forward support.
- Enabled SM120 Pack-GQA backward execution with improved M-splitting and split-KV handling.
Improvements
- Enhanced local/sliding-window masking support and SM120 routing behaviors.
- Improved numerical stability for edge cases (including sink/empty-row LSE handling).
Documentation
- Updated CuTe DSL FlashAttention-4 README with SM120 capabilities and limitations.
Tests
- Added/expanded SM120 regression, fp8 decode, Pack-GQA, local window, paged-KV, and mask/score-mod aux-scalar coverage.

`FlashAttentionForwardSm80.__call__` sets self.use_tma_O = self.arch >= Arch.sm_90 but the base class kernel never constructs a TMA atom for O — it passes None as `tma_atom_O` to `self.epilogue` (line ~1069). The check exists because the file once intended to support a Hopper-style TMA-O path that was never wired up here. On Hopper / SM_100 hardware this is dead code because those archs use their own forward classes (`FlashAttentionForwardSm90` / `FlashAttentionForwardSm100`) with their own `__call__`. But `FlashAttentionForwardSm120` inherits from this class, and `FlashAttentionForwardBase.__init__` reads `self.arch` from the DSL, which is `Arch.sm_120` on consumer Blackwell. The epilogue then takes the TMA-output branch and crashes inside `quack.copy_utils.tma_get_copy_fn` -> `cpasync.tma_partition` with `AttributeError: 'NoneType' object has no attribute '_trait'`. Static `arch = 80` on `FlashAttentionForwardSm120` was intended to prevent this but is overwritten by `__init__`. Force `use_tma_O = False` here; SM90 and SM100 are unaffected because they have their own `__call__`. Reproduced on RTX 5090 (SM_120, cuTeDSL 4.4.2, torch 2.10.0+cu128). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

`FlashAttentionForwardSm80.epilogue` chose the rmem->smem store atom via `get_smem_store_atom(self.arch.major*10 + self.arch.minor, ...)`, which returns - `CopyUniversalOp` for arch < 90 (or non-16-bit data), and - `StMatrix8x8x16bOp(num_matrices=4)` (Hopper `stmatrix`) for arch >= 90 on 16-bit data. `stmatrix` is hardware-paired with WGMMA's output register layout. The SM80 base class uses `mma.sync.aligned.m16n8k16` whose output register layout is *not* what `stmatrix` consumes. With WGMMA-output the atom permutes bytes from a fixed pattern of threads/registers; feeding it SM80-MMA-output silently scrambles values across nearby register lanes during the store. On native SM_80 hardware this branch never fires because the DSL arch is sm_80 < sm_90. The bug only surfaces when this class is reused on SM_120 via `FlashAttentionForwardSm120`, where `self.arch` is read from the DSL as `sm_120` and the >= 90 branch picks `stmatrix`. Symptom: the kernel completes without error and returns the correct output shape and a roughly correct output norm (each scrambled value is replaced by a same-magnitude neighbour), but element-wise diffs vs fp32 SDPA are 0.5-1.2 (non-causal) and 3.4-3.9 (causal), versus SDPA-bf16's own ~0.003 and ~0.008. Determinism still holds and error scales linearly with input magnitude — the precision/permutation signature, not a logic bug. Fix: pass a fixed `80` to `get_smem_store_atom` here so the SM80 base class always takes the universal-copy path, matching its actual MMA output layout. Verified on RTX 5090 (SM_120, cuTeDSL 4.4.2, torch 2.10.0+cu128): 240/240 correctness configs pass against fp32 SDPA reference across {fp16, bf16} x {causal, full} x B in {1,2} x S in {128..4096} x {MHA, GQA, MQA} x D in {64, 128}, with max abs diff matching SDPA-bf16's own (~0.012 worst case, ~0.0027 mean). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

`pack_gqa.compute_ptr` calls utils.elem_pointer(tensor, ((h_idx, m_idx),)) with a tensor whose layout is supposed to keep a composite `(qhead_per_kvhead, seqlen_q)` first mode (created by `pack_gqa_layout`). The slice `mO[None, 0]` that lands in `compute_ptr` is meant to preserve that compositeness so the rank-2 coord matches. On SM_120 with `cuTeDSL==4.4.2` the slice collapses the composite mode into a rank-1 layout. `cute.crd2idx` then refuses the rank-2 coord and raises at trace time with unable to compute crd2idx with '!cute.layout<"(?):(?{i64 div=8})">' and '!cute.coord<"((?,?))">' resulting in a `ValueError: Operation creation failed` before the kernel can run. Every default-policy GQA / MQA shape on consumer Blackwell hits this. The non-packed GQA path is numerically identical (pack_gqa is a perf-only optimization for the GQA Q-load / O-store), so flipping the auto-default to False on SM_120 makes GQA / MQA work out of the box while a deeper cuTeDSL fix is investigated. Explicit `pack_gqa=True` from the caller is still honoured. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Add FlashAttentionForwardSm120Tma class that uses TMA (cp.async.bulk) for Q/K/V loads with 1 DMA warp + 4 MMA warps, enabling producer/consumer overlap via PipelineTmaAsync with mbarrier synchronization. Key design: - TMA-compatible SMEM swizzle: Swizzle(B, 4, 3) instead of (B, 3, 3) - KV double-buffering (kv_stages=2), 160 threads (5 warps), 99KB SMEM - All pipeline operations inlined in the mainloop (not delegated to a separate @cute.jit method), which avoids CuTe DSL compiler hangs when pipeline states flow through method boundaries - is_first=False with pre-reset softmax state eliminates the need for a compile-time is_first flag in the single-loop mainloop - Dispatch: TMA default for SM120 non-paged, non-varlen. Falls back to CpAsync for paged KV and varlen (TMA addressing constraints). Validated on SM121a (DGX Spark): - 8/8 configs pass: non-causal + causal, B=1/2, Sq=64/128/256, Sk=128/256/512, H=4/8, D=128 - All diffs 0.002-0.008 vs reference Contributed by Second Nature Computing (https://joinsecondnature.com) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

@2imi9

Override self.arch = Arch.sm_80 after parent __init__ to prevent base class code paths from seeing the runtime arch (12.x) and enabling SM90+ features. The parent __init__ overwrites the class-level arch=80 attribute with the actual GPU arch. This was found by @2imi9 in Dao-AILab#2420 for the CpAsync kernel — same bug applies here. Add can_implement() check before TMA dispatch in interface.py so that configs exceeding SM120's 99KB SMEM (e.g. hdim=192 with kv_stages=2) fall back to the CpAsync kernel instead of failing at instantiation. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Blake Ledden <blake@secondnaturecomputing.com>

The base FlashAttentionForwardSm80.__call__ and FlashAttentionForwardSm100.__call__ both keep `stream` as the final parameter, with a comment: "Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI)". cute.compile binds arguments positionally against the compile_args list in interface.py, which ends with `current_stream`. The TMA kernel had `stream` at position 7 (right after softmax_scale). On this branch alone the kernel still works as advertised, but the mismatch breaks when composed with PRs that append further positional arguments to the compile path — most visibly when combined with Dao-AILab#2348's paged-KV plumbing or Dao-AILab#2439's dropout seeds, where the extra positions push `current_stream` onto a parameter that no longer exists or has the wrong type. Aligning with the base-class convention is mechanical and preserves correctness in isolation: Validation on SM121a (DGX Spark GB10), causal ∈ {False, True}, dtype ∈ {bf16, fp16}, B=1 S=256 H=4 D=64: causal=False bf16: max_diff=0.0020 PASS causal=False fp16: max_diff=0.0002 PASS causal=True bf16: max_diff=0.0078 PASS causal=True fp16: max_diff=0.0010 PASS Signed-off-by: Blake Ledden <blake@secondnaturecomputing.com>

Block-sparse attention processes only the KV blocks specified by block_sparse_tensors rather than the full KV sequence. Two block types are supported: mask_blocks (partially masked, apply mask_mod per element) and full_blocks (fully unmasked, skip masking entirely). Design follows the same mma_one_n_block callback pattern as SM90/SM100. The SM80 base class gets a new mma_one_n_block_bs method (load K, load V, wait, GEMM QK, score_mod, mask, softmax, GEMM PV) and a corresponding run_block_sparse_mainloop_sm80 utility in block_sparse_utils.py that iterates mask blocks then full blocks, mirroring consume_block_sparse_loads. Key implementation details: - run_block_sparse_mainloop_sm80: iterate mask_blocks first (highest n), then full_blocks. First full block always gets mask_seqlen=True since it may be at a higher n position than any mask block. - mma_one_n_block_bs: no pipeline overlap (block address unknown ahead of time), load K then V with separate cp_async_wait_group(1)/wait_group(0). - SM120 inherits SM80 base class and gets block sparsity for free. - FlashAttentionForwardSm120.__init__ forces self.arch = Arch.sm_80 to prevent the SM80 epilogue from using TMA-O (which would crash on SM121a since tma_atom_O is None in this kernel variant). - SM120: num_splits clamped to 1 in interface.py (no SplitKV support yet). - Block sparsity assert removed from SM120 interface path. Validated on SM121a (DGX Spark GB10): - test_block_sparsity.py: 4621 passed, 40 skipped - causal/non-causal, various head dims and sequence lengths Contributed by Second Nature Computing (https://joinsecondnature.com) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Replace inline blocksparse_tensors[0]/[2] index access with the get_total_block_count() utility from block_sparse_utils.py. This keeps variable naming consistent with the rest of the block-sparse codebase (which unpacks by name, not by index position). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

@jayhshah

…b#2572) (Dao-AILab#2590) * Fix bwd postprocess 2CTA gating to include sm_11x The 2CTA gating in flash_bwd_postprocess.py used `arch // 10 == 10`, which only matches SM 10.x (B100/B200/B300) and misses SM 11.x (Thor). The rest of the codebase (e.g. interface.py:549, 563, 834) consistently gates Blackwell-family 2CTA features as `arch // 10 in [10, 11]`. Bring the two postprocess sites in line with that convention. Flagged by @jayhshah in Dao-AILab#2572 follow-up discussion. * Include sm_110 in interface.py Blackwell-family heuristics Three sites in interface.py gate Blackwell-family behavior using `arch // 10 == 10`, which appears inconsistent with the rest of the file's `arch // 10 in [10, 11]` convention (used at lines 549, 563, 834, 974, 1035, etc.): - L533: `q_stage` heuristic for Blackwell forward - L579: `use_dedicated_hd256_kernel` (forward) - L1335: `use_dedicated_hd256_kernel` (backward) The dispatch in `_flash_attn_fwd` already routes both sm_10x and sm_11x through the same `FlashAttentionForwardSm100` / MLA classes, so these gates likely should treat them the same. NOTE FOR REVIEWERS: I'm not certain these are all oversight vs. intentional SM100-only paths. If any of them is intentional, please flag so I can revert just that hunk. The FP8 assert at L480 is left untouched on purpose — its error message reads as deliberate. * Apply ruff format to flash_bwd_sm100.py Pre-existing format drift surfaced by pre-commit. Not in the cute_exclude pattern, so it gets auto-fixed when other files in flash_attn/cute/ are touched in the same commit chain.

@ocss884

* Use is_family_of for sm_90 and sm_103 arch checks Follow-up to Dao-AILab#2572 — apply the same is_family_of pattern to the two remaining range-style arch checks for consistency: - flash_fwd_sm90.py:69 (SM 9.x assert) - flash_fwd_sm100.py:195 (is_sm103 flag) Same semantic narrowing as Dao-AILab#2572: bare-base SMs (sm_90, sm_103) are excluded. These kernels rely on wgmma / UMMA / 2CTA paths that require the a/f PTX variant anyway, so bare-base targets could not compile. * Clarify is_sm103 forward-inclusive semantics is_family_of(sm_103f) also matches any future sm_10x with x > 3, not just sm_103a/f. This was raised in PR review (@ocss884) — adding an inline comment clarifying that this forward-inclusive behavior is intentional: the flag gates ex2 emulation, sm_103 (B300) has fast hardware ex2, and later Blackwell variants in the same family are assumed to inherit it. No code-behavior change.

* Bump aiter submodule commit Co-authored-by: sstamenk <170634954+sstamenk@users.noreply.github.com> * Bump aiter submodule to 3b2e6f48ce97e1d494e8b3f1af5c65f74e304b28 (#2) Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: sstamenk <170634954+sstamenk@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: sstamenk <170634954+sstamenk@users.noreply.github.com>

…ao-AILab#2594) * Clamp kv_stage to avoid SMEM overflow for small head_dims on SM100 Fixes Dao-AILab#2591. The unbounded formula at flash_fwd_sm100.py:335 ignores per-stage state (mbarriers, sScale, pipeline counters) and yields kv_stage values that overflow the sm_100a 227 KB SMEM cap when head_dim_padded=16 (head_dim in {8, ..., 16}). Repro: hd=8/16 + seqlen >= 256 + bf16 fails with cudaErrorInvalidValue ("launch shared memory exceeds current GPU arch sm_100a allowed. Allocated: 233472 bytes. Max: 232448 bytes."). Clamp kv_stage at 32. Surgical to the broken case: the unbounded formula maxes at 26 stages for head_dim_padded >= 32, and the 2CTA gate at interface.py:572 restricts 2CTA to hd_padded in {128, 192} (both no-op), so the clamp only fires at hd_padded in {8, 16}. Verified across 24 configs (hd in {8,16,32,64,96,128} x causal in {T,F} x seqlen in {128,2048}) on B200 with max_err vs torch SDPA <= 0.0078. * Add test_flash_attn_small_head_dim regression test The main test_flash_attn_output parametrizes d over {64, 96, 128, 192, 256} and never exercises head_dim < 64, even though _validate_head_dims accepts head_dim >= 8 for sm_100/110. That coverage gap let the SMEM-overflow bug in Dao-AILab#2591 slip through. This focused test covers d in {8, 16, 32} x causal x seqlen in {128, 2048}. The seqlen=2048 cases push q_stage 1->2 (the actual bug trigger); the seqlen=128 cases also exercise the q_stage=1 boundary that fits on main today but is structurally adjacent. d=32 serves as a canary against any future tighter kv_stage clamp regressing it.

…#2595) apply_exp2_convert selected the exp2 implementation based on mask_fn presence: hardware ex2.approx.ftz for causal-masked tiles, polynomial emulation for unmasked tiles. Different q_stage values (1 for decode, 2 for prefill) compute different m_block for the same logical Q row, shifting which tiles are processed with vs without mask_fn. The same K tile could receive different exp2 methods across variants. Fix: always pass self.ex2_emu_freq regardless of mask_fn presence. Add regression test for decode↔prefill bitwise consistency on MLA (192,128) shapes.

…AILab#2605) cutlass 4.5.2 is safe to update, and quack 0.5.0 has been published, so bump the FA4 (flash_attn/cute) requirement floors to match. Updates the dependencies and the cu13 extra in pyproject.toml, and the documented versions in CLAUDE.md. Verified on NVIDIA GB300 (SM100, CUDA 13.2): deps resolve cleanly (nvidia-cutlass-dsl 4.5.2 base+cu13, quack-kernels 0.5.0), imports OK, and a representative GPU sample of tests/cute/test_flash_attn.py passes (6 passed / 6 skipped / 0 failed across hd 64/96/128/192, causal, mha/gqa/mqa, fwd+bwd).

…ILab#2558) * refactor mla sm100 forward * add benchmark; address deprecation warnings; tweak ptx gemm dispatch * update interface and tests

…Lab#2617) v0.2.30 only ships URLs up to CUDA 13.1.0; bumping to v0.2.35 adds 13.1.1, 13.2.0, and the matching aarch64 SBSA installers. Signed-off-by: oliver könig <okoenig@nvidia.com>

* graph capture fix * rm env flag

…ILab#2557)

stack-info: PR: Dao-AILab#2625, branch: drisspg/stack/42

ruff format flagged flash_attn/cute/flash_bwd_sm100.py (trailing whitespace in a comment and an over-split call). It was missed by the lint sweep in Dao-AILab#2625.

Passing weights_only=False (the pre-2.4 default) to torch.load allows arbitrary Python object deserialization from the checkpoint file. A malicious .pt/.pth file can execute arbitrary code on the machine loading it — a well-known PyTorch deserialization vector (CWE-502). Four call sites updated: training/src/utils/checkpoint.py load_checkpoint() training/src/eval.py eval checkpoint loader flash_attn/utils/pretrained.py partial(torch.load, ...) loader flash_attn/models/llama.py state_dicts_from_checkpoint() weights_only=True restricts deserialization to tensors, dicts, lists, tuples, and other primitive types — no arbitrary Python objects. Requires PyTorch >= 1.13; FA4's CuTeDSL dependency already requires a modern PyTorch 2.x build, so no compatibility regression. Fixes Dao-AILab#2583

…ao-AILab#2629)

…ressions (Dao-AILab#2616) stack-info: PR: Dao-AILab#2616, branch: drisspg/stack/41

…k-GQA, correctness fixes Builds on the SM120 forward foundation (TMA + cp.async paths): per-shape tile selection, LPT scheduling, SplitKV, paged-KV plumbing, pack-GQA folding (incl. the pack-gqa-aware SplitKV partial epilogue), and the forward correctness fixes (varlen guard, pack_gqa+SplitKV, block-sparse WAR race, strict is_sm120 gating, PDL/combine). All sm120-gated.

… dispatch arch==120-gated backward: dQ-accum scatter, pack-GQA M-split policy, sm120 postprocess, local-window fix. Real SM80/SM90/SM100 backward unchanged.

…rrectness fixes Local/sliding-window + block-sparse + mask_mod masking, block-sparse runtime utils, paged-KV manager, varlen seqlen clamps, softmax sink (sm120-gated), pack-GQA, shared utils. Shared-file changes are sm120-gated or behavior-identical for other archs.

Public-API dispatch for sm120: tile selection, backward M-split policy, SplitKV/paged routing, compile-cache keying, q_subtile_factor, and fp8 KV-cache decode routing with a cutlass-dsl version guard. sm120-only branches gated on arch//10==12.

Quantized-KV decode kernel: fp8 K/V cache (half the bytes) with per-(batch,head) descale, bf16 compute. Memory-bandwidth win for GQA decode; auto-routes when fp8 K/V is supplied.

* Fix CuTe SM120 compile-time argument handling * clean up * guard empty SM120 local backward tiles --------- Co-authored-by: Kevin-Li-2025 <2242139@qq.com> Co-authored-by: drisspg <drisspguessous@gmail.com>

…tKV fallback) (Dao-AILab#2656) * [CuTe,Fwd,sm120] Fix use_tma_O crash on SM120 (issue Dao-AILab#2649) On SM120 (Blackwell GeForce / RTX PRO 6000 / DGX Spark) the forward kernel set `use_tma_O = self.arch >= Arch.sm_90`, enabling the TMA-based O-store epilogue. But SM120 does not build the TMA store atom (tma_atom_O is None), so any forward call crashes in cpasync.tma_partition with: AttributeError: 'NoneType' object has no attribute '_trait' This makes the CuTe-DSL forward unusable on every SM120 GPU. Restrict the TMA O-store to sm_90..sm_119, which is where the WGMMA-era epilogue path is actually available: self.use_tma_O = Arch.sm_90 <= self.arch < Arch.sm_120 SM120 falls back to the non-TMA register->gmem O store (already used for the SM80 path), which is correct and what the CpAsync SM120 kernel expects. Verified on RTX PRO 6000 Blackwell (sm_120, cc 12.0), torch 2.12.0+cu130, nvidia-cutlass-dsl 4.5.2: forward now runs and matches PyTorch SDPA reference for hdim 64/96/128, causal and non-causal (max abs err <= 8e-3 in bf16). Before this fix every SM120 forward call raised the AttributeError above. * [CuTe,Fwd,sm120] Implement Pack-GQA on SM120; graceful SplitKV fallback Pack-GQA was only half-wired in the SM80/SM120 CpAsync forward: the epilogue referenced PackGQA.store_O/store_LSE, but the Q-load and head-indexing used the plain (unpacked) path. So pack_gqa=True crashed in pack_gqa.store_O (crd2idx on a packed (h_idx, m_idx) coordinate against an unpacked mO layout). This implements Pack-GQA end to end on SM120 (and SM80), mirroring the SM90 path: - Reshape mQ/mO (head_idx=2) and mLSE (head_idx=1) via pack_gqa_layout so qhead_per_kvhead folds into the seqlen mode ((qhead, seqlen)). - Scheduler args use cute.size(mQ.shape[0]) (packed total rows) and seqlen_q_static = mQ.shape[0][1] (logical seqlen), so causal/mask q_idx stay correct. - Kernel head-indexing: when pack_gqa, num_head from the scheduler already indexes the KV head (mQ/mK share nheads_kv); no division. - Q-load: gather rows via PackGQA.load_Q (per-row (h_idx, m_idx) gmem pointers) instead of the contiguous local_tile path. SplitKV (num_splits>1) is an SM100-only feature (SM80/SM90 also assert it unsupported); SM120 has no forward+combine path. Fall back to num_splits=1, which is numerically correct, instead of crashing in _check_type on the fp32 partials. Verified on RTX PRO 6000 Blackwell (sm_120): pack_gqa=True matches PyTorch SDPA GQA/MQA reference (err <= 8.4e-3 bf16) AND is bit-identical to the unpacked path (max |packed - unpacked| = 0.0) across MHA/GQA/MQA, causal/non-causal, hd 64/128, seqlen 512-2048. num_splits=3 falls back and matches reference (err 6.8e-4). Stacked on the SM120 use_tma_O fix (Dao-AILab#2649). * re-enable SM120 pack-gqa after rebase * clean up SM120 pack-gqa split handling * fix SM120 varlen pack-gqa offset --------- Co-authored-by: drisspg <drisspguessous@gmail.com>

stack-info: PR: Dao-AILab#2696, branch: drisspg/stack/47

…n is a tensor (Dao-AILab#2507) * Fix backward compile key instability when max_seqlen is a tensor When max_seqlen_q/max_seqlen_k are passed as torch.Tensor (e.g. by HuggingFace Transformers _prepare_from_posids), the arithmetic in _flash_attn_bwd produces tensor results that leak into the compile key tuple. Since pickle.dumps(torch.Tensor) produces a unique hash per object, every backward call generates a new compile key, causing infinite kernel recompilation and filling the persistent JIT cache with identical .o files. Cast max_seqlen_q/k to int() before they enter the seqlen_q/k computation path, ensuring the compile key contains only Python scalars. * Replace int() with host-scalar guard to avoid CPU-GPU sync

… symlink escape (Dao-AILab#2702) * hopper/setup.py: harden tarfile extraction against path traversal and symlink escape download_and_copy() extracted NVIDIA toolchain archives with a bare tarfile.extractall() into the predictable ~/.flashattn/nvidia/<name> cache, allowing arbitrary file write at build time via a pre-planted symlink or a malicious archive member (issue Dao-AILab#2637). - Add safe_extractall(): use the PEP 706 data filter when available (3.12, backported to 3.10.12/3.11.4), else fall back to per-member path containment and link rejection (stream-safe, single pass). - Refuse extraction into a symlinked cache path, closing the primary pre-planted-symlink vector on all Python versions. Signed-off-by: Aryan Putta <aryansputta@gmail.com> * hopper/setup.py: allow in-destination links in extractall fallback Address review on Dao-AILab#2702: 1. The no-data-filter fallback rejected every link member, which regressed real builds: the cuda_nvcc archives ship intra-package symlinks (e.g. libnvvm.so -> libnvvm.so.4) that the data filter permits. Allow links whose resolved target stays inside the extract dir instead, matching the data-filter behavior, and keep rejecting escaping and absolute-target links. 2. Harden the cache-path check: os.path.islink only inspects the leaf, so also require the fully resolved tmp_path to stay under the cache root, catching a symlinked parent directory. Signed-off-by: Aryan <aryansputta@gmail.com> --------- Signed-off-by: Aryan Putta <aryansputta@gmail.com> Signed-off-by: Aryan <aryansputta@gmail.com>

) The split-KV kernel (compute_attn_1rowblock_splitkv) indexes block_table[n_block * kBlockN / page_block_size], bounded only by actual_seqlen_k. In the kvcache path actual_seqlen_k is seqlens_k[b] + seqlen_knew, but block_table only has max_num_blocks_per_seq columns per sequence. If a caller passes a cache_seqlens (or appends new keys) exceeding max_num_blocks_per_seq * page_block_size, the kernel reads block_table out of bounds with no in-kernel check (see issue Dao-AILab#2709). Validate the caller contract host-side and raise a clear error instead. The .max().item() sync is only paid on the paged-KV path. Add test_flash_attn_kvcache_paged_block_table_bounds covering both the cache-length overflow and the appended-new-keys overflow, plus a positive control exactly at capacity. Co-authored-by: yunweili3 <yunweili3@users.noreply.github.com>

Resolve conflicts in 6 files. Pull out SM120 work now upstreamed independently; keep the more complete SM120 implementation where it is entangled with paged-KV / SplitKV / varlen-GQA that upstream's base path lacks. Pulled out (took upstream's version): - flash_fwd.py use_tma_O: replaced hardcoded False (my ead67d3) with upstream Dao-AILab#2656 'Arch.sm_90 <= self.arch < Arch.sm_120' (evaluates identically for SM80/SM120; re-added the dropped Arch import). - flash_bwd.py softmax_scale_log2: replaced my inline compute with upstream Dao-AILab#2671 'softmax_scale_log2, _ = compute_softmax_scale_log2(...)' (equivalent; discards the None-adjusted scale via _). Kept (my implementation; upstream's is a subset): - flash_fwd.py pack-GQA fold / num_head_kv / prologue: is_sm120-gated, carries paged-KV, SplitKV, and the varlen-GQA offset fix absent from upstream Dao-AILab#2656's base __call__. Upstream's cute.size() composite-shape wrapping auto-merged and is compatible. - flash_bwd.py local-window backward mainloop: my version handles pack_gqa packed-rows, pack_gqa_m_splits, r2p gating and the skip_full_causal_mask two-phase loop that upstream Dao-AILab#2671's BlockInfo.get_m_block_min_max path does not. - interface.py backward: adopted upstream Dao-AILab#2621 sparse-MLA (qv is not None) dispatch, re-added pack_gqa=ctx.pack_gqa to the dense _flash_attn_bwd calls. Complementary (both kept, no conflict): - interface.py max_seqlen: my host-int coercion (SM120 skip_full_causal const_expr) + upstream Dao-AILab#2507 single_block tensor guard fix different symptoms of the same root cause.

… call The _flash_attn_bwd dispatch passed 'local' both as the 12th positional arg AND as the is_local=local keyword to FlashAttentionBackward{Sm80,Sm120}. Since __init__ takes is_local as a trailing keyword (position 11 is is_causal, 12 is SdP_swapAB), the positional 'local' shifted every following arg by one and collided with V_in_regs: TypeError: __init__() got multiple values for argument 'V_in_regs' so every SM120 backward call raised. This was latent in the branch (a leftover from an earlier main-merge that added is_local positionally upstream while this branch keeps it as a trailing kwarg); it only fires when the backward actually runs. Drop the redundant positional; is_local is still passed via the keyword. Validated on RTX PRO 6000 (sm_120, torch 2.12+cu130, cutlass-dsl 4.6.0.dev0): fwd matches SDPA (maxerr<=4e-3 bf16) and bwd yields finite gradients for dense / causal / GQA / D128 / local(causal+bidirectional).

The seqlen_q<=8 decode auto-split (interface.py:~1313) rewrites num_splits 1->0 to request the SplitKV heuristic, but the SM120 guard immediately did 'assert num_splits == 1', so every small-seqlen (decode) SM120 forward crashed: AssertionError: SM120 forward only supports num_splits=1 Route the auto sentinel (num_splits < 1) through num_splits_heuristic on SM120 as well; it returns 1 when the grid is already filled and >1 only for underfilled decode/small-batch shapes, so the SplitKV forward path runs only where it helps. An explicit user-requested num_splits > 1 is still rejected (test_flash_attn_sm120_rejects_splitkv stays green). This was latent in the branch (the auto-split predates this guard); it only fires for seqlen_q<=8, which is why the seqlen>=64 suite was green. Validated on RTX PRO 6000 (sm_120): decode S=1/3/4 MHA+GQA match SDPA (maxerr 0); explicit num_splits=3 still raises.

…void P saturation (Dao-AILab#2717) * [CuTe, SM100] Make FP8 max_offset dtype-aware to avoid e4m3 P saturation With rescale_threshold=4 the online-softmax row max can be stale by up to 4 (in log2 units), so P reaches 2^(max_offset + 4). max_offset=8 puts that at 4096, past e4m3fn's 448 ceiling: the largest probabilities saturate on the f32->fp8 satfinite convert and e4m3 accuracy degrades below e5m2 (up to 1.6x worse rel_l2, growing with seqlen). Cap max_offset at 4 for e4m3 so the worst case is 2^8 = 256 <= 448; e5m2 keeps 8 (57344 ceiling absorbs the overshoot). B200: restores e4m3 to ~2x lower error than e5m2 across seqlen 256-4096, uniform and peaked softmax, matching quantization-only emulation; LSE consistent; fwd timing unchanged (0.387 vs 0.390 ms, hd128 s4096). Related: Dao-AILab#2716 * [CuTe, Tests] Unrot the FP8 dtype path in test_flash_attn_output Running the suite with dtype=float8_e4m3fn has bit-rotted: - the test sets requires_grad on fp8 tensors, which the interface now rejects (FP8 is forward-only); gate it on non-fp8 dtypes. - it generates random descales and applies them in attention_ref, but the flash_attn_func call site has no descale kwargs (only _flash_attn_fwd takes them), so kernel and reference disagreed by construction; stop generating them. With these, the fp8 sweep runs cleanly (378 cases on SM100 with the e4m3 max_offset fix; 190 of them fail without it). fp8 stays out of the default dtype parametrize. Related: Dao-AILab#2716

* allow varlen score mod in backward; add tests and examples * add recompute fastdiv_mods to sm90 bwd * remove softcap != 0 limitation in test * fix linter error * guard use 2cta against softcap in bwd * undo formatting in test_flash_attn.py * reset test_flash_attn * update tests for score mod varlen bwd, guard blocksparse varlen bwd * aux_tensors -> aux_data; unpack args in test * aux_tensors -> aux_data for sm90 backward * make_fragment -> make_rmem_tensor in score_mod_definitions * predicate on aux_data.tensors, not aux_data * relax test tolerance in vectorized score mod tests - bitwise equality failing on sm103 though within tolerance * revert erroneous test reformatting to main

Dao-AILab#2669) * Expand FLASHATTENTION_DISABLE_DROPOUT to not bring in unneeded headers Summary: Previously, using the FLASHATTENTION_DISABLE_DROPOUT flag still pulled in unneed dependencies from ATen for at::Generator and Philox related headers. This change sets up the codebase so that using the flag will not pull in these unnecessary headers. There are two major changes of note: 1. We remove needing an RNG gen in the schema--the Python frontend always passed in None so this should not be BC breaking to most users. 2. Instead of referencing the PhiloxState directly, in order to detach dependencies when dropout is not needed, we introduce an opaque buffer that will hold the philox state when dropout is desired. Test Plan: pytest tests/test_flash_attn.py::test_flash_attn_output -k "113-203-64 and dtype0 and mha" pytest tests/test_flash_attn.py::test_flash_attn_varlen_output -k "113-203-64 and dtype0 and mha" g++ -c -O1 -std=c++17 -D_GLIBCXX_USE_CXX11_ABI=1 <torch+cutlass+cuda -I flags> \ csrc/flash_attn/flash_api.cpp -o /tmp/fa.o nm -C /tmp/fa.o | grep -E 'mha_(fwd|bwd|varlen)\(' | grep -c Generator returns 0 g++ -E -DFLASHATTENTION_DISABLE_DROPOUT <same -I flags> csrc/flash_attn/flash_api.cpp \ | grep -c 'CUDAGeneratorImpl.h\|philox_unpack.cuh' also returns 0 Reviewers: Subscribers: Tasks: Tags: Add trivially copyable assert Add back gen * use mark.skipIf

…Lab#2692) * fix combine kernel bug for full cudagraph * linearzie kernel

stack-info: PR: Dao-AILab#2731, branch: drisspg/stack/49

…Lab#2559) * add dynamicpersistentvarlenscheduler to flash_fwd_sm100 and prepare kernel * mild refactor to tile scheduler protocol, guard num_m_blocks_ptr for sm100, update tests to use scheduler metadata * rename varlen_batch_idx -> virtual_batch_idx, because it is relevant for non-varlen blocksparse batch sorting * split out VarlenSchedulerBase to share code between SingleTile and DynamicPersistent schedulers * add benchmark script for varlen dynamic persistent scheduler * minor clean up * updates to has_work logic, tile scheduler selection, and varlen test suite * fix tile scheduler dispatch logic * integrate binary batch search for single tile varlen * refactor tile scheduler for compositionality * work PR 2520 into interface and kernels * fix linter errors * wip: modify scheduler metadata public api * clean up scheduler metadata API; add docstrings; split out _get_fwd_config method; remove cluster_size==1 restriction; guard architectures against unused scheduler metadata args * address driss' comments * fix compute_tile_cumsum guards in interface * simplfiy benchmark, guard against _compute_tile_cumsum with small batch size * fix compute_tile_cumsum guard * update to 4.6.0 * fix linter error * fix rebase bug * add cu_blocks_kernel to replace _compute_tile_cumsum * fix linter error * address comments on PR * add seqlen_k_per_split, add single tile varlen scheduler to combine kernel * add blocks to batch idx O(1) lookup path to varlen scheduler * fix lint errors * various legibility improvements and bug fixes

* add methodology doc * revise method * concision pass

…ces (Dao-AILab#2755) * [CuTe,Sm100] Sparse MLA bwd: skip dK/dV scatter at -1 sentinel indices The sparse-MLA (gather_kv_indices) backward scatter epilogues atomically accumulated dV/dK at row indices read straight from gather_kv_indices with no validity guard, while every load path treats -1 (the documented sentinel for invalid top-k slots, which any causal top-k index tensor contains as padding) as invalid and predicates the gather. The masked math is correct -- p/dS are exactly 0.0 at sentinel slots -- but the atomic itself is destructive: index -1 addresses one row before the (batch-sliced) buffer base, and red.add.f32 flushes subnormal destinations to +0.0 and canonicalizes NaN payloads even when the addend is 0.0. For batch 0 this lands out of bounds in whatever tensor the caching allocator placed before dv/dk (silently corrupting e.g. int32 tensors, whose small values are all subnormal fp32 bit patterns); for later batches it lands in the previous batch's last row. Symptoms depend purely on allocation layout: bitwise-correct results, silently wrong grads, or IMA. Fix: skip the atomic when the index is negative, mirroring the load-side guard. The skipped contribution is mathematically 0.0, so numerics for valid slots are unchanged. Also fix _flash_attn_bwd_sparse_mla discarding caller-supplied dq=/dk= buffers (dq = dk = None right after recording prealloc_dq/dk, after which the reallocation is skipped because prealloc is set, so passing dq=/dk= crashed). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * [CuTe,Sm100] Test sparse MLA bwd with -1-padded gather_kv_indices Adds test_flash_attn_mla_sparse_bwd_sentinel and a varlen counterpart (causal x shared_kv, seqlen 512/1024 non-varlen, packed docs [512, 4, 1024] varlen): builds causal top-k indices with -1 tail padding, checks out/lse/grads against attention_ref through the public autograd path, then reruns the backward with preallocated dk/dv buffers surrounded by int32 canaries (values 1..N, all subnormal fp32 bit patterns, so one misdirected red.add.f32 -- even of +0.0 -- flushes them to zero) and asserts the canaries are untouched. The varlen kernels are separate compile-time specializations, and the dK epilogue guard must apply to the doc-relative index before seqlen_k_offset is added; the varlen test pins that down (doc 0's row -1 is the canary-visible case) and includes a doc shorter than topk_len whose index rows are almost entirely sentinels. Fails deterministically without the sentinel-scatter guard; existing sparse-MLA tests never hit the bug because they generate gather_kv_indices as full argsort permutations with no -1 slots. Also makes attention_ref's top-k mask sentinel-aware: its scatter_ used to route -1 indices into key 0 (unmasking it) and trip scatter's bounds check; out-of-range indices now also map to the dummy column, and the mask applies regardless of topk_len vs seqlen_k (equivalent for permutation indices). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

…, matching the dense path (Dao-AILab#2756) Co-authored-by: qaf <qaf@example.org>

…ab#2706) * Support learnable sink backward * Simplify learnable sink backward plumbing * Format learnable sink postprocess * Support learnable sink with hd256 and frozen QKV * Format standalone sink reduction * Tighten learnable sink backward checks * Keep learnable sink scope lean * Remove standalone learnable sink varlen test * Always test learnable sink backward * Simplify learnable sink test setup * Detach learnable sink test tensors * Use Tuple return type for backward * Remove sink backward tensor wrapper * Clarify sink reduction CTA selection * Refine learnable sink backward handling * Handle sink-only rows in backward * Address learnable sink review feedback * Separate learnable sink dtype coverage * Use dtype-aware learnable sink gradient tolerance * Handle empty learnable sink backward * Address learnable sink follow-up feedback

…ao-AILab#2705) * [CuTe, FA4] Preserve first-tile flag during scheduler reconstruction * Pin nvidia-cutlass-dsl to 4.7.0 * Keep scheduler fix separate from DSL upgrade

…Dao-AILab#2761) When scheduler metadata provides per-batch dynamic num_splits, the varlen tile schedulers pack num_splits into the top 16 bits of split_idx. The dense path unpacks it inside BlockInfo.get_n_block_min_max, but the SM100 block-sparse paths pass split_idx to the block-sparse helpers as is: the load and MMA warps passed the packed value, for which split_block_range yields an empty block range, while the softmax and correction warps passed the unpacked value, yielding a non-empty range. The warps then disagree on whether a tile has work, the softmax/correction/MMA mbarrier handshake never completes, and the kernel spins forever. Every varlen + block-sparse + SplitKV forward hangs this way. Also pass the dynamic per-batch num_splits (instead of the static maximum) to the block-sparse helpers so the split ranges cover the whole block list, matching the dense path.

Fix forward issues exposed by dynamic-shape and layout fuzzing. Canonicalize unaligned inputs, distinguish static broadcast and auxiliary tensor ABIs in the compile cache, and compile SplitKV combine optional operands exactly as called. Use target-SKU SM metadata during fake selection. Size aliased SM100 K/V shared memory for the larger staged layout and ceil-divide non-TMA paged-loader entries so partial row waves receive page pointers. Add one focused regression for each underlying bug.

PR Dao-AILab#2559 added static- and dynamic-persistent dispatch for varlen SM100 kernels, but test_clc_fuzz still required the pre-change single-tile scheduler, so the scheduler assertions failed before numerical validation could run. Validate each scheduler class against its scheduling mode, account for dynamic SplitKV dispatch, and update the two static fallback expectations. STATIC mode accepts both StaticPersistentTileScheduler and SingleTileVarlenScheduler for varlen because the dispatch depends on whether every batch fits in a single m-block. Verification: the full tests/cute/test_clc_fuzz.py suite passes on SM100 (B200, CC 10.0) — 189 passed — both on this branch's base (1cc7ff6) and cherry-picked onto c68c592. The GQA + SplitKV varlen cases were confirmed to still select the CLC SingleTileVarlenScheduler, matching the assertion precedence.

Brings the branch up to date with 35 upstream commits. 22 conflict hunks across 7 files. Redundant SM120 fixes dropped in favour of upstream's: - softmax.py: upstream Dao-AILab#2706 implements the all-masked-row (row_max == -inf) sink guard unconditionally and with a proper max-shift, so the is_sm120-gated version and its plumbing are removed (3 call sites in flash_fwd.py). - paged_kv.py: upstream landed the same page_entry_per_thread ceil-div. - interface.py: dropped the local maybe_contiguous, which was shadowing upstream's alignment-aware version everywhere in the file. Kept, with upstream's additions merged in: - flash_fwd.py / flash_bwd.py: TileSchedulerArguments gained cu_total_m_blocks alongside the existing is_split_kv. - flash_fwd_combine.py: upstream rewrote it for linearized scheduling (Dao-AILab#2692); the SM120 use_pdl gate (griddepcontrol.wait is illegal on sm_120) is re-applied on top. - bwd postprocess: compile key carries both upstream's cu_total_m_blocks / learnable_sink_dtype and the SM120 pack_gqa pair, in signature order. Forward config refactor: upstream extracted tile selection into _get_fwd_config(), now shared with the public get_scheduler_metadata(). The SM120 per-shape tuning moved into _get_fwd_config_sm120() behind an arch-12 branch, with an Sm120FwdContext for the extra inputs and FwdConfig carrying the SM120 outputs (num_threads, cp.async stages, Q-in-regs). Verified dispatch-identical to the pre-merge branch over 1080 shapes. Fixes required by the merge: - flash_fwd_sm120_tma.py: upstream added mCuTotalMBlocks/mCuTotalSplitsMBlocks to the arch 8/9/12 launch args, but the SM120 TMA kernel had no such parameters, so every d64 TMA forward failed to bind. Wired through to its SingleTileVarlenScheduler. - interface.py: current_stream moved into upstream's compile-miss block; the fp8 decode path compiles earlier and now makes its own. - interface.py: learnable-sink backward. Upstream Dao-AILab#2706 asserts SM90/SM100/110 whenever a sink is present, which the autograd backward always passes. That removed working SM120 behaviour: dSink is a pure side-output of the dQ postprocess and dq/dk/dv receive the sink only through LSE. The assert is now gated on compute_dsink, derived from needs_input_grad, so a frozen sink still backprops on SM120 and only a real dSink request errors out. Pre-existing bug fixed: _bwd_postprocess_dkv_sm120 unpacked 14 values from make_fake_bwd_tensors, which returns 15 (mScaleP). Tested on RTX PRO 6000 (sm_120): 2652 passed / 246 skipped across the SM120 suites, the combine suite and a 350-case random sample of test_flash_attn.py; forward and backward match SDPA for causal, GQA, D256, sliding-window, non-square and decode shapes.

Supersedes the compute_dsink approach in the merge commit. Same behaviour, but _flash_attn_bwd is now byte-identical to origin/main again: instead of teaching it to skip dSink, the autograd Functions simply don't hand it the sink when the sink is frozen. dSink is a pure side-output of the dQ postprocess (only SM90/SM100/SM110 implement the reduction); dq/dk/dv never read sink_tensors and already receive the sink's contribution through LSE. So passing learnable_sink=None for a frozen sink yields identical gradients while avoiding upstream Dao-AILab#2706's arch assert, which the autograd backward would otherwise trip on SM120 for any sink-using model. An actual dSink request still raises upstream's original error. Verified on RTX PRO 6000: frozen-sink fwd/bwd match SDPA (1.6e-3 / 3.9e-3, same as before), dSink request errors loudly, 256 passed / 246 skipped across the SM120 suites and the 350-case sample.

jganbar and others added 30 commits May 24, 2026 00:31

replace deprecated apis (Dao-AILab#2602)

eaf806d

[CuTe,Fwd,Sm100] refactor mla sm100 forward and add page table (Dao-A…

6dba037

…ILab#2558) * refactor mla sm100 forward * add benchmark; address deprecation warnings; tweak ptx gemm dispatch * update interface and tests

ci: bump Jimver/cuda-toolkit to v0.2.35 for CUDA 13.2 support (Dao-AI…

c184200

…Lab#2617) v0.2.30 only ships URLs up to CUDA 13.1.0; bumping to v0.2.35 adds 13.1.1, 13.2.0, and the matching aarch64 SBSA installers. Signed-off-by: oliver könig <okoenig@nvidia.com>

[ROCm] Bump Triton to >=3.6.0 and aiter submodule (Dao-AILab#2614)

b02b07e

[Triton] Fix graph capture issues and env var (Dao-AILab#2620)

fefa96a

* graph capture fix * rm env flag

[CuTe,Bwd,Sm100] allow 2cta with score mod and mask mod in bwd (Dao-A…

d80a771

…ILab#2557)

[CuTe] Fix lint failures (Dao-AILab#2625)

22a7223

stack-info: PR: Dao-AILab#2625, branch: drisspg/stack/42

[CuTe] Fix lint failure in flash_bwd_sm100.py (Dao-AILab#2627)

766ed20

ruff format flagged flash_attn/cute/flash_bwd_sm100.py (trailing whitespace in a comment and an over-split call). It was missed by the lint sweep in Dao-AILab#2625.

use correction warps if not tma store; remove outdated packgqa guard (D…

98688fd

…ao-AILab#2629)

Add aux-scalars to interface to enable dynamic ints and floats in exp…

bc58abc

…ressions (Dao-AILab#2616) stack-info: PR: Dao-AILab#2616, branch: drisspg/stack/41

sm120: backward — SM80-base extensions, postprocess, pack-GQA M-split…

8163761

… dispatch arch==120-gated backward: dQ-accum scatter, pack-GQA M-split policy, sm120 postprocess, local-window fix. Real SM80/SM90/SM100 backward unchanged.

sm120: fp8 (e4m3/e5m2) KV-cache decode kernel

3190fa4

Quantized-KV decode kernel: fp8 K/V cache (half the bytes) with per-(batch,head) descale, bf16 compute. Memory-bandwidth win for GQA decode; auto-routes when fp8 K/V is supplied.

yinli-systems and others added 30 commits July 6, 2026 19:48

Fix CuTe SM120 compile-time argument handling (Dao-AILab#2671)

6e646e0

* Fix CuTe SM120 compile-time argument handling * clean up * guard empty SM120 local backward tiles --------- Co-authored-by: Kevin-Li-2025 <2242139@qq.com> Co-authored-by: drisspg <drisspguessous@gmail.com>

ad tcgen.ld.red support to sm103a arch (Dao-AILab#2696)

af04967

stack-info: PR: Dao-AILab#2696, branch: drisspg/stack/47

Enable 2CTA for SM100 block-sparse backward (Dao-AILab#2661)

2402cb0

add linearize scheduling to combine kernel for full cudagraph (Dao-AI…

14c3779

…Lab#2692) * fix combine kernel bug for full cudagraph * linearzie kernel

Numeric tweaks to fp8 (Dao-AILab#2731)

849f660

stack-info: PR: Dao-AILab#2731, branch: drisspg/stack/49

Remove SM100 Functions from Hopper (Dao-AILab#2746)

c75d019

[AI] Add doc on debug methodology (Dao-AILab#2753)

df61ab6

* add methodology doc * revise method * concision pass

[ROCm] Fix CK varlen_fwd binding argument mismatch (Dao-AILab#2742)

4a948e9

[CuTe,Bwd,Sm90] Fix: wait for bwd_preprocess on the block-sparse path…

4460ebb

…, matching the dense path (Dao-AILab#2756) Co-authored-by: qaf <qaf@example.org>

[CuTe, FA4] Preserve first-tile flag during scheduler reconstruction (D…

7a08d7a

…ao-AILab#2705) * [CuTe, FA4] Preserve first-tile flag during scheduler reconstruction * Pin nvidia-cutlass-dsl to 4.7.0 * Keep scheduler fix separate from DSL upgrade

Fix duplicated word in layer norm comment (Dao-AILab#2744)

d7e4dba

[CuTe, Fwd] Stabilize tensor max_seqlen compile key (Dao-AILab#2762)

1cc7ff6

Navigation Menu

FA4 consumer Blackwell (sm_120) integration — squashed history - #3