FlashAttention-4: How Together AI Sped Up Attention on Blackwell GPUs

FlashAttention-4 is a rework of the attention algorithm built specifically for modern GPUs, where tensor cores are growing in performance much faster than memory and other resources.

Why the standard approach no longer works

At first glance, attention performance is governed by the speed of two matrix multiplications: S = Q × K^T and O = P × V. But analysis of the Blackwell B200 reveals something unexpected: the bottleneck is not in the tensor cores, but in the special function units (SFU) that compute the exponential in softmax (forward pass) and in shared memory traffic (backward pass). From Hopper (H100) to Blackwell (B200), BF16 tensor core performance grew from 1 to 2.25 petaflops, while the number of SFUs and shared memory bandwidth stayed the same. This asymmetry breaks the standard optimization approach: you cannot simply assume the tensor cores fully determine performance. In practice, the units get in each other's way.

How Together AI solved the problem

The research team, together with engineers from NVIDIA, Meta* and Princeton, proposed three key ideas:

• New pipelining: software pipelines that overlap the work of the tensor cores, SFU and memory as much as possible, without idle time.
• Emulating the exponential with polynomials: instead of the slow SFU, the forward pass computes the exponential on fast FMA units (fused multiply-add).
• TMEM and 2-CTA MMA: use of the new tensor memory (256 KB per SM) plus a mode in which two thread blocks work on a single matrix operation, reducing shared memory traffic.

Blackwell features that made this possible

Blackwell added several features specifically for optimizations like this:

Tensor memory (TMEM) is fast on-chip storage (256 KB per SM) wired directly into the tensor cores. Intermediate results can stay in TMEM without a trip to the slower shared memory, which sharply cuts access latency.

Asynchronous 5th-generation tensor cores: each operation is launched by a single thread and accumulates results in TMEM. The maximum tile for BF16 is 128×256×16 (roughly twice that of Hopper), which allows deeper pipelining without overflowing registers.

2-CTA MMA is a new mode in which two thread blocks work on one matrix operation at the same time. It halves the load on shared memory and reduces the number of atomic operations.

Numbers and results

On a Blackwell B200 with BF16, FlashAttention-4 reaches 1605 TFLOPs/s (71% utilization). That is 1.3× faster than cuDNN 9.13 and 2.7× faster than Triton. For context: 1605 TFLOPs is almost half of Blackwell's peak performance, yet it is squeezed out of a complex attention kernel.

What this means

FlashAttention-4 shows how to work in the era of asymmetric GPU scaling: instead of tweaking old algorithms, rebuild them together with the new capabilities of the hardware. The result looks like an achievement, but it is only the start of adapting to the new reality of hardware.

*Meta is designated as an extremist organization and banned in Russia.