FlashAttention-2

The Problem

The original FlashAttention algorithm (2022) was a massive leap forward because it solved the memory bandwidth bottleneck by keeping intermediate attention results in fast SRAM. However, it still wasn't utilizing the GPU's compute (FLOPs) to its absolute maximum theoretical potential. It topped out at roughly 25-40% of the maximum theoretical FLOPs on A100 GPUs.

The inefficiency stemmed from how the workload was partitioned. The original algorithm split the workload across GPU thread blocks based on the batch size and number of attention heads. If you had a small batch size and a small number of heads, but a very long sequence length, the GPU's many streaming multiprocessors (SMs) were mostly sitting idle, waiting for a few thread blocks to do all the heavy lifting.

The Idea

Tri Dao rewrote the algorithm to optimize the low-level CUDA thread execution. The core idea of FlashAttention-2 is to partition the workload along the sequence length dimension, not just the batch and head dimensions.

How It Works

1. Parallelism over Sequence Length: Instead of assigning an entire sequence to one thread block, FlashAttention-2 breaks the sequence length itself into smaller chunks. Multiple thread blocks can now work on different parts of the same sequence concurrently, drastically increasing GPU occupancy.
2. Reduced Non-Matmul FLOPs: The original algorithm spent a surprising amount of time doing non-matrix-multiplication math (like scaling and masking). FlashAttention-2 aggressively refactored these operations, moving scaling operations to happen just once at the end, rather than inside the inner loop.
3. Optimized Thread Block Geometry: It tweaked the sizes of the tiles loaded into SRAM to better match the specific hardware geometry of modern Nvidia GPUs (like the A100 and H100).

Why It Mattered

FlashAttention-2 doubled the speed of the original FlashAttention, achieving up to 73% of the theoretical maximum FLOPs on A100 GPUs.

This pure engineering optimization effectively acted as a 2x free compute multiplier for the entire AI industry. It made training long-context models dramatically cheaper and faster. Because it was integrated directly into PyTorch's scaled_dot_product_attention, almost every researcher and company automatically benefited from the upgrade without changing their model architecture.

What Came After

FlashAttention-2 became the universal standard. As hardware architecture evolved to the Nvidia Hopper generation (H100), Tri Dao and team released FlashAttention-3 (2024), which utilizes Hopper-specific features like asynchronous memory transfers (TMA) and the Tensor Memory Accelerator to push performance even further.