Why architecture matters here
Flash Attention matters because attention is the compute-and-memory workhorse of transformers. Standard attention materializes the full N×N QK^T matrix in HBM — this is O(N²) memory, which becomes untenable at long context. More importantly, moving that matrix in and out of HBM dominates run time; the GPU spends more time on memory bandwidth than on math.
Cost impact is direct. Training a large model with Flash Attention runs 2-3x faster than standard attention at the same context length. Inference with long prompts sees similar wins. The alternative — buying more or bigger GPUs — is far more expensive.
Long context is the second lever. Standard attention can't reach 128k tokens on any reasonable GPU because of the memory blowup. Flash Attention keeps memory linear in N, which is what makes 128k, 200k, and million-token contexts economically viable.
The architecture: every step explained
Walk the diagram from top to bottom.
Q / K / V tensors. The three inputs to attention. Sharded by attention head so heads run in parallel. Each head has its own N×d Q, K, V matrices.
Standard attention. Compute S = Q · K^T (N×N matrix), then softmax(S), then P · V. The N×N intermediate lives in HBM. For N=8k, d=128, this is roughly 256 MB per head per layer — moved between HBM and SRAM twice.
IO cost. O(N²·d) bytes moved to and from HBM per attention call. HBM bandwidth on H100 is ~3 TB/s; at long context this dominates. Attention becomes memory-bound.
Flash Attention: tile. Split Q into blocks of rows (say, 128 rows); split K, V into blocks of columns. Load one Q block into SRAM. Loop over K, V blocks: load each into SRAM.
Online softmax. The core trick. Instead of computing softmax over the full row, maintain a running max and running denominator as you process K, V blocks. On each new block, adjust prior partial output with the running rescale factor.
SRAM tile compute. Each block-level QK^T is small enough to fit in SRAM (~128 KB per SM on H100). The GPU does all the math on-chip. HBM traffic drops dramatically.
Rescale + accumulate. On each K/V block, compute local partial output. Rescale prior accumulated output using the running softmax max. Add the new partial. Continue.
Causal mask fusion. For decoder attention, upper-triangular positions have no effect. Flash Attention 2 skips those blocks entirely.
Backward pass. The forward pass doesn't store the full attention matrix; the backward pass recomputes what's needed on the fly. Trade compute for memory — beneficial because HBM was the bottleneck.
Flash Attention 2 / 3. FA2 (Dao 2023) improved warp partitioning and non-matmul FLOP efficiency, delivering ~2x over FA1. FA3 (2024) added FP8 support and Hopper-specific optimizations, delivering another 1.5-2x on H100.
Integration. PyTorch's scaled_dot_product_attention picks Flash Attention when available. vLLM, TensorRT-LLM, and every serving stack use it. xFormers and Triton kernels expose it directly.
End-to-end forward pass
Trace a forward pass. Model processes a batch of sequences with N=8192 tokens. In the attention layer, Q/K/V are computed as usual: shape [B, H, N, d].
For each head in each batch, Flash Attention runs. It picks tile sizes: Br (query rows per tile) = 128; Bc (key columns per tile) = 128. Total tiles: 64 Q blocks × 64 K/V blocks = 4096 tiles per head. Each tile is 128 rows × 128 cols × d = ~64 KB, fits in SRAM.
The outer loop iterates over Q blocks. For each Q block, load into SRAM. Initialize running max m = -infty, denominator l = 0, accumulator O = 0.
Inner loop over K, V blocks. Load Kj, Vj into SRAM. Compute Sij = Qi · Kj^T (small matrix in SRAM). Compute mij = max(Sij), Pij = exp(Sij - mij).
Update running: new_m = max(m, mij); rescale prior l and O by exp(m - new_m); add exp(mij - new_m) · Pij · Vj. Continue.
After all K, V blocks, output for Qi is O / l. Write back to HBM. Move to next Q block.
Result: same math as standard attention, computed with a fraction of HBM traffic. Speedup: 2-3x on typical H100 workloads. No accuracy loss.