FlashAttention
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré
The paper in one sentence
FlashAttention speeds up exact attention by minimizing expensive reads and writes between accelerator memory levels.
What the paper does
FlashAttention is an exact attention algorithm designed around the memory hierarchy of modern accelerators. It tiles the computation so intermediate attention matrices do not need to be repeatedly materialized in slower high-bandwidth memory, reducing data movement while preserving the mathematical result.
Why it matters to engineers
FLOP counts alone do not predict model performance. This paper is a useful systems lesson: an algorithm can be mathematically equivalent yet much faster because it is designed for real hardware memory traffic.
Key claims from the paper
- Attention can be computed exactly without materializing the full intermediate attention matrix in high-bandwidth memory.
- IO-aware tiling reduces memory traffic and improves wall-clock performance.
- The method enables longer contexts within the same memory budget.
Architecture, step by step
- 1
Tile inputs
Query, key, and value matrices are divided into blocks sized for fast on-chip SRAM.
- 2
Compute locally
Attention contributions are calculated for one block at a time close to the compute units.
- 3
Online normalization
Running softmax statistics combine block results without storing the complete score matrix.
- 4
Write compact output
Only the necessary output and normalization state move back to high-bandwidth memory.
Production takeaways
- Profile memory movement as well as arithmetic utilization when optimizing accelerator workloads.
- Kernel availability, tensor shapes, precision, and hardware generation determine the realized speedup.
- Exact numerical equivalence at the algorithm level still requires testing under reduced precision.
- Longer context support changes batching and capacity economics; it does not make attention cost disappear.
Limitations and caveats
- Benefits depend on hardware, kernel implementation, sequence shapes, and surrounding model operations.
- The method improves exact attention execution but does not change its quadratic arithmetic complexity.
- End-to-end speedups can be smaller when attention is not the dominant workload.
Interview questions to test understanding
- Why can reducing memory traffic matter more than reducing FLOPs?
- How does online softmax avoid storing the entire attention matrix?
- Which workload shapes would reduce the benefit of a specialized attention kernel?
- What should you benchmark before enabling a new kernel in production?