NeurIPS 2022 (Outstanding Paper)

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

By Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré · Stanford University & SUNY Buffalo · Published June 2022

TL;DR

"Standard self-attention is bottlenecked by memory read/writes (IO) to GPU High Bandwidth Memory, scaling quadratically with sequence length. FlashAttention uses tiling and online softmax tracking to compute exact attention in fast GPU SRAM blocks, reducing memory accesses from quadratic to linear and unlocking up to 7.6x speedups."

IO-Aware Algorithm SRAM Tiling Gradient Recomputation 100% Exact Math

Why This Matters

Transformers have revolutionized machine learning, but they suffer from a fundamental weakness: the core Self-Attention mechanism scales quadratically $O(N^2)$ with the sequence length $N$. If you double your sequence length, the memory and compute requirements quadruple.

To solve this, researchers designed dozens of "approximate" attention methods (like Linformer, Reformer, or Performer) that reduce computational complexity to linear $O(N)$. Yet, in practice, most of these approximate methods run slower than standard attention.

Why? Because they focus purely on reducing floating-point operations (FLOPs) while ignoring the real hardware bottleneck of modern GPUs: memory bandwidth. GPU compute speed has drastically outpaced memory speed. The time spent loading and storing data to relatively slow High Bandwidth Memory (HBM) dominates the actual matrix multiplication time.

The Big Idea

Instead of changing the math of attention to reduce FLOPs, FlashAttention changes how the math is executed on hardware. It is an IO-aware exact attention algorithm.

"We must avoid reading and writing the intermediate $N \times N$ attention matrix to and from relatively slow HBM. We do this by loading inputs in blocks to fast on-chip SRAM, computing attention locally, and updating the output incrementally."

FlashAttention achieves this through two key techniques:

Tiling: We restructure the softmax operation so that it can be computed incrementally on small blocks of the input, without needing to see the entire matrix at once.
Recomputation: Instead of saving the massive $N \times N$ attention matrix in HBM during the forward pass to use for backpropagation, we only store the softmax normalization statistics. During the backward pass, we recompute the attention matrix block-by-block in fast SRAM, which is significantly faster than reading it from HBM.

How It Works

Step-by-step breakdown of exact block-based attention computation.

In standard attention, we compute: $$S = QK^T \in \mathbb{R}^{N \times N}, \quad P = \text{softmax}(S) \in \mathbb{R}^{N \times N}, \quad O = PV \in \mathbb{R}^{N \times d}$$ Writing the $N \times N$ matrices $S$ and $P$ to HBM consumes massive memory and bandwidth. FlashAttention splits $Q, K, V$ into blocks, loads them into fast on-chip SRAM, and computes attention incrementally.

Interactive Tiling Simulator

Step through the block-wise incremental softmax execution.

Q (Query)

Block 0

Block 1

Block 2

Block 3

K (Key)

Blk 0

Blk 1

Blk 2

Blk 3

SRAM (Softmax)

→

O (Output)

O[0]

O[1]

O[2]

O[3]

Step 0 / 16

Initialization

Click "Next Step" to begin the block-wise attention loop. We will load blocks of Q, K, and V into fast SRAM memory.

Active Row (i): -

Active Col (j): -

Row Max ($m_i$): -

Row Sum ($\ell_i$): -

Understanding GPU Memory Bottlenecks

Drag the slider to adjust the sequence length and see how memory accesses explode in Standard Attention compared to FlashAttention.

Sequence Length ($N$) 1024

Standard Attention (HBM IO)

4.19 MB

Requires writing/reading $O(N^2)$ intermediate matrices back and forth to slow HBM.

FlashAttention (HBM IO)

0.52 MB

Optimal $O(N)$ memory access. Fuses operations inside fast SRAM blocks.

SRAM vs HBM Speed Ratio

SRAM is ~12.6x faster than HBM (19 TB/s vs 1.5 TB/s)

HBM Read/Write Footprint Scaling

Lower is Better

Standard

FlashAttention

As $N$ grows, Standard Attention's memory IO scales quadratically, completely saturating GPU memory bandwidth.

Results & Benchmarks

FlashAttention is both faster and more memory-efficient than any other exact or approximate attention method.

BERT-Large Training

15%

End-to-end wall-clock speedup compared to the highly optimized NVIDIA MLPerf 1.1 record.

GPT-2 Speedup

3.0x

Speedup in end-to-end training compared to HuggingFace implementations (sequence length 1K).

Path-X Breakthrough

61.4%

First Transformer to achieve better-than-chance accuracy on Path-X (seq length 16K).

Speedup & Memory Scaling Benchmarks

Actual GPU benchmark numbers (measured on A100 GPU with FP16 precision).

Standard PyTorch Attention

FlashAttention

Block-Sparse FlashAttention

Limitations & Open Questions

While FlashAttention represents a massive step forward, the authors highlight several limitations and areas for future exploration:

Writing CUDA Kernels is Hard: FlashAttention requires writing custom CUDA kernels to achieve fine-grained control over memory accesses. This is highly non-trivial and hard to transfer across different GPU architectures (e.g., AMD or custom TPUs).
Compiling to High-Level Languages: There is a critical need for high-level deep learning compilers (like Halide or TVM) that can automatically compile standard attention definitions into optimal, IO-aware CUDA implementations without manual engineering effort.
Multi-GPU IO Optimization: FlashAttention is currently optimized for a single GPU. Scaling IO-awareness to multi-GPU setups—where data transfer across NVLink or PCIe is the next bottleneck—is an open challenge.

Glossary

Self-Attention

A mechanism in deep learning that relates different positions of a single sequence in order to compute a representation of the sequence. It calculates a weight for each pair of elements in the sequence, resulting in $O(N^2)$ computational complexity.

FLOPs (Floating Point Operations)

A measure of computer performance, useful in fields of scientific computations that require floating-point calculations. In deep learning, reducing FLOPs does not always translate to faster runtime if the hardware is bottlenecked by memory bandwidth.

High Bandwidth Memory (HBM)

A high-speed computer memory interface for 3D-stacked SDRAM from AMD and Hynix. On GPUs (like the A100), HBM acts as the main global memory. It is massive (40GB-80GB) but relatively slow compared to on-chip SRAM cache.

SRAM (Static Random-Access Memory)

Extremely fast, on-chip memory situated directly next to the GPU processor cores. It is highly limited in size (e.g., 192KB per Streaming Multiprocessor on an A100) but provides bandwidth up to 19 TB/s.

Citation

@inproceedings{dao2022flashattention,
  title={Flashattention: Fast and memory-efficient exact attention with io-awareness},
  author={Dao, Tri and Fu, Daniel Y and Ermon, Stefano and Rudra, Atri and R{\'e}, Christopher},
  booktitle={Advances in Neural Information Processing Systems},
  volume={35},
  pages={16344--16359},
  year={2022}
}