What FlashAttention Does Instead of Storing S and P

FlashAttention never writes the $N \times N$ score matrix to HBM during the forward pass. The backward pass still needs $S$ (the scores) and $P$ (the softmax probabilities) to compute gradients.

For one attention layer with $N = 8192$ , 32 heads, fp16, what would storing $S$ and $P$ cost, and what does FlashAttention do instead?

About 8.6 GB, avoided by sparsifying the score matrix before the backward pass, which makes the gradients approximate but close enough in practice.

Nothing extra: the forward pass already wrote $S$ to HBM, so the backward pass simply reads it back.

About 8.6 GB per layer, two 8192 x 8192 fp16 matrices per head across 32 heads. FlashAttention stores only the output and the softmax statistics $(m, \ell)$ , which are $O(N)$ per head, and recomputes $S$ and $P$ from Q, K, V tiles in SRAM.

About 8.6 GB, avoided by storing $S$ and $P$ in fp8 instead, which is where the memory saving comes from.

About 4.3 GB, and recomputation is a net loss in wall-clock time, since it repeats the QK matmul that the forward pass already paid for.