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

FlashAttention made attention fast by never writing the full attention matrix to GPU memory: it tiles queries and keys, computes softmax on the fly ("online softmax") and keeps everything on chip. Each new GPU generation needs the kernel rebuilt around its hardware: FlashAttention-2 for Ampere, FlashAttention-3 for Hopper.

Blackwell (B200, GB200) changed the balance. Per SM per clock, BF16 tensor-core throughput doubled to 8,192 operations, but the exponential unit stayed at 16 operations and shared memory at 128 bytes. The paper's roofline for a 128×128 tile shows the result: in the forward pass, exponentials (1,024 cycles) now take as long as the matrix multiplies (1,024); in the backward pass, shared-memory traffic (3,328 cycles) exceeds them (2,560) by 30%. Making the matrix math faster no longer helps unless the rest keeps up.

FlashAttention-4 attacks the non-matmul work:

• Pipeline for asynchronous tensor cores. Blackwell's tensor cores write results to a new 256 KB tensor memory per SM without blocking registers. FA4 processes two query tiles in a ping-pong so one tile's softmax runs while the other's matrix multiplies, and moves output rescaling to a separate warpgroup.
• Software exponentials. For 10% to 25% of softmax entries, 2^x is computed on the general FMA units: split x into integer and fraction, evaluate a cubic polynomial on the fraction and add the integer to the float's exponent bits. After rounding to BF16, it matches the hardware unit within one unit in the last place on 99% of inputs.
• Lazy rescaling. Online softmax rescales the running output whenever a new maximum appears. FA4 rescales only when the maximum grows by more than 8 in log2 units (a factor of 256), and fixes up at the end. The result stays exact.
• Paired-CTA backward pass. Blackwell can run one matrix multiply across two thread blocks, each holding half an operand. That cuts shared-memory traffic for operands and halves the global atomic adds for dQ.

It also orders work tiles longest first for causal masks and variable lengths, offers a deterministic backward mode, and is written in CuTe-DSL in Python. On B200 with BF16, forward attention reached up to 1,613 TFLOPS (71% of peak), 1.1-1.3x cuDNN 9.13 and 2.1-2.7x Triton.

Customers buying Blackwell expect their models to get faster. For long-context inference and training, attention is often the largest single cost, and its speed depends on whether the software stack has kernels written for the new hardware.

Check the attention kernel version, not just the GPU. Running Hopper-era kernels (or generic fallbacks) on Blackwell leaves much of the hardware idle. Serving and training frameworks pick up FA4-class kernels on their own schedules; confirm what a deployment actually runs.

Long sequences gain most. The paper's speedups are largest from 4K tokens up and with causal masks, the regime of long-document, codebase and reasoning workloads.

Faster hardware shifts bottlenecks rather than removing them. The paper's roofline argument (matrix units doubling while exponentials and memory stay flat) applies to any kernel. When a new GPU underdelivers on a workload, look at the non-matmul work.