FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
FlashAttention: an IO-aware exact attention algorithm using tiling to cut GPU memory traffic, speeding Transformer training and enabling longer context.
Self-attention scales quadratically with sequence length, making Transformers slow on long sequences, and approximate methods cut compute but often lose quality without speedups. The authors argue the missing principle is IO-awareness, counting reads and writes between GPU memory levels. FlashAttention is an IO-aware exact attention algorithm that uses tiling to cut transfers between GPU HBM and on-chip SRAM, needing fewer HBM accesses. It trains Transformers faster (15% on BERT-large, 3x on GPT-2, 2.4x on long-range arena) and enables longer context and new capabilities.
Based on: FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness · Neural Information Processing Systems