How FlashAttention computes softmax one block at a time
One query's attention scores, split into two blocks and worked through by hand, to see how FlashAttention gets the softmax without having all the scores at once.
One query's attention scores, split into two blocks and worked through by hand, to see how FlashAttention gets the softmax without having all the scores at once.
A transformer redoes the same attention projections for every past token at each decoding step. The KV cache stores those keys and values so they get reused instead of recomputed, and the one cost it adds is memory that grows with the sequence.