F LASH A TTENTION : Fast and Memory-Efﬁcient 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 IO-aware — accounting for reads and writes between levels of GPU memory.
We propose
F LASH A TTENTION , 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 F LASH A TTENTION , showing that it requires fewer HBM accesses than standard attention, and is optimal for a range of SRAM sizes.
We also extend F LASH A TTENTION to block-sparse attention, yielding an approximate attention algorithm
that is faster than any existing approximate attention method.
F LASH A TTENTION 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).
F LASH A TTENTION and block-sparse F LASH A TTENTION 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 classiﬁcation) and entirely new capabilities: the ﬁrst 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).
Introduction
Transformer models [86] have emerged as the most widely used architecture in applications such as natural language processing and image classiﬁcation.
Transformers have grown larger [5] and deeper [87], but equipping them with longer context remains difﬁcult [83], since the self-attention module at their heart has time and memory complexity quadratic in sequence length.
An important question is whether making attention faster and more memory-efﬁcient can help Transformer models address their runtime and memory challenges for long sequences.
Many approximate attention methods have aimed to reduce the compute and memory requirements of attention.
These methods range from sparse-approximation [53, 77] to low-rank approximation [13, 52, 88], and their combinations [3, 9, 96].
Although these methods reduce the compute requirements to linear or near-linear in sequence length, many of them do not display wall-clock speedup against standard attention and have not gained wide adoption.
One main reason is that they focus on FLOP reduction (which may not correlate with wall-clock speed) and tend to ignore overheads from memory access (IO).
In this paper, we argue that a missing principle is making attention algorithms IO-aware [1]—that is, carefully accounting for reads and writes to different levels of fast and slow memory (e.g., between fast GPU on-chip SRAM and relatively slow GPU high bandwidth memory, or HBM [47], Figure
FlashAttention
Memory Hierarchy with Bandwidth & Memory Size
Attention on GPTFlashAttention PyTorch
T i m e ( m s )
Matmul
Mask
Softmax
Dropout
Matmul
FusedKernel
Q: N x d
V: N X d
K T : d x N
Q K T : N x N sm (QK T )V: N x d
Outer Loop
Copy Block to SRAM
Copy
O u t e r L oo p
Copy
I nn
e r
L oo p
Compute Block on SRAM
Output to HBM
Inner Loop
I nn e r L oo p
Outer Loop
[Figure]
left).
On modern GPUs, compute speed has out-paced memory speed [64–66], and most operations in Transformers are bottlenecked by memory accesses [45].
IO-aware algorithms have been critical for similar memory-bound operations, when reading and writing data can account for a large portion of the runtime—such as database joins [74], image processing [73], numerical linear algebra [4], and more [42, 89].
However, common Python interfaces to deep learning such as PyTorch and Tensorﬂow do not allow ﬁne-grained control of memory access.
We propose F LASH A TTENTION , a new attention algorithm that computes exact attention with far fewer memory accesses.
Our main goal is to avoid reading and writing the attention matrix to and from HBM.
This requires (i) computing the softmax reduction without access to the whole input (ii) not storing the large intermediate attention matrix for the backward pass.
We apply two well-established techniques to address these challenges.
(i)
We restructure the attention computation to split the input into blocks and make several passes over input blocks, thus incrementally performing the softmax reduction (also known as tiling ).
(ii) We store the softmax normalization factor from the forward pass to quickly recompute attention on-chip in the backward pass, which is faster than the standard approach of reading the intermediate attention matrix from HBM.
We implement
F LASH A TTENTION in CUDA to achieve ﬁne-grained control over memory access and fuse all the attention operations into one GPU kernel.
Even with the increased FLOPs due to recomputation, our algorithm both runs faster (up to 7.6x on GPT-2 [70], Figure 1 right) and uses less memory —linear in sequence length—than standard attention, thanks to the massively reduced amount of HBM access.
We analyze the IO complexity [1] of F LASH A TTENTION , proving that it requires 𝑂 ( 𝑁 2 𝑑 2 𝑀 − 1 ) HBM accesses where 𝑑 is the head dimension and 𝑀 is the size of SRAM, as compared to Ω ( 𝑁𝑑 + 𝑁 2 ) of standard attention.
For typical values of 𝑑 and 𝑀 , F LASH A TTENTION requires many times fewer
HBM accesses compared to standard attention (up to 9 × fewer, as shown in Fig. 2).
Moreover, we provide a lower bound, showing that no exact attention algorithm can asymptotically improve on the number of HBM accesses over all SRAM sizes.
We also show that F LASH A TTENTION can serve as a useful primitive for realizing the potential of approximate attention algorithms by overcoming their issues with memory access overhead.
As a proof of concept, we implement block-sparse F LASH A TTENTION , a sparse attention algorithm that is 2-4 × faster than even F LASH A TTENTION , scaling up to sequence length of 64k.
We prove that block-sparse F LASH A TTENTION has better IO complexity than F LASH A TTENTION by a factor proportional to the sparsity ratio.
We discuss further extensions to other operations (attention on multi-GPU, kernel regression, block-sparse matrix multiply) in Section 5.
We open-source
F LASH A TTENTION to make it easier to build on this primitive 1 .
We empirically validate that F LASH A TTENTION speeds up model training and improves model quality by modeling longer context.
We also benchmark the runtime and memory footprint of F LASH A TTENTION and block-sparse F LASH A TTENTION compared to prior attention implementations.
• Faster Model Training.
F LASH A TTENTION trains Transformer models faster in wall-clock time.
We train BERT-large (seq. length 512) 15% faster than the training speed record in MLPerf 1.1 [60], GPT2 (seq.
length 1K)
3 × faster than baseline implementations from HuggingFace [91] and Megatron-LM [80], and long-range arena (seq.
length 1K-4K) 2.4 × faster than baselines. • Higher Quality Models.
F LASH A TTENTION scales Transformers to longer sequences, which improves their quality and enables new capabilities.
We observe a 0.7 improvement in perplexity on GPT-2 and 6.4 points of lift from modeling longer sequences on long-document classiﬁcation [14]. F LASH A TTENTION enables the ﬁrst Transformer that can achieve better-than-chance performance on the Path-X [83] challenge, solely from using a longer sequence length (16K).
Block-sparse F LASH A TTENTION enables a Transformer to scale to even longer sequences (64K), resulting in the ﬁrst model that can achieve better-than-chance performance on Path-256. • Benchmarking Attention.
F LASH A TTENTION is up to 3 × faster than the standard attention implementation across common sequence lengths from 128 to 2K and scales up to 64K. Up to sequence length of 512, F LASH A TTENTION is both faster and more memory-efﬁcient than any existing attention method, whereas for sequence length beyond 1K, some approximate attention methods (e.g., Linformer) start to become faster.
On the other hand, block-sparse F LASH A TTENTION is faster than all existing approximate attention methods that we know of.
Background
We provide some background on the performance characteristics of common deep learning operations on modern hardware (GPUs).
We also describe the standard implementation of attention.
2.1 Hardware Performance
We focus here on GPUs.
Performance on other hardware accelerators are similar [48, 50].
GPUMemoryHierarchy.
The GPU memory hierarchy (Fig. 1 left) comprises multiple forms of memory of different sizes and speeds, with smaller memory being faster.
As an example, the A100 GPU has 40-80GB of high bandwidth memory (HBM) with bandwidth 1.5-2.0TB/s and 192KB of on-chip
SRAM per each of 108 streaming multiprocessors with bandwidth estimated around 19TB/s [46, 47].
The on-chip SRAM is an order of magnitude faster than HBM but many orders of magnitude smaller in size.
As compute has gotten faster relative to memory speed [64–66], operations are increasingly bottlenecked by memory (HBM) accesses.
Thus exploiting fast SRAM becomes more important.
Execution Model.
GPUs have a massive number of threads to execute an operation (called a kernel).
Each kernel loads inputs from HBM to registers and SRAM, computes, then writes outputs to HBM.
Performance characteristics.
Depending on the balance of computation and memory accesses, operations can be classiﬁed as either compute-bound or memory-bound.
This is commonly measured by the arithmetic intensity [89], which is the number of arithmetic operations per byte of memory access.
Compute-bound: the time taken by the operation is determined by how many arithmetic operations there are, while time accessing HBM is much smaller.
Typical examples are matrix multiply with large inner dimension
and convolution with large number of channels.
Memory-bound: the time taken by the operation is determined by the number of memory accesses, while time spent in computation is much smaller.
Examples include most other operations: elementwise (e
g., activation, dropout), and reduction (e
g., sum, softmax, batch norm, layer norm).
Kernel fusion.
The most common approach to accelerate memory-bound operations is kernel fusion: if there are multiple operations applied to the same input, the input can be loaded once from HBM, instead of multiple times for each operation.
Compilers can automatically fuse many elementwise operations [55, 68, 78].
However, in the context of model training, the intermediate values still need to be written to HBM to save for the backward pass, reducing the effectiveness of naive kernel fusion.
2.2 Standard Attention Implementation
Given input sequences Q , K , V ∈ R 𝑁 × 𝑑 where 𝑁 is the sequence length and 𝑑 is the head dimension, we want to compute the attention output O ∈ R 𝑁 × 𝑑 :
[Equation]
where softmax is applied row-wise.
Standard attention implementations materialize the matrices S and P to HBM, which takes 𝑂 ( 𝑁 2 ) memory.
Often 𝑁 (cid:29) 𝑑
(e.g., for GPT2, 𝑁 = 1024 and 𝑑 = 64 ).
We describe the standard attention implementation in Algorithm 0.
As some or most of the operations are memory-bound (e.g., softmax), the large number of memory accesses translates to slow wall-clock time.
This problem is exacerbated by other elementwise operations applied to the attention matrix, such as masking applied to S or dropout applied to P .
As a result, there have been many attempts to fuse several elementwise operations, such as fusing masking with softmax [80].
In Section 3.2, we will show that the standard attention implementation performs HBM accesses quadratic in the sequence length 𝑁 .
We also compare the number of FLOPs and number of HBM accesses of standard attention and of our method (F LASH A TTENTION ).
Require: Matrices Q , K , V ∈ R 𝑁 × 𝑑 in HBM.
F LASH A TTENTION : Algorithm, Analysis, and Extensions
We show how to compute exact attention with fewer HBM reads/writes and without storing large intermediate matrices for the backward pass.
This yields an attention algorithm that is both memory efﬁcient and faster in wall-clock time.
We analyze its IO complexity, showing that our method requires much fewer HBM accesses compared to standard attention.
We further show that
F LASH A TTENTION can serve as a useful primitive by extending it to handle block-sparse attention.
We focus here on the forward pass for ease of exposition; Appendix B contains details for the backward.
3.1 An Efﬁcient Attention Algorithm With Tiling and Recomputation
Given the inputs Q , K , V ∈ R 𝑁 × 𝑑 in HBM, we aim to compute the attention output O ∈ R 𝑁 × 𝑑 and write it to HBM.
Our goal is to reduce the amount of HBM accesses (to sub-quadratic in 𝑁 ).
We apply two established techniques (tiling, recomputation) to overcome the technical challenge of computing exact attention in sub-quadratic HBM accesses.
We describe this in Algorithm 1.
The main idea is that we split the inputs Q , K , V into blocks, load them from slow HBM to fast SRAM, then compute the attention output with respect to those blocks.
By scaling the output of each block by the right normalization factor before adding them up, we get the correct result at the end.
Tiling.
We compute attention by blocks.
Softmax couples columns of K , so we decompose the large softmax with scaling [53, 62, 69].
For numerical stability, the softmax of vector 𝑥 ∈ R 𝐵 is computed: 𝑚
[Equation]
For vectors 𝑥 ( 1 ) ,𝑥 ( 2 ) ∈
R 𝐵 , we can decompose the softmax of the concatenated 𝑥 = (cid:2) 𝑥 ( 1 ) 𝑥 ( 2 ) (cid:3)
∈ R 2 𝐵 as: 𝑚 ( 𝑥 ) =
𝑚
( (cid:2) 𝑥 ( 1 ) 𝑥 (
[Equation]
) (cid:3) )
=
max
[Equation]
𝑚 ( 𝑥 (
[Equation]
𝑚
[Equation]
𝑥 ( [Equation] [Equation] Therefore if we keep track of some extra statistics ( 𝑚 ( 𝑥 ) ,ℓ ( 𝑥 ) ), we can compute softmax one block at a time.
2 We thus split the inputs Q , K , V into blocks (Algorithm 1 line 1), compute the softmax values along with extra statistics (Algorithm 1 line 1), and combine the results (Algorithm 1 line 1).
Recomputation.
One of our goals is to not store 𝑂 ( 𝑁 2 ) intermediate values for the backward pass.
The backward pass typically requires the matrices S , P ∈ R 𝑁 × 𝑁 to compute the gradients with respect to Q , K , V .
However, by storing the output O and the softmax normalization statistics ( 𝑚,ℓ ) , we can recompute the attention matrix S and P easily in the backward pass from blocks of Q , K , V in
SRAM.
This can be seen as a form of selective gradient checkpointing [10, 36].
While gradient 2 This style of aggregation is called algebraic aggregation [35].
checkpointing has been suggested to reduce the maximum amount of memory required [69], all implementations (that we know off) have to trade speed for memory.
In contrast, even with more
FLOPs, our recomputation speeds up the backward pass due to reduced HBM accesses (Fig. 2).
The full backward pass description is in Appendix B.
Implementation details: Kernel fusion.
Tiling enables us to implement our algorithm in one CUDA kernel, loading input from HBM, performing all the computation steps (matrix multiply, softmax, optionally masking and dropout, matrix multiply), then write the result back to HBM (masking and dropout in Appendix B).
This avoids repeatedly reading and writing of inputs and outputs from and to HBM.
Require: Matrices Q , K , V ∈ R 𝑁 × in HBM, on-chip SRAM of size
Load Q 𝑖 , O 𝑖 ,ℓ 𝑖 ,𝑚 𝑖 from HBM to on-chip SRAM. On chip, compute S 𝑖𝑗 =
Q 𝑖 K 𝑇𝑗 ∈ R 𝐵 𝑟 × 𝐵 𝑐 .
On chip, compute ˜ 𝑚 𝑖𝑗 = rowmax ( S 𝑖𝑗 ) ∈ R 𝐵 𝑟 , ˜ P
𝑖𝑗 = exp ( S 𝑖𝑗 − ˜ 𝑚 𝑖𝑗 )
∈ R 𝐵 𝑟 × 𝐵 𝑐 (pointwise), ˜ ℓ 𝑖𝑗 = rowsum ( ˜ P 𝑖𝑗 )
∈ R 𝐵 𝑟 .
On chip, compute 𝑚 new 𝑖 = max ( 𝑚 𝑖 , ˜ 𝑚
𝑖𝑗 ) ∈ R 𝐵 𝑟 , ℓ new 𝑖 = 𝑒 𝑚 𝑖 − 𝑚 new 𝑖 ℓ 𝑖 + 𝑒 ˜ 𝑚
𝑖𝑗 − 𝑚 new 𝑖 ˜ ℓ 𝑖𝑗 ∈ R 𝐵 𝑟 .
Write O 𝑖 ← diag ( ℓ new 𝑖 ) − 1 ( diag ( ℓ 𝑖 ) 𝑒 𝑚 𝑖 − 𝑚 new 𝑖 O 𝑖 + 𝑒 ˜ 𝑚
𝑖𝑗 − 𝑚 new 𝑖 ˜ P 𝑖𝑗 V 𝑗 ) to HBM.
Write ℓ 𝑖 ← ℓ new 𝑖 , 𝑚 𝑖 ← 𝑚 new 𝑖 to HBM.
14: end for 15: end for 16:
Return O .
We show F LASH A TTENTION ’s correctness, runtime, and memory requirement (proof in Appendix C).
Theorem 1.
Algorithm 1 returns O = softmax ( QK (cid:62) ) V with 𝑂 ( 𝑁 2 𝑑 ) FLOPs and requires 𝑂 ( 𝑁 ) additional memory beyond inputs and output.
3.2 Analysis: IO Complexity of F LASH A TTENTION
We analyze the IO complexity of F LASH A TTENTION , showing signiﬁcant reduction in HBM accesses compared to standard attention.
We also provide a lower bound, proving that no exact attention algorithm can asymptotically improve on HBM accesses over all SRAM sizes.
Proofs are in Appendix C.
Theorem 2.
Let 𝑁 be the sequence length, 𝑑 be the head dimension, and 𝑀 be size of SRAM with 𝑑 ≤ 𝑀 ≤ 𝑁𝑑 .
Standard attention (Algorithm 0) requires Θ ( 𝑁𝑑 + 𝑁 2 ) HBM accesses, while
F LASH A TTENTION (Algorithm 1) requires Θ ( 𝑁 2 𝑑 2 𝑀 − 1 ) HBM accesses.
For typical values of 𝑑 (64-128) and 𝑀 (around 100KB), 𝑑 2 is many times smaller than 𝑀 , and thus
F LASH A TTENTION requires many times fewer HBM accesses than standard implementation.
This leads to both faster execution and lower memory footprint, which we validate in Section 4.3.
The main idea of the proof is that given the SRAM size of 𝑀 , we can load blocks of K , V of size Θ ( 𝑀 ) each (Algorithm 1 line 1).
For each block of K and V , we iterate over all blocks of Q (Algorithm 1 line 1) to compute the intermediate values, resulting in Θ ( 𝑁𝑑𝑀 − 1 ) passes over Q .
Each pass loads Θ ( 𝑁𝑑 ) elements, which amounts to Θ ( 𝑁 2 𝑑 2 𝑀 − 1 ) HBM accesses.
We similarly prove that the backward pass of standard attention requires Θ ( 𝑁𝑑 + 𝑁 2 ) HBM accesses while the backward pass of F LASH A TTENTION requires Θ ( 𝑁 2 𝑑 2 𝑀 − 1 ) HBM accesses (Appendix B).
We prove a lower-bound: one cannot asymptotically improve on the number of HBM accesses for all values of 𝑀 (the SRAM size) when computing exact attention. [Figure]
Proposition 3.
Let 𝑁 be the sequence length, 𝑑 be the head dimension, and 𝑀 be size of SRAM with 𝑑 ≤ 𝑀 ≤ 𝑁𝑑 .
There does not exist an algorithm to compute exact attention with 𝑜 ( 𝑁 2 𝑑 2 𝑀 − 1 ) HBM accesses for all 𝑀 in the range [ 𝑑,𝑁𝑑 ] .
The proof relies on the fact that for 𝑀 =
Θ ( 𝑁𝑑 ) any algorithm must perform Ω ( 𝑁 2 𝑑 2 𝑀 − 1 )
=Ω ( 𝑁𝑑 )
HBM accesses.
This type of lower bound over a subrange of 𝑀 is common in the streaming algorithms literature [92].
We leave proving parameterized complexity [29] lower bounds in terms of 𝑀 as exciting future work.
We validate that the number of HBM accesses is the main determining factor of attention run-time.
In Fig.
2 (left), we see that even though F LASH A TTENTION has higher FLOP count compared to standard attention (due to recomputation in the backward pass), it has much fewer HBM accesses, resulting in much faster runtime.
In Fig.
2 (middle), we vary the block size 𝐵 𝑐 of F LASH A TTENTION , which results in different amounts of HBM accesses, and measure the runtime of the forward pass.
As block size increases, the number of HBM accesses decreases (as we make fewer passes over the input), and runtime decreases.
For large enough block size (beyond 256), the runtime is then bottlenecked by other factors (e.g., arithmetic operations).
Moreover, larger block size will not ﬁt into the small SRAM size.
3.3 Extension: Block-Sparse F LASH A TTENTION
We extend F LASH A TTENTION to approximate attention: we propose block-sparse F LASH A TTENTION , whose IO complexity is smaller than F LASH A TTENTION by a factor proportional to the sparsity.
Given inputs Q , K , V ∈ R 𝑁 × 𝑑 and a mask matrix ˜ M ∈ { 0 , 1 } 𝑁 × 𝑁 , we want to compute:
[Equation]
where ( S (cid:12) 𝟙 ˜ M ) 𝑘𝑙 = S 𝑘𝑙 if ˜ M 𝑘𝑙 = 1 and −∞ if M 𝑘𝑙 = 0 .
We require ˜
M to have block form: for some block sizes 𝐵 𝑟 ,𝐵 𝑐 , for all 𝑘,𝑙 ,
˜ M 𝑘,𝑙 =
M 𝑖𝑗 with 𝑖 = (cid:98) 𝑘 / 𝐵 𝑟 (cid:99) ,
𝑗 = (cid:98) 𝑙 / 𝐵 𝑐 (cid:99) for some M ∈ { 0 , 1 } 𝑁 / 𝐵 𝑟 × 𝑁 / 𝐵 𝑐 .
Given a predeﬁned block sparsity mask M ∈ { 0 , 1 } 𝑁 / 𝐵 𝑟 × 𝑁 / 𝐵 𝑐 we can easily adapt Algorithm 1 to only compute the nonzero blocks of the attention matrix.
The algorithm is identical to Algorithm 1, except we skip zero blocks.
We reproduce the algorithm description in Algorithm 5 in Appendix B.
We also analyze the IO complexity of block-sparse F LASH A TTENTION .
Proposition 4.
Let 𝑁 be the sequence length, 𝑑 be the head dimension, and 𝑀 be size of SRAM with 𝑑 ≤ 𝑀 ≤ 𝑁𝑑 .
Block-sparse F LASH A TTENTION (Algorithm 5) requires Θ ( 𝑁𝑑 + 𝑁 2 𝑑 2 𝑀 − 1 𝑠 ) HBM accesses where 𝑠 is the fraction of nonzero blocks in the block-sparsity mask.
We see that applying block-sparsity yields a direct improvement by the sparsity to the larger term in the IO complexity.
For large sequence lengths 𝑁 , 𝑠 is often set to 𝑁 − 1 / 2 [12] or 𝑁 − 1 log 𝑁 [3, 18, 96], resulting in Θ ( 𝑁 √ 𝑁 ) or Θ ( 𝑁 log 𝑁 ) IO complexity.
For downstream experiments, we use the ﬁxed butterﬂy sparsity pattern [18], which has been shown to be able to approximate arbitrary sparsity [17].
In Fig.
2 (right), we validate that as the sparsity increases, the runtime of block-sparse F LASH A TTENTION improves proportionally.
On the LRA benchmark, block-sparse F LASH A TTENTION achieves 2.8 × speedup, while performing on par with standard attention (Section 4).
Experiments
We evaluate the impact of using F LASH A TTENTION to train Transformer models.
We validate two claims about training time and model accuracy, and report attention runtime and memory benchmarks.
• Training Speed.
F LASH A TTENTION outperforms the MLPerf 1.1 [60] speed record for BERT by 15%, and speeds up GPT-2 up to 3 × over HuggingFace [91] and 1 . 8 × over Megatron [80] over standard Transformers.
F LASH A TTENTION speeds up the long-range arena (LRA) benchmark 2.4 × .
• Quality.
F LASH A TTENTION scales Transformers to longer sequences, yielding higher quality.
F LASH A TTENTION trains GPT-2 with context length 4K faster than Megatron trains GPT-2 with context length 1K, while achieving 0.7 better perplexity.
Modeling longer sequences yields 6. points of lift on two long-document classiﬁcation tasks.
Finally, F LASH A TTENTION yields the ﬁrst Transformer that can achieve better-than-random performance on the challenging Path-X task (sequence length 16K), and block-sparse F LASH A TTENTION yields the ﬁrst sequence model that we know of that can achieve better-than-random performance on Path-256 (sequence length 64K). • Benchmarking Attention.
We measure the runtime and memory performance of F LASH A TTENTION and block-sparse F LASH A TTENTION based on sequence length.
We conﬁrm that the memory footprint of F LASH A TTENTION scales linearly with seq.
length and is up to 3 × faster than standard attention for common seq.
lengths (up to 2K).
We conﬁrm that runtime of block-sparse F LASH A TTENTION scales linearly in seq.
length and is faster than all existing approximate attention baselines.
Additional experiment details are in Appendix E.
4.1 Faster Models with F LASH A TTENTION BERT.
F LASH A TTENTION yields the fastest single-node BERT training speed that we know of.
We train a BERT-large [24] model with F LASH A TTENTION on Wikipedia.
Table 1 compares our training time to the implementation from Nvidia that set the training speed record for MLPerf 1.1 [60, 63].
Our implementation is 15% faster.
[Table]
F LASH A TTENTION yields faster training times for GPT-2 [70] on the large OpenWebtext dataset [34] than the widely used HuggingFace [91] and Megatron-LM [80] implementations.
Table 2 shows up to 3 × end-to-end speedup compared to Huggingface and 1.7 × speedup compared to Megatron-LM.
F LASH A TTENTION achieves the same perplexity as the other two implementations, as we do not change the model deﬁnition.
Appendix E includes plots of the validation perplexity throughout training, conﬁrming that F LASH A TTENTION is as numerically stable as the baselines and produces the same training / validation curves
[Table]
[Table]
Long-range Arena.
We compare vanilla Transformer (with either standard implementation or F LASH A TTENTION ) on the long-range arena (LRA [83]) benchmark.
We measure accuracy, throughput, and training time of all models.
Each task has a different sequence length varying between 1024 and 4096.
We follow the implementation and experimental setting in Tay et al.
[83]and Xiong et al.
3 Table 3 shows that F LASH A TTENTION achieves up 2.4 × speed-up compared to standard attention.
Block-sparse F LASH A TTENTION is faster than all of the approximate attention methods that we have tested.
4.2 Better Models with Longer Sequences
Language Modeling with Long Context.
The runtime and memory-efﬁciency of F LASH A TTENTION allow us to increase the context length of GPT-2 by 4 × while still running faster than the optimized implementation from Megatron-LM.
Table 4 shows that that GPT-2 with
[Table]
F LASH A TTENTION and context length 4K is still 30% faster than GPT-2 from Megatron with context length 1K, while achieving 0.7 better perplexity.
[Table]
Long Document Classiﬁcation.
Training Transformers with longer sequences with F LASH A TTENTION improves performance on the MIMIC-III [49] and ECtHR [6, 7] datasets.
MIMIC-III contains intensive care unit patient discharge summaries, each annotated with multiple labels.
ECtHR contains legal cases from the European Court of Human Rights, each of which is mapped to articles of the Convention of Human Rights that were allegedly violaged.
Both of these datasets contain very long text documents; the average number of tokens in MIMIC is 2,395 tokens, and the longest document contains 14,562 tokens, while the average and longest numbers in ECtHR are 2,197 and 49,392, respectively.
We evaluate lift from increasing the sequence length of a pretrained RoBERTa model [58] (we repeat the positional embeddings, as in Beltagy et al.
Table 5 shows that sequence length 16K outperforms length 512 by 4.3 points on MIMIC, and that length 8K outperforms length 512 by 8.5 points on ECtHR.
The discrepancies may be due to subtle distribution shifts: MIMIC-III contains specialized medical text and thus may be more susceptible to a distribution shift in the document length, whereas ECtHR contains general language.
[Table] [Table]
[Table]
Path-X and Path-256.
The Path-X and Path-256 benchmarks are challenging tasks from the long-range arena benchmark designed to test long context.
The task is to classify whether two points in a black and white 128 × 128 (or 256 × 256) image have a path connecting them, and the images are fed to the transformer one pixel at a time.
In prior work, all transformer models have either run out of memory, or only achieved random performance [83].
There has been a search for alternative architectures that can model such long context [39].
We present here the ﬁrst result of Transformer models being able to solve Path-X and Path-
(Table 6).
We pretrain a transformer on Path-64, and then transfer to Path-X by spatially interpolating the positional embeddings.
F LASH A TTENTION achieves 61.4 accuracy on Path-X. Additionally, block-sparse F LASH A TTENTION enables the
Transformers to scale to sequence length 64K, achieving 63.1 accuracy 4 on Path-256. [Figure] u n t i m e ( m s )
M e m o r y F oo t p r
i n t ( G B )
[Figure] 4.3 Benchmarking Attention
We vary sequence length and measure runtime and memory usage of F LASH A TTENTION and block-sparse F LASH A TTENTION against various attention baselines on one A100 GPU with 40 GB
HBM, with dropout and a padding mask.
We compare against reference implementations for exact attention, approximate attention, and sparse attention.
We report a subset of baselines in the main body; Appendix E contains more baselines and full details.
Runtime.
Figure 3 (left) reports the runtime in milliseconds of the forward + backward pass of F LASH A TTENTION and block-sparse F LASH A TTENTION compared to the baselines in exact, approximate, and sparse attention (exact numbers in Appendix E).
Runtime grows quadratically with sequence length, but F LASH A TTENTION runs signiﬁcantly faster than exact attention baselines, up to 3 × faster than the PyTorch implementation.
The runtimes of many approximate/sparse attention mechanisms grow linearly with sequence length, but F LASH A TTENTION still runs faster than approximate and sparse attention for short sequences due to fewer memory accesses.
The approximate attention runtimes begin to cross over with F LASH A TTENTION at sequences between 512 and 1024.
On the other hand, block-sparse F LASH A TTENTION is faster than all implementations of exact, sparse, and approximate attention that we know of, across all sequence lengths.
Memory Footprint.
Figure 3 (right) shows the memory footprint of F LASH A TTENTION and block-sparse F LASH A TTENTION compared to various exact, approximate, and sparse attention baselines.
F LASH A TTENTION and block-sparse F LASH A TTENTION have the same memory footprint, which grows linearly with sequence length.
F LASH A TTENTION is up to 20 × more memory efﬁcient than exact attention baselines, and is more memory-efﬁcient than the approximate attention baselines.
All other algorithms except for Linformer run out of memory on an A100 GPU before 64K, and F LASH A TTENTION is still 2 × more efﬁcient than Linformer.
Limitations and Future Directions
We discuss limitations of our approach and future directions.
Related work is given in Appendix A.
Compiling to CUDA.
Our current approach to building IO-aware implementations of attention requires writing a new CUDA kernel for each new attention implementation.
This requires writing the attention algorithm in a considerably lower-level language than PyTorch, and requires signiﬁcant engineering effort.
Implementations may also not be transferrable across GPU architectures.
These limitations suggest the need for a method that supports writing attention algorithms in a high-level language (e.g., PyTorch), and compiling to IO-aware implementations in CUDA—similar to efforts such as Halide in image processing [73].
IO-Aware Deep Learning.
We believe that the IO-aware approach can extend beyond attention.
Attention is the most memory-intensive computation in Transformers, but every layer in a deep network touches GPU HBM.
We hope our work inspires IO-aware implementations of additional modules.
We discuss these potential extensions in Appendix D.
Multi-GPU IO-Aware Methods.
Our IO-aware implementation of attention is optimal within constants for computing attention on a single GPU.
However, the attention computation may be par- allelizable across multiple GPUs [75].
Using multiple GPUs adds an additional layer to IO analysis— accounting for data transfer between GPUs.
We hope our work inspires future work in this direction.
Societal Impacts.
As Transformer-based foundation models grow in size and data, our work seeks to understand how to train these large models more efﬁciently.
This may allow a general community with limited access to computational resources to train and understand those foundation models.
Our method is applicable to all Transformer-based models, which have a variety of applications, both positive and negative.
For example, language modeling may make it easier to spread misinformation, while image classiﬁcation models may make automatic surveillance easier.
Alleviating these risks requires addressing application-speciﬁc issues such as privacy, bias, and discrimination.
Acknowledgments
Our implementation uses Apex’s FMHA code ( https://github.com/NVIDIA/apex/tree/ master/apex/contrib/csrc/fmha ) as a starting point.
We thank Young-Jun Ko for the in-depth explanation of his FMHA implementation and for his thoughtful answers to our questions about
CUDA.
We thank Sabri Eyuboglu, Megan Leszczynski, Laurel Orr, Yuhuai Wu, Beidi Chen, and
Xun Huang for their constructive feedback and suggestions on early drafts of the paper.
We thank
Markus Rabe and Charles Staats for helpful discussion of their attention algorithm.
We gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos.
{+CCF
+}
(Beyond Sparsity), CCF1563078 (Volume to Velocity), and(RTML); ARL under No.
W911NF-21-2-0251 (Interactive Human-AI Teaming);
ONR under
{+No. N
+}
(Unifying Weak Supervision)
; ONR N00014-20-1-2480: Understanding and Applying Non-Euclidean Geometry in Machine Learning; N000142012275 (NEPTUNE);
NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF,
Accenture, Ericsson, Qualcomm, Analog Devices, Google Cloud, Salesforce, Total, the HAI-GCP & HAI-Azure Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI), Department of Defense (DoD) through the National Defense Science and Engineering Graduate Fellowship (NDSEG) Program, and members of the Stanford DAWN project: Facebook, Google and VMWare.
The U S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, ﬁndings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reﬂect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, or the U S. Government.
Atri Rudra’s research is supported by NSF grant CCF-1763481.
A Related Work
IO-Aware Runtime Optimization.
The broad concept of optimizing for reading and writing to fast/slow memory has a long history in computer science and has been known by many names.
We draw the most direct connection to the literature of analyzing I/O complexity in this work [1], but concepts of memory hierarchies are fundamental and has appeared in many forms, from the working set model [23], to data locality [90], to the Rooﬂine model of arithmetic intensity [89], to analyses of scalability [61], to standard textbook treatments of computer architecture [42].
We hope that this work encourages the community to adopt these ideas in more parts of the deep learning stack.
Efﬁcient ML Models with Structured Matrices.
Matrix multiply is the core computational bottleneck of most machine learning models.
To reduce the computational complexity, there have been numerous approaches to learn over a more efﬁcient set of matrices.
These matrices are called structured matrices , which have subquadratic ( 𝑜 ( 𝑛 2 ) for dimension 𝑛 × 𝑛 ) number of parameters and runtime.
Most common examples of structured matrices are sparse and low-rank matrices, along with fast transforms commonly encountered in signal processing (Fourier, Chebyshev, sine/cosine, orthogonal polynomials).
There have been several more general classes of structured matrices proposed in machine learning: Toeplitz-like [81], low-displacement rank [51], quasi-separable [27]).
The butterﬂy pattern we use for our block-sparse attention is motivated by the fact that butterﬂy matrices [16, 67] and their products have been shown to be able to express any structured matrices with almost optimal runtime and number of parameters [17, 21].
However, even though structured matrices are efﬁcient in theory, they have not seen wide adoption since it is hard to translate their efﬁciency to wall-clock speedup since dense unconstrained matrix multiply has very optimize implementation, a phenomenon known as the hardware lottery [43].
Extensions of butterﬂy matrices [18, 19] aimed to make butterﬂy matrices more hardware-friendly.
Sparse Training.
Our block-sparse F LASH A TTENTION can be seen as a step towards making sparse model training more efﬁcient.
Sparse models have seen success in compressing models for inference (pruning) by sparsifying the weight matrices [25, 40, 41, 57, 79].
For model training, the lottery tickets hypothesis [30–32] suggests that there are a set of small sub-networks derived from a larger dense network that performs as well as the original dense network.
Out block-sparse
F LASH A TTENTION can also be seen as a ﬁxed lottery ticket in the context of attention: we ﬁx the sparsity pattern to be the butterﬂy pattern through training, and observe that it performs almost as well as the (dense) F LASH A TTENTION on the Long-range Arena tasks.
Efﬁcient Transformer.
Transformer-based models have become the most widely-used architecture in natural language processing [24] and computer vision [26, 95].
However, one of their computational bottlenecks is that their time and memory scales quadratic in the sequence length.
There are numerous approaches to overcome this bottleneck, including approximation with hashing (i.e., sparse) such as Reformer [53] and Smyrf [20] and with low-rank approximation such as Performer [13, 56].
One can even combine sparse and low-rank approximation for better accuracy (e.g., Longformer [3], BigBird [96], Scatterbrain [9], Long-short transformer [98], Combiner [76]).
Other approaches include compressing along the sequence dimension to attend to multiple tokens at once [54, 59, 82, 93].
One can also attend over the states from previous sequences to help lengthen the context (e.g., TransformerXL [15] and Compressive Transformer [72]).
We recommend the survey [84] for more details.
There are several lines of work on developing other modules instead of attention to model longer context.
HiPPO [37] and its extensions, most notably S4 [33, 38, 39] projects the history on a polynomial basis, allowing accurate reconstruction of the history through state-space models.
They combine the strengths of CNNs (efﬁcient training), RNNs (efﬁcient inference), and continuous models (robust to change in sampling rates).
LambdaNetworks [2], AFT [97] and FLASH [44] are other attempts at replacing attention in the context of image classiﬁcation and language modeling.
B Algorithm Details
We ﬁrst derive the forward and backward passes of attention and show that they can be computed in a memory-efﬁcient manner (requiring extra memory linear instead of quadratic in the sequence length).
Though they reduce the amount of extra memory required, naively they still incur quadratic HBM accesses, resulting in slower execution speed.
We describe the F LASH A TTENTION algorithm to implement both the forward and the backward passes on GPUs that reduces HBM accesses, leading to both faster runtime and smaller memory footprint.
B.1 Memory-efﬁcient forward pass
The main challenge in making attention memory-efﬁcient is the softmax that couples the columns of K (and columns of V ).
Our approach is to compute the softmax normalization constant separately to decouple the columns.
This technique [62] has been used in the literature [53, 69] to show that attention computation does not need quadratic extra memory (though the number of HBM accesses is still quadratic, resulting in slow run-time).
For simplicity, we omit here the max-shifting step during softmax.
The full algorithm in Appendix B. contains all the steps.
Recall that given input sequences Q , K , V ∈ R 𝑁 × 𝑑 , we want to compute the attention output O ∈ R 𝑁 × 𝑑 :
[Equation]
We have that 𝑆 𝑖𝑗 =
𝑞
𝑇𝑖 𝑘 𝑗 where 𝑞 𝑖 and 𝑘 𝑗 are the 𝑖 -th and 𝑗 -th columns of Q and K respectively.
Deﬁne the normalization constants of softmax: 𝐿
[Equation]
Let 𝑣 𝑗 be the 𝑗 -th column of V , then the 𝑖 -th columns of the output is
[Equation]
We see that once 𝐿 𝑖 is computed, we can compute 𝑜 𝑖 without extra memory by repeatedly summing 𝑒 𝑞𝑇𝑖 𝑘𝑗 𝐿 𝑖 𝑣 𝑗 .
Therefore the forward pass can be computed with 𝑂 ( 𝑛 ) extra memory:
Compute 𝐿 𝑖 for all 𝑖 according to Eq.
(1), which takes 𝑂 ( 𝑛 ) extra memory.
Compute 𝑜 𝑖 for all 𝑖 according to Eq.
(2), which takes 𝑂 ( 𝑑 ) extra memory.
B.2 Memory-efﬁcient backward pass
We derive the backward pass of attention and show that it can also be computed with linear memory.
Rabe and Staats [69] suggests that the backward pass can be done without quadratic extra memory by applying gradient checkpointing to the memory-efﬁcient forward pass.
We instead derive the backward pass explicitly and show how it can be computed in a memory-efﬁcient manner.
Suppose that there is a scalar loss function 𝜙 , and let the output gradient be dO ∈ R 𝑛 × 𝑑 (where dO denotes 𝜕𝜙𝜕 O ).
We want to compute the input gradients dQ , dK , dV ∈ R 𝑛 × 𝑑 (where dQ , dK , dV denote 𝜕𝜙𝜕 Q , 𝜕𝜙𝜕 K , 𝜕𝜙𝜕 V respectively).
The gradient dV is easy to see.
Applying reverse-mode autodiff by hand (aka the chain rule), we obtain (in matrix notation) dV
=
P 𝑇 dO .
Thus:
[Equation]
Since we already computed 𝐿 𝑖 , 𝑑𝑣 𝑗 can be computed without extra memory by repeated summing.
The gradients dQ and dK are a little more complicated.
We go through the gradients dP and dS ﬁrst.
From Eq.
(2), we have that dP = dOV 𝑇 , and so:
[Equation]
Recall that 𝑃 𝑖 : =
softmax ( 𝑆 𝑖 : ) .
Using the fact that the Jacobian of 𝑦 = softmax ( 𝑥 ) is diag ( 𝑦 )− 𝑦𝑦 𝑇 , we have that
[Equation]
where ◦ denotes pointwise multiplication.
Deﬁne
[Equation]
then
[Equation]
Hence 𝑑𝑆
[Equation]
Now we can get the gradients dQ and dK .
Recall that 𝑆 𝑖𝑗 = 𝑞
𝑇𝑖 𝑘 𝑗 , so
[Equation]
Similarly,
[Equation]
Therefore the backward pass can also be computed with 𝑂 ( 𝑛 ) extra memory:
Compute 𝑑𝑣 𝑗 for all 𝑗 according to Eq. which takes
𝑂 ( 𝑑 ) extra memory
Compute 𝐷 𝑖 for all 𝑖 according to Eq.
(4), which takes 𝑂 ( 𝑛 ) extra memory.
Compute 𝑑𝑞 𝑖 for all 𝑖 according to Eq.
(5), which takes 𝑂 ( 𝑑 ) extra memory.
Compute 𝑑𝑘 𝑗 for all 𝑗 according to Eq.
(6), which takes 𝑂 ( 𝑑 ) extra memory.
B. F LASH A TTENTION : Forward Pass
We describe the full details of F LASH A TTENTION forward pass.
Given input sequences
Q , K , V ∈ R 𝑁 × 𝑑 , we want to compute the attention output O ∈ R 𝑁 × 𝑑 :
[Equation]
[Equation] where 𝜏 ∈ R is some softmax scaling (typically 1 √ 𝑑 ), MASK is some masking function that sets some entries of the input to −∞ and keep other entries the same (e.g., key padding mask when sequences in the batch don’t have the same lengths and are padded), and dropout ( 𝑥,𝑝 ) applies dropout to 𝑥 elementwise (i.e., output 𝑥 1 − 𝑝 with probability 1 − 𝑝 and output 0 with probability 𝑝 for each element 𝑥 ).
The full algorithm is in Algorithm 2.
We save the output O , the softmax statistics ℓ and 𝑚 , and the pseudo-random number generator state R for the backward pass.
) ∈ R 𝐵 𝑟 , ˜ P 𝑖𝑗 = exp ( S masked 𝑖𝑗 − ˜ 𝑚
𝑖𝑗 ) ∈
R 𝐵 𝑟 × 𝐵 𝑐 (pointwise), ˜ ℓ 𝑖𝑗 = rowsum ( ˜ P 𝑖𝑗 )
∈ R 𝐵 𝑟 .
On chip, compute 𝑚 new 𝑖 = max ( 𝑚 𝑖 , ˜ 𝑚
𝑖𝑗 ) ∈ R 𝐵 𝑟 , ℓ new 𝑖 = 𝑒 𝑚 𝑖 − 𝑚 new 𝑖 ℓ 𝑖 + 𝑒 ˜ 𝑚
𝑖𝑗 − 𝑚 new 𝑖 ˜ ℓ 𝑖𝑗 ∈ R 𝐵 𝑟 .
On chip, compute ˜ P dropped 𝑖𝑗 = dropout ( ˜ P 𝑖𝑗 ,𝑝 drop ) .
Write O 𝑖 ← diag ( ℓ new 𝑖 ) − 1 ( diag ( ℓ 𝑖 ) 𝑒 𝑚 𝑖 − 𝑚 new 𝑖 O 𝑖 + 𝑒 ˜ 𝑚
𝑖𝑗 − 𝑚 new 𝑖 ˜ P dropped 𝑖𝑗
V 𝑗 ) to HBM.
Write ℓ 𝑖 ← ℓ new 𝑖 , 𝑚 𝑖 ← 𝑚 new 𝑖 to HBM.
17: end for 18: end for 19: Return O ,ℓ,𝑚, R .
B. F LASH A TTENTION : Backward Pass
We describe the full details of F LASH A TTENTION backward pass.
Given input sequences
Q , K , V ∈ R 𝑁 × 𝑑 , the output O ∈ R 𝑁 × 𝑑 , and the output gradient dO , we want to compute the input gradients dQ , dK , dV ∈ R 𝑁 × 𝑑 .
We ﬁrst describe the standard attention backward pass in Algorithm 3 for completeness. Algorithm
Require: Matrices Q , K , V , dO ∈ R
𝑁 × 𝑑 , P ∈ R 𝑁 × 𝑁 in HBM.
1: Load P , dO by blocks from HBM, compute dV
= P (cid:62) dO ∈ R
𝑁 × 𝑑 , write dV to HBM. 2: Load dO , V by blocks from HBM, compute dP = dOV (cid:62) ∈ R 𝑁 × 𝑁 , write dP to HBM. 3: Read P , dP from HBM, compute dS ∈ R 𝑁 × 𝑁 where 𝑑𝑆 𝑖𝑗 = 𝑃 𝑖𝑗 ( 𝑑𝑃 𝑖𝑗 − (cid:205) 𝑙 𝑃 𝑖𝑙 𝑑𝑃 𝑖𝑙 ) , write dS to HBM.
4: Load dS and K by blocks from HBM, compute dQ = dSK , write dQ to HBM. 5: Load dS and Q by blocks from HBM, compute dK = dS (cid:62) Q , write dK to HBM.
6: Return dQ , dK , dV . We now make two observations about F LASH A TTENTION backward pass:
We do not need to store the dropout mask of size 𝑂 ( 𝑁 2 ) from the forward pass.
Instead, we can save the pseudo-random number generator states from the forward pass and re-generate the dropout mask in the backward pass.
This allows us to only use 𝑂 ( 𝑁 ) extra memory.
When computing the softmax gradient, we use Eq.
(4) to compute 𝐷 𝑖 = 𝑃 (cid:62) 𝑖 : 𝑑𝑃 𝑖 : without reducing over 𝑃 𝑖 : and 𝑑𝑃 𝑖 : of size 𝑁 (they might not ﬁt into SRAM).
Instead we can rewrite 𝐷 𝑖 = 𝑑𝑜 (cid:62) 𝑖 𝑜 𝑖 and compute the dot product between vectors of size 𝑑 .
The full F LASH A TTENTION backward pass algorithm is in Algorithm 4.
Conceptually it is just a block version of the derivation in Appendix B.2.
We see that similar to the forward pass, the backward pass performs 𝑂 ( 𝑁 2 ) FLOPs and only requires 𝑂 ( 𝑁 ) extra memory beyond inputs, output, output gradient, and input gradients.
We analyze the IO-complexity of the backward pass, similar to the forward pass (Theorem 2).
Theorem 5.
Let 𝑁 be the sequence length, 𝑑 be the head dimension, and 𝑀 be size of SRAM with 𝑑 ≤ 𝑀 ≤ 𝑁𝑑 .
Standard attention (Algorithm 0) backward pass requires Θ ( 𝑁𝑑 + 𝑁 2 ) HBM accesses, while F LASH A TTENTION backward pass (Algorithm 4) requires Θ ( 𝑁 2 𝑑 2 𝑀 − 1 ) HBM accesses.
The proof is in Appendix C.
B.5 Comparison with Rabe and Staats [69]
We describe here some similarities and differences between our F LASH A TTENTION algorithm and the algorithm of Rabe and Staats [69].
Conceptually, both F LASH A TTENTION and Rabe and Staats [69] operate on blocks of the attention matrix using the well-established technique of tiling (or softmax scaling) [53, 62].
To reduce the memory footprint, both methods avoid storing the large attention matrix in the forward pass and recompute it in the backward pass.
The ﬁrst major difference is that Rabe and Staats [69] focuses on the reducing the total memory footprint (maximum amount of GPU memory required) while F LASH A TTENTION focuses on reducing memory accesses (the number of memory reads/writes).
As mentioned in Section 2, the amount of memory access is the primary determining factor of runtime.
Reducing memory accesses also necessarily reduces the total amount of memory required (e.g., if an operation incurs 𝐴 memory accesses, then its total memory requirement is at most 𝐴 ).
As a result, F LASH A TTENTION is faster than standard attention (2-4 × ) while Rabe and Staats [69] is around the same speed or slightly slower than standard attention.
In terms of total memory required, both methods offer substantial memory saving.
The second difference between the two methods is the way information is summarized from each block to pass to the next block.
Rabe and Staats [69] summarizes each block with its temporary output along with the softmax normalization statistics.
At the end of the forward pass, the temporary outputs of all the blocks are combined using the statistics to produce the ﬁnal output.
= ( 0 ) 𝐵 𝑐 × 𝑑 on SRAM.
9: for 1 ≤ 𝑖 ≤ 𝑇 𝑟 do 10:
Load Q 𝑖 , O 𝑖 , dO 𝑖 , dQ 𝑖 ,ℓ 𝑖 ,𝑚 𝑖 from HBM to on-chip SRAM.
On chip, compute S 𝑖𝑗 =
𝜏 Q 𝑖 K 𝑇𝑗 ∈ R 𝐵 𝑟 × 𝐵 𝑐 .
On chip, compute S masked 𝑖𝑗 = MASK
( S 𝑖𝑗 ) .
On chip, compute P 𝑖𝑗 = diag ( 𝑙 𝑖 ) − 1 exp ( S masked 𝑖𝑗 −
𝑚 𝑖 ) ∈ R 𝐵 𝑟 × 𝐵 𝑐 .
On chip, compute dropout mask Z 𝑖𝑗 ∈ R 𝐵 𝑟 × 𝐵 𝑐 where each entry has value 1 1 − 𝑝 drop with probability 1 − 𝑝 drop and value 0 with probability 𝑝 drop .
On chip, compute P dropped 𝑖𝑗
=
P
𝑖𝑗 ◦ Z 𝑖𝑗 (pointwise multiply).
On chip, compute ˜ dV
𝑗 ← ˜ dV
𝑗 +( P dropped 𝑖𝑗 ) (cid:62) dO 𝑖 ∈ R 𝐵 𝑐 × 𝑑 .
On chip, compute dP dropped 𝑖𝑗
= dO 𝑖 V (cid:62) 𝑗 ∈ R 𝐵 𝑟 × 𝐵 𝑐 .
On chip, compute dP 𝑖𝑗 = dP dropped 𝑖𝑗 ◦ Z 𝑖𝑗 (pointwise multiply).
On chip, compute 𝐷 𝑖 = rowsum ( dO 𝑖 ◦ O 𝑖
) ∈ R 𝐵 𝑟 .
On chip, compute dS 𝑖𝑗 = P 𝑖𝑗 ◦( dP 𝑖𝑗 − 𝐷 𝑖 ) ∈
R 𝐵 𝑟 × 𝐵 𝑐 .
Write dQ 𝑖 ← dQ 𝑖 + 𝜏 dS 𝑖𝑗 K 𝑗 ∈ R 𝐵 𝑟 × 𝑑 to HBM.
On chip, compute ˜ dK
𝑗 ← ˜ dK
𝑗 + 𝜏 dS (cid:62) 𝑖𝑗 Q 𝑖 ∈ R 𝐵 𝑐 × 𝑑 .
23: end for 24:
Write dK 𝑗 ← ˜ dK
𝑗 , dV 𝑗 ← ˜ dV 𝑗 to HBM.
25: end for 26: Return dQ , dK , dV .
instead incrementally updates the output (Algorithm 1 line 1) after processing each block, so only one copy of the output is needed (instead of 𝐾 copies for 𝐾 blocks).
This means that F LASH A TTENTION has smaller total memory requirement compared to Rabe and Staats [69].
The ﬁnal major difference is the way the backward pass is computed.
Rabe and Staats [69] uses gradient checkpointing to recompute the attention matrix and the temporary output of each block.
F LASH A TTENTION instead simpliﬁes the backward pass analytically (Appendices B.2 and B.4).
It only recomputes the attention matrix and does not recompute the temporary output of each block.
This reduces the memory requirement for the backward pass and yields speedup.
C Proofs
Proof of Theorem 1.
We ﬁrst count the number of FLOPs and extra memory required.
The dominating FLOPs are from matrix multiplication.
In the inner loop, (Algorithm 1 line 1), we compute Q 𝑖 K (cid:62) 𝑗 ∈ R 𝐵 𝑟 × 𝐵 𝑐 for Q 𝑖 ∈ R 𝐵 𝑟 × 𝑑 and K 𝑗 ∈ R 𝐵 𝑐 × 𝑑 , which takes 𝑂 ( 𝐵 𝑟 𝐵 𝑐 𝑑 ) FLOPs.
We also compute (Algorithm 1 line 1) ˜ P
𝑖𝑗 V 𝑗 ∈ R 𝐵 𝑟 × 𝑑 for ˜ P 𝑖𝑗 ∈ R 𝐵 𝑟 × 𝐵 𝑐 and V 𝑗 ∈ R 𝐵 𝑐 × 𝑑 , which takes 𝑂 ( 𝐵 𝑟 𝐵 𝑐 𝑑 )
FLOPs.
We execute the inner loops 𝑇 𝑐 𝑇 𝑟
[Equation]
times.
Therefore the total number of FLOPs is
[Equation]
In terms of extra memory required, we see that we need 𝑂 ( 𝑁 ) memory to store the statistics ( ℓ,𝑚 ) .
We now prove the algorithm’s correctness by induction on 𝑗 for 0 ≤ 𝑗 ≤ 𝑇 𝑐 .
Let K :
𝑗 ∈ R 𝑗𝐵 𝑐 × 𝑑 be the ﬁrst 𝑗𝐵 𝑐 rows of K , and similarly V : 𝑗 ∈ R 𝑗𝐵 𝑐
×
𝑑 the the ﬁrst 𝑗𝐵 𝑐 rows of V .
Let S : , : 𝑗 = QK (cid:62) : 𝑗 ∈ R 𝑁 × 𝑗𝐵 𝑐 , and P : , : 𝑗 = softmax ( S : , : 𝑗 ) ∈ R 𝑁 × 𝑗𝐵 𝑐 (softmax applied row-wise).
Let 𝑚 𝑗 ,ℓ ( 𝑗 ) , O ( 𝑗 ) be the values of 𝑚,ℓ, O in HBM after the 𝑗 -th iteration of the outer loop (Algorithm 1 line 1).
(Note that these values of 𝑚,ℓ, O are updated after each iteration of the outer loop.)
We want to show that after the 𝑗 -th iteration of the outer loop, we have computed in HBM: 𝑚 ( [Equation] = rowmax ( S : , :
[Equation]
) ∈ [Equation] 𝑁 ,
[Equation]
Based on our initialization (Algorithm 1 line 1), this claim is true for 𝑗 = 0 (i.e., before the any iteration of the outer loop is executed).
Suppose that the claim holds for some 𝑗 = 0 ,...,𝑇 𝑐 − 1 .
We want to show that the claim also holds for 𝑗 + 1 .
Indeed, when we update the statistics in the inner loop (Algorithmline 1) on the ( 𝑗 + 1 ) -th iteration of the outer loop, we update 𝑚 ( 𝑗 + 1 ) =
max ( 𝑚 ( 𝑗 ) , ˜ 𝑚 ) where ˜ 𝑚 ∈ R 𝑁 is the row-max of S : ,𝑗 : 𝑗 + 1 , the slice of S from column 𝑗𝐵 𝑐 to column ( 𝑗 + 1 )
𝐵 𝑐 − 1 .
This implies that 𝑚 ( 𝑗 + 1 ) =
rowmax ( S : , : 𝑗 + 1 ) ∈ R 𝑁 .
Similarly, we update
[Equation]
where ˜ ℓ = rowsum ( exp ( S : ,𝑗 : 𝑗 + 1 − ˜ 𝑚 ))
∈ R 𝑁 .
By the same algebraic manipulation in Section 3.1, we obtain:
[Equation]
Let V 𝑗 : 𝑗 + 1 be the slice of V from column 𝑗𝐵 𝑐 to column ( 𝑗 + 1 ) 𝐵 𝑐 − 1 , we also update: [Equation] 𝑗 + 1 ) = diag ( ℓ ( 𝑗 + 1 ) )
[Equation]
[Equation]
exp [Equation] S : , : 𝑗 [Equation]
𝑚
[Equation]
𝑗
[Equation]
V : 𝑗
[Equation]
exp
[Equation]
S 𝑗 : 𝑗 +
[Equation]
𝑚 ( 𝑗 + 1 ) )
V 𝑗 :
[Equation]
) =
[Equation]
[Equation]
= softmax ( S : 𝑗 + 1 ) V : 𝑗 + 1 .
We then see that the claim is also true for 𝑗 + 1 .
By induction, the claim is true for all 𝑗 = 0 ,...,𝑇 𝑐 .
When 𝑗 = 𝑇 𝑐 , we conclude that the ﬁnal value of O in HBM is softmax ( S ) V = softmax ( QK (cid:62) ) V .
(cid:3)
Proof of Theorem 2.
We ﬁrst analyze the IO complexity of standard attention implementation.
The inputs Q , K , V ∈ R 𝑁 × 𝑑 reside in HBM, and the at the end of the algorithm the output O ∈ R 𝑁 × 𝑑 is written to HBM.
In the ﬁrst step of computing the matrix multiply S =
QK (cid:62) , the inputs Q , K are read from HBM and the output S ∈ R
𝑁 × 𝑁 is written to HBM (Algorithm 0 line 0).
This incurs Θ ( 𝑁𝑑 + 𝑁 2 ) HBM accesses.
In the second step of computing P = softmax ( S ) , the input S is read from HBM and the output P is written to HBM (Algorithm 0 line 0).
This incurs Θ ( 𝑁 2 )
HBM accesses.
In the last step of computing O = PV , the inputs P , V are read from global memory and the output
O is written to HBM (Algorithm 0 line 0).
This incurs Θ ( 𝑁𝑑 + 𝑁 2 ) HBM accesses.
Overall, standard attention implementation requires Θ ( 𝑁𝑑 + 𝑁 2 ) global memory accesses.
We now analyze the IO complexity of streaming attention.
Following Algorithm 1, we see that each element of K and V is loaded from HBM once (Algorithmline 1).
We make 𝑇 𝑐 passes over Q and O , each pass loading all of Q and all of O to HBM (Algorithm 1 line 1).
Therefore the number of HBM accesses is Θ ( 𝑁𝑑
+ 𝑁𝑑𝑇 𝑐 ) =
Θ ( 𝑁𝑑𝑇 𝑐 ) .
We derive the conditions on the block sizes 𝐵 𝑐 and 𝐵 𝑟 .
We need the blocks K 𝑗 and V 𝑗 of size 𝐵 𝑐 × 𝑑 to ﬁt into on-chip memory, which translates to:
[Equation]
Similarly, we need the blocks Q 𝑖 , O 𝑖 of size 𝐵 𝑟 × 𝑑 to ﬁt into on-chip memory, which translates to:
[Equation]
Finally, we need the block S 𝑖𝑗 of size 𝐵 𝑟 × 𝐵 𝑐 to ﬁt into on-chip memory, which translates to: 𝐵 [Equation] 𝐵 𝑐 = 𝑂
( 𝑀 ) .
We therefore set:
[Equation]
We then have:
[Equation]
As a result, the number of HBM accesses is:
[Equation]
Proof of Proposition 3.
For contradiction, suppose that there exists an algorithm that computes exact attention where the number for HBM access for all 𝑀 ∈ [ 𝑑,𝑁𝑑 ] is
[Equation]
In the regime of 𝑀 =
Θ ( 𝑁𝑑 ) , this results in the number of HBM accesses:
[Equation]
However, the input to attention (matrices Q , K , V ) and the output O have size 𝑁𝑑 and they start out being in HBM, so if the algorithm computes exact attention it must incur at least Ω ( 𝑁𝑑 ) HBM accesses.
This is a contradiction.
[Equation]
Proof of Theorem 5.
The IO complexity of the attention backward is very similar to the IO complexity of the attention forward (Theorem 2).
Here we provide a sketch of the proof.
We ﬁrst analyze the IO complexity of standard attention backward pass.
The inputs Q , K , V , dO ∈ R 𝑁 × 𝑑 reside in HBM, and the at the end of the algorithm the outputs dQ , dK , dV ∈ R 𝑁 × 𝑑 are written to HBM.
At each step of the standard attention backward pass, one needs to load inputs of size 𝑁𝑑 or 𝑁from HBM, and needs to write the outputs of size 𝑁 2 or 𝑁𝑑 to HBM.
This incurs Θ ( 𝑁𝑑 + 𝑁 2 ) HBM accesses.
We now analyze the IO complexity of F LASH A TTENTION backward pass.
Similar to Theorem 2, we see that each element of K and V is loaded from HBM once.
Each element of dK and dV is only written to HBM once.
We make 𝑇 𝑐 passes over Q , O , dO , each pass loading all of Q , O , dO to HBM.
We also make 𝑇 𝑐 passes over dQ , each pass reading/writing all of dQ from/to
HBM.
Therefore the number of HBM accesses is Θ ( 𝑁𝑑
+ 𝑁𝑑𝑇 𝑐 ) =
Θ ( 𝑁𝑑𝑇 𝑐 ) .
As in the proof of Theorem 2, the constraints on the block sizes are that:
[Equation]
We then have:
[Equation]
As a result, the number of HBM accesses is:
[Equation]
D Extension Details D.1 Block-sparse F LASH A TTENTION
We describe the full block-sparse F LASH A TTENTION algorithm in Algorithm 5.
The algorithm is identical to Algorithm 2, except that we skip zero blocks.
Require: Matrices Q
K V ∈ R HBM chip SRAM size , softmax scaling constant 𝜏 ∈
R , masking function MASK , dropout probability 𝑝 drop , block sizes 𝐵 𝑐 = (cid:6)
𝑀 4 𝑑 (cid:7) ,𝐵 𝑟 =
min (cid:0)(cid:6) 𝑀 4 𝑑 (cid:7) ,𝑑 (cid:1) , block sparsity mask 𝑀 ∈ { 0 , 1 } 𝑁 / 𝐵 𝑟 × 𝑁 / 𝐵 𝑐 ..
R 𝐵 𝑟 × 𝐵 𝑐 .
On chip, compute S masked 𝑖𝑗 = MASK
( S 𝑖𝑗 ) .
On chip, compute ˜ 𝑚
𝑖𝑗 = rowmax ( S masked 𝑖𝑗 )
∈ R 𝐵 𝑟 , ˜ P
𝑖𝑗 = exp
( S masked 𝑖𝑗 − ˜ 𝑚
𝑖𝑗 ) ∈
R 𝐵 𝑟 × 𝐵 𝑐 (pointwise), ˜ ℓ 𝑖𝑗 = rowsum ( ˜ P 𝑖𝑗 )
∈ R 𝐵 𝑟 .
On chip, compute 𝑚 new 𝑖 = max ( 𝑚 𝑖 , ˜ 𝑚
𝑖𝑗 ) ∈ R 𝐵 𝑟 , ℓ new 𝑖 = 𝑒 𝑚 𝑖 − 𝑚 new 𝑖 ℓ 𝑖 + 𝑒 ˜ 𝑚
𝑖𝑗 − 𝑚 new 𝑖 ˜ ℓ 𝑖𝑗 ∈ R 𝐵 𝑟 .
On chip, compute ˜ P dropped 𝑖𝑗 = dropout ( ˜ P 𝑖𝑗 ,𝑝 drop ) .
Write O 𝑖 ← diag ( ℓ new 𝑖 ) − 1 ( diag ( ℓ 𝑖 ) 𝑒 𝑚 𝑖 − 𝑚 new 𝑖 O 𝑖 + 𝑒 ˜ 𝑚
𝑖𝑗 − 𝑚 new 𝑖 ˜ P dropped 𝑖𝑗
V 𝑗 ) to HBM.
Write ℓ 𝑖 ← ℓ new 𝑖 , 𝑚 𝑖 ← 𝑚 new 𝑖 to HBM.
17: end if 18: end for 19: end for 20: Return O ,ℓ,𝑚, R .
We prove the IO-complexity of block-sparse F LASH A TTENTION .
Proof of Proposition 4.
The proof is very similar to the proof of Theorem 2.
For the block-sparse case, notice that we only need to load blocks corresponding to nonzero blocks.
As a result, the number of HBM accesses are scaled by 𝑠 , the fraction of nonzero blocks in the block-sparsity mask.
However, for small values of 𝑠 , we would still need to write the result O ∈ R 𝑁 × 𝑑 .
Therefore the number of HBM accesses is
[Equation]
D.2 Potential Extensions
We discuss here a few potential extensions of the IO-aware approach to speed up deep learning training.
Multi-GPU Attention.
Large language models are trained on hundreds or thousands of GPUs, and one typically splits the attention computation between 4-8 GPUs on the same node [80].
This introduces another level of memory hierarchy: beside GPU SRAM and GPU HBM, we also have the
HBM of other GPUs.
For very long sequences, the different GPUs on the same node can cooperate to compute attention by taking into account the asymmetry of different levels of memory hierarchy.
Sparse MLP layers.
Typical dense MLP layers are compute-bound and not memory-bound.
To improve their efﬁciency, MLP layers with sparse weight matrices can be used [18].
However, many sparse MLP layers are instead memory-bound, and their speedup is often not proportional to the sparsity.
We believe that an IO-aware implementation can alleviate this issue and realize the beneﬁts of sparsity.
We are excited about future work in this direction, to reduce the computational requirement of large models and improve their wall-block runtime.
Kernel machine learning.
Our approach in F LASH A TTENTION relies on the fact that the 𝑁 × 𝑁 attention matrix is a function of a low-rank matrix QK (cid:62) (of rank 𝑑 (cid:28) 𝑁 ).
As a result, we can repeatedly load the inputs Q , K and recompute the block of the attention matrix that we need, signiﬁcantly reducing HBM access.
As similar scenario happens in kernel machine learning: each element 𝐾 𝑖𝑗 of the 𝑁 × 𝑁 kernel matrix K is a function of two vectors of size 𝑑 (cid:28) 𝑁 , as it measures the similarity between two datapoints 𝑥 𝑖 and 𝑥 𝑗 .
The KeOps library [8, 28] is a successful example of how reducing memory reads/writes can speed up kernel operations.
We hope that this will motivate kernel methods that focus more on reducing IOs instead of just FLOPs.
E Full Experimental Results E.1 BERT
We train BERT-large following the training procedure and hyperparameters of the reference MLPerf 1.1 implementation.
In particular, we use the LAMB optimizer with learning rate 3.75e-3, with batch size 448, trained for at most 7100 steps.
The training is stopped once the validation accuracy (for masked language modeling) reaches the target 72.0%, and the wall-clock run-time is measured.
We train with FP16 precision using Apex AMP (with O2 optimization level).
We compare our results with the reported training speed from Nvidia that was submitted to MLPerf 1.1 (Table 1).
We use the same train / validation data split provided by MLPerf 1.1 reference implementation.
In particular, we evaluate on the same 10000 validation examples as the baseline from Nvidia.
We train the model on 8 × A100-80GB GPUs.
Each training run takes between 16 and 19 minutes, and we average the results of 10 runs.
We see a memory saving of 1.8 𝑡𝑖𝑚𝑒𝑠 (from 58GB to 32GB) for the same batch size.
In Table 7, we additionally compare against the commonly used Huggingface implementation.
F LASH A TTENTION is 3.2 × faster than this implementation.
[Table]
E.2 GPT-
We use the standard implementations of GPT-2 [70] from Huggingface transformers library and from Nvidia’s Megatron-LM repo.
We follow the training recipe of the Megatron-LM repo.
We use an effective batch size of 512, and use gradient accumulation to ﬁt into available GPU memory.
We use the AdamW optimizer, with learning rate 6e-4 for GPT-2 small and 1.5e-4 for
GPT-2 medium, and weight decay of 0.1.
All models are trained with the same hyperparameters for 400K steps.
We run all implementations with mixed-precision training (PyTorch AMP).
We use the Openwebtext dataset, with the GPT-2 BPE tokenizer.
We randomly select 0.5% of the dataset as the validation set, with the rest being used as training set.
This random selection of validation set is done once, and all models are evaluated on the same validation set.
We train the model on 8 × A100-40GB GPUs, and we measure the wall-clock training time.
Training
GPT-2 small takes between 2.7-9.5 days, and training GPT-2 medium takes between 6.9-21.0 days (Table 2).
For GPT-2 small, we see a memory saving of 3.5 𝑡𝑖𝑚𝑒𝑠 (from 39GB to 11GB) for the same batch size of 16 (which means we could run F LASH A TTENTION with 4 × larger device batch size while keeping the global batch size of 512 the same).
In Fig.
4, we plot of the validation perplexity throughout training of GPT-2 small/medium, using either HuggingFace implementation or our F LASH A TTENTION implementation.
We see that
F LASH A TTENTION behaves the same as the baseline implementation and the validation perplexity curves of the two implementations almost lie on top of each other.
[Figure]
i t [Figure] We additionally compare the speedup of FlashAttention as we scale the number of GPUs fromto 8, for GPT-2 Small.
All the training hyperparameters are kept the same.
As we change the number of GPUs, we change the number of gradient accumulations to keep the global batch size the same (512).
We use PyTorch DistributedDataParallel when there are more than 1 GPUs.
Table 8 shows that the speedup is consistent, around 3 .
7 × .
[Figure]
[Table] [Figure] [Table]
[Figure]
[Table]
Long Document Classiﬁcation.
For MIMIC-III and ECtHR, we follow the hyperparameters of Dai et al.
E.3 LRA details
We follow the hyperparameters from the Long-range arena paper [83], the Long-range arena repo ( https://github.com/google-research/long-range-arena ), and the Nyströmformer reproduction [94].
To be generous to the baseline methods, if we are unable to reproduce the performance of any baseline for any of the ﬁve tasks, we report the better performance from Tay et al.
[83] or Xiong et al.
[94] for that baseline on that task.
After hyperparameter tuning, almost all of the attention methods achieve similar accuracy on all of the ﬁve LRA tasks.
We run all methods with mixed-precision training, except for Performer (not stable with mixed precision) and Local Attention (implementation does not support FP16).
To calculate the overall wallclock-time speedup, we take the geometric mean of the wallclock-time speedup of each of the ﬁve tasks.
Path-X For Path-X and Path-256, we follow the hyperparameters from the PathFinderexperiments from the long-range arena paper[83].
For both, we ﬁrst pretrain a model on Path-64.
We take the checkpoint after 200 epochs, upsample its positional embedding (we duplicate the positional embeddings gridwise in space), and ﬁne-tune it on the downstream task for 200 epochs with one epoch of linear warmup, and cosine decay of the learning rate.
For Path-X, we take the best performing checkpoint (according to val accuracy), and additionally ﬁne-tune it for 200 epochs with the same warmup and learning rate (this adds roughly 4 points of accuracy to F LASH A TTENTION for Path-X, but the model starts overﬁtting afterwards).
E.4 Faster Vision Transformer with F LASH A TTENTION on ImageNet
On the popular vision benchmark, ImageNet [22], we show that F LASH A TTENTION can also speedup
Vision Transformers (ViT) [26] by 1.5 𝑡𝑖𝑚𝑒𝑠 , where the sequence length is 196 (patch size 16 ×for 224 × 224 images).
For longer sequences, F LASH A TTENTION yields up to 3.5 × speedup.
We use the ViT-base implementation from the widely-used library timm , and replace the standard attention implementation with F LASH A TTENTION .
We follow the same training recipe as that of
DeiT [85], which improves on the original training recipe of ViT.
We measure accuracy and training time of both models (for 300 epochs) on 8 × A100s.
Table 9 shows that F LASH A TTENTION achieves up 1.5 × speed-up compared to standard attention.
[Table]
We also compare ViT-Large with smaller patch sizes (i.e., longer sequence lengths).
Table 10 shows that F LASH A TTENTION yields 3.5 × speedup and saves up to 3.6 𝑥 memory, compared to standard attention. [Table]
[Table]
E.5 Comparison with Automatic Fusion
We compare FlashAttention with automatic fusion methods: NVFuser from Pytorch 1.12 (newest version at the time of writing), AOT compiler from Functorch [Table] and TVM [11].
For context, we also include the runtime of standard implementation from Pytorch and the more optimized implementation from Megatron-LM [80].
We benchmark for batch size 16, [Table] and 64, sequence length 1024, 16 heads, head dimension 64, with key-padding mask (and no dropout).
The runtime is measured on an A100-SXM4-40GB GPU
[Table] Table 11 shows that F LASH A TTENTION is about 2-3 × faster than these methods [Table]
[Table]
E.6 Comparison with Apex FMHA
We compare our method/implementation with Apex FMHA ( https://github.com/NVIDIA/ apex/tree/master/apex/contrib/csrc/fmha ).
[Table]
When we started this project, Apex FMHA was the fastest implementation of attention (that we knew of), tailored for short sequences of length at most 512.
In fact, almost all MLPerf submissions for
BERT training benchmark running on Nvidia GPUs use FMHA for their model code, as of MLPerf 1.1 [60].
Since FMHA targets BERT models, it only supports head dimension 64, and only runs on A100 GPUs.
FMHA fuses the attention computation dropout ( softmax ( MASK ( QK (cid:62) )))
V into one
CUDA kernel.
In the forward pass, it stores the attention matrix softmax ( MASK ( QK 𝑇 )) to HBM to be used in gradient computation.
As a result, it does not offer substantial memory saving (though for shorter sequences memory footprint is often not a primary concern).
We use FMHA code as a starting point, and apply two well-established techniques (tiling and recomputation) to deal with long sequences and to save memory as mentioned in Section 3.
As a result, we can support much longer sequences (e.g., up to length 64K).
We also support more head dimensions (16, 32, 64, 128) and broader GPU types (all Turing and Ampere GPUs at the time of writing).
In Table 12, we compare the performance of F LASH A TTENTION and Apex FMHA for short sequences (as FMHA only supports sequence length at most 512).
Generally F LASH A TTENTION is slightly faster than FMHA in the forward pass and slightly slower than FMHA in the backward pass.
This is because we do not store the attention matrix in the forward pass and recompute it in the backward pass.
Compared to FMHA, the overall runtime of F LASH A TTENTION is about 4% slower for sequence length 128, 8% faster for sequence length 256, and 5% faster for sequence length 512.
E.7 Rooﬂine analysis
In Fig.
5, we include a rooﬂine analysis of the F LASH A TTENTION forward pass, taken from Nvidia
Nsight Compute (batch size 16, seqlen 512, 16 heads, head dimension 64) on an A100-SXM4-40GB
GPU.
While
TTENTION substan- speeds up attention, there is still some potential headroom to gain further speedup
E.8 Speedup On Different Hardware and Conﬁgurations
Speedup varies between different types of GPU types and generations depending on HBM bandwidth and SRAM size.
In this section, we proﬁle F LASH A TTENTION speedup on different GPUs and conﬁgurations.
A100 Figure 6 shows speedup on an A100 GPU with batch size 8, head dimension 64, andattention heads, across different sequence lengths.
We generally see 2-4 × speedup, and we see more speedup when using dropout and masking due to kernel fusion.
A100, Head Dimension
Speedup also changes when we increase the head dimension.
Each block requires more memory, so we need to use smaller block sizes to ﬁt into SRAM.
Figureshows speedup with head dimension 128 on an A100 (batch size 16, 12 heads).
We see less speedup overall—but we can still see signiﬁcant speedup (up to 3 × ) with a causal mask, where half the blocks are masked out.
RTX 3090 Figure 8 shows speedup on an RTX 3090 GPU.
Here, we use batch size 12 withattention heads.
We observe slightly higher speedups on the RTX 3090 (between 2.5-4.5 × ), since the memory bandwidth on an RTX 3090 is lower than on an A100 (roughly 900 GB/s vs. 1.5 TB/s).
T4 Figure 9 shows speedup on a T4 GPU.
T4 SRAM is smaller than A100, so we need to make the block sizes smaller in F LASH A TTENTION .
As a result, we observe less speedup on T4, which matches the IO complexity analysis in Section 3.2.
T4 GPUs are commonly used for inference, so we also report speedup on the forward pass only.
E.9 Full Benchmarking Results
We report the full benchmarking results and experimental details on A100.
[Table]
[Table] [Table]
[Table] [Table]
[Table]
Table
Baselines We compare against reference implementations for exact attention from PyTorch/HuggingFace and Megatron, approximate attention, and sparse attention.
For approximate attention, we compare against reference implementations of Reformer [53], Local Attention [71],
Linformer Attention [88], Smyrf [20], and LongShortFormer (LSFormer) [98].
For sparse attention, we compare against reference implementations of Block-Sparse Attention form OpenAI [12],
Longformer[3], and BigBird Attention [96].
For the approximate and sparse attention, we use a compression ratio of 1/8, or a compressed sequence length of 256, whichever is smaller.
Setup We measure runtime and memory usage of the attention computation with 8 heads of dimension 64, and batch size 16 on a machine with one A100 GPU with 40 GB of GPU HBM.
We vary sequence length in our experiments.
We compute attention on random vectors for Q , K , and V (we do not measure the projection from the hidden layer).
For dropout, we use dropout 0.1; for masking, we use a padding mask with uniformly-random mask lengths between the total sequence length and the total sequence length minus 20.
To measure runtime, we take the average of 100 measurements of the attention call.
We only measure memory footprint once, since it does not vary between runs.
We report timing results on the forward pass, backward pass, and combined forward + backward pass.
We measure each method with and without dropout, masking, or both—except for Block Sparse,
Longformer, and BigBird.
These methods did not successfully run the backward pass with masking due to a bug in external libraries, so we measured them without masking to be generous.
We use
FP16 for all measurements, except for Local Attention, whose implementation only supports FP32.
For each baseline, we increase sequence length until it runs out of memory on the GPU, except for the following exceptions: The Megatron implementation does not support sequence lengths longer than 2048.
Block-Sparse (OpenAI) does not support sequence lengths longer than 4096.
Longformer and BigBird do not support sequence lengths longer than 8092.
We measure memory usage on the combined forward + backward pass, without dropout or masking.
Results Table 13 summarizes all the experimental conﬁgurations and contains pointers to the results tables.
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]