FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Authors: Tri Dao
Daniel Y Fu
Stefano Ermon
Atri Rudra
Christopher Re
Venue: NeurIPS 2022 Accept (2022)
Keywords: Attention
GPUs
Hardware-efficient model
Long context
IO complexity

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 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, yielding an approximate attention algorithm that is faster than any existing approximate attention method. FlashAttention, 3x speedup on GPT-2 (seq. length 1K), and 2.4x speedup on long-range arena (seq. length 1K-4K). FlashAttention, 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).

FLASHATTENTION: Fast and Memory-Eficient Exact Attention with IO-Awareness
Tri Dao", Daniel Y. Fu 1, Stefano Ermon 1, Atri Rudra 3, Christopher Ré'
Department of Computer Science, Stanford University $ Department of Computer Science and Engineering, University at Buffalo, SUNY (trid, danf @stanford. edu, ermon@stanford.cc edu, atri@buffalo. edu, chrismre@cs. stanford. edu
Abstract
Transformers are slow and memory-hunge 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 qual- ity reduce the compute complexity, but often do not achieve wall-clock speedup. We argue that a missing principle is making attention algorithms 10-aware accounting for reads and writes between levels of GPU memory. We propose FLASHATTENTION,, an 10-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 FL.AHHAT ATTENTION, 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 at- tention, 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, 3x speedup on GPT-2 (seq. length IK), and 2.4x speedup on long-range arena (seq. length IK-4K). FLASIIAT- TENTION and block-sparse PLASSIATTENTION enable longer context in Transform- ers, yielding higher quality models (0.7 better perplexity on GPT-2 and 6.4 points lift on long-document classification) and entirely new capabilities: the first Trans- formers to achieve getter-than-chame performance on the Path X challenge (seq. length 16K, 61.4% accuracy) and Path-256 (seq. length 64K, 63. % accuracy).
Introduction
Transformer models [86] have emerged as the most widely used architecture in applications such natural language processing and image classification. Transformers have grown larger (3] and deeper [87), but equipping them with longer context remains difficult [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-efficien 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 aprroxaoppooiiiaiionn
53 to low-rank approxima- and their combinations 3 g 96]. Although these methods reduce the compute
tion 13 requirements to linear or near-linear sequence length, many of them 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 overhcads from memory access (IO).
52 88
In this paper, we argue that missing principle is making attention algorithms 10-aware 1-th is, carefully accounting for reads and writes 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
36th Conference Neural Information Processing Systems (NeurIPS 2022).

---

Outer Loop
KtdxN
Attention GPT-2
Copy to SRAM
Q:Nxd
Outer Loop
V.N
15-
Matmul
GPG
SRAM: TB/s (20 MB)
QR: N
SRAM
Dropout
GPU
HBM: 1.5 TB/s (40 GB)
Inner Loop
Copy
Outerr Loog
Time (m3)
10+
HBM
Compute Block
Softmax
Main Memory
DRAM: 12.8 GB/s
Copy
on SRAM
Inner Loop
Fused
(CPU DRAM)
(>1TB)
Mask
Kernel
Matmul
Memory Hierarchy with Bandwidth & Memory Size
Output to HBM
PyTorch
FlashAttention
smidK]V:Nxd
Inner Loop
FlashAttention
Figure 1: Left: FLASHATTENTION uses tiling to prevent materialization of the large NxN attention matrix (dotted box) on (relatively) slow GPU HBM. In the outer loop (red arrows). FLASHATTENTION loops through blocks of the K and V matrices and loads them to fast on-chip SRAM. In each block, FTTSHATTENTION loops over blocks of Q matrix (blue arrows), loading them to SRAM, and writing the output of the attention computation back to HBM. Right: Speedup over the PyTorch implementation of attention on GPT-2. ATTENTION does not read and write the large N NxN attention matrix to HBM, resulting in 7.6x speedup on the attention computation.
left). On modern GPUs, compute speed has out-paced memory speed (64)66], and most operations in Transformers are bottlenecked by memory accesses [45). 10-aware algorithms have been critical for similar memory-bound operations, when reading and writing data can account for large portion of the runtime- -scch as database joins (74,, image processing (73,, numerical linear algebra 4 and more 42 [8]] However, common Python interfaces to deep learning such as Py PyTorch and Tensorflow do not allow fine-grained control of memory access.
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 address these challenges.
We propose FLASHATTENTION
(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 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 FLASIATTENTION in CUDA to achieve fine-grained control over memory access and fuse all the attention operations into one GPU kernel. Even with the increased FLOPs due to recomputation, our right) and uses less memory- in sequence length-than standard attention, thanks to the massively reduced amount of HBM access.
algorithm both runs faster (up to 7.6x on GPT-2 77),, Figure D
We analyze the 10 complexity [ of FLASHATTENTION, proving that it requires of HBM accesses where is the head dimension and M the size of SRAM, as compared ((Nd+N3) of standard attention. For typical values of d and M, FLASHATTENTION requires many times fewer HBM accesses compared standard attention (up to 9x fewer, as shown in Fig. 2 Moreover, 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 FLASHATTENTION approximate attention algorithms by overcoming their issucs with memory access overhead. As proof of concept, implement block-sparse FLASHATTENTION that is 2-4x faster than even FLASHATTENTION, scaling up to sequence length of 64k. We prove that block-sparse FLASHATTENTION has better I0 complexity than FLASHATTENTION by 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 FLASHATTENTION to make it easier to build on this primitive
useful primitive for realizing the potential of
sparse attention algorithm
FLASHATTENTION code is available at httppss: htt.ttp://ggccuu...o//////////// http:///htttaccooo//////aaaaaaaa

---

We empirically validate that FLASHATTENTION speeds up model training and improves model quality by modeling longer context. We also benchmark the runtime and memory footprint of FLASHAT- TENTION and block-sparse FLASHATTENTION compared to prior attention implementations.
Faster Model Training. FLASHATTENTION trains Transformer models faster in wall-clock time. We train BERT-arge (seq. length 512) 15% faster than the training speed record in MI.Perf 1.1 [60]. GPT2 (seq. length 1K) 3x faster than baseline implementations from HuggingFace [91 and Megatron-I.M [80], and long-range arena (seq. length 1K-4K) 2.4x faster than baselines.
Higher Quality Models. FLASHATTENTION scales Transformers 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 lift from modeling longer sequences on long-documng docemetttnntt classification [14. FLASHATTENTION enables the first Transformer that can achieve better-than siiter-than-chaaaee chaanc performance on the Path-X 83 challenge, solely from using longer sequence length (16K). Block-sparse FLASHATTENTION enables Transformer to scale to even longer sequences (64K), resulting in the first model that can achieve peiter-than-change performance on Path-256.
Benchmarking Attention. FLASHATTENTION is up to 3x 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, FLASHÁTTENTION is both faster and more memory-efficiet 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 FLASHATTENTION is faster than all existing approximate attention methods that we know of.
2 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
GPU Memory Hierarehy. The GPU memory hierarchy (Fig. leftt comprises multiple forms of mem- ory 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 operations are increasingly bottlenecked by memory (HBM) accesses. Thus exploiting fast SRAM becomes more important.
Execution Model. GPUs have a massive number of threads execute operation (called 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, oper- ations can be classified as either compute-bound or memory-bound This is commonly measured by the arithmetic intensity g9 which the number of arithmetic operations per byte of memory access.
1. Compute-bounde the time taken by the operation 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- operations 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 HBM to save for the backward pass, reducing the effectiveness of naive kernel fusion.
2.2 Standard Attention Implementation
Given input sequences Q.K. Nxd where N is the sequence length and d is the head dimension, we want to compute the attention output OE Nxd,
S=QKT NxN (= ERN
O=PVERNxd
3

---

where softmax is applied row-wise.
Standard attention implementations materialize the matrices S and P to HBM, which takes O(N2) memory. Often N > d (e.g., for GPT2, N = 1024 and =d64), We describe the standard attention implementation in Algorithm [ As some or most of the operations are memory- (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 clementwise operations, such as fusing masking with softmax [80].
In Section 3.2
we will show that the standard attention implementation performs HBM accesses quadratic the sequence length N. We also compare the number of FLOPS and number of HBM accesses of standard attention and of our method (FL.ASHATTENTION)
Algorithm 0 Standard Attention Implementation
Require: Matrices Q.K, V.K.VRNAA Nxd in HBM.
1: Load Q.K by blocks from HBM, compute S=QKT, write S to HBM.
2: Read S from HBM, compute P = softmax(SS), write P to HBM.
3: Load and V by blocks from HBM, compute O=PV, write O to HBM.
4: Return O.
3 FLASHATTENTION: Algorithm, Analysis, and Extensions
We show how compute exact attention with fewer HBM reads/writes and without storing large intermediate matrices for the backward pass. This yields an attention algorithm that both memory efficient and faster in wall-clock time. We analyze its I0 complexity. showing that our method requires much fewer HBM accesses compared to standard attention. We further show that FLASHÂTTENTION can serve useful primitive by extending it to handle block-sparse attention.
We focus here on the forward pass for ease of exposition; Appendix backward.
B contains details for the
3.1 An Efficient Attention Algorithm With Tiling and Recomputation
Given the inputs Q.K. V E Nxd in HBM, we aim compute the attention output O R and write it to HBM. Our goal is to reduce the amount of HBM accesses (to aub-quadratic in N).
We apply two established techniques (tiling, recomputation) to overcome the technical challenge of computing exact attention in sub- sub-quadratic HBM accesses. We deseribe this in Algorithm T 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 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 X is computed:
f(x) = Σ softmax(x) f(x)
m(x) = max Xis
fxx = [===<>i-mii (m(())
)xn-m(x))
((x)
For vectors x (1) (2)) ER we can decompose the softmax of the concatenated ( E as:
m(x)=m( (x(1) =ax mmmx(())) f(x)= (L)==((m<xx))) (m)(()))mm(()) f(x(1) ) ((x(2)).
((x) ((x==(([(()) x22) m(x""")> ) -m(()) ((x +e )-m()) ((x softmax(x) f(x)
((x)
Therefore if we keep track of some extra statistics (m))), we can compute softmax one block time We thus split the inputs Q.K. V into blocks (Algorithm I line L). compute the softmax values along with extra statistics (Algorithm line D). and combine the results (Algorithm [line !
Recomputation. One of our goals is to not store (NV2) intermediate values for the backward pass. The backward pass typically requires the matrices NXN to compute the gradients with respect to Q.K.V V. However, by storing the output O and the softmax normalization statistics (m, 6), we 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 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. contrast, even with more FLOPs, our recomputation speeds up the backward pass due to reduced HBM accesses (Fig, 2 The full backward pass description in Appendix B
implement our algorithm in 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 Appendix B This avoids repeatedly reading and writing of inputs and outputs from and
Implementation details: Kernel fusion. Tiling enables
to HBM.
Algorithm 1 FLASHATTENTION
Require: Matrices Q.K, VER in HBM, on-chip SRAM of size M.
1: Set block sizes B
M B, m
4d I.d).
M
ad
2: Initialize 0= (0)
Nxd ((= (xxxxd(((())))) (0)x ER _== ER in HBM.
3: Divide Q into 7,
N blocks
Qr, of size B, x each, and divide K.V in to T, Re
N
B
blocks KI.... Kr, and I..... Vr. -r..00 of size Bexd each.
4: Divide O into T, blocks Oi....Or, of size B, xd each, divide é into T, blocks C,,...,,, of size B, each, divide m into T, blocks mI.....y,, of size B, each.
5: for i TT do
6: Load v, from HBM to on-chip SRAM.
7:
for 1 IsisT, TT do
8:
Load O, from HBM on-chip SRAM.
9:
On chip, compute S,j Q E BB
On chip, compute mij rowmax Si,j) E R. P.j = exp(Sii -mij) R.R.BBBB (pointwise), =y==rowuum(P;) B
11:
On chip, compute m max (m; ERR, Cnew em=mn li (,j ER.
12:
Write O; - diag - (diag( ( mu-mi Ou PijV,) to HBM.
Write f e:ee mi m. new to HBM.
14:
end for
end for
16:
Return O.
We show FLASHATTENTION correctness, runtime, and memory requirement (proof in Appendix C
Theorem 1. Algorithm 1 returns (aftmaxx(KK)) with O(N-d) FLOPs and requires O
additional memory beyond inputs and output.
3.2 Analysis: 10 Complexity of FLASHAFTENTION
We analyze the 10 complexity of FLASHATTENTION, showing significant reduction in HBM accesses compared to standard attention. We also provide a lower bound, proving that no exact attention algo- rithm can asymptotically improve HBM accesses
all SRAM sizes. Proofs in Appendix C
Theorem 2. Let N be the sequence length, d be the head dimension, and M be size of SRAM with d & M & Nd. Standard attention (Algorithm 0h requires O(Nd+N3) HBM accesses, while FLASHATTENTION (Algorithm ! requires @ d M HBM accesses.
For typical values (64-128) and M (around 100KB), d² is many times smaller than M, and thus FLASHATTENTION 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 that given the SRAM size of M, we can load blocks of K, V of
size ((M) each (Algorithm] [ line T.. For each block of K and V, we iterate over all blocks of Q (Algorithm I line I to compute the intermediate values, resulting in O(NdM-1) passes over Q. Each pass loads (NN)) elements, which amounts to @ df prove that the backward pass of standard attention requires O(N Nd+ N2) HBM accesses while the backward pass of FLASHATTENTION requires df M-!) HBM accesses (Appendix B)
HBM accesses. We similarly
We prove a lower-bound: one cannot asymptotically improve on the number of HBM accesses for all values of M (the SRAM size) when computing exact attention.
5

---

Effect of Block Size
Sparsity Speedup
Centimme Fwo
Fuid + Bwdd (msi)
Attention
Standard
FLASHATTENTION 75.2
HBM Accesses
150
Dense FlashAttantion
HBM R/W (GB) 35.3
GFLOPs
66.6
4.4
HBM Accesses
Runtime
100
Block. Sparse
FlashAtention
Runtime (ms) 35.1
11.7
s0
128
256
20
Block Size
%6 Non-Zero Blocks
Figure 2: Left: Forward + backward runtime of standard attention and FLASHATTENTION for seq. length 1024, head dim. 64, 16 heads, batch size 64, key-padding mask and no dropout on A100 GPU. HBM access is one of the primary factors affecting runtime. Middle: Forward runtime of FLASHATTENTION (seq. length 1024, head dim. 64, 16 heads, batch size 64) on GPU. Fewer HBM accesses result in faster runtime, up to point. Right: The runtime (for seq. length 4K) of block sparse FLASHATTENTION faster than FLASHATTENTION by factor proportional to the sparsity.
Proposition 3. Let N be the sequence length, be the head dimension, and M be size of SRAM with d M Nd. There does not exist an algorithm to compute exact atlention with of (N22d² M-1) HBM accesses for all M in the range [d.Nd].
The proof relies on the fact that for (NN)) any algorithm must perform ( HBM accesses. This type of lower bound over a subrange of M is common in the streaming algorithms literature [92]. We leave proving parameterized complexity (29 lower bounds in terms of M as exciting future work.
We validate that the number of HBM accesses is the main determining factor of attention run-time. In Fig. [ (left), we see that even though FLASHATTENTION has higher FLOP count compared to stan- dard 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 B, of
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 make fewer passes 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 fit into the small SRAM size.
the input), and
3.3 Extension: Block-Sparse FLASHATTENTION
We extend FLASHATTENTION to approximate attention: propose block-sparse FLASHATTEN- TION, whose I0 complexity is smaller than FLASUITTENTION by a factor proportional to the sparsity.
Given inputs Q.K. and a mask matrix Me (0.1) NxN, we want to compute:
S=QK² NRN P s=sotmmxxSS kx) ERN O=PVER-xd.
where (So = if Mę 1 and if Mę/=0. We require M to have block form: for some block sizes B, for all k,/, M with i= B= for some Me (0,1) /B.
Given a predefined block sparsity mask Me (0.1) /B, we can easily adapt Algorithm to
only compute the nonzero blocks of the attention matrix. The algorithm is identical to Algorithm except we skip zero blocks. We reproduce the algorithm description in Algorithm 3 in Appendix B
We also analyze the 10 complexity of block sparse FLASHATTENTION.
Proposition 4. Let N be the sequence length, be the head dimension, and M be size of SRAM with d & M S Nd. Block-sparse FLASHATTENTION (Algorithm) 5 requires O(Nd+ + & HBM accesses where is the fruction of nonzero blocks in the block sprccksysrtttyyyoooo 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 N, is often set to -11 12 or No [8, 96 resulting in ONN or (NNlggN IO complexity. For downstream experiments, we use the fixed butterfly sparsity pattern [8] which has been shown to be able to approximate arbitrary sparsity 17
In Fig. 1 (right), we validate that as the sparsity increases, the runtime of block-sparse FLASHATTEN- TION improves proportionally. On the LRA benchmark, block-sparse FLASHÄTTENTION achieves 2.8x speedup, while performing on par with standard attention (Section 4..
4 Experiments
We evaluate the impact of using FLASHATTENTION to train Transformer models. We validate two claims about training time and model accuracy, and report attention runtime and memory benchmarks.

---

Training Speed. FLASHATTENTION outperforms the MLPerf 1.1 [60] speed record for BERT by 15%, and speeds up GPT-2 up to 3x over HuggingFace [9]] and 1.8x over Megatron 80 over standard Transformers. FLASHATTENTION speeds up the long-range arena (LRA) benchmark 2.4x.
Quality. FLASHATTENTION scales Transformers to longer sequences, yielding higher quality. FLASHATTENTION trains GPT-2 with context length 4K faster than Megatron trains GPT-2 with context length IK, while achieving 0.7 better perplexity. Modeling longer sequences yields 6.4 points of lift on two cong-document classification tasks. Finally, FLASHATTENTION yields the first performance on the challenging Path-X X task (sequence length 16K), and -parse ATTENTITNN yields the first sequence model that we know of that can achieve better-than-randon performance on Path-256 (sequence length 64K).
Transformer that can achieve better-
Benchmarking Attention. We measure the runtime and memory performance of FLASHATTEN- TION and block-sparse FLASHATTENIIN based on sequence length. We confirm that the memory footprint of FLASIJATTENTION scales linearly with seq. length and is up to 3x faster than standard attention for common seq. lengths (up to 2K). We confirm that runtime of block-sparse FLASIIA-- TENTION 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 ATTENTONN
BERT. FLASHATTENTION yields the fastest single-node BERT training speed that we know of. We train a BaRT-large 24 model with FLASHATTENTION on Wikipedia. Table [ compares our training time to the implementation from Nvidia that set the training speed record for MLPerf 1.1 660 63). Our implementation is 15% faster.
Table Training time BERT-large, starting from the same initialization provided by the MI.Perf benchmark, to reach the target accuracy of 72.0% on masked language modeling. Averaged over 10 runs on 8xA GPUs. BERT Implementation
Training time (minutes)
Nvidia MI.Perf 1.I 66
20.0 1.5
FLASHATTETIOO (ours)
17.4 + 1.4
GPT-2. FLASHATTENTION yields faster training times for GPT-2 770 on the large Open Webtext
dataset 334 than the widely used HuggingFace (9]] and Megatron-LM (80) implementations. Table shows up to 3x end-to-end speedup compared to Huggingface and 1.7x speedup compared
Megatron-LM. FLASUIATTENIOON achieves the same perplexity as the other two implementations, as we do not change the model definition. Appendix [ inccluess plots of the validation perplexity throughout training, confirming that FLASIIATTENTION is as numerically stable as the baselines and produces the same training/ validation curves.
Table 2: GPT-2 small and medium using FLASHÁTTENTION achieve up to 3x speed up compared to Huggingface implementation and up to 1.7x compared Megatron- Training time reported on BxA100s GPUs. WeenWebText (ppl) Training time (speedup)
Model implementations
GPT-2 small Huggingtace [91
18.2
9.5 days (1.0x)
GPT-2 small Mcgatron-LM 80
18.2
4.7 days (2.0x)
GPT-2 small FLASHATTENTION
18.2
2.7 days (3.5x)
GPT-2 medium medium- Huggingtace 91
14.2
21.0 days (1.0x)
GPT-2 medium MMgatronn 80
14.2
11.5 days (1.8x)
GPT-2 medium FLASHATTENTION
14.2
6.9 days (3.0x)
Long-range Arena. We compare vanilla Transformer (with either standard implementation or FLASHATTENTION) 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 Tay al.
[8] Xiong et al. [94] Table 3 shows that FLASHATTENTION achieves up 2.4x speed-up compared to standard attention. Block-sparse FLASHATTENTION 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-efficiency of FLASHAT- TENTION allow to increase the context length of GPT-2 by 4x while still running faster than the optimized implementation from Megatron-LM. Table shows that that GPT-2 with
3 LLR accuracy results are known to be highly dependent on the tuning procedure [94] Our reproduced baselines perform better than reported in the original comparison (83).
7

---

Table 3: The performance of standard attention, FLASHATTENTION, block-sparse FLASHATTENTION, and
approximate atlention baselines on the Long-Range-Arre benchmarks.
Models
Ls
Text
Retrieval
Image
Pathfinder
Avg
Speedup
Transformer
36.0
63.6
81.6
42.3
72.7
59.3
CASSSATTENTION
37.6
63.9
81.4
43,5
72.7
59.8
2.4x
Block-sparse FLASHATTENTIOX
37.0
63.0
81.3
43.6
73.3
59.6
2.8%
Lintormer 88
35.6
55.9
77.7
37.8
67.6
54.9
2.5X
Linear Attention
38.8
63.2
80.7
42.6
72.5
59.6
2.3x
Performer 13
36.8
63.6
82.2
42.1
69.9
58.9
L88
Local Attention 83
36.1
60.2
76.7
40.6
66.6
56.0
1.7x
Reformer 55
36.5
63.8
78.5
39.6
69.4
57.6
1.3x
Smyrf 20
36.1
64.1
79.0
39.6
70.5
57.9
1.7x
FLASHATTENTION and context length 4K is still 30% faster than GPT-2 from Megatron with context length 1K, while achieving 0.7 better perplexity.
Table 4; GPT-2 small with FLASH ATTENTION, with 4x larger context length compared to Megatron-L.M, is still 30% faster while achieving 0.7 better perplexity. Training time on 8xA100 GPUs is reported.
Model implementations
Context length
Open WebText (ppl)
Training time (speedup)
GPT-2 small Megatron- LM
1k
18.2
4.7 days (LUx)
GPT-2 small FLASHATTEN
18.2
2.7 days (1.7x)
GPT-2 small FLASHATTENTION
2k
17.7
3.0 days (L.6x)
GPT-2 small FLASHÄTTENTION
4k
17.2
3.6 days (1.3%)
Long Document Classification. Training Transformers with longer sequences with FLASHAAT- TENTION improves performance on the MIMIC-III [49 and ECIHR 66] 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 ECUHR 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, in Beltagy et (3).
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 ECIHR. The discrepancies may be due to subtle distribution shifts: MIMIC-III contains specialized medical text and thus may be more susceptible to distribution shift in the document length, whereas ECtHR contains general language.
Table 6: We report the first Transformer model that achieve non-random performance on Path-X and Path-256.
Table 5: Long Document performance (micro Fi) at different sequence lengths using FLASHATTEN-
Model
Path-
P5th-256
Transtormer
X
X
TION.
Linformer
X
512
1024
2048
4096
8192
16384
Lincar Attention 32]
Performer [3
X
X
MIMICIII
ECHHR 6
52.8
72.2
50.7
74.3
77.1
51.7
78.6
54.6
80.7
56.
79.2
Local Atteation [3
Reformer 53
X
SMYRF 20
X
FLASHATTENTION
61.4
Block sparse FLASHATTENNION
56.0
63.1
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 classify whether two points in a black and white 128x 128 (or 256x256) image have path connecting them, and the images are fed to the transformer one pixel at 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 (9). We present here the first result of Transformer models being able to solve Path-X X and P5th-256 (Table We pretrain a transformer on Path-644 -6, and then transfer to Path-X by spatially interpolating the positional embeddings. FLASHATTENTION achieves 61.4 accuracy on Path-X. Additionally, block-sparse FLASHATTENTION enables the Transformers to scale to sequence length 64K, achieving 63.1 accuracy on Path-256.
8

---

Attention Runtime (Fwd Pass Bwd Pass)
Attention Memory Usage
Memory Footprint (GB)
10a
Runtime (ms)
Crossover Points
20
10'
10
20x
10
128
256
512
1024
2048
4096
8192
256
8K
16K
32K
64K
Sequence Length
Sequence Length
FlashAttention
PyTorch Attention
Linformer Attention
Spark-Sparse FlashAttention
Megatron Attention
OpenAl Sparse Attention
Figure 3: Left: runtime of forward pass + backward pass. Right: attention memory usage.
4.3
Benchmarking Attention
We vary sequence length and measure runtime and memory usage of FLASHATTENTION and block-sparse FLASHATTENTION against various attention baselines on one A A100 GPU with 40 GB HBM, with dropout and padding mask. We compare against reference implementations for exact attention, approximate attention, and sparse attention. We report a subset baselines in the main body; Appendix
E contains more baselines and full details.
Runtime. Figure (left) reports the runtime in milliseconds of the forward + backward pass of FLASUITTTNTION and block-sparse FLASIATTENTION compared to the baselines in exact, approximate, and sparse attention (exact numbers in Appendix E.. Runtime grows quadratically with sequence length, but FLASHATTENTION runs significantly faster than exact attention baselines, up to 3x faster than the PyTorch implementation. The runtimes of many approximate/sparse attention mechanisms grow linearly with sequence length, but FLTENTTTNNTIO still runs faster than approximate and sparse attention for short sequences due to fewer memory accesses. The approximate attention runtimes begin cross over with FLASHATTENTION at sequences between 512 and 1024. On the other hand, block-sparse FLASHATTENTION is faster than all implementations of exact, sparse, and approximate attention that we know of, across all sequence lengths.
Memory Footprint Figure (right) shows the memory footprint of FLASIIATTENTION and block-sparse FLASIIATTENTION compared to various exact, approximate, and sparse attention baselines. FLASIIATTENTION and block-sparse FLASIATTENTION have the same memory footprint, which grows linearly with sequence length. FLASTATTENTION is up to 20x more memory efficient than exact attention baselines, and is more memory-efficient than the approximate attention basclines. All other algorithms except for Linformer run out of memory on an A100 GPU before 64K, and FLASHÄTTENTION is still 2x more efficient than Linformer.
5 Limitations and Future Directions
We discuss limitations of our approach and future directions. Related work is given in Appendix A
Compiling CUDA. Our current approach building IO-aware implementations of attention requires writing a new CUDA kernel for each new attention implementation. This requires writing the attention algorithm in considerably lower-level language than PyTorch, and requires significant engineering effort. Implementations may also not be transferrable across GPU architectures. These limitations suggest the need for method that supports writing attention algorithms in high-level language (c.g., Py Torch), and compiling to 1O-aware implementations in CUDA CUDA-similar such Halide in image processing [73
similar to efforts
10-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 deep network touches GPU HBM. We hope our work inspires 10-aware implementations of additional modules. We discuss these potential extensions in Appendix[]
Multi-GPU 10-Aware Methods. Our IO-aware implementation of attention optimal within constants for computing attention on a single GPU. However, the attention computation may be par-
*Path-255 requires longer sequences but has relatively shorter paths than Path- X, so it is easier to obtain higher accuracy.

---

allelizable across multiple GPUs [75]. Using multiple GPUs adds an additional layer to I0 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 efficiently. This may allow a general community with limited access computational resources to train and understand those foundation models. Our method applicable to all Transformer-basse models, which have a variety of applications, both positive and negative. For example, language modeling may make it easier spread misinformation, while image classification models may make automatic surveillance easier. Alleviating these risks requires addressing application-specilee issues such as privacy, bias, and discrimination.
Acknowledgments
Our implementation uses Apex's FMHA code (https://githu) CO.VVIIIII apex/tree/ tree/ master/ /app//ape///// contrib/carc/fmmae 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. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); ARL under No. W911NF-21-2-0255 (Interactive Human-AI Teaming); ONR under No. N000141712266 (Unifying Weak Supervision); ONR N00014-20-1-2488 Understanding and Applying Non-Euclidean Geometry in Machine Learning: N000142012275 (NEPTUNE); NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASE, 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, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, the U.S. Government. Atri Rudra's ' research is supported by NSF grant CCF- (CF-1773481.
References
[1] Alok Aggarwal and S Vitter, Jeffrey. The input/output complexity of sorting and related problems. Communications of the ACM, 31(9):1 1111111611100 127, 1988.
[2] Irwan Bello. LambdaNetworks: Modeling long-range interactions without attention. arXiv preprint arXiv:2102.08660) 2021.
[3] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long- docemeett transformer. arXiv preprint arXiv:2004.05223 2020.
[4] Susan Blackford, Antoine Petitet, Roldan Pozo, Karin Remington, Clint Whaley, James Demmel, Jack Dongarra, lain Duff, Sven Hammarling, Greg Henry, et al. An updated set of basic linear algebra subprograms (blas). ACM Transactions on Mathematical Softwure, 28(2): 135-151, 2002.
[5] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[6] Ilias Chalkidis, Ion Androutsopoulos, and Nikolaos Aletras. Neural legal judgment prediction in English. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 1317-4323, Florence, Italy, 2019. Association for Computational Linguistics. doi: 10.18653/v1/P194 1424. URL http:::/w/wwwwwwbbbbb/p wwwwwb..wbbbb.... www.s/ebb.rrg..... http:a/www....o//attooooo/////-/ http://toooooo/1-11//0/1 1424
[7] Ilias Chalkidis, Manos Fergadiotis, Dimitrios Tsarapatsanis, Nikolaos Aletras, Ion An- droutsopoulos, and Prodromos Malakasiotis. Paragraph-level rationale extraction through regularization: A case study curopean court of human rights cases. In Proceedings of
10

---

the Anmual Conference the North American Chapter of the Association for Computational Linguistics, Mexico City, Mexico, 2021. Association for Computational Linguistics.
[8] Benjamin Charlier, Jean Feydy, Joan Alexis Glaunès, François-David Collin, and Ghislain Durif. Kernel operations on the gpu, with autodiff, without memory overflows. Journal of Machine Learning Research, 22(74): 1-6, 2021. URL.htpp://mmlzz http://jiii...orrppppess//////// pwtppppppps///////////// 20- 275 html
[9] Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[10] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin, Training deep nets with sublinear memory cost. arXiv preprint arXiv:/604.00 2016.
[11] Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, et al. (TVM): An automated (End-to-Ene -oo--dd-nnn optimizing compiler for deep learning. In 13th USENIX Symposium on Operating Systems Design and Implementation ( (OSDI 18), pages 578-594, 2018.
[12) Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
[13] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos. Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. International Conference on Learning Representations (ICLR), 2020.
[14] Xiang Dai, Ilias Chalkidis, Sune Darkner, and Desmond Elliott. Revisiting transformer- based models for long document classification. arXiv preprint (=XXi::220400683,, 2022.
[15] Zihang Dai, Zhilin Yang, Yiming Yang. Jaime G Carbonell, Quoc Le, and Ruslan Salakhutdinov, Transformer-Xe XL: Attentive language models beyond fixed-length context. In Proceedings of the Annual Meeting of the Association for Computational Linguistics, pages 2978-228888 2988, 2019.
[16] Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré. Learning fast algorithms for linear transforms using butterfly factorizations. In International Conference on Machine Learning (ICML), 2019.
[17] Tri Dao, Nimit Sohoni, Albert Gu, Matthew Eichhorn, Amit Blonder, Megan Leszczynski, Atri Rudra, and Christopher Ré. Kaleidoscope: An efficient, learnable representation for all structured linear maps. In International Conference on Learning Representations (ICLR), 2020.
[18] Tri Dao, Beidi Chen, Kaizhao Liang, Jiaming Yang. Zhao Song, Atri Rudra, and Christopher Ré. Pixelated butterfly: Simple and efficient sparse training for neural network models. In International Conference on Learning Representations (ICLR), 2022.
[19] Tri Dao, Beidi Chen, Nimit Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive structured matrices for efficient and accurate training. In International Conference Machine Learning (ICML), 2022.
[20] Giannis Daras, Nikita Kitaev, Augustus Odena, and Alexandros G Dimakis. Smyrf-efficient attention using asymmetric clustering. Advances in Neural Information Processing Systems, 33:6476-6498 2020.
[21] Christopher De Sa, Albert Gu, Rohan Puttagunta, Christopher Ré, and Atri Rudra. A two- pronged progress in structured dense matrix vector multiplication. In Proceedings of the Twenty- Ninth Annual ACM-SIAM Symposium on Discrete Algorithms, pages 1060-1079. SIAM, 2018
[22] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248 255.. Ieee, 2009.
[23] Peter Denning. The working set model for program behavior. Communications of the ACM, 11(5):323-333, 1968.

---

[24] Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. 2019.
[25] Xin Dong, Shangyu Chen, and Sinno Jialin Pan. Learning prune deep neural networks via 2017.
layer-wise optimal brain surgeon. arXiv preprint arXiv:
[26] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition at scale. In International Conference Learning Representations, 2020.
[27] Y Eidelman and Gohberg. On new class of structured matrices. Integral Equations and Operator Theory, 34(3):293-322 1999.
[28] Jean Feydy, Joan Glaunès, Benjamin Charlier, and Michael Bronstein. Fast geometric learning with symbolic matrices. Advances in Neural Information Processing Systems, 33, 2020.
[29] Jörg Flum and Martin Grohe. Parameterized Complexity Theory. Springer, 2006.
[30] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference on Leurning Representations, 2018
[31] Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M Roy, and Michael Carbin. Stabilizing the lottery ticket hypothesis. arXiv preprint wwtXiv:1903.01611 2019.
[32] Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. In International Conference on Machine Learning. pages 3259-3269. PMLR, 2020.
[33] Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It's raw! audio generation with state-space models. In International Conference on Machine Learning (ICML), 2022. 2022..
[34] Aaron Gokaslan, Vanya Cohen, Pavlick Ellie, and Stefanie Tellex. Openwebtext corpus, 2019.
[35] Jim Gray, Surajit Chaudhuri, Adam Bosworth, Andrew Layman, Don Reichart, Murali Venka- trao, Frank Pellow, and Hamid Pirahesh. Data cube: A relational aggregation operator generaliz- ing group-by, cross-tab, and sub-totals. Data mining and knowledge discovery, 1(1):29-53, 53, 1997.
[36] Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and techniques of algorithmic differentiation. SIAM, 2008
[37] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: with optimal polynomial projections. In Advances in neural information processing systems (NeurIPS), 2020.
Recurrent memory
[38] Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances Neural Information Processing Systems, 34, 2021.
[39] Albert Gu Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The International Conference on Learning Representations (ICLR), 2022.
[40] Song Han, Jeff Pool, John Tran, and William J Dally. Learning both weights and connections for efficient neural networks, arXiv preprint arXiv: 1506. 02626, 2015.
[41] Song Han, Huizi Mao, and William Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. In International Conference on Learning Representations, 2016.
[42] John Hennessy and David Patterson. Memory hierarchy design. Computer Architecture: A Quantitative Approach, pages 390-525, 2003.
[43] Sara Hooker. The hardware lottery. arXiv preprint urXiv:2009.06-48 2020.
[44] Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoe V Le. Transformer quality in linear time. arXiv preprint 204:2202..1044, 2022.
12

---

[45] Andrei Ivanov, Nikoli Dryden, Tal Ben-Nun, Shigang Li, and Torsten Hoefler. Data movement all you need: A case study on optimizing transformers. Proceedings of Machine Learning and Systems, 3.711-732, 2021.
[46] Zhe Jia and Peter Van Sandt. Dissecting the Ampere GPU architecture via microbenchmarking. GPU Technology Conference, 2021.
[47] Zhe Jia, Marco Maggioni, Benjamin Staiger, and Daniele P Scarpazza. Dissecting the nvidia 2018.
Volta GPU architecture via microbenchmarking. arXiv preprint
[48] Zhe Jia, Blake Tillman, Marco Maggioni, and Daniele Paolo Scarpazza. Dissecting the graphcore IPU architecture via microbenchmarking, arXiv preprint
2019.
[49] Alistair EW Johnson, Tom J Pollard, Lu Shen, Li-wei H Lehman, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo Authony Celi, and Roger G Mark. Mimic-iii, a freely accessible critical care database. Scientific data, 3(1):1-9, 2016.
[50] Norman P Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, et al. In-datacenter performance analysis of tensor processing unit. In Proceedings of the 44th annual international symposium on computer architecture, pages 1-12, 2017.
[51] Thomas Kailath, Sun- Yuan Kung, and Martin Morf. Displacement ranks of matrices and linear equations. Journal of Mathematical Analysis and Applications, 68(2):395-407, 1979.
[52] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages 5156-5165. PMLR, 2020.
[53] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In The International Conference Machine Learning (ICML), 2020.
[54] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite BEDRT for self-supervised learning of language representations. In The International Conference on Learning Representations (ICLR), 2020.
[55] Mingzhen Li, Yi Liu, Xiaoyan Liu, Qingxiao Sun, Xin You, Hailong Yang, Zhongzhi Luan, Lin Gan, Guangwen Yang, and Depei Qian. The deep learning compiler: A comprehensive survey. IEEE Transactions on Parallel and Distributed Systems, 32(3):708-725 2020.
[56] Valerii Likhosherstov, Krzysztof Choromanski, Jared Davis, Xingyou Song. and Adrian Weller. Sub-linear memory: How make performers slim. arXiv preprint arXh::2012.11346, 2020.
[57] Ji Lin, Yongming Rao, Jiwen Lu, and Jie Zhou. Runtime neural pruning. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[58] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint 1rXiv:1907.116999 2019.
[59] Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer. Luna: Linear unified nested attention. Advances in Neural Information Processing Systems, 34, 2021.
[60] Peter Mattson, Christine Cheng, Gregory Diamos, Cody Coleman, Paulius Micikevicius, David Patterson, Hanlin Tang, YuuYeon Wei, Peter Bailis, Victor Bittorf, et al. MIperf training benchmark. Proceedings of Machine Learning and Systems, 2:336-349 2020.
[61] Frank McSherry, Michael Isard, and Derek G Murray, Scalability! but at what (COST)? In 15th Workshop on Hot Topics in Operating Systems (HotOS XV). 2015.
[62] Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv preprint
2018.
13

---

[63] MLCommons. Miperf 1.1 training results, 2021. URL https://mlcomoons.com
training-nowie moonoooomlllll 11/
[64] NVIDIA. Nvidia Tesla V100 GPU architecture, 2017.
[65] NVIDIA. Nvidia A100 tensor core GPU architecture, 2020.
[66] NVIDIA. Nvidia H100 tensor core GPU architecture, 2022.
[67] D Stott Parker. Random butterfly transformations with applications in computational linear algebra. 1995.
[68] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[69] Markus N Rabe and Charles Staats. Self-attention does not need O(n²) memory. arXiv preprint 1rXiv:211..05682 2021.
[70] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[71] Jack Rae and Ali Razavi. Do transformers need deep long-range memory? In Proceedings of the 58th Annual Meeting of the Association for Computational Lin- guisties. Online, July 2020. Association for Computational Linguistics. URL httpp:://ww/ aclweb. org/ aatp//nthhooooy//a 2020
wet.0....1--1... acl-mclamiii 672
[72] Jack Rac, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. In The International Conference on Learning Representations (ICLR), 2020.
[73] Jonathan Bagan-Kelley, Connelly Barnes, Andrew Adams, Sylvain Paris, Frédo Durand, and Saman Amarasinghe. Halide: language and compiler for optimizing parallelism, locality, and recomputation image processing pipelines. Acm Sigplan Notices, 48(6):519-530, 2013.
[74] Raghu Ramakrishnan, Johannes Gehrke, and Johannes Gehrke. Database management systems, volume 3. McGraw-Hii New York, 2003.
[75] Benjamin Recht and Christopher Ré. Parallel stochastic gradient algorithms for large-scale matrix completion. Mathematical Programming Computation, 5(2):201-226 2013.
[76] Hongyu Ren, Hanjun Dai, Zihang Dai, Mengjiao Yang, Jure Leskovec, Dale Schuurmans, and Bo Dai. Combiner: Full attention transformer with sparse computation cost. Advances in Neural Information Processing Systems, 34, 2021.
[77] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53-68, 2021.
[78] Amit Sabne. XLA: Compiling machine learning for peak performance. 2020.
[79] Victor Sanh, Thomas Wolf, and Alexander M Rush. Movement pruning: Adaptive sparsity by fine-tuning. arXiv preprint arXiv:2005. 07683, 2020.
[80] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-LM: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv: 090500003, 2019.
[81] Vikas Sindhwani, Tara Sainath, and Sanjiv Kumar. Structured transforms for small-footprint deep learning. In Advances in Neural Information Processing Systems, pages 0188-3096, 2015.
[82] Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers. In Proceedings of the Annual Meeting of the Association for Computational Linguisties, 2019.
14

---

[83] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: benchmark for efficient transformers, In International Conference on Learning Representations, 2020.
[84] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey, arXiv preprint N00Xi::20090663322 2020.
[85] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning, pages 10347-10357. PMLR, 2021.
[86] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[87] Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang. Dongdong Zhang, and Furu Wei. Deepnet: Scaling transformers to 1,000 layers. arXiv preprint arXivv 2022.
[88] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint Ntt1:20006.4768, 2020.
[89] Samuel Williams, Andrew Waterman, and David Patterson. Roofline: an insightful visual performance model for multicore architectures. Communications of ACM, 52(4))65-76, 2009.
[90] Michael E Wolf and Monica S Lam. A data locality optimizing algorithm. In Proceedings the ACM SIGPLAN 1991 conference on Programming language design and implementation. pages 30-44, 1991.
[91] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer. Patrick von Platen, Clara Ma, Yacine Jernite. Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariana Drame, Quentin Lhoest, and Alexander M. Rush. thet--rrt natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations. pages 38-45, Online, October 2020. Association for Computational Linguistics. URL
Transformers:
https:/ http::/wwwwwwwwww..w/... awwwwb.wwbb....... org/ anthology/ 2020. 2020.emmllpp.....mmmmmmmmmm.2mmp emnlp-demos.. 6
[92] David P Woodruff. Optimal space lower bounds for all frequency moments. In SODA, volume pages 167-175. Citeseer, 2004.
[93] Felix Wu, Angela Fan, Alexei Baevski, Yann N Dauphin, and Michael Auli. Pay less attention with lightweight and dynamic convolutions. In The International Conference on Learning Representations (ICLR), 2019.
[94] Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: nystöm-based algorithm for approximating self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence. AAAI Conference on Artificial Intelligence, volume 35, page 14138, 2021.
[95] Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Zi-Hang Jiang, Francis EH Tay, Jiashi Feng, and Shuicheng Yan. Tokens-to-token vit: Training vision transformers from scratch on imagenet. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 558-567, 2021.
[96] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santi- ago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang. al. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33, 2020.
[97] Shuangfei Zhai, Walter Talbott, Nitish Srivastava, Chen Huang, Hanlin Goh, Ruixiang Zhang, and Josh Susskind. An attention free transformer. arXiv preprint 2105, 2021.
[98] Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long-short transformer: Efficient transformers for language and vision. Advances in Neural Information Processing Systems, 34, 2021.
15

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section 5
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See
Section[5]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full of assumptions of all theoretical results? [Yes] See Section 3.2
(b) Did you include complete proofs of all theoretical results? [Yes] See Appendix C
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material as URL)? [Yes] See Appendix E
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Appendix
(c) Did you report error bars (e.g.. with respect to the random seed after running experiments multiple times)? [Yes] See Section 4
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes] See Appendix
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] See Section
and Appendix E
(b) Did you mention the license of the assets? [Yes] See Appendix
(c) Did you include any new assets either in the supplemental material URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation? [N/A]
16


Reviews:
[{"review_id": "yBQDZgiqlbE", "reviewer": "Reviewer_jjqW", "summary": "This paper studies an important topic for current Transformer architectures: accelerating the speed and memory consumption for self attention operations. The proposed method leverages tiling and recomputation to reduce the memory access of HBM, avoiding the memory bound issues in standard self attention kernels. The authors also extend the proposed Fast-Attention with block-wise sparsity, and  is the fastest implementation among exsiting approximate attention methods. The experiments show the superior performance w.r.t. memory and speed against existing solutions. The paper is quite solid with technical details, yet it is still well-written and friendly to the general audience. The idea to exploit IO awareness should be inspiring for the community to develop more efficient algorithms and kernels. ", "questions": "\n- What are the separative gains of reduced HBM access and tiling? In Algorithm 0, the Q,K,V matrices are loaded to HBM by blocks. I wonder whether similar techniques (tiling) are applied in standard attention implementation?\n\n- The baselines in Table 2 and 3 do not fuse the attention kernel. I wonder if there are any public baselines with fused kernels for fair comparisons?\n\n- As a suggestion, it is a bit miss-leading to show that $\\Theta(N^2 d^2 M^{-1})$ requires less memory access than $\\Theta(Nd + N^2)$, which may be attributed to tiling (due to the presence of $M^{-1}$). The main reduction of HMB access should come from not storing the attention score $S$, which takes $\\Theta(N^2)$ accesses. \n", "limitations": "The authors have listed and discussed the potention directions for the limitations. ", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "- The paper studies an important research question: designing fast and memory-cheap attention kernels for the general transformer architectures. The authors conduct extensive experiments and achieve impressive results w.r.t. training time and memory consumption across multiple tasks and network architectures.\n\n- The proposed Fast-Attention can be readily combined with block-sparsity to allow longer sequences, and the results are still faster than existing approximate attention designs. \n\n- The writing is clear and self-contained, rigorous and yet friendly to the general audience. Though I am not an expert in low-level designs on CUDA but I find it rather pleasing to read the aritcle.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "gzO66mRp8hl", "reviewer": "Reviewer_UuW7", "summary": "The paper leverages fusion and tiling techniques to reduce the overall memory consumption of attention mechanism. The results corroborate with some of the existing work on how fusion/tiling would help to reduce the memory footprint and runtime.", "questions": "(1) Fusion and tiling are well-established techniques in the literature, which is also highlighted in the paper. The block-level Softmax calculation is also not new. What are the main contributions of FlashAttention over existing methods?\n\n(2) There are recent proposals to leverage fusing opportunities in transformers, such as [FLAT](https://arxiv.org/abs/2107.06419) and [Self-Attention Does Not Need O(n$^2$) Memory](https://arxiv.org/abs/2112.05682), how is FlashAttention different from these work?\n\n(3) Which version of PyTorch did you use in your implementations? I am wondering if newer versions of PyTorch have already added some supports for fusion.\n\n(4) I understand the limitation to study the scalability of the FlashAttention method over larger number of GPUs. However, I suspect that as we increase the number of GPUs the benefit of FlashAttention over existing methods may plateau. Do you have any intuition on the performance projection of FlashAttention as number of GPUs increase?\n\n(5) What is the overhead of recompute step compared to the end-to-end training runtime? \n\n(6) In Figure 2-Middle, I suspect if you include the data for larger block sizes the number of HBM accesses would increase due to the limited on-chip resources. Is my understanding correct? \n\n(7) In addition, in the most recent release of [MLPerf](https://mlcommons.org/en/training-normal-20/), the training runtime has been improved across all the models. I understand that the NeurIPS submission deadline is before the most recent MLPerf release, but I am wondering if you could provide some insights on whether FlashAttention would outperform the SOTA implementations?\n\n(8) This is more a general comment about transformer models and the contribution of each layer to the overall runtime. Even for reasonably large sequence lengths, FF layers are still the main contributor to the end-to-end runtime. I am wondering if we are targeting the right problem to reduce the runtime of transformer models (I understand that memory consumption is still a large part of self-attention layers, but I doubt that this statement is correct for their runtime).\n\n**Minor**\n\n- There is a missing reference in p. 23. \n- In some of the tables, the accuracy of FlashAttention is lower than SOTA. Nothing major, but would be good to make sure that in all cases, the final accuracy matches the SOTA implementation. This would provide a better head-to-head comparison.", "limitations": "The authors summarizes some of the main limitations of this work in a dedicates section. As authors clearly mentioned, this is not a one-fit-all solutions for all the platforms and models. In this current state, the engineering efforts to adopt such approach may not be transferrable to other platforms and ML libraries. ", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n- The paper consider both forward pass and backward pass in the results that makes it more interesting compared to recent fusion/tiling approaches.\n\n- The paper compares with the hand-tuned MLPerf results (although it is old by now) and show up to 14% improvement for one transformer model.\n\n**Weaknesses**\n\n- Fusion/tiling is not new and has been extensively explored in the literature. It is not clear what the contributions of this work are over previous methods.\n\n- The paper seems to merely provide an engineering solution, merging multiple known approaches (fusion, tiling, algebraic aggregation) --- some already applied to transformer models.\n\n- The paper needs to perform head-to-head comparison with existing fusion techniques for transformers.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "YZxfMlFlZwK", "reviewer": "Reviewer_uahY", "summary": "This paper puts forward the hardware performance perspective in the design of attention algorithms. Specifically, it proposes to use established techniques such as tiling and re-computation to reduce the amount of high-bandwidth memory accesses for exact and approximate attention in transformers which is shown to be the main bottleneck in existing methods. Evaluation demonstrates that the proposed optimization enables training BERT and GPT-2 style architectures faster and achieves new state-of-the-art in the long range benchmark. Apart from training efficiency, the proposed attention improves performance in language modeling, long-document classification and it is the first result that is better than random in Path-X challenge.", "questions": "What is the runtime memory and decoding speed during training/inference for the models and baselines in Tables 1, 2, 4 and 5? An investigation beyond the runtime and memory of forward/backward pass of an individual attention layer is missing.  \n\nThe introduction does not cover the complete view of prior work. Some emphasized the fact that intermediate results are often written into GPU memory and proposed unbiased approximations for iterative computation of attention e.g. kernel-based ones to avoid that and make attention compute-bound instead of memory-bound. Also, reducing the objective to reduce quadratic compute complexity is still important for scaling to extremely long sequences as this paper shows (e.g. Table 13) which is not reflected well in the text.  \n\nDo the authors plan to provide implementations for GPT-2 and BERT or the provided kernel would work?  Since it was highlighted that implementing a custom kernel is not straightforward this might be a concern reproducibility. \n\nThe algorithm appears to be applicable to pretrained models with non-optimized exact attention and does not necessarily require training the model from scratch. Is that the case and is there any evidence to support it?  If not please clarify this in the text. \n\nFor GPT-2 style models, could the authors explain how the causal masking is handled? It would be useful to include a description in section 3 for completeness and extra clarity. ", "limitations": "Yes. Limitations require a bit more discussion (see comments above). ", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- In terms of quality, the exposition of ideas was clear and provided sufficient background for modern hardware which was necessary for the reader to follow. Writing overall was of high quality and the experimentation/proofs thorough. \n- The paper brings awareness about reducing the number of memory accesses for computing exact attention which has been considered not scalable because of its quadratic complexity and memory access handling of existing implementations.\n- The key idea of the algorithm is to compute attention sequentially by loading into on-chip memory blocks of Q, K, V and store the partial result along with easy-to-store statistics required to compute the final result into standard accelerator memory. This is based on the observation that modern accelerators have improved a lot in speed and access to standard memory can become a bottleneck for several standard operations used in deep learning (softmax, batch norm, etc). This kind of optimization is applicable to approximate attention as well and could help improve the efficiency of other operators in the future. \n-  Experiments are well designed and provide evidence of the superiority of the proposed algorithm for both exact and approximate attention methods. It is quite important that this is the first attention that is exact and can extract information from such long contexts in high resolution and improves over all existing alternatives in tasks that require longer context. Should next generation of deep learning frameworks adopt such optimizations the impact can be big for scaling transformers to very long contexts without trading off performance. \n\n**Weaknesses** \n- Framing and motivation presents previous efforts as unaware of hardware performance and that they fail to speed up over existing transformers in terms of wall-clock time which is not universally true. \n-  The main concern is how much hardware/deep learning framework/experiment-specific the low-level optimization results are. Also, the memory and speed numbers are thoroughly reported for attention layer but they are mostly missing for full models (except training speed).\n- The code is provided which is helpful but implementations for big pretrained models such as BERT and GPT-2 are not available which might be an issue for reproducing results.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "NMXby1GMsJn", "reviewer": "Reviewer_5Whx", "summary": "The paper proposes an algorithm to compute exact attention that reduces memory data transfers. \nThey tile computation of softmax, which is a major bottleneck in running transformers. \nAnd they recompute attention matrix to avoid storing it in memory.\nThey also propose block-sparse attention in which they eliminate zero blocks in model with high sparsity.\nThis improves training time and produces more accurate models because it can model longer sequence.\nThey provide IO-complexity for the algorithm.\nPrevious methods that change the attention module and use approximate function of softmax. \nThe proposed method was evaluated on transformer models (e.g: BERT, GPT2). They show improved training throughput and ability to use longer sequence lengths and showcase this in Path-X and Path-256 challenges.\n", "questions": "Some suggestions and questions that could improve the paper are:\n\nRemains to be seem if the method is relevant and if can be applied to mobile variation of transformers (mobile-ViT) and Vision transformers.\n\nThe effectiveness of the method is bounded by hardware constraints (SRAM size). It would interesting to see a roofline plot to demonstrate the compute-bound and memory access trade-off with and without flash-attention.  \n\nEvaluation or comment on other memory technologies other than HBM and GPUs to expand coverage of this work\n\nCompilers (e.g: TVM) have mechanisms for tiling and adding passes. It may be interesting to see if other optimizations in TVM could be applied here and further improve Flashattention", "limitations": "Yes. the paper appropriately discuss limitations and social impact", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros:\nThe paper tackles an important problem and demonstrates algorithm optimizations to improve training runtime and long-range accuracy of relevant transformer models.\n\nThe claims are supported by evaluations\n\nOverall the paper is clear\n\nCons:\nThe techniques: tiling and recomputation are known algorithms that is applied into the new transformers workload. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "KRQfEiBB7mu", "reviewer": "Reviewer_ucdo", "summary": "## Summary:\nProblem: The paper addresses the problem of large memory requirements with high compute complexities of transformer models.\nSolution: the paper attributes the above problem to the missing principle of IO-awareness, where the auditing is missing between different levels of GPU memory hierarchy. Which is a fair finding since the transformers are not designed with that in mind and moreover the GPUs at that time may not be having such complex memory hierarchies either or they may be evolving by then. The proposed solution is FLASHATTENTION\nResults: The paper reports some exciting improvements to the speedup while also helping to improve the quality of models with the ability to process longer sequences. \n\n## Distinction from state-of-the-art:\nThe paper clearly distinguishes from the state-of-the-art methods in the sub-domain and positions within the body of the knowledge. \n", "questions": "Please refer to the strengths and weaknesses above.", "limitations": "The paper describes the limitations to a satisfactory level.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Pros:\n- The paper is written well, easy to understand and follow the concepts.\n\n- The novelty of the paper is in identifying the underlying hardware causes of long time consumption by transformers and proposing the optimizations in the form of efficient data movement.\n\n- The paper is a good fit for machine learning hardware software codesign in many respects.\n\n## Cons:\n\n- Line 54, typo, “out algorithm “ should be “our algorithm”\n\n- One of the recent improvements are warp level matrix multiplications named as wmma in tensor core api. Since transformers involve a number of cascaded matrix multiplications starting from the input layer, wondering if you can compare the proposed approach with some of those mma/wmma based optimizations.\n\n- I understand the focus of the paper is to perform efficient computations on GPU architectures. On the contrary, what if we  are operating on a resource constraint device (only in terms of memory but not that much on compute) where the SRAM size is not as big as that on the GPUs? Then, this approach looks to perform worse than the standard attention due to the N^2d^2. However, this can be out of the scope but can be listed as a limitation.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept