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. 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 GPU on-chip SRAM. We analyze the IO complexity of FlashAttention, showing that it requires fewer HBM accesses than standard attention, and is optimal for a range of SRAM sizes. We also extend FlashAttention to block-sparse attention, yielding an approximate attention algorithm that is faster than any existing approximate attention method. FlashAttention trains Transformers faster than existing baselines: 15% end-to-end wall-clock speedup on BERT-large (seq. length 512) compared to the MLPerf 1.1 training speed record, 3$\times$ speedup on GPT-2 (seq. length 1K), and 2.4$\times$ speedup on long-range arena (seq. length 1K-4K). FlashAttention and block-sparse FlashAttention enable longer context in Transformers, yielding higher quality models (0.7 better perplexity on GPT-2 and 6.4 points of lift on long-document classification) and entirely new capabilities: the first Transformers to achieve better-than-chance performance on the Path-X challenge (seq. length 16K, 61.4% accuracy) and Path-256 (seq. length 64K, 63.1% accuracy).
核心贡献 · Key contributions
提出 FlashAttention，一种 IO 感知的精确注意力算法，通过分块减少 GPU 高带宽内存访问。 Proposes FlashAttention, an IO-aware exact attention algorithm using tiling to reduce GPU HBM accesses.
实现显著的端到端加速：BERT-large 提升 15%，GPT-2 提升 3 倍，长距离竞技场提升 2.4 倍。 Achieves significant wall-clock speedups: 15% on BERT-large, 3x on GPT-2, 2.4x on long-range arena.
支持 Transformer 更长上下文，GPT-2 困惑度降低 0.7，长文档分类提升 6.4 个百分点。 Enables longer context in Transformers, improving perplexity by 0.7 on GPT-2 and 6.4 points on long-document classification.
首个在 Path-X（61.4%）和 Path-256（63.1%）上超越随机水平的 Transformer。 First Transformer to achieve better-than-chance on Path-X (61.4%) and Path-256 (63.1%) via longer sequences.
扩展至块稀疏 FlashAttention，比所有现有近似注意力方法更快。 Extends to block-sparse FlashAttention, faster than all existing approximate attention methods.
提供 IO 复杂度分析，并证明在多种 SRAM 大小下精确注意力的最优性。 Provides IO complexity analysis and proves optimality for exact attention over a range of SRAM sizes.
局限 · Limitations
需要自定义 CUDA 内核，限制了跨 GPU 架构的可移植性和易用性。 Requires custom CUDA kernel, limiting portability and ease of use across GPU architectures.
仅对单 GPU 最优；多 GPU 的 IO 分析和并行化尚未探索。 Optimal only for single GPU; multi-GPU IO analysis and parallelization remain unexplored.
块稀疏注意力依赖预定义稀疏模式，可能不适用于所有任务。 Block-sparse attention relies on predefined sparsity patterns, which may not suit all tasks.
反向传播中的重计算增加 FLOPs，尽管被减少的 HBM 访问所抵消。 Recomputation in backward pass increases FLOPs, though offset by reduced HBM accesses.
对于极短序列，性能提升减弱，因为内存访问不再是主要瓶颈。 Performance gains diminish for very short sequences where memory access is less dominant.
论文章节 · Sections（共 14）
摘要Abstract
1 引言1 Introduction
2 背景2 Background
2.1 硬件性能2.1 Hardware Performance
2.2 标准注意力实现2.2 Standard Attention Implementation
3 FlashAttention：算法、分析与扩展3 FlashAttention: Algorithm, Analysis, and Extensions
3.1 一种结合分块和重计算的高效注意力算法3.1 An Efficient Attention Algorithm With Tiling and Recomputation
3.2 分析：FlashAttention 的 IO 复杂度3.2 Analysis: IO Complexity of FlashAttention