You train with 16K-token sequences. Each head's attention score matrix then has about 270 million entries, over half a gigabyte in half precision. Standard attention writes the scores to GPU HBM, reads them back for the softmax, writes the probabilities and stores them for the backward pass, which reads them again.
FlashAttention holds a query block in on-chip SRAM (about 200 KB per streaming multiprocessor), streams key/value blocks past it with a running max and exponential sum per query, rescaling the partial output as it goes, and writes back only the exact output and one softmax statistic per query. Its backward pass stores no score matrix and recomputes the scores block by block from and .