ai-agent-book/chapter5/paper-to-ppt/paper/sample_paper.md

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

**Authors**: Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré (Stanford University)

**Venue**: NeurIPS 2022

## Abstract

Transformers are slow and memory-hungry on long sequences, since the time and
memory complexity of self-attention are quadratic in sequence length. We argue
that a missing principle is making attention algorithms *IO-aware* — accounting
for reads and writes between levels of GPU memory. We propose **FlashAttention**,
an IO-aware exact attention algorithm that uses tiling to reduce the number of
memory reads/writes between GPU high-bandwidth memory (HBM) and on-chip SRAM.
FlashAttention trains Transformers faster than existing baselines: 15% end-to-end
speedup on BERT-large, 3x speedup on GPT-2, and enables longer context, yielding
higher quality models (0.7 better perplexity on GPT-2) and entirely new
capabilities (the first Transformers to achieve better-than-chance performance
on Path-X with 16K sequence length).

## 1. Introduction

Transformer models have grown larger and deeper, but equipping them with longer
context remains difficult, because the self-attention module at their heart has
time and memory complexity **quadratic** in sequence length. A key question is
whether making attention faster and more memory-efficient can help Transformer
models address their runtime and memory challenges for long sequences.

Many approximate attention methods (sparse, low-rank) aim to reduce the compute
and memory requirements. Although these methods reduce the FLOP count, they often
do not achieve wall-clock speedup, mainly because they focus on FLOP reduction
and ignore overheads from **memory access (IO)**.

Our main observation is that the principal missing ingredient is making attention
algorithms **IO-aware** — carefully accounting for reads and writes to different
levels of fast and slow memory (e.g., between fast on-chip SRAM and relatively
slow HBM on a GPU).

## 2. Background: GPU Memory Hierarchy

A GPU has a memory hierarchy with different bandwidth/size trade-offs:

| Memory level | Size (A100) | Bandwidth |
|--------------|-------------|-----------|
| SRAM (on-chip) | 20 MB | 19 TB/s |
| HBM (main GPU) | 40-80 GB | 1.5-2.0 TB/s |
| CPU DRAM | > 1 TB | 12.8 GB/s |

Standard attention implementations materialize the large N x N attention matrix S
and P to HBM, incurring O(N^2) HBM accesses. Since HBM is much slower than SRAM,
this memory traffic dominates the runtime for long sequences.

## 3. Method: FlashAttention

FlashAttention computes exact attention with far fewer HBM accesses via two
classical techniques adapted to attention:

- **Tiling**: split the inputs Q, K, V into blocks, load them from slow HBM to
  fast SRAM, compute attention per block, and accumulate the output. We never
  materialize the full N x N attention matrix in HBM.
- **Softmax rescaling / recomputation**: the softmax normalization is computed
  incrementally across blocks using online softmax; in the backward pass, the
  attention matrix is recomputed on-chip from stored statistics rather than
  read from HBM.

The result is an algorithm with **O(N^2 d)** FLOPs but only **O(N^2 d^2 / M)**
HBM accesses, where M is the SRAM size and d is the head dimension. For typical
values, this is many times fewer HBM accesses than standard attention.

## 4. Experiments and Results

FlashAttention delivers strong end-to-end and micro-benchmark results:

- **BERT-large** (seq. 512): 15% faster training than the MLPerf 1.1 record.
- **GPT-2** (seq. 1K): up to **3x** end-to-end speedup over HuggingFace and
  Megatron-LM implementations, with identical model quality.
- **Long-range Arena** (1K-4K): 2.4x speedup vs standard attention.
- **Long context quality**: enabling 4K context on GPT-2 gives **0.7 better
  perplexity**; a Path-X task with 16K sequence becomes solvable for the first
  time (61.4% accuracy), and block-sparse FlashAttention solves Path-256 (63.1%).

Memory usage scales **linearly** in sequence length (versus quadratic for
standard attention), up to 20x memory savings at long sequence lengths.

## 5. Conclusion

By treating attention as an IO-bound problem and minimizing HBM accesses through
tiling and recomputation, FlashAttention computes exact attention faster and with
a smaller memory footprint than approximate methods, while remaining numerically
exact. IO-awareness is a broadly useful principle: we hope FlashAttention
inspires IO-aware implementations of more deep learning primitives.