FlashAttention-4

Abstract

Rendered from the arXiv abstract page (fetched 2026-09-11):

Attention is a computational bottleneck in transformer language models, especially at long context. Prior work optimised attention for Hopper through asynchrony and warp specialisation; Blackwell presents a different problem because its scaling is asymmetric — tensor-core throughput roughly doubles while shared-memory bandwidth and the exponential units scale minimally or not at all. The paper introduces three techniques: (1) redesigned pipelines using fully asynchronous matrix operations and larger tiles, (2) software-emulated exponentials and conditional softmax rescaling to minimise non-matmul overhead, and (3) tensor memory plus 2-CTA MMA mode to cut shared-memory traffic and atomics in the backward pass. It reaches up to 1.3× over cuDNN 9.13 and 2.7× over Triton on B200 with BF16, at 1,613 TFLOPs/s (71% utilisation), implemented in CuTe-DSL within Python with 20–30× faster compilation than C++ template approaches while preserving full functionality.

Key Contributions — and which bound moved

The finding that earns this rung's §10 test:

For typical Blackwell attention workloads, shared-memory traffic and exponential operations now dominate execution time, exceeding MMA compute by 25–60%.

That is the whole argument. A generation ago the kernel was matmul-bound and the job was to feed the tensor cores. On Blackwell the tensor cores are not the constraint — the softmax path and the memory movement around it are. Every technique in the paper follows from that inversion.

1. Pipeline redesign for asynchronous overlap. Exploits Blackwell's fully asynchronous MMA and 128×128 tiles (Hopper used 64×128), with accumulators held in tensor memory rather than registers. That decouples the softmax computation from the tensor-core work so the two genuinely overlap instead of taking turns.

2. Attacking the exponential unit. The MUFU hardware unit is a fixed resource that did not scale, so the exponential is also computed in software as a degree-3 polynomial on the FMA units, running in parallel with MUFU — spreading one bottlenecked operation across two functional units. Degree 3 is stated as sufficient accuracy for BF16. Added to that, conditional softmax rescaling: skip the renormalisation when the new running maximum does not exceed a threshold, typically log₂(256) = 8.0.

3. Cutting shared-memory traffic in the backward pass. Blackwell's 2-CTA MMA mode pairs CTAs to execute cooperatively with partitioned accumulator tiles, halving shared-memory bandwidth for operand B, and a restructured dQ computation cuts global atomic adds by 50%.

Results

Limitations

• Blackwell-specific by construction. The authors say the implementation is optimised for Blackwell; the principle (optimise for whichever unit did not scale) generalises, the kernel does not.
• The deterministic backward mode costs performance — up to 75% of the non-deterministic variant's speed, because the reductions are serialised through semaphore locks. Reproducible training pays for reproducibility.
• The polynomial exponential is an accuracy trade priced for BF16; the paper does not claim it at higher precision.
• 1.3× is stated as up to, against one cuDNN version (9.13). Baselines in kernel work age in weeks.
• 71% utilisation is measured against theoretical peak, which is itself a vendor number.

Why this is on the compilers rung

compilers held zero sources. Its scope is the layer between a model graph and the silicon, and the rung's standing test is whether a source names which bound moved. This one does, with a number: 25–60%, non-matmul over matmul. It is upstream of every serving-throughput figure the site quotes, and it is the reason "more FLOPs" stopped translating into "faster attention".

It pairs with C1 (unpicked, 2026-09-05), which makes the same argument at batch-1 on a different hardware class.

No live market call rests on this rung — brand and teaching value, plus it underwrites the serving numbers.

Source: FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling by Ted Zadouri, Markus Hoehnerbach, Jay Shah, Timmy Liu, Vijay Thakkar and Tri Dao — arXiv:2603.05451, submitted 2026-03-05.