FlashAttention

The Problem

The standard self-attention mechanism requires computing an $N \times N$ attention matrix (where $N$ is the sequence length).

The problem is not just the FLOPs (math operations) required to compute it, but the memory bandwidth. GPUs have a massive but slow main memory (HBM) and a tiny but incredibly fast on-chip memory (SRAM). Standard implementations compute parts of the attention matrix in SRAM, write the intermediate results back to HBM, read them back to apply the softmax, write them back again, and read them again to multiply with the Values.

For long sequences, this constant reading and writing to slow HBM (the "memory wall") absolutely chokes the GPU. The GPU spends most of its time waiting for data to arrive from memory, rather than doing math.

The Idea

Tri Dao and his team realized that attention algorithms were not "IO-aware"—they were ignoring the physical memory hierarchy of the GPU.

They designed FlashAttention to compute the exact same mathematical result as standard attention, but by aggressively minimizing reads and writes to HBM.

How It Works

FlashAttention uses a technique called tiling.

1. It loads blocks (tiles) of the Query, Key, and Value matrices from HBM into the fast SRAM.
2. It computes the attention scores and multiplies them by the Values entirely within SRAM for that block.
3. Crucially, it uses an online softmax calculation trick that allows it to compute the softmax incrementally block-by-block, without ever needing access to the full un-normalized row.
4. It writes the final output directly back to HBM, bypassing the need to ever instantiate or write the massive $N \times N$ intermediate attention matrix to HBM.

Additionally, to save memory during the backward pass (gradient calculation), FlashAttention does not store the intermediate attention matrix; instead, it rapidly recomputes it on the fly during the backward pass. Recomputing is actually faster than reading the massive matrix from HBM.

Why It Mattered

FlashAttention was a software engineering triumph that fundamentally changed the economics of LLM training. It achieved 2x-4x speedups in training time and reduced memory usage from $O(N^2)$ to $O(N)$.

Because memory usage now scaled linearly instead of quadratically, researchers could suddenly train models on vastly longer context windows (e.g., jumping from 2K to 16K or 32K) without running out of GPU memory.

What Came After

FlashAttention was rapidly integrated into PyTorch and adopted by virtually every AI lab. It was quickly superseded by its own author with FlashAttention-2 (2023), which optimized the work partitioning between GPU thread blocks to achieve even closer to theoretical maximum FLOPs, and later FlashAttention-3 (2024) for Hopper architecture GPUs.