FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling
Together AI
FlashAttention-4 is a new attention algorithm and kernel implementation that addresses asymmetric hardware scaling on Blackwell GPUs by pipelining tensor cores, special function units, and memory operations to reduce bottlenecks. The implementation achieves 1605 TFLOPs/s with 71% utilization on BF16, delivering 1.3× speedup over cuDNN and 2.7× over Triton by using software-emulated exponentials, tensor memory for intermediate storage, and 2-CTA MMA modes to reduce shared memory traffic. The design enables more efficient training and inference for large language models on next-generation hardware through careful co-optimization of algorithm and kernel implementation.
Why it matters
As GPU throughput outpaces memory bandwidth, kernels must evolve. We introduce FlashAttention-4, featuring new pipelining for maximum overlap, 2-CTA MMA modes to reduce shared memory traffic, and a hardware-software hybrid approach to softmax exponentials.
