FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling

TL;DR

FlashAttention-4 is a re-design of the FlashAttention kernel for NVIDIA Blackwell (B200 / GB200) GPUs, motivated by the observation that Blackwell doubles BF16 tensor-core throughput vs. Hopper (2.25 vs. 1 PFLOPS) while shared-memory bandwidth and the MUFU exponential unit stay flat — so on B200 the attention bottleneck has shifted off the matmul units onto softmax and SMEM traffic. The paper introduces three co-designed remedies: (1) a new ping-pong pipeline that uses Blackwell’s fully-asynchronous MMA + 128×128 tiles + tensor memory (TMEM) and moves output rescaling into a separate “correction” warpgroup off the critical path, (2) software-emulated 2^x via Cody-Waite range reduction + Sollya-tuned polynomial on FMA units, run on 10–25% of softmax entries to use MUFU and FMA in parallel, plus conditional softmax rescaling that skips renormalization when m_j − m_{j−1} ≤ log2(256), and (3) backward-pass changes that use the 2-CTA MMA mode so each CTA stages only half of operand B and halve the atomic adds in dQ. End result on B200 BF16: up to 1.3× over cuDNN 9.13, 2.7× over Triton, ~1613 TFLOPs/s (71% peak). The kernel ships in CuTe-DSL embedded in Python, with reported 20–30× faster compile times than the C++ template path.

Key claims

Blackwell B200 doubles BF16 tensor-core throughput vs. Hopper H100 (2.25 vs. 1 PFLOPS) while SMEM read bandwidth (128 B/clock/SM) and MUFU exponential throughput (16 ops/clock/SM) remain the same, shifting the attention bottleneck to non-matmul resources [§1, §2.2].
Roofline for forward attention with tile $M{=}256, N{=}d{=}128$ : MMA compute and exponential unit each cost ~2048 cycles, SMEM ~1536 cycles — i.e. exp throughput is co-dominant with matmul, not negligible as on Hopper [Table 1, §3.1.1].
Forward pipeline uses two softmax warpgroups + one correction warpgroup + one tensor-core/TMA warpgroup; ping-pong over two output tiles per CTA with two S tiles overlapping with P in TMEM (chosen over the alternative of one S + two P tiles to start the pipeline by immediately computing two S tiles) [§3.1.2, Fig. 1].
Software-emulated $2^x$ via Cody-Waite range reduction + Horner-evaluated polynomial (degree 3 sufficient; degree 5 matches hardware to within 2× max-rel-err) runs on FMA units in parallel with MUFU; emulation is applied to only 10–25% of softmax entries, tuned per tile to balance register pressure [§3.1.3, Table 2].
After rounding to BF16, the degree-3 polynomial’s quantization error matches MUFU.EX2 to within 1 BF16 ULP on 99% of inputs — i.e. emulation cost in accuracy is dominated by BF16 quantization itself [Table 2, §3.1.3].
Conditional online-softmax rescaling skips the e^{m_{j−1}−m_j} correction whenever $m_j - m_{j-1} \le \tau$ (default $\tau = \log_2 256 = 8.0$ ), recovering the exact result only at the final normalization step; rescaling is gated at the warp granularity to avoid divergence [§3.1.4, Eq. 6].
Backward pass exploits the Blackwell 2-CTA MMA mode so each CTA in a pair stages only half of operand B in SMEM (the hardware consumes the combined tile via TMEM across the pair), and restructures the dQ step to halve the number of atomic reductions [§3.2 intro, §2.2].
Backward roofline (5 MMAs + 2 elementwise ops, $M{=}N{=}d{=}128$ ): MMA compute scales as $10MNd / 8192$ cycles; SMEM and exp are pushed below MMA, making MMA the binding constraint [§3.2.1, Eq. 7–8, Fig. 2].
A deterministic execution mode is supported with “minimal performance overhead” for reproducible RL training [§1].
Headline benchmarks (B200, BF16): up to 1.3× over cuDNN 9.13, 2.7× over Triton, peak ~1613 TFLOPs/s = 71% of theoretical 2.25 PFLOPS [Abstract, §1].
Implementation is entirely in CuTe-DSL embedded in Python with reported 20–30× faster compile times than C++ template-based equivalents; code released under permissive license at Dao-AILab/flash-attention/.../flash_attn/cute [§1].

Method

The paper is a kernel co-design: pick the bottleneck via a roofline that explicitly accounts for SMEM read traffic, exponential-unit cycles, and MMA compute on Blackwell; then move work between hardware units until the bottleneck is rebalanced. Three primary algorithmic levers. (1) Pipeline. Forward: two output tiles per CTA, two softmax warpgroups (each thread owns a full 128-element row, eliminating inter-warp shuffles for row-max), one correction warpgroup that does the O_{j-1} rescale off the critical path, and one warpgroup driving TMA + MMA. Because MMA accumulators live in TMEM (not registers as on Hopper), S and P can be staged in TMEM, with two S tiles overlapping P. Backward similarly redistributes the 5 MMAs across producer/consumer warps. (2) Exponential emulation. Decompose $2^x = 2^{\lfloor x \rfloor} \cdot 2^{x_{\mathrm{frac}}}$ ; integer part via IEEE-754 bit manipulation, fractional part via Sollya-tuned polynomial evaluated with FMA instructions. Apply only to a tunable fraction of entries (10–25%) to avoid register spills; the rest go through MUFU.EX2 in parallel. Conditional rescaling (τ = 8) skips renormalization steps whose multiplier is bounded, with the algebraic invariant that the final $1/\ell_{\mathrm{final}}$ normalization recovers the exact softmax-times-V output. (3) Backward 2-CTA mode. Use Blackwell’s 2-CTA MMA so a pair of CTAs cooperatively executes a single MMA over M = 256, partitioning operand B across the two CTAs’ SMEM (each holds half) — this halves SMEM traffic for the large operands and lets the authors restructure dQ to halve atomic adds. Implementation is in CuTe-DSL (Python), with the paper claiming the 20–30× compile-time speedup vs. C++ templates is the load-bearing productivity win that made this much architecture-specific tuning feasible.

Results

On B200 with BF16, FA4 reaches up to ~1613 TFLOPs/s (71% of the 2.25-PFLOPS peak), up to 1.3× over cuDNN 9.13, and 2.7× over the Triton attention implementation; the speedup grows with sequence length, and BF16 is the headline configuration. The paper reports the speedup is achieved “on the shifted bottleneck resources” — i.e. near-peak utilization of the SMEM + exp budget on Blackwell, not just matmul utilization. Backward-pass details, including the 2-CTA dQ atomic-adds reduction and the full roofline, are in §3.2 but truncated in the available extract. B300/GB300 (which double MUFU throughput to 32 ops/clock/SM) are mentioned as a future regime where the exp-emulation lever would matter less.

Why it’s interesting

For any internal team running attention-heavy workloads on B200/GB200 — large-context training, long-rollout video DiTs, agent rollouts with long KV caches — FA4 is the new baseline, and the design choices it forces are about which non-matmul resource is binding now. The framing “tensor cores scale faster than everything else, so attention is no longer matmul-bound” is the same observation that motivates the structured-attention work in MonarchRT: Efficient Attention for Real-Time Video Generation (which beats FA-2/3 on RTX 5090 at 95% sparsity by avoiding the dense N×N matrix entirely) — these are complementary axes: FA4 makes dense attention faster on the bottleneck resources, MonarchRT changes the attention algorithm so the bottleneck shifts. The conditional softmax rescaling (skipping e^{m_{j-1}-m_j} updates when the rescale factor is bounded by 256) is the same algebraic trick that Inference-Time Hyper-Scaling with KV Cache Compression-style approximations rely on for KV-cache compression, but here applied without lossiness. Also worth flagging: the CuTe-DSL-in-Python compile-time claim (20–30× faster than C++) is the methodological subtext — Tri Dao et al. are arguing that architecture-specific kernels need Python-speed iteration loops to be tractable.