Explain Why FlashAttention Is Faster Despite Doing More FLOPs

A reviewer says: "FlashAttention recomputes attention scores in the backward pass and adds online-softmax rescaling, so it does more arithmetic than standard attention, yet your benchmark shows it is 2 to 3 times faster and lets you train at 32k context. Either the benchmark or the claim is wrong."

Setting: 32 heads, head dimension 128, sequence length 8192, bf16, A100-80GB (about 312 TFLOPs bf16 peak, about 2 TB/s HBM; spec values).

1. Compute the size of the score matrix, the HBM traffic of standard attention's four score passes and the matmul FLOPs for one layer and one sequence. Is standard attention compute-bound or memory-bound?
2. Explain the tiling and online-softmax mechanism, including the update rule that makes exact softmax possible block by block.
3. State what FlashAttention stores for the backward pass and its memory complexity versus standard attention.
4. Why is the benchmark not contradictory? Name what is and is not approximated, and what FlashAttention-2 and -3 change.