FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
Authors: Tri Dao
Venue: ICLR 2024 poster (2024)
Keywords: Attention
GPUs
LLM training
LLM inference
Primary area: infrastructure, software libraries, hardware, etc.
TL;DR: By reducing the non-matmul FLOPs and better work partitioning, we speed up FlashAttention by 2x, allowing us to train language models with 2x longer sequences.

Abstract: Scaling Transformers to longer sequence lengths has been a major problem in the last several years, promising to improve performance in language modeling and high-resolution image understanding, as well as to unlock new applications in code, audio, and video generation. The attention layer is the main bottleneck in scaling to longer sequences, as its runtime and memory increase quadratically in the sequence length. FlashAttention [5] exploits the asymmetric GPU memory hierarchy to bring significant memory saving (linear instead of quadratic) and runtime speedup (2-4× compared to optimized baselines), with no approximation. However, FlashAttention is still not nearly as fast as optimized matrix-multiply (GEMM) operations, reaching only 25-40% of the theoretical maximum FLOPs/s. We observe that the inefficiency is due to suboptimal work partitioning between different thread blocks and warps on the GPU, causing either low-occupancy or unnecessary shared memory reads/writes. We propose FlashAttention-2, with better work partitioning to address these issues. In particular, we (1) tweak the algorithm to reduce the number of non-matmul FLOPs (2) parallelize the attention computation, even for a single head, across different thread blocks to increase occupancy, and (3) within each thread block, distribute the work between warps to reduce communication through shared memory. These yield around 2× speedup compared to FlashAttention, reaching 50-73% of the theoretical maximum FLOPs/s on A100 and getting close to the efficiency of GEMM operations. We empirically validate that when used end-to-end to train GPT-style models, FlashAttention-2 reaches training speed of up to 225 TFLOPs/s per A100 GPU (72% model FLOPs utilization).

FLASHATTENTION-2: FASTER ATTENTION WITH BETTER PARALLELISM AND WORK PARTITIONING
Tri Dao!2 12
! 'Department of Computer Science, Princeton University Department of Computer Science, Stanford University tr tridao@pri oceton. edu
ABSTRACT
Scaling Transformers to longer sequence lengths has been a major problem in the last several years, promising to improve performance in language modeling and high- resolution image understanding, as well as to unlock new applications in code, audio, and video generation. The attention layer is the main bottleneck in scaling to longer sequences, as its runtime and memory increase quadratically in the sequence length. FLASHATTENTION (Dao et al. 2022) exploits the asymmetric GPU memory hier- archy to bring significant memory saving (linear instead of quadratic) and runtime speedup (2-4x compared to optimized baselines), with no approximation. However, FLASHATTENTION is still not nearly as fast as optimized matrix- multiply (GEMM) operations, reaching only 25-40% of the theoretical maximum FLOPs/s. We observe that the inefficiency due to suboptimal work partitioning between different thread blocks and warps on the GPU, causing either low- low-occypancy or unnecessary shared memory reads/writes. We propose FLASHATTENTION- 2, with better work partition- ing to address these issues. In particular, we (1) tweak the algorithm to reduce the number of non- matmul FLOPs (2) parallelize the attention computation, even for a single head, across different thread blocks to increase occupancy, and (3) within each thread block, distribute the work between warps to reduce communication through shared memory. These yield around 2X speedup compared to FLASHATTENTION reaching 50-73% of the thcoretical maximum FLOPs/s on 1 and getting close to the efficiency of GEMM operations. We empirically validate that when used end-to-end train GPT-style models, PLASHAATEENTIO-22 reaches training speed of up to 225 TFLOPs/s per A100 GPU (72% model FLOPs utilization).
INTRODUCTION
Scaling up the context length of Transformers (Vaswani et al. 2017) is challenge, since the attention layer at their heart has runtime and memory requirements quadratic in the input sequence length. Ideally, we would like to go beyond the standard 2k sequence length limit to train models to understand books, high resolution images, and long-form videos. Just within the last year, there have been several language models with much longer context than before: GPT-4 (0penA][2023) with context length 32k, MosaicML's MPT with context length 65k, and Anthropic's Claude with context length 100k. Emerging use cases such as long document querying and story writing have demonstrated need for models with such long context.
To reduce the computational requirement of attention such long context, there have been numerous methods proposed to approximate attention (Kitaev et al., 2020 [Roy et al. 2021 Wang et al. 2020
Katharopoulos al. 2020 Choromanski al.) [2200] Beltagy et al. [020 Zaheer et al. 2020 Chen
et al. 2021). Though these methods have seen some use cases, as far as we know, most large- scale training runs still use standard attention. Motivated by this, Dao et al. (2022) proposed reorder the attention computation and leverages classical techniques (tiling, recomputation) to significantly speed it up and reduce memory usage from quadratic to linear in sequence length. This yields 2-4x wall-clock time speedup over optimized baselines, up to 10-20x memory saving, with no approximation, and as a result FLASHArTENTION has seen wide adoption in large-scale training and inference of Transformers.
However, context length increases even more, FLASHATTENTION is still not nearly as efficient as other primitives such as matrix-multiply (GEMM). In particular, while FLASHATTENTION is already 2-4x faster than standard attention implementation, the forward pass only reaches 30-50% of the

---

theoretical maximum FLOPs/s of the device (Fig. 6 while the backward pass is even challenging, reaching only 25-35% % maximum throughput on A 100 GPU (Fig. F....7.. In contrast, optimized GEMM reach up to 80-90% of the theoretical maximum device throughput. Through careful profiling, still has suboptimal work partitioning between different thread blocks and warps the GPU, causing either low-occupancy unnecessary shared memory reads/writes.
observe that
Building on ATTENTIN,, we propose FLASHATTENTION- with better parallelism and work partitioning to address these challenges.
1. In Section 3.1 we tweak the algorithms to reduce the number of non-matmul FLOPs while not changing the output. While the non-matmul FLOPs only account for small fraction of the total FLOPs, they take longer to perform as GPUs have specialized units for matrix multiply, and as result the matmul throughput can be up to 16x higher than non-man matmul throughput. It is thus important to reduce non mon-matmul FLOPs and spend as much time as possible doing matmul FLOPs.
2. We propose to parallelize both the forward pass and backward pass along the sequence length dimen- sion, in addition to the batch and number of heads dimension. This increases occupancy (utilization of GPU resources) in the case where the sequences are long (and hence batch size is often small),
3. Even within one block of attention computation, we partition the work between different warps of a thread block to reduce communication and shared memory reads/writes.
In Section 4 we empirically validate that TLASHATTENTTIN
yields significant speedup compared to even FLASHTTTENIION. Benchmarks on different settings (with or without causal mask, different head dimensions) show that FLASHATTENIIO achieves around 2x speedup TION, reaching up to 73% of the theoretical max throughput in the forward pass, and up to 63% of the theoretical max throughput in the backward pass. During LLM inference, FLASHATTENTION-22 kernel is up to 7x faster than the attention kernel from PasterTransforeet When used end and to train GPT-style models, reach training speed of up 225 TFLOPs/s per A100 GPU.
FLASHATTEN-
2 BACKGROUND
We provide some background on the performance characteristics and execution model of GPUs. We also describe the standard implementation of attention, as well as FLASHATTENTION.
2.1 HARDWARE CHARACTERISTICS
GPU performance characteristics. The GPU consists of compute elements (e.g.. floating point arith- metic units) and memory hierarchy. Most modern GPUs contain specialized units to accelerate matrix multiply in cow-prccision (c.g.. Tensor Cores on Nvidia GPUs for FP16/BFII 16 matrix multiply). The memory hierarchy comprise of high bandwidth memory (HBM), and on-chip SRAM (aka shared mem- ory), As an example, the A 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 (Jia et al. 2018 Jia and Van Sandi 2021)) ). As the L2 cache is not directly controllable by the programmer, we focus on the HBM and SRAM for the purpose of this discussion.
Execution Model. GPUs have a massive number of threads execute an operation (called kernel). Threads are organized into thread blocks, which are scheduled to run on streaming multiprocessors (SMs). Within each thread blocks, threads are grouped into warps (a group of 32 threads). Threads within a warp can communicate by fast shuffle instructions or cooperate to perform matrix multiply. Warps within thread block can communicate by reading from writing shared memory. Each kernel loads inputs from HBM to registers and SRAM, computes, then writes outputs to HBM.
2.2 STANDARD ATTENTION IMPLEMENTATION
Given input sequences Q.K, E Nxd where is the sequence length and is the head dimension, we want to compute the attention output O Nxd.
S=QK e R NxN
P=
O=PVERNKC Nxd
where softmax is applied row-wise!! For multi-head attention (MHA), this same computation is performed in parallel across many heads, and parallel over the batch dimension (number of input sequences in a batch).
I For clarity of exposition, omit the scaling of QKT (typically by 1/d). and optionally elementwise masking S and/or dropout applicd to
2

---

The backward pass of attention proceeds as follows. Let d0 E be the gradient of O with respect to some loss function. Then by the chain rule (aka backpropagation):
=V=Pa P= Xxd
dP=dOV ERNNN
dS = ddoottmaaxdPP
dQ=dSKERW>d dk=dSTQER
where dsoftmax is the gradient (backward pass) of softmax applied row-wise. One can work out that if softmax(s) s) for some vector and then with output gradient dp, the input gradient ds (dia(()))ppp))e
Standard attention implementations materialize the matrices and P to HBM, which takes O(N²) memory. Often N > (typically N is on the order of 1k-8k and d is around 64-128). The standard attention implementation (1) calls the matrix multiply (GEMM) subroutine to multiply writes the result to HBM, then (2) loads S from HBM to compute softmax and write the result P to HBM, and finally (3) calls GEMM to get O=PV. As most of the operations are bounded by memory bandwidth, the large number of memory accesses translates to slow wall-clock time. Moreover, the required memory is O N2) due to having to materialize and P. Moreover, one has to save P R. NNN for the backward pass to compute the gradients.
QK'
2.3 FLASHATTENTION
To speed up attention hardware accelerators such as GPU, (Dao et al. (022) proposes an algorithm to reduce the memory reads/writes while maintaining the same output (without approximation).
2.3.1 FORWARD PASS
FLASHATTENTION applies the classical technique of tiling to reduce memory IOs, by (1) loading blocks of inputs from HBM to SRAM, (2) computing attention with respect that block. and then (3) updating the output without writing the large intermediate matrices and to HBM. As the softmax couples entire rows blocks of row, online softmax (Milakov and Gimelshein 2018 Rabe and Staats 202I) can split the attention computation into blocks, and rescale the output each block to finally get the right result (with no approximation). By significantly reducing the amount of memory reads/writes, FLASHIATTENIION yields 2-4x wall-clock speedup over optimized baseline attention implementations.
We describe the online softmax technique (Milakov and Gimelshein, 2018) and how it is used in attention (Rabe and Staats 2021). For simplicity, consider just row block of the attention matrix S, of the form [$()) ((2) for some matrices {(1), ,S(2) e RBBBB. where B, and Be are the row and column block sizes. We want to compute softmax of this row block and multiply with the value, of
(1)
the form
for some matrices ((1) (() R Bexd Standard softmax would compute:
(2)
m = max "ttttmmaax("...o.noaaas rowmax(((
(roowumm e ((1) 3m +rowsuowmmmmmmmmm -m) ER.
P=
(p())
p(2) =diag(() ( S E x2B.
(
0=
(P())
2(2
= ddaa(())
( (1
+e
m mm B.xd
(
Online softmax instead computes "local" softmax with respect to each block and rescale get the right output at the end:
m (1) rowmax(S(2)) ER
( rrossmm (t")mm" ERB.
(()) diag(e(U)) ( ()
ERB, B.B.
O''p(DD()) () =dia(( mmmmm
mm v(1)
BR
m (2) max(m (rowma((S()))
(2) e mll) mm!)m(!))) ()) ((1) +rowsum( S -m rowsum (")) (
p diag (( 2m
0 () diag(e 10)0 =daa (01mm +diag(
We show how FLASH uses online softmax to enable tiling (Fig. to reduce memory reads/writes.
3

---

- Loooo
gs o (
Output
ymu
gm
- (iitsat A
g = pppspi
Feecial sornec decombiator
Figure 1: Diagram of how FLASHARTENTION forward pass is performed, when the key K is partitioned into two blocks and the value V is also partitioned into two blocks. By computing attention with respect to each block and rescaling the output, get the right answer at the end, while avoiding expensive memory reads/ writes of the intermediate matrices and P. We simplify the diagram, omitting the step in softmax that subtracts each element by the row-wise max.
2.3.2 BACKWARD PASS
In the backward pass, by re-computing the values of the attention matrices S and P once blocks of inputs Q, K. V are already loaded to SRAM, FLASHATTENNTIO avoids having to store large intermediate values. By not having save the large matrices S and P of size NxN, FLASHATTENTION yields 10-20x memory saving depending on sequence length (memory required linear sequence length N instead of quadratic). The backward pass also achieves 2-4x wall-clock speedup due to reduce memory reads/writes.
The backward pass applies tiling to the equations in Section 2.2 Though the backward pass is simpler than the forward pass conceptually (there is softmax rescaling). the implementation is significantly more involved. This is because there are more values to be kept in SRAM to perform matrix multiples in the backward pass, compared to just 2 matrix multiples in the forward pass.
FLASHATTENTION-2
ALGORITHM, PARALLELISM, AND WORK PARTITIONING
We describe the FLASHATTENTION- algorithm, which includes several tweaks to FLASHATTENTION to reduce the number of non- matmul FLOPs. We then describe how to parallelize the computation on different thread blocks to make full use the GPU resources. Finally we describe we partition the work between different warps within one thread block to reduce the amount of shared memory access. These improvements lead to 2-3x speedup as validated in Section
3.1 ALGORITHM
We tweak the algorithm from FLASHATTENTION to reduce the number of non-matmul FLOPs. This is because modern GPUs have specialized compute units (e.g.. Tensor Cores on Nvidia GPUs) that makes matmul much faster. As an example, the A 100 GPU has a max theoretical throughput of 312 TFLOPs/s of FP16/BF16 matmul, but only 19.5 TFLOPs/s of non-matmul FP32. Another way to think about this is that each non- non-matmul FLOP 16x more expensive than matmul FLOP. To maintain high throughput (c.g., more than 50% of the maximum theoretical TFLOPs/s), we want to spend as much time on matmul FLOPs as possible.
3.1.1 FORWARD PASS
We revisit the online softmax trick as shown in Section 2.3 and make two minor tweaks to reduce non-matmul FLOPs:
1. We do not have to rescale both terms of the output update by diag(
-12 8(2) (m (2) 0 (2) =1) (2(2) (0()))))))))))))))))))) +diag( (1))) (12))
4

---

We can instead maintain "un-scaled" version of ((2) and keep around the statistics ((2);
(2) d) +e - (
Only at the every end the loop do we scale the final Ô (last) by (liag(((aan))) (last) to get the right output.
2.
We do not have to save both the max m() and the sum of exponentials e for the backward pass.
We only need to store the logsumexp L =m==m()) + +log( (()))) (())..
In the simple case of 2 blocks in Section
2.3 the online softmax trick now becomes: rowsum(e 8(D) )ERR
m (1) =rowmax($ E
0 = ( -m! ((1) R
Bd
mmaxmm (S
((2) = -m 22) ((1) +rowsum (12) rowsum e 8 +rowsum( (1))
( diag( e -m()
( diag(e
e m1))) -m 10 ((2) v = (1l ( (m
0 () diag( ((2) (2) =0
We describe the full FLASHATTENTION- forward pass in Algorithm
Algorithm 1 FLASHATTTENTINN forward pass
Require:
Matrices V.K.VER in HBM, block sizes Be, Br. blocks Q2...., of size B, x each, and divide K, V in to T. R2
M 1: Divide Q into T, =
Br
blocks KI....,y,, and 1 Vi........ of size B.xd each.
2: Divide the output O Nxd into T, blocks Oi....,yy, size B, x d each, and divide the logsumexp
L into T, blocks LI, of size B, each.
3: for Ti do
4: Load Q from HBM to on-chip SRAM.
5: On chip, initialize 0,00 (0) B, e. =()=, R m (0) = =(-00)B, ERS
6:
for jorjsje STc do
7:
Load V, from HBM to on-chip SRAM.
On chip, compute S."= K ER BBB.
8:
On chip, compute m; ( max m (j-1) rrowmax
(()) ((
p100 = exp(s()) -m
9:
R B, BB (pointwise), (1ع) = m! -my c (111)) +rowsum (P()) ER?
10:
On chip, compute 0 = diag(e" (
in (-1)) "
end for
On chip, compute O, (
12:
13: On chip, compute L; = (Te) +log( (
(TTe) (TT)).
14: Write O to HBM as the i-th block of O.
15: Write L, to HBM as the i-th block of L.
16: end for
17: Return the output O and the logsumexp L.
Causal masking.
One common use case of attention is in auto-regressive language modeling, where we need to apply causal mask to the attention matrix S (i.c., any entry Siy with is set to -oo).
1. As FLASHATTENTIO and FLTETTTTEN already operate by blocks, for any blocks where all the column indices are more than the row indices (approximately half of the blocks for large sequence length), we skip the computation of that block. This leads to around 1.7-1.8x speedup compared to attention without the causal mask.
We do not need to apply the causal mask for blocks whose row indices are guaranteed to be strictly less than the column indices. This means that for each row, only need apply causal mask to I block (assuming square block).
Correctness, runtime, and memory requirement. As with FLASHATTENTION, Algorithm returns the correct output O ==softmax (aaimmaax))a)(((( V (with no approximation), using (NV2d FLOPs and requires
5

---

(N) additional memory beyond inputs and output (to store the logsumexp L). The proof is almost the same as the proof of Dao et al. (2022) Theorem 1), so we omit here.
3.1.2 BACKWARD PASS
The backward pass of FLASIAATTENTION-2 is almost the same as that of FLASIIATTENTION. We make minor tweak to only use the row-wise logsumexp L instead of both the row-wise max and row-wise sum of exponentials the softmax. We include the backward pass description in Algorithm for completeness.
Multi-query attention and grouped- query attention. Multi-query attention (MQA) (Shazeer 2019) and grouped-quers attention (GQA) (Ainslie et al. 2023) are variants of attention where multiple heads of query attend to the same head of key and value, in order to reduce the size of KV cache during inference. Instead of having to duplicate the key and value heads for the computation, we implicitly manipulate the indices into the head to perform the same computation. In the backward pass, need to sum the gradients dk and dV across different heads that were implicitly duplicated.
3.2 PARALLELISM
The first version of FLASHATTENTIO parallelizes over batch size and number of heads. We use thread block to process one attention head, and there are overall batch surbenumber of heads thread blocks. Each thread block is scheduled to run streaming multiprocessor (SM), and there are 108 of these SMs on an A100 GPU for example. This scheduling efficient when this number is large (say2 80), since
effectively almost all of the compute resources the GPU.
In the case of long sequences (which usually means small batch sizes or small number of heads), make better use of the multiprocessors on the GPU, we now additionally parallelize over the sequence length dimension. This results in significant speedup for this regime.
Forward pass. We see that the outer loop (over sequence length) embarrassingly parallel, and schedule them on different thread blocks that do not need to communicate with each other. We also parallelize over the batch dimension and number of heads dimension, done in FLASHATTENTION. The increased parallelism over sequence length helps improve occupancy (fraction of GPU resources being used) when the batch size and number of heads small, leading speedup in this case.
Backwar ward pass. Notice that the only shared computation between different column blocks is in update dQ in Algorithm 2 where we need to load dQ, from HBM to SRAM, then on chip, update
dQ; - dQ, +ds () j, and write back HBM. We thus parallelize over the sequence length dimension as well, and schedule thread block for each column block of the backward pass. We use atomic adds communicate between different thread blocks update dQ.
We describe the parallclization scheme Fig
Forward pass
Backward pass
Worker
Worker 2
Worker
Worker 4
Worker
Worker 1
Worker 2
Worker 3
Worker 4
Worker 5
Figure In the forward pass (left), we parallelize the workers (thread blocks) where each worker takes care of a block of rows of the attention matrix. In the backward pass (right), each worker takes care of block of columns of the attention matrix.
Decoding, During LLM inference, most of the time is spent on iterative decoding, where one token predicted at time. The bottleneck for the attention operation during decoding different from that during training or prefill (prompt processing), because the query length is very short (often query length 1 since only the new extra token is attending all the previous tokens, stored in the KV cache). As

---

result, the bottleneck is no longer the read/write of intermediate matrices (the scores QK" and attention probabilities softmax( QK') )) Instead, the bottleneck is to load the KV cache as quickly as possible.
To accommodate this setting, we split the KV cache loading among different thread blocks, to increase occupancy and saturate the HBM bandwidth. However, since the thread blocks cannot easily communicate with each other, we write intermediate results to HBM, then call a separate kernel to reduce the results and produce final output.
3.3
WORK PARTITIONING BETWEEN WARPS
As Section describe how we schedule thread blocks, even within each thread block, we also have to decide how to partition the work between different warps. We typically use 4 or 8 warps per thread block, and the partitioning is described in Fig. 3
Forward pass. For each block, PLASHATTENTION splits K and V across warps while keeping Q accessible by all warps. Each warp multiplies to get slice of QKT, then they need to multiply with a slice of and communicate to add up the result. This is referred to as the "split-K" scheme. However, this is inefficient since all warps need to write their intermediate results out to shared memory, synchronize, then add up the intermediate results. These shared memory reads/writer wrrtte slow down the forward pass in FLASHATTENTION,
we instead split Q across 4 warps while keeping K and V accessiile by all warps. After each warp performs matrix multiply to get a slice of QK', they just need to multiply with their shared slice of V to get their corresponding slice of the output. There is no need for communication between warps. The reduction in shared memory reads/writes yields speedup (Section 4
In FLASIIATTEN ATTIAATTTOONN
Worn
Wars
Merp
Wap3
(a) FLASHATTENTION
(b) FLASHATTENTION--
Figure 3: Work partitioning between different warps in the forward pass Backward pass. Similarly for the backward pass, choose partition the warps to avoid the "split-K" scheme. However, it still requires some synchronization due to the more complicated dependency between all the different inputs and gradients Q.K, V,O.dO,ddl V.O.dd,.d..de.. ddV.d..VV. Nevertheless, avoiding "split-K" reduces shared memory reads/writes and again yields speedup (Section
Tuning block sizes Increasing block sizes generally reduces shared memory loads/stores. but increases the number of registers required and the total amount of shared memory. Past certain block size, register spilling causes significant slowdown, or the amount of shared memory required is larger than what the GPU has available, and the kernel cannot (64,128) x (64,128), depending on the head dimension d and the device shared memory size.
all. Typically we choose blocks of size
We manually tune for each head dimensions since there are essentially only 4 choices for block sizes, but this could benefit from auto-tuning avoid this manual labor. We leave this to future work.
EMPIRICAL VALIDATION
We evaluate the impact of using FLASHATTENTION to train Transformer models.
Benchmarking attention. We measure the runtime of FLASHATTENTION- across different sequence lengths and compare it to a standard implementation in Py Torch, PLASHATTENTION and FLASHATTENTION in Triton. We confirm that FLASHATTENTION- is 1.7-3.0x faster than FLASHATTENTION, 1.3-2.5x faster than FLASHATTENTION in Triton, and 3-10x faster than a standard attention implementation. FLASHATTENTION-2 reaches up to 230 TFLOPs/s. 73% of the theoretical maximum TFLOPs/s on 1 GPUs
End-to-end training speed When used end-to-end to train GPT-style models size 1.3B and 2.7B on sequence lengths either 2k or 8k, FLASHATTENTION yields up to 1.3x speedup compared
7

---

FLASHÄTTENTION and 2.8x speedup compared to baseline without FLASHATTENTION, FLASHATTENTION-2 reaches up to 225 TFLOPs/s (72% model FLOPs utilizzation) per A GPU.
BENCHMARKING ATTENTION FOR TRAINING
We measure the runtime of different attention methods on an 1 80GB SXM4 GPU for different settings (without / with causal mask, head dimension 64 or 128). We report the results in Fig. Fig. and Fig. 7 showing that PLASHATTENTION-2 is around 2x faster than FLASHATTENTION and FLASHÂTTENTION in xformers (the "cutlass" implementation) 1.3-1.5x faster than FLASHATTENTION in Triton the forward pass and around 2x faster in the backward pass. Compared to standard attention implementation in PyTorch, FLASILATTENTION-2 can be up to 10x faster.
PLASHÄTTENTION- is around
Benchmark setting: we vary the sequence length from 512, Ik, Ik,.., 16k, and set batch size so that the total number of tokens is 16k. We set hidden dimension to 2048, and head dimension be either 64 or 128 (i.e., 32 heads or 16 heads). To calculate the FLOPs the forward pass, we use:
4-seqlen" head dimension number of heads.
With causal mask, we divide this number by 2 to account for the fact that approximately only half of the entries are calculated. To get the FLOPs of the backward pass, multiply the forward pass FLOPs by 2.5 (since there are 2 matmuls in the forward pass and 5 matmuls in the backward pass, due to recomputation).
forward backward speed (A100 SXM4)
forward backward speed (A100 SXM4)
Pytorch
Pytorch
Speed (TFLOPS/S)
200
PlashAttantion
PlashAttentinn22
Flashictentio Thitan
aftrmers
262
Speed (TFLOPSIS)
200
Flashittantion Trito
Flasivittationn
xftemers
232
322
3.00
32
300
303
300
07
KAP 7ST4
co1
Sequence length
Sequence length
(a) Without causal mask, head dimension 64
(b) Without causal mask, head dimension 128
Attention forward backward speed (A100 80GB 5XM4)
Attention forward backward speed (A100 80GB 5X044
Pytorch
125
Speed (TFLOP)/S)
Flashlctention F
315
Speed (TFLOPS/S)
Flashiitontion
172
2
250
250
130
:10
30
50
10
14
512
A fk
16k
512
2k
4t
fk
14k
Sequence length
Sequence length
With causal mask, head dimension
(d) With causal mask, head dimension
Figure 4: Attention forward backward speed on A00 GPU
Just running the same implementation on H100 GPUs (using no special instructions to make use of new features such TMA and 4th-gen Tensor Cores), we obtain up 335 TFLOPs/s (Fig. 8 We expect that by using new instructions, we can obtain another 1.5x-2x speedup on H100 GPUs. We leave that future work.
4.2 BENCHMARKING ATTENTION FOR INFERENCE
We benchmark the attention kernel during decoding for the of multi-query attention, where the bot- tleneck is loading the KV cache. In Fig. 5 we see that the attention kernel from ATAAATTTNN up 28x faster than naive implementation in PyTorch, and up to 7x faster than implementation
8

---

from Faster Traktferaasfrmerr This is thanks to better work partitioning where multiple thread blocks loading the KV cache at the same time to saturate HBM bandwidth.
Attention decoding time
Hechiitenion2
Time
(61, 145
(12,
116.
04, 10k)
124
6115
(Batch Sequence length)
Figure 5: Attention decoding time A100 80GB, with hidden dimension 2048 and multi-query atten- tion. The attention kernel from FLTETIATTNNION-22 up to 7x faster than that of PasterTransforme and 28x faster than naive implementation in PyTorch.
4.3 END-TO-END PERFORMANCE
We measure the training throughput of GPT-style models with either 1.3B or 2.7B parameters, 8xA 100 80GB SXM4. As shown in Table FLASHATTENTION-2 yields 2.8x speedup compared
to baseline without FLASHATTENTION and 1.3x speedup compared to FLASHATTENTION, reaching up to 225 TFLOPs/s per A100 GPU.
Note that we calculate the FLOPs by the formula, following Megatron-LM (Shoeybi et al. many other papers and libraries):
2019) (and
6 seqlen number of params+ 12- 22-number of layers- hinden dim dim-seqlen".
The first term accounts for the FLOPs due to weight-input multiplication, and the second term accounts for the FLOPs due to attention. However, one can argue that the second term should be halved, as with causal mask only need compute approximately half the number of elements in attention. We choose to follow the formula from the literature (without dividing the attention FLOPSs by 2) for consistency.
Table 1: Training speed (TFLOPs/s/PPU) of GPT-style models on 8xA1000 GPUs, FLASHATTENTION-
2 reaches up to 225 TFLOPs/s (72% model FLOPs utilization). We compare against baseline running without FLASIIATTENTION.
Model
Without FLASIATTENTION
FLASUATTENTION
FLASILATTENTION-2
GPT3-1.3B 2k context
142 TFLOPs/s
189 TFLOPs/s
196 TFLOPs/s
GPT3-1.3B 8k context
72 TFLOPS/s
170 TFLOPs/s
220 TFLOPs/s
GPT3-2.7B 2k context
149 TFLOPs/s
189 TFLOPs/s
205 TFLOPs/s
GPT3-2.7B 8k context
TFLOPs/s
175 TFLOPs/s
225 TFLOPs/s
DISCUSSION AND FUTURE DIRECTIONS
FLASHATTENTION- is 2x faster than FLASHATTENTIO: which means that train models with 16k longer context for the same price as previously training 8k context model, for the same number of tokens. We excited about how this be used understand long books and reports, high resolution images, audio and video. FLASHATTENTION-2 will also speed up training, finetuning, and inference of existing models.
In the near future, we plan collaborate with researchers and engineers make FlashAttention widely applicable in different kinds of devices (e.g., H100 GPUs, AMD GPUs), as well as new data types such FP8. As an immediate next step, plan to optimize Attentonnn for H100 GPUs to use new hardware features (TMA, 4th-gen Tensor Cores, fp8). Combining the low-level optimizations in Flash Attention-2 with high-level algorithmic changes (e.g., local, dilated, block-sparse attention) could allow us to train AI models with much longer context. We are also excited to work with compiler researchers make these optimization techniques easily programmable.

---

ACKNOWLEDGMENTS
We thank Phil Tillet and Daniel Haziza, who have implemented versions of FLASIIATTENTION in Triton (Tillet al. 2019) and the xformers library (Lefaudeux al. [022] FLASIIATTENTION-2 was motivated by exchange ideas between different ways that attention could implemented. We are grateful to the Nvidia CUTLASS team (especially Vijay Thakkar, Cris Cecka, Haicheng Wu, and Andrew Kerr) for their CUTLASS library, particular the CUTLASS 3.x release, which provides clean abstractions and powerful building blocks for the implementation of FLASHATTENTION-2. We thank Driss Guessous for integrating ATTENTTNNNTION to PyTorch. FLASSHATENNTIO-2 has benefited from helpful discussions with Phil Wang, Markus Rabe, James Bradbury, Young- Joung-Jun Ko, Julien Launay, Daniel Hesslow, Michaöl Benesty, Horace He, Ashish Vaswani, and Erich Elsen. Thanks to Stanford CRFM and Stanford NLP for the compute support. We thank Dan Fu and Christopher Ré for their collaboration, constructive feedback, and constant encouragement on this line of work of designing serdwware-fffiieee algorithms. We thank Albert Gu and Beidi Chen for their helpful suggestions on early drafts of this paper.
REFERENCES
Joshua Ainslie, James Lee- cerph, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305. 3245, 2023.
Iz Beltagy, Matthew Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint (oXi::2004.0555555 2020.
Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. Advances in Neural Information Processing Systems (NeurIPS), 2021.
Krzysztof Marcin Choromanski, Valerii Likhosherstov. David Dohan, Xingyou Song, Andreca Gane, Tamas Sarlos, Peter Hawkins, Jared Quiney Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In International Conference on Learning Representations (ICLR), 2020.
Tri Dao, Daniel Y. Fu. Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with 10-awareness. In Advances in Neural Information Processing Systems, 2022.
Zhe Jia and Peter Van Sandt. Dissecting the Ampere GPU architecture via microhenchmarkingg GPU Technology Conference, 2021.
Zhe Jia, Marco Maggioni, Benjamin Staiger, and Daniele Scarpazza. Dissecting the nvidia Volta GPU architecture via microbenchmarking, arXiv preprint /004.00800000266 2018.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers
RNNs: Fast autoregressive transformers with linear attention. International Conference on Machine Learning, pages 5156-5165. PMLR, 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In The International Conference on Machine Learning (ICML), 2020.
Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, and Daniel Haziza. xformers: A modular and hackable transformer modelling library. https::/ wttp:/oooooccoeeeeeeeeeeeeeeeeee www.e/ooooreaarcc//ooooooooo
2022.
Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv preprins arXin:1805.02887.. 2018.
OpenAI. Gpt-4 technical report. ArXiv, abss23333//03 2023.
Markus N Rabe and Charles Staats. Self-attention does not need 0(z2) memory. arXiv preprint arXiv:2/ 12.06622, 2021.
10

---

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.
Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150, 2019.
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 1 1909008099,,,,,, 2019.
Philippe Tillet, Hsiang-Tssin Tuung Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPIAN International Workshop on Machine Learning and Programming Languages, pages 10-19, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Iliia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Sinong Wang, Belinda Zii, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006. 04768, 2020.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33, 2020.

---

A FLASHATTENTION-2 BACKWARD PASS
Algorithm 2 FLASHATTENTION Backward Pass
Require: Matrices Q.K. exx Nxd HBM, vector LERN in HBM, block sizes Be, Br blocks Qu....Quu Qr, of size B, x each, and divide K, V in to to B.
1: Divide Q into T, - B,
N blocks Ky. and I...., Vr., size x. each.
N
2: Divide into T, blocks Or, of size x, each, divide dO into T, blocks do,,...,dOr, of size B, x each, and divide into T, blocks Li....Lr, size By each.
3: Initialize dQ= NO) in HBM and divide it into T, blocks dQ......drr, of size B, x d each.
Divide dK,dV e R Nxd in to T. blocks dK,,,., and dV T of size Be xd each.
4: Compute D D=rowsun( dOo 0) ER" (pointwise multiply), write D to HBM and divide it into T, blocks D Dr, of size B, each.
5: forjjjs ST, do
6: Load K,.V, from HBM to on-chip SRAM.
7: Initialize (V
SRAM.
8: for ST
9:
Load d0 from HBM on-chip SRAM.
On chip, compute 5 = BRB.
11:
On chip, compute P/ =ex((s; L;) ER BR..H
On chip, compute dV dP Td,, ER Bexd
On chip, compute dP /) =do, _ ER
13:
On chip, compute ds (1 (dP ED)
14:
15:
Load dQ, from HBM to SRAM, then on chip, update dQ, - dQ, K, E Bxd and write back HBM.
On chip, compute dk dk dds+d(()) Q ER
17:
end for
18:
Write dKj.dV, to HBM.
19:
end for
20:
Return dQ.dK,dV.
B BENCHMARKING ATTENTION ON A 100 AND H100
12

---

speed (A1 8OGB 5XM44
speed (A100 5X044)
Pytsees
Pytorm
Plachkctentios
Plahnittattonn
Speed (TFLOP%/5)
200
Fisshictention
Speed (TFLOPs/S)
200
xieemers
Flashettnntin
FlastActaniios
2.50
3.00
10%
100
200
6,72
c0o
coo
Sequence length
Sequence length
(a) Without causal mask, head dimension 64
(b) Without causal mask, head dimension 128
Attantion forward speed (A100 BOGB SXM4)
Attantion forward speed (A100 BOGB SXM4)
Flashictention
Plashittontion
130
Speed (TFL))s
Flachéetention Thiton FlashAttentin22
mers
267
100
Speed (TFLOPs/5)
Flashletention Thhton
xfommers
PlsttAttzntion 2
3.5
344
3190
200
200
12
512
1k
4k
8
16x
512
4k
8k
16k
Sequence length
Sequence length
(c) With causal mask, head dimension 64
(d) With causal mask, head dimension 128
Figure 6: Attention forward speed on A A100 GPU
Attention speed SXM4)
Attention speed SXM4)
Pylsech
Pytorch
Plashiittntion
PlasMittettion
Speed (TFLOPs/S)
200
FlachAttention Thilon
aformers
170
Speed (TFLOPs/S)
200
xfoemers
Flachiitention
FisshAttattonn
FissMettntion 2
3.5
3150
385
100
100
co2
com
Sequence length
Sequence length
Without causal mask, head dimension 64
(b) Without causal mask, head dimension 128
Attention backward speed (A300 80EB SKM4)
Attention backword sp speed (A100 80GB SXM4)
Pytoreh
Speed (PL))))))
Flasháttention22
FlashAttentian
mers
310
106
Speed (FL))))
Fashhlttnntion
Flachiitention
350
2:50
111
ESRY
37gh
58
2000
6
50-
50
41
512
8k
15k
512
5k
15k
Sequence length
Sequence length
(c) With causal mask, head dimension 64
(d) With causal mask, head dimension 128
Figure 7: Attention backward speed A100 GPU
13

---

speed (4100 SXMS)
Attention speed (H100 SXMS)
Pytoech
Pytorch
Flachirttntion
FlachVetention
332
Flashittontion2
FlachAttintion22
Speed (TFLOPsS)
300
206
Speed LTFLOSSS
300
100
1.2
200
200
00W
004
Sequence length
Sequence length
Without causal mask, head dimension 64
(b) Without causal mask, head dimension 128
Attention forward backward speed (H100 8OGB SXMS)
Attention forward backward speed (H160 8OGB 5XM5)
F
Flashittartion
Speed (FLOPs)S)
300
Speed (TFLOPSIS)
300
205
200
200
123
130
500
300
300
12
512
1x
2x
4k
36c
512
2k
4k
16k
Sequence
length
Sequence
length
(c) With causal mask, head dimension 64
(d) With causal mask, head dimension 128
Figure 8: Attention forward backward speed on H100 GPU
14


Reviews:
[{"review_id": "QeaOL3CvbI", "reviewer": "Reviewer_AK2B", "summary": "Scaling Transformers for longer sequences holds the promise of enhancing language modeling and understanding complex inputs, but is hindered by the attention layer's quadratic scaling in memory and runtime. FlashAttention has mitigated this by bringing linear memory usage and considerable runtime speedup, yet it still lags behind the efficiency of optimized matrix multiplication operations. To address this, FlashAttention-2 is introduced with improved work partitioning, yielding a significant speedup and reaching closer to the efficiency of matrix multiply (GEMM) operations. Empirical validation shows that FlashAttention-2 significantly increases the training speed of GPT-style models on both A100 and H100 GPUs.", "questions": "Thank you for submitting to ICLR 2024. FlashAttention-2 is a very useful artifact that has the potential to benefit both researchers and practitioners, and the proposed optimization techniques appear sound.\n\nHere are my questions I would like the authors to answer:\n* Perhaps the most significant omission in this paper is the lack of an ablation study. This makes it challenging to discern the contributions of individual optimizations. Among the proposed optimizations, which one has the highest impact?\n* In Section 3.1, is the technique of skipping blocks for \"causal masking\" also applied to FlashAttention? As the authors mention, this technique can be applied to both FlashAttention and FlashAttention-2, and I am curious about how the application of this technique would affect the performance gap between FlashAttention and FlashAttention-2 if it had not been applied to FlashAttention.\n* In Section 3.3, what is the performance impact of \"tuning block sizes\"? Was the same level of parameter tuning effort applied to FlashAttention? My question concerns the extent to which the performance gains over FlashAttention can be attributed to algorithmic improvements versus additional tuning effort, with the latter being less significant.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "* The proposed platform-specific optimizations are clever and sound.\n* The resulting software artifacts are useful and have has the potential to benefit both researchers and practitioners.", "weaknesses": "* The work is mostly engineering-focused, with several \"tweaks\" made to FlashAttention.\n* The performance gains are relatively marginal, especially when compared to those of the original FlashAttention over the baseline.\n* The absence of an ablation study makes it difficult to pinpoint the exact sources of efficiency.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "RchnThk41v", "reviewer": "Reviewer_RLZf", "summary": "This paper presents improvements to FlashAttention (Dao et al., 2022), a established method for efficiently computing attention through fused ops. These improvements are designed specifically for better parallelism and work partitioning in GPUs, resulting in the development of FlashAttention v2. Performance benchmarks have been conducted for both training and inference phases. Additionally, the authors provide comprehensive results from training end-to-end GPT-style models with 1.3 billion and 2.7 billion parameters and 2k and 8k context sizes.", "questions": "Which GPU architectures currently support FlashAttention v2?\nWhat are the minimum requirements for its use?\nWhat modifications are necessary to adapt FlashAttention v2 for use with relative positional encoding methods (e.g., RoPE and ALiBi)?\nDoes FlashAttention v2 offer compatibility with sparse block masks (as in v1)?", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "There are many things to like in this paper, such as:\n\n1. The paper is well-written\n2. The improvements are well-explained and justified\n3. The paper covers both training and inference time optimization\n4. The results are encouraging\n\nIn summary, I expect widespread adoption of FlashAttention v2 within the community. Furthermore, the methods proposed and utilized in this paper could inspire the creation of more efficient components in machine learning.", "weaknesses": "FlashAttention v2 has a notable limitation: it relies on recent, specialized GPU architectures like the A100 (and H100). Additionally, the requirement for custom CUDA kernels adds a layer of complexity. \n\nA small critique is that Figure 3 could benefit from a more descriptive caption.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "UmS7FpjK6j", "reviewer": "Reviewer_hPzR", "summary": "The authors present a new algorithm, FlashAttention-2, which builds in FlashAttention to improve the efficiency of the attention algorithm when executed on GPUs. The authors focus, in particular, on maximizing the amount of time spent in \"matrix-multiply\" FLOPs, that is, computation that is using matrix multiplication units, which are better lent to GPU hardware given particular division of work amongst warps and loose requirements around shared memory accesses.\n\nThese optimizations include:\n1. Deferring scaling of values in the online softmax computation to further reduce HBM utilization\n2. Save a logsumexp for online softmax backwards rather than max and sums to reduce the memory usage.\n3. The bulk of the changes present in FlashAttention-2 are related to scheduling. Parallelization over the sequence dimension results in better warp occupancy in cases where there are few attention heads or a low batch size. Further, the authors change how the backward pass shares computation in the query derivative update, which also reduce HBM utilization. Splitting the KV cache amongst thread blocks also helps to saturate memory bandwidth. Better partitioning amongst warps, overall, drives better utilization.", "questions": "- Why are the gaps between FA/FA-2 different with Cutlass versus in Triton, if the authors were to speculate? Further, the comparison in section 4.1 can be clarified -- is the assertion that the performance of the vanilla CUDA implementation of FlashAttention and the Cutlass FlashAttention implementation in xformers are congruent?\n- The authors might also consider mentioning in the manuscript what they think trends in changing GPU hardware will mean for FA-2's general direction. Given that HBM bandwidth is not improving as quickly as SM arithmetic latency and that the amount of vram available on GPUs is not increasing, how will the approaches used in FA-2 change in relevance over time?\n- Do you think a compiler could realistically generate FA-2? What sort of cost models might be required?\n- How does FA-2 function when there is no explicit sequence length dimensions? In many training setups, tokens are padded without sequence boundaries, and models learn end-of-sequence tokens implicitly. What is the default behavior in this regime?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The proposed approach shows promising improvements in a performance-critical parts of end-to-end transformer computation.\n- The approach is symmetrically applicable for both training and inference with broad applications in both research and production settings.\n- It is important that engineering contributions such as FlashAttention-2 (which I will refer to FA-2) are part of conference literature, and the attention to detail therein is what drives the solid impact of this work rather than a special \"algorithmic contribution\" in a classical sense.\n- The devil is in the details: careful analysis of scheduling and an intuitive approach to laying out computation drives the approach's strong results.\n- The work builds on an already-strong baseline, FlashAttention.\n- The paper is well-written and organized and clearly lays out the authors' contributions; in particular, the paper is quite accessible to those without a low-level background in machine learning computation or GPU programming.", "weaknesses": "- The baselines benchmarked in the paper can be stronger.\n  - For latency benchmarks in Figure 5, the only baselines are a FasterTransformer and PyTorch. The authors do not consider compilers\n  - Do the PyTorch benchmarks in Figure 5 use CUDA Graphs? PyTorch has significant framework overhead, and CUDA Graphs can give an order of magnitude speedup for some workloads, especially latency-sensitive ones.\n- The above applies more generally to the other evaluation in Section 4; FlashAttention-2 is compared to PyTorch, then implementations with Triton and Cutlass, but not with any other frameworks capable of code generation. For example: while XLA may not be memory-bandwidth-aware by default, it can still generate kernels with fused operators that significantly reduce total memory I/O.\n- While not needing to resort to approximations is a significant advantage of FlashAttention-2, this could be highlighted much more in the manuscript. Section 1 discusses many alternative attention approximations -- even speculation about why these aren't used (i.e. they are riskier when researchers have limited resources and don't adapt as easily) would strengthen the exactness boon of the authors' approach.\n\nSeveral improvements to writing might improve the paper:\n- The constant and equation in general in Sections 4.1 (i.e. 4) is not adequately explained (why is the sequence length squared? why 4?). Clarifying these might help new readers.\n- The usage of \"major problem\" in the first sentence of the abstract is unclear -- it's clear that scaling sequence lengths is difficult; are the authors suggesting the problem is difficult, significant, or both?\n- Section 1: \"However, context length increases\" <-- is missing \"**as** context length increases\"\n- Text in all of the provided diagrams can be made clearer, and the diagrams can be rendered more clearly. It is difficult to read them as is.\n- Section 3.2 might more clearly explain \"prefill\" and \"KV cache\" to readers. While somewhat ubiquitously understood amongst people doing performance engineering for large-scale transformers, some clarification would help the paper flow and increase its accessibility.\n- Figure 2 can be clearer with respect to rows and columns -- this is the attention matrix -- what is its size/can the axes be labeled?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}, {"review_id": "67qRa7n4yf", "reviewer": "Reviewer_Jn9J", "summary": "This paper describes FlashAttention-2 which improves upon FlashAttention by introducing \"tweaks\" to improve performance on GPUs.   The paper claims the tweaks improve performance by increasing occupancy and use of matrix-multiply hardware (tensor cores).  The paper reports a bit under $1.3\\times$ wall clock speedup versus FlashAttention.", "questions": "Regarding the equation at the top of Page 5, I am unclear \"$\\mbox{diag}(l^{(1})^{-1}$\" is to the power -1.  Comparing to the prior equation seems like exponent of -1 should be 1.\n\nI think it would help some readers (like me) understand the contribution a bit better if the paper briefly summarized the key changes in the six (unnumbered) equations on Page 5 that are described as the \"online softmax trick\" versus the six on Page 3.  \n\nHow do the \"tweaks\" in Section 4.1.1 help reduce non-matrixmul FLOPs?  I know a fair amount about tensor cores, but it wasn't obvious to me.\n\nThe paper claims occupancy is increased on Page 6 but it was unclear: (i) what definition of occupancy is being used (GPU resources could mean many things and occupancy often just refers to number of warps that can concurrently run versus max number supported by hardware ); and (ii) whether any measurement has been made to confirm the claimed improvement (e.g., using NVIDIA Parallel Nsight or similar approaches for collecting performance counters).\n\nMuch of Algorithm 1 seems similar to the original FlashAttention.  It may help summarizing which lines are different.  It would also help the reader if there was a summary of which lines lead to the reduction in non-matrixmul FLOPs and improved occupancy.\n\n\"Only at the every end of the\" - typo.\n\nFor the backward pass (Section 3.1.2): It was unclear what the relevance of the paragraph on MQA and GQA is to the changes in FlashAttention-2 versus FlashAttention.  \n\nIn Figure 2, does an uncolored square mean no computation?  Does the backward pass for a given worker start right away or do workers need to synchronize between forward and backward pass?  Do you not need to compute the combined result for the forward pass before you can start the backward pass?    If you do need to wait, then how can one achieve greater than 50% use of peak performance if roughly half the compute cycles are spent waiting for the longest running forward/backward pass thread block to complete?   If you don't need to wait, why not?\n\nI'm not sure how to relate Figure 3 to Algorithm 1 (i.e., which lines it is meant to illustrate).  From the two paragraphs above Figure 3 I get it there are two potential sources of reduced execution time: fewer shared memory accesses and fewer synchronizations (__syncthreads, I assume).  Unclear which of those matters most and why given that shared memory accesses proceed about as fast as register file accesses and synchronization with a thread block is low overhead.  \n\nWhy is FlashAttention (version 1) missing in Figure 5?\n\nAs someone who knows GPUs well, I would have liked to see more performance counter data to backup the claims of the sources of performance improvements.   I understand space is limited in the main text, but in checking the supplemental material, while it is great to see all the code, there appeared to be no PDF providing additional data or details.  Including one might have helped.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Improving training speed of LLMs is of great interest to many.", "weaknesses": "Could do a better good job explaining how a given \"tweak\" helps achieve a given improvement (occupancy, use of tensor cores).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)