← All questions
Inference & ServingHard· FlashAttention· Memory Optimization· Backward Pass

FlashAttention does not store the full attention weight...

Because the attention score matrix has an arithmetic intensity (AI) of approximately 1, it is memory-bound. Instead of storing the massive attention weight matrix, the backward pass recomputes the scores (S = QKT / sqrt(d)) tile-by-tile in SRAM using the stored Q, K, output O, and gradient dO. This approach leverages the excess compute supply to eliminate the O(T^2) HBM reads that would otherwise dominate the backward pass time and memory usage.

Code playground (Python)