FlashAttention

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré · NeurIPS 2022

Read paper · arXiv:2205.14135

Accumulate exact attention in blocks with online softmax without materializing the full attention matrix.

What you will implement

Implement blockwise recurrence and rescaling. This CPU algorithm exercise does not benchmark GPU kernels.

• 24. Online Softmax and Block Attention

  Merge score tiles with online softmax to recover exact attention and log-sum-exp.

  Implement only online_softmax_update; query/key tiling and final normalization are supplied. This CPU teaching implementation does not reproduce GPU acceleration.