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

Abstract

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 IOaware-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× speedup on GPT-2 (seq. length 1K), and 2.4× 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).

OpenFold: Retraining AlphaFold2 yields new insights into its learning mechanisms and capacity for generalization

“…These optimizations create trade-offs between memory consumption and speed that can be tuned differently for training and inference. They include advanced implementations of neural network attention mechanisms (Vaswani et al 2017) with favorable properties for unusually short and long sequences (Rabe and Staats 2021, Dao et al 2022), module refactoring for lower memory usage, optional approximations of certain computations that reduce the memory burden, and specialized low-level code customized for GPU hardware. For technical details see appendices F.1 and F.2.…”

“…FlashAttention: We incorporate FlashAttention (Dao et al 2022), an efficient fused attention implementation that tiles computation in order to reduce data movement between different levels of GPU memory, greatly improving peak memory usage and runtime in the process. We find it to be particularly effective for short sequences with 1,000 residues or less, on which it contributes to an OpenFold speedup of up to 15% despite only being compatible with a small number of the attention modules in the network.…”

“…Apart from new training code and data, OpenFold has several advantages over AlphaFold2: (i) it runs up to three times faster on most proteins, (ii) it uses less memory, allowing prediction of extremely long proteins and multi-protein complexes on a single GPU, and (iii) it is implemented in PyTorch (Paszke et al 2019), the most widely used machine learning framework (AlphaFold2 uses Google’s JAX (Bradbury et al 2018)). As such, OpenFold can be readily used by the widest community of developers and interfaces with a rich ecosystem of existing machine learning software (Rasley et al 2020, Charlier et al 2021, Falcon et al 2019, Charlier et al 2021, Dao et al 2022).…”

scFormer: A Universal Representation Learning Approach for Single-Cell Data Using Transformers

“…This greatly exceeds the input length of common transformers used in NLM. Efficient self-attention techniques can be used (Katharopoulos et al, 2020;Wang et al, 2020;Dao et al, 2022). Also, since the order of the genes is not sequential in scRNA-seq data, and the transformer computation is agnostic to the order, we can dynamically sample subsets of the input.…”

ByteTransformer: A High-Performance Transformer Boosted for Variable-Length Inputs

“…Among which, PyTorch provides a standard implementation of MHA [28]; NVIDIA TensorRT provides fused MHA for short sequences whose lengths are smaller than 512 [29]. To scale the fused MHA to long sequences, Stanford researchers propose FlashAttention [30], which assumes identical shapes of inputs and assigns the workload of a whole attention unit to a single CTA. However, FlashAttention brings significant wasted computations if input sequence lengths are variable.…”