Member-only story
Standard transformer attention materializes an N×N matrix of attention scores for every layer and head. At N=8192 tokens and d=64 per head, that is 268 MB of float32 just for the score matrix — per head, per layer, per batch item. Flash Attention (Dao et al., 2022) eliminates this matrix entirely by tiling the computation into blocks that fit in fast on-chip memory (SRAM, on-chip static RAM that sits next to each Streaming Multiprocessor), streaming through K and V without ever assembling the full N×N matrix.
This article builds both from scratch in NumPy, verifies they produce numerically equivalent output (float64 round-off, ~1e-15, well below the 1e-10 PASS threshold), and shows the memory savings clearly: 8128× at N=8192.
Disclaimer: The opinions expressed in this article are my own and do not represent the views of Google. This content is based solely on publicly available information.
Why Standard Attention Has a Memory Problem #
To understand why Flash Attention exists, you first need to see precisely where the standard formulation breaks down. The memory cost grows quadratically with sequence length, and the numbers become unworkable faster than most practitioners expect.