Why does FlashAttention save memory during training, no...

Standard attention requires storing the full N x N probability matrix P to compute gradients during the backward pass, which consumes massive amounts of memory at long context lengths. FlashAttention avoids this by storing only the log-sum-exp (Li) values per head. During the backward pass, it recomputes the necessary tiles of P. This reduces the activation memory footprint from O(N^2) to O(N), which is often the difference between being able to run a training job or encountering an out-of-memory error.