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

FlashAttention-4 addresses an imbalance in NVIDIA's Blackwell GPUs: matrix multiplication accelerated more than some operations surrounding it. Attention must also normalize scores and move intermediate data, so faster multiplication alone leaves bottlenecks. Together AI's technical account describes overlapping these operations, sharing exponential calculations across hardware resources and keeping selected intermediate results in tensor memory. The report connects each change to a specific bottleneck.

For forward attention on B200 GPUs using BF16 arithmetic, the authors report speeds up to 1.3 times cuDNN 9.13 and 2.7 times the tested Triton implementation. These are attention-kernel comparisons; their effect on a complete model depends on how much execution time attention occupies. The article separately compares a newer cuDNN version.

“From the Hopper H100 to the Blackwell B200 , for instance, BF16 tensor core throughput increases from 1 to 2.25 PFLOPs, while both the SFU count and shared memory bandwidth remains unchanged.”

Together AI

“However, doing a “feeds and speeds” analysis for B200 in fact shows the opposite: the main performance bottleneck lies not in how fast the tensor cores can do MMA, but rather (a) in the SFU units for softmax exponential during the FWD computation, and (b) in the shared-memory traffic during the BWD computation.”

Together AI

“On B200 with BF16, it reaches up to 1605 TFLOPs/s (71% utilization), up to 1.3× faster than cuDNN version 9.13 and 2.7× faster than Triton.”

Together AI

“The critical detail is that exp is the bottlenecked section. We explicitly synchronize the two softmax warpgroups so they do not evaluate exp at the same time, thereby reducing MUFU contention.”

Together AI

“In practice, deterministic backward reaches up to about 85-90% of the nondeterministic throughput in our benchmarks.”

Together AI