FlashAttention Kernel & Online Softmax

FlashAttention is a GPU-memory-optimized attention kernel: by tiling $Q/K/V$ into SRAM and merging partial softmaxes via Online Softmax, HBM traffic drops from $O(N^2)$ to $O(N^2 / M)$ (≈ $O(N)$ for large $N$; $M$ = SRAM block capacity). Standard attention computes $S = \frac{QK^\top}{\sqrt{d}}$, softmax, then $O = SV$ — the $N \times N$ matrix $S$ is written to and read back from HBM, an $O(N^2)$ memory bottleneck. FlashAttention keeps blocks of $Q/K/V$ resident in SRAM, computes local softmax per block, and merges incrementally with a running max and rescaling: for element $x_i$, $m_{new} = \max(m, x_i)$, $l_{new} = l \cdot e^{m - m_{new}} + e^{x_i - m_{new}}$, rescaling the output by $e^{m - m_{new}}$ and normalizing by $\frac{1}{l}$ at the end — numerically identical to standard softmax.

the default kernel for long-context attention (4K/32K/1M tokens) in both training and inference; interview favorites: why FlashAttention is fast, what Online Softmax is, and FA-1/2/3 differences; it also underpins inference optimizations like PD separation and KV quantization.

standard attention is dominated by HBM bandwidth — the $O(N^2)$ $S$-matrix round-trip dominates runtime at long sequences, and $O(N^2)$ memory does not fit. FlashAttention-1 (2022) combines tiling, Online Softmax and backward recomputation (no $S$ stored) to cut HBM traffic by an order of magnitude and memory from $O(N^2)$ to $O(N)$ — 2-4× faster than eager PyTorch at 64K sequences; FA-2 (2023) parallelizes across heads and moves the $N^2$ softmax overhead out of the parallel region, roughly doubling again; FA-3 (2024) exploits Hopper warp specialization, TMA asynchronous copies, and adds FP8.

Derive standard attention's HBM bottleneck: why the $N \times N$ round-trip between $S = \frac{QK^\top}{\sqrt{d}}$ and softmax costs $O(N^2)$ traffic; why bandwidth dominates at long sequences (e.g. 128K)?

Write the Online Softmax updates: $m_{new} = \max(m, x_i)$, $l_{new} = l \cdot e^{m - m_{new}} + e^{x_i - m_{new}}$; why do you need a running max with rescaling instead of directly accumulating $e^{x_i}$?

Tiling details: what limits block size (e.g. ~228KB SRAM per SM on H100)? How does HBM traffic drop from $O(N^2)$ to $O(N^2 / M)$; how is the causal mask handled inside the kernel?

FA-1/2/3 evolution: why backward recomputation cuts memory from $O(N^2)$ to $O(N)$; FA-2's head parallelism and upcast optimization; what FA-3's warp specialization/TMA async and FP8 solve?

Relation between FlashAttention and KV cache/PagedAttention: why kernel-level optimization composes with engine-level optimization; why the decode bandwidth bottleneck (reading weights per token) cannot be fixed by FA alone?