WishWiki

FlashAttention

FlashAttention is an Algorithm for computing exact Self-attention in Transformer models while dramatically reducing the memory traffic that ordinarily makes attention slow on long sequences. Introduced in 2022 by Tri Dao and collaborators at Stanford, it belongs to a small class of ideas that feel almost magical: it changes nothing about the mathematics of attention, yet makes it run many times faster and fit sequences many times longer, purely by rethinking where data lives while the computation happens.

The Problem It Solves

Ever since Attention Is All You Need introduced the Attention mechanism that powers modern large language models, engineers have wrestled with an ugly fact: attention costs grow quadratically with sequence length, both in compute and in memory. Earlier fixes tried approximate attention — sparsifying or low-ranking the attention matrix — trading fidelity for speed. Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem by trading off model quality to reduce the compute complexity, but often do not achieve wall-clock speedup.

The IO-Aware Insight

FlashAttention's breakthrough was recognizing that the real bottleneck wasn't raw arithmetic but the movement of data between a GPU's large, slow high-bandwidth memory (HBM) and its tiny, blazing-fast on-chip SRAM. 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 and GPU on-chip SRAM. By splitting the input into blocks — a technique called tiling — and computing softmax incrementally without ever materializing the full attention matrix in slow memory, the algorithm achieves speed that scales with hardware reality rather than theoretical FLOP counts alone. The result: identical outputs to standard attention, but with far fewer trips across the memory hierarchy, plus tangible downstream gains such as improved perplexity on GPT-2 and better handling of long documents.

Evolution: FlashAttention-2 and -3

The idea proved fertile ground for iteration. FlashAttention-2 improved parallelism and work partitioning across GPU thread blocks. FlashAttention-3 pushed further by exploiting the asynchronous execution capabilities of newer Hopper-architecture GPUs, overlapping computation with data movement and using low-precision FP8 arithmetic. FlashAttention-3 achieves speedup on H100 GPUs by 1.5-2.0X with BF16 reaching up to 840 TFLOPs/s (85% utilization), and with FP8 reaching 1.3 PFLOPs/s.

Why It Matters

FlashAttention has become foundational infrastructure across the machine learning ecosystem, quietly underpinning training and inference for models built on T5-like and GPT-3-like architectures, and complementing efficiency techniques like Low-rank adaptation and Mixture of experts. Because it is exact rather than approximate, it lets researchers extend context windows without sacrificing the correctness that downstream methods like Reinforcement learning from human feedback and InstructGPT-style fine-tuning depend upon. It stands as a reminder that in Networked systems and hardware alike, sometimes the fastest path forward isn't a cleverer formula but a cleverer sense of place — knowing exactly where your data should sit at every moment.

Related

Sources

  1. [2205.14135] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
  2. [2205.14135] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
  3. fast and memory-efficient exact attention with IO-awareness
  4. [PDF] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness | Semantic Scholar
  5. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
  6. GitHub - Dao-AILab/flash-attention: Fast and memory-efficient exact attention · GitHub
  7. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
  8. Storage usage
  9. FlashAttention-2: Faster Attention with Better Parallelism ...
  10. FLASHATTENTION: Fast and Memory-Efficient Exact Attention with IO-Awareness
Wishing…
your wish is being written

✨ Wish for a new page

👁 Wish for another view of this page

Sign in to WishWiki

Keep your wishes together, see your activity — and later, get your own private wiki space.

⏱ Page history