Why Flash Attention Is Fast: IO-Awareness Explained

Why Flash Attention Is Fast: How IO-Awareness Changed the Attention Computation

When asked where the bottleneck is when running a Transformer on a GPU, most people immediately say "attention's O(N²) complexity." That's not wrong, but you need to go one level deeper to understand what Flash Attention's design is actually solving. The real latency cost doesn't come from the O(N²) computation itself being slow — it comes from the round-trip cost of writing the N×N intermediate values to GPU HBM and reading them back.

Tri Dao et al.'s original FlashAttention paper was the first to quantify this. The framing of memory IO count — not FLOP count — as the bottleneck is exactly why they called it "IO-Aware Attention."

The GPU Memory Hierarchy Is Where the Problem Starts

The A100 memory hierarchy has three levels.

HBM has high capacity but is relatively slow. SRAM is only 192 KB per SM, but its bandwidth is roughly 10× that of HBM. The problem is that standard attention repeatedly writes and reads intermediate values to and from HBM — values too large to fit in SRAM.

Standard Attention's Real Bottleneck: HBM Access Count

Implementing standard attention in PyTorch breaks down into three steps:

# Conceptual flow (what happens when kernels are separate)
S = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d)   # N×N → written to HBM
P = torch.softmax(S, dim=-1)                                # read from HBM → written back
O = torch.matmul(P, V)                                      # read from HBM → output

Computing S produces an N×N matrix that gets written to HBM. Computing softmax requires reading it back. The result P is written back to HBM, then read again for the PV multiplication. Just counting the HBM round-trips for S and P: 2 writes + 2 reads, moving 4N² elements. At fp16, that's 8N² bytes.

Every time N doubles, the HBM traffic for intermediates quadruples. As Horace He's roofline analysis shows, standard attention is not compute-bound — the GPU compute units aren't waiting on arithmetic — it's memory-bandwidth-bound, waiting on data to come back from HBM. This is why reducing FLOPs alone doesn't improve speed. Running matmul → softmax → matmul as three separate kernels forces an HBM round-trip at each step.

Tiling: Processing the N² Matrix Entirely in SRAM

The core idea behind Flash Attention is to never write the N×N matrix to HBM. Instead, Q, K, and V are split into small tiles, loaded into SRAM, and the full attention output is computed entirely within those tiles.

With tile size Br=Bc=64, d=64, and fp16, the SRAM footprint is:

Q tile    : 64 × 64 × 2B = 8 KB
K tile    : 64 × 64 × 2B = 8 KB
V tile    : 64 × 64 × 2B = 8 KB
S tile    : 64 × 64 × 2B = 8 KB  (lives in SRAM only, never written to HBM)
O accumulator : 64 × 64 × 2B = 8 KB
──────────────────────────────
Total     : 40 KB

This uses only 40 KB of the 192 KB SRAM per SM. At N=1024, there are Tr=16 tiles along the Q dimension and Tc=16 tiles along the K·V dimension. For each tile pair, the S block is computed in SRAM, softmax is applied, O is updated, and S is simply discarded. Only the final O is ever written to HBM.

The problem is that numerically stable softmax requires knowing the maximum value across the entire row — but values from tiles not yet processed are unavailable. Flash Attention solves this with online softmax. With each tile, the running maximum m and denominator sum l are accumulated, and the O accumulator is corrected as each new tile arrives:

m_new = max(m_old, max(S_tile))
l_new = exp(m_old - m_new) × l_old + sum(exp(S_tile - m_new))
O_new = (exp(m_old - m_new) × l_old × O_old + exp(S_tile - m_new) × V_tile) / l_new

This correction produces exactly the same softmax result without ever holding the full N×N matrix in memory. It's not an approximation — it's exact attention. That's why the paper includes "Exact Attention" in its framing.

Tile size is tuned based on SRAM capacity and GPU architecture. Larger Br and Bc reduce HBM accesses but increase SRAM usage, which can cause shared memory contention with other kernels. Recommended block sizes for each GPU architecture are documented in the flash-attn GitHub repository.

Recomputation: Deliberately Spending FLOPs to Save Memory

During training, the backward pass needs intermediate values produced in the forward pass. Standard attention stores S and P for gradient computation — both of size N². As sequence length grows, these intermediates eat into GPU memory.

Flash Attention doesn't store S and P. Instead, it recomputes them during the backward pass. The only values saved are the online softmax statistics m and l, which are O(N). This is obviously more FLOPs, but it's a rational choice because on GPUs, bandwidth is more expensive than arithmetic. The cost of storing N²-sized tensors in HBM and reading them back exceeds the cost of recomputing those values in SRAM. It's a trade: sell memory, buy FLOPs.

Inference has no backward pass. Recomputation is a design decision for gradient computation, so the benefit of Flash Attention during inference comes purely from the HBM IO reduction that tiling provides. Conflating training memory savings with inference latency improvements leads to the wrong root-cause analysis.

FA2 and FA3: What Each Version Targeted Next

FA1 addressed HBM IO. FA2 and FA3 tackled the remaining bottlenecks in order.

The key contribution of the FA2 paper is switching the outer loop direction. In FA1, iterating over K·V in the outer loop means multiple thread blocks share responsibility for updating the same O block, requiring synchronization to merge results. FA2 switches Q to the outer loop so each thread block independently produces a complete O block. This improves occupancy, and FA2 also reduces non-matmul FLOPs in the softmax rescaling step, boosting Tensor Core utilization. Where FA1 reached 25–40% of theoretical peak, FA2 achieves 50–73%.

FA3 exploits capabilities unique to the Hopper architecture (H100). Warp specialization assigns matmul and softmax to separate warps and overlaps them asynchronously. In FP8 low-precision mode, it reaches approximately 1.2 PFLOPs/s on H100.

When Flash Attention Delivers Less Speedup Than Expected in Serving

The fact that the benefit scales with N is frequently overlooked in serving workloads. The savings Flash Attention achieves relative to standard attention's N² intermediate traffic grow proportionally to N/d. With d=64 fixed, that's a 64× reduction at N=4096 but only 8× at N=512.

At short sequence lengths (128–512 tokens) with batch=1 on A100, the latency difference between FA2 and naive attention is minimal, and in some cases the tiling setup overhead can actually reverse the advantage. There's simply less N² traffic to eliminate. Even the FA1 paper's benchmarks show speedups dropping below 2× at N≤256.

With large batches like batch=32, GPU occupancy naturally improves, making Flash Attention's tiling effect comparatively smaller. Flash Attention still helps in this regime, but it's worth distinguishing whether the primary driver is IO reduction or occupancy improvement.

For models using GQA (Grouped Query Attention) or MQA, fewer heads mean fewer opportunities for parallelism. Running batch=1 decoding on a model like Llama 3 70B with 8 KV heads means Flash Attention processes fewer tiles to begin with, and GPU occupancy issues can outweigh the benefit. In these conditions, decoding-specialized kernels like FlashDecoding or FlashInfer are a better fit.

The conditions where Flash Attention provides dramatic gains are long sequences during the prefill phase. Applied mechanically to short-sequence decoding, the benefit nearly vanishes. This is the first thing to check before drawing conclusions about Flash Attention's impact in a serving environment.

Positioning Relative to PagedAttention and Chunked Prefill

vLLM and SGLang use Flash Attention at the kernel level and implement PagedAttention on top of it. These operate at different layers. Flash Attention is a kernel that takes contiguous Q·K·V tensors and computes attention in a tiled fashion. PagedAttention is a higher-level structure that manages the KV cache in non-contiguous memory blocks.

Flash Attention fundamentally assumes contiguous K·V memory layout. Handling PagedAttention's non-contiguous blocks requires either gathering into a contiguous buffer first, or using a variant that natively supports paged access — such as the flash_attn_with_kvcache API in the flash-attn repo. vLLM implements additional custom CUDA kernels that integrate paged KV layout with Flash Attention directly.

The same applies to Chunked Prefill. Long prompts are split into chunks and prefilled sequentially, with Flash Attention kernels handling the attention computation within each chunk. The KV cache connection across chunks is managed by the layer above — Flash Attention on either side has no awareness of chunk or page boundaries. It only sees a contiguous tensor view.