FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
Authors: Jay Shah
Ganesh Bikshandi
Ying Zhang
Vijay Thakkar
Pradeep Ramani
Tri Dao
Venue: NeurIPS 2024 spotlight (2024)
Keywords: attention
hardware-aware algorithms
H100
Primary area: infrastructure
TL;DR: We speed up FlashAttention on modern GPUs (Hopper) with asynchrony and low-precision

Abstract: Attention, as a core layer of the ubiquitous Transformer architecture, is the bottleneck for large language models and long-context applications. elaborated an approach to speed up attention on GPUs through minimizing memory reads/writes. However, it has yet to take advantage of new capabilities present in recent hardware, with FlashAttention-2 achieving only 35% utilization on the H100 GPU.
We develop three main techniques to speed up attention on Hopper GPUs: exploiting asynchrony of the Tensor Cores and TMA to (1) overlap overall computation and data movement via warp-specialization and (2) interleave block-wise matmul and softmax operations, and (3) block quantization and incoherent processing that leverages hardware support for FP8 low-precision. We demonstrate that our method, FlashAttention-3, achieves speedup on H100 GPUs by 1.5-2.0$\times$ with BF16 reaching up to 840 TFLOPs/s (85\% utilization), and with FP8 reaching 1.3 PFLOPs/s. We validate that FP8 FlashAttention-3 achieves 2.6$\times$ lower numerical error than a baseline FP8 attention.

FlashAttention-33 Fast and Accurate Attention with Asynchrony and Low-precision
Jay Shah"! Ganesh Bikshandi"
Ying Zhang 2, Vijay Thakkar 3,4 Pradeep Ramani Tri D
Colfax Research, Meta, NVIDIA, Georgia Institute of Technology Princeton University, Together AI
ganesh] -intl. com, yingzemeta.....m com.
(vithakkar, ppraman)errrr @nvidiiiiiiiiiii com. triêtridao me
Abstract
Attention, core layer of the ubiquitous Transformer architectrr,,, is the bottleneck for large language models and long-context applications. FLASIATTTETTONN elaborated an approach to speed up attention on GPUs through minimizing memory reads/writes. However, it has yet take advantage of new capabilities present in recent hardware, with FLASHATTENTION- achieving only 35% utilization on the H100 GPU. We develop three main techniques to speed up attention on Hopper GPUs: exploiting asynchrony of the Tensor Cores and TMA to (1) overlap overall computation and data movement via warp sarp-specialization and (2) interleave block-wise matmul and softmax operations, and (3) block quantization and incoherent processing that leverages hardware support for FP8 low-precision. We demonstrate that our method, FLASHATTENTION -3, achieves speedup on H100 GPUs by 1.5-2.0x with BF16 reaching up to 840 TFLOPs/s (85% utilization), and with FP8 reaching 1.3 PFLOPs/s. We validate that FP8 FLASHATTENTION-3 achieves 2.6x lower numerical error than baseline FP8 attention.
1 Introduction
For the Transformer architecture (58), the attention mechanism constitutes the primary computational bottleneck, since computing the ceff-attention scores of queries and keys has quadratic scaling in the sequence length. Scaling attention to longer context will unlock new capabilities (modeling and reasoning over multiple long documents [24
43
49
and files in large codebases 30
modalities (high-resolution images [10], audio 23 video 25 and new applications (user interaction
with long history [52), agent workflow with long horizon 61 This has generated significant interest
in making attention faster in the long-context regime, including by approximation 13 55 and
software optimization ([16] 29 45)
even alternative architectures (22 42 54].
In this work, we build on the work of Dao et [66] on developing exact-attention algorithms that integrate knowledge of the GPU's execution model and hardware characteristics into their high-level design. [16), Dao introduced ITLASHATTENTION a novel tiling strategy for parallelizing atten- tion that climinates intermediate reads/writes to slow global memory through fusing all of the attention operations into a single GPU kernel. Dao [4] restructured the algorithm as FLASHATTENTION also parallelize over the sequence length dimension and perform the inner loop of the forward pass over blocks of the key and value matrices, thus improving the occupancy and distribution of work on the GPU. However, we observe that FLASHATTEN INTENTION22 nonetheless achieves poor utilization newer GPUs relative to optimized matrix multiiliialpiiiionn (GEMM) kernels, such as 35% vs. 80-85% on the Hopper H000 GPU. Partially, this may be attributed implementation-leve differences, such not using sppccfc-- instructions in place of Ampere ones when targeting the Tensor Cores. Several
to
'Equal contribution
38th Conference Neural Information Processing Systems (NeurIPS 2024).

---

work such as ThunkerKitten [51] and cuDNN 9 [39 has shown that with Hopper-speciffe instructions and tile-based abstractions, one can speedup attention computation and simplify the implementation.
More fundamentally, FLASHATTENTION 2's a algorithm adheres to a simplified synchronous model and makes no explicit use of asynchrony and low-precision in its design. Asynchrony is result of hardware specialization to accelerate the most important operations in a ML workload: specific hardware units performing matrix multiplication (Tensor Cores) memory loading (Tensor Memory Accelerator- TMMA, separate from the rest of the CUDA cores performing logic, integer, and floating point computation. Low precision such FP8 in Hopper and FP4 in Blackwell, continuing the trend of FP16 (Pascal in 2017) and BF16 (Ampere in 2020), is a proven technique to get double or quadruple throughput for the same power and chip area. We review the capabilities afforded by Hopper in these directions in $2. The technical challenge is to redesign FLASHATTENTIO
to make use of these hardware features: asynchrony requires overlapping computation between matmul and softmax even though one depends on the output of the other, and low- low-precision requires care to minimize quantization error, especially in the case of outlier features in LLMs (20 53).
To this end, we propose FLASHATTENTION 3, which contributes and synthesizes three new ideas to further improve performance newer GPU architectures
1. Producer-Consumee asynchrony: We define a warp-specialized software pipelining scheme that exploits the asynchronous execution of data movement and Tensor Cores by splitting producers and consumers of data into separate warps, thereby extending the algorithm'ss 's ability to hide memory and instruction issue latencies.
2. Hiding softmax under asynchronous block-wise GEMMs:
We overlap the comparatively low- throughput non-GEn GEMM operations involved in softmax, such as floating point multiply-add and exponential, with the asynchronous WGMMA instructions for GEMM. As part of this, we rework the FLASHATTENTION 2 algorithm to circumvent certain sequential dependencies between softmax and the GEMMs. For example, in the 2-stage version of our algorithm, while softmax executes on one block of the scores matrix, WGMMA executes in the asynchronous proxy to compute the next block.
3. hardwarr-aaceleratee low-precision GEMM: We adapt the forward pass algorithm to allow for targeting the FP8 Tensor Cores for GEMM, nearly doubling the measured TFLOPs/s. This requires bridging the different layout conformance requirements of WGMMA in terms of how blocks of FP32 accumulator and FP8 operand matrices are assumed to be laid out in memory. We use the techniques of block quantization and incoherent processing to mitigate the loss of accuracy that results from moving FP8 precision.
To validate our method empirically, we benchmark FLASHATTENTION on the H100 SXM5 GPU over range of parameters and show that (1) BF16 achieves 1.5-2.0x speedup over ATTENTIN.. 2 in the forward pass (reaching up to 840 TFLOPs/s) and 1.5- 1.7-1.75x in the backward pass, (2) FP8 achieves 1.3 PFLOPs/s. and (3) for large sequence length, BF16 outperforms and FPS is par compared to the ofate-of-the-aar implementation of attention from NVIDIA's cuDNN library. We also yields the same numerical error FLASHATTENTION-2 and
validate that 6P16 EEL.SHATTENTION is better than the standard attention implementation as intermediate results (e.g., softmax rescaling) are kept in FP32. Moreover, FP8 PLASHIATTENTION- with block quantization and incoherent processing is 2.6x more accurate than standard attention with per-tensor quantization in cases with outlier features.
We open-source PTASHATTENTOON with a permissive license and plan to integrate it with PyTorh to benefit the largest number of researchers and developers.
2 Background: Multi-Head Attention and GPU Characteristics
2.1 Multi-Head Attention
Let Q.K.V Nxd be the query, key and value input sequences associated to single head, where N is the sequence length and is the head dimension. Then the attention output O is computed as:
S=æQK NxN
O=PVERNE
where softmax is applied row-wise and one typically sets = 1/ Vd the scaling factor. In practice, we subtract rowmax(S) from S to prevent numerical instability with the exponential function. For
We deserihe our results in the context of NVIDIA's Hopper architecture. However, our algorithm is operative for any GPU architecture with sufficiently robust asynchronous execution and low capabilities. "FLASHATTENTION 3 is available at //tp://b//ttttt/// wow/ AlLab/ attention

---

multi-head attention (MHA), each head has its own set of query, key and value projections, and this computation parallelizes across multiple heads and batches produce the full output tensor.
Now let & be a scalar loss function and let d(-) output gradient dO e Nxd, compute dQ. dK, and dV according to the chain rule as follows:
( be notation for the gradient. Given the
dV = P dO Nxd
dP do=
ds = dsoftmax ddsoottmaax(dPPERR (RP)
dQ=adSKER
dK=adS' aK=adS"QERM>
Here, we have that ds (diag()) p dp for p ==softmax((1) function of vector s. and write dsoftmax (dP) for this formula applied row-wise. Finally, this computation again parallelizes across the number of heads and batches for the backward pass of MHA.
2.2 GPU hardware characteristics and execution model
We describe the aspects of the GPU's execution model relevant for FLASHATTENTION-3 with a focus on the NVIDIA Hopper architecture concrete instantiation of this model.
Memory hierarchy: The GPU's memories organized as hierarchy of data locales, with capacity inversely related to bandwidth (Table 1 Global memory (GMEM), also known as HBM, is the off- chip DRAM accessible to all streaming multiprocessors (SMs). Data from GMEM gets transparently cached into an on-chip L2 cache. Next, each SM contains a small on-chip, programmer programmer-managee highly banked cache called shared memory (SMEM). Lastly, there is the register file within each SM.
Thread hierarchy: The GPU's programming model is organized around logical groupings of execution units called threads. From the finest to coarsest level, the thread hierarchy is comprised of threads, warps (32 threads), warpgroups (4 contiguous warps), threadblocks (i.e., cooperative thread arrays or CTAs), threadblock clusters (in Hopper), and grids.
These two hierarchies are closely interlinked. Threads in the same CTA are co-scheduled on the same SM, and CTAs in the same cluster are co-scheduled on the same GPC. SMEM is directly addressable by all threads within CTA, whereas each thread has at most 256 registers (RMEM) private to itself.
Table I: Thread-Memon hierarchy for the NVIDIA Hopper H100 SXM5 GPU,
Hardware Level
Parallel Agent
Data Locale
Capacity @ Bandwidth
Chip
Grid
GMEM
80 GiB 3.35 TB/s
GPC
Threadblock Clusters
L2
50 MiB @ 12 TB/s
SM
Threadblock (CTA)
SMEM
228 KiB SM, 31TB/s per GPU
Thread
Thread
RMEM
256 KiB per SM
Asynchrony and warp-specialization: GPUs are throughput processors that rely concurrency and asynchrony to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the Tensor Memory Accelerator (TMA) a dedicated hardware unit [38 $7.29.. Furthermore, unlike prior architectures such as Ampere, the Tensor Core of Hopper, exposed via the warpgroup-wide WGMMA instruction [40] 19.7.14], is also asynchronous and can source its inputs directly from shared memory.
Hardware support for asynchrony allows for warp-specialized kernels, where the warps of a CTA are divided into producer or consumer roles that only ever issue either data movement or computation. Generically, this improves the compiler's s ability to generate optimal instruction schedules 4.. In addition, Hopper supports the dynamic reallocation of registers between warpgroups via setmaxnreg 440) $9.7.17.1], so those warps doing MMAs can obtain a larger share of RMEM than those just issuing TMA (for which only a single thread is needed).
Modern GPUs have specialized hardware units for accelerating low-precision computation. For example, the WGMMA instruction can target the FP8 Tensor Cores on Hopper to deliver 2x the throughput per SM when compared to FP16 or BF16.
Low-precision number formats:
* *Luo et al. (34] reports shared memory bandwidth of 128 bytes per clock cycle per SM, and we multiply that by 132 SMs and the boost clock of 1830 MHz.
3

---

However, correctly invoking FP8 WGMMA entails understanding the layout constraints on its operands. Given a GEMM call to multiply AXB for an M K-marrix and an N' K K-matrix B, we say that the A or B operand is mn-major if it is contiguous in the outer M or N dimension, and k-major if is instead contiguous in the inner K-dimension. Then for FP16 WGMMA, both mn-major and k-major input operands are accepted for operands in SMEM, but for FP8 WGMMA, only the k-major format is supported. Moreover, situations such as attention where one wants to fuse GEMMs in a single kernel, clashing FP32 accumulator and FP8 operand layouts pose an obstacle to invoking dependent FP8 WGMMAs.
In the context of attention, these layout restrictions entail certain modifications to the design of an FP8 algorithm, which we describe in 3.3
2.3 Standard Attention and Flash Attention
Following Dao et al. [16]. we let standard attention denote an implementation of attention on the GPU that materializes the intermediate matrices S and P to HBM. The main idea of FLASHATTENTIN to leverage local version of the softmax reduction avoid these expensive intermediate reads/writes and fuse attention into a single kernel. Local softmax corresponds to lines [8 I of the consumer mainloop in Algorithm Algorithm[togethee ttoethher with the rescalings of blocks of O. The simple derivation that this procedure indeed computes O can be found in [14 $2.3.1].
was
3 FlashAttention: Algorithm
In this section, we describe the FLASHATTENTION-- algorithm. For simplicity. we focus on the forward pass, with the backward pass algorithm described in Appendix B.T We first indicate how to integrate warp-specialization with a circular SMEM buffer into the base algorithm of FLASHATTENTION-2 We then explain how to exploit asynchrony of WGMMA to define an overlapped GEMM-softmmn 2-stage pipeline. Finally, we describe the modifications needed for FP8, both in terms of layout conformance and accuracy via block quantization and incoherent processing.
3.1 Producer-Consumer asynchrony through varp-specialization and pingpong scheduling
Warp-specialization As with FLASHATTENTION-2 2, the forward pass of PLLASHATTENTION-3 is embarrassingly parallel in the batch size, number of heads, and query sequence length. Thus, will suffice to give a CTA-level view of the algorithm, which operates on a tile Q; of the query matrix compute the corresponding tile O of the output. To simplify the description, give the warp-specialization scheme with a circular SMEM buffer that does not have in addition the GEMM-sofimax overlapping. Let be the head dimension, N the sequence length, and fix query block size B, to divide Q into T, = B, blocks Qr....r,.
first
For our implementation of Algorithm on Hopper, we use setmaxnreg for (de)allocations, TMA for loads of Q and VK,Vyoosj-aa and WGMMA to execute the GEMMs in the consumer mainloop. with the SS or RS prefix indicating whether the first operand is sourced from shared memory or register file. For interpreting the execution flow of Algorithm| | note that issuing TMA loads does not stall on the completion of other loads due to asynchrony. Moreover, in the producer mainloop, no waits will be issued for the first iterations as the buffer gets filled.
Pingpong scheduling The asynchronous nature of WGMMA and TMA, along with warp- specialization, opens up the opportunity to overlap the softmax computation of one warpgroup with the GEMM of another warpgroup. To motivate this, notice that non-matmul operations have much lower throughput than matmul operations on modern hardware accelerators. As example, the H100 SXM5 GPÚ has 989 TFLOPS of FP16 matmul but only 3.9 TFLOPS of special functions such as exponentia (necessary for softmax). For the attention forward pass in FP16 with head dimension 128, there are 512x more matmul FLOPS compared to exponential operations, but the exponential has 256x lower throughput, so exponential can take 50% of the cycle compared matmul. The situation is even worse with FP8, where the matmul throughput doubles but the exponential throughput stays the same.
'The CUDA programming guide specifies that 16 operations of special functions can be performed per streaming multiprocessor (SM) per clock cycle. We multiply 16 by 132 SMs and 1830 MHz clock speed to get 3.9 TFLOPS of special functions.

---

Algorithm 1 FLASHATTENTION- forward pass without intra-consumer overlapping CTA view
Require: Matrices Q BR and K, E.R in HBM, key block size B. with T. - N
1: Initialize pipeline object to manage barrier synchronization with s-stage circular SMEM buffer.
2: in'nn producer warpgroup then
3: Deallocate predetermined number of registers.
4: Issue load Q from HBM to shared memory.
5:
Upon completion, commit notify consumer of the load Q
6:
for j <Tc do
7:
Wait for the (j% stage of the buffer be consumed. Issue loads of K,.V, J from HBM to shared memory at the (J'os)thh stage of the buffer. Upon completion, commit to notify consumers of the loads of Kj.V V.
8:
9:
10:
end for
else
12: Reallocate predetermined number of registers as function of number of consumer warps.
On-chip, initialize O, = (0) & Rd and m, (0)))
14:
Wait for Q; to be loaded in shared memory. for 0 <T do
16:
Wait for K to be loaded in shared memory. Compute S Kj (SS-GEMM), Commit and wait.
17:
18:
Store oldd = = and compute m, = =mx mith rrowmax(S)
Compute
= = (eeeep((""m))) -m;) and 4 mn)
20:
Wait for V to be loaded in shared memory.
21:
Compute 0, dixiie(eepp(m)
m, - (RS-GEMM). Commit and wait.
22:
Release the (5%s))th stage of the buffer for the producer.
23:
end for
24:
Compute O;=diag(66 (-)))))))00)) -0 and L; Llgg))
25:
Write O, and L; HBM as the ith block of O and
26:
end if
Since the exponential performed by a separate hardware unit (the multi-function unit), ideally we' d want the exponential calculation be scheduled when the Tensor Cores are performing the matmul. To do so, we use synchronization barriers (bar. . synn instructions) to force the GEMMs (GEMM1 of one iteration, and GEMMO QKT of the next iteration) of warpgroup to be scheduled before the GEMMs of warpgroup As result, the softmax of warpgroup will be scheduled while warpgroup 2 is performing its GEMMs. Then the roles swap, with warpgroup 2 doing softmax while warpgroup 1 doing GEMMs (hence, "pingpong" scheduling). This is illustrated in Fig. T Though in practice the pingpong scheduling not clean as depicted in the figure, we generally find this improve performance (e.g., from 570 TFLOPS to 620-640 TFLOPS for FP16 forward with head dimension 1 128 and sequence length 8192).
Werpgroup
Gcvse
GCM01
Sofrren
CZMM
CCABO
Softee
Werpgroup
Schmu
devrn
(BNAU)
time
Figure 1: Pingpong scheduling for 2 warpgroups overlap softmax and GEMMs: the softmax of warpgroup should scheduled when the GEMMs of another warpgroup are running. The same color denotes the same iteration.
Attention variants For multi-query attention (50] and grouped query attention 3 we follow the approach in FLASHATTENTION- and adjust the tensor indexing to avoid duplicating K and V in HBM.
5

---

3.2 Intra-warpgrou overlapping GEMMs and softmax
Even within one warpgroup, overlap some instructions in the softmax with instructions in the GEMMs. We describe one technique to do so.
In the attention algorithm, operations within the inner loop (main loop) have sequential dependencies that impede parallelization within a single iteration. For example, (local) softmax (lines) 18 to 19) relies on the output 8 of the first GEMM, while the second GEMM takes its result P as an operand. Indeed, the wait statements lines 17 and of Algorithm| serialize the execution of softmax and GEMMs. However, we can break these dependencies by pipelining across iterations through additional buffers in registers. Pursuing this idea, we propose the following two-stage GEMM-softmax pipelining algorithm:
WGMMA0
2
N-1
Softmax
0
2
N-1
WGMMA1
N-2
N-1
time
Figure 2-stage WGMMA-softmax pipelining
Algorithm 2 FLASHATTENTION- consumer warpgroup forward pass
Require: Matrices Q; BR and K, V in HBM, key block size B. with T. [ N R. 1.
1: Reallocate predetermined number of registers as function of number of consumer warps.
2: On-chip, initialize O, (0) eR Rd and f, = (0), (-o0) ER
3: Wait for Q, and Ko to be loaded in shared memory.
4: Compute Seu Q using WGMMA. Commit and wait.
5: Release the Oth stage of the buffer for K.
6: Compute m; Pour and l, based on Scurs and rescale O,-
7: for & j < 1 do
8: Wait for K, to loaded in shared memory.
9: Compute Snext = K using WGMMA. Commit but do not wait.
Wait for jj-1 be loaded shared memory.
11:
Compute 0, j-1 using WGMMA. Commit but not wait.
Wait for the WGMMA K,K.
13: Compute m, Paext and based on Spatt
14:
Wait for the WGMMA Peur V and then rescale O;
15: Release the (/%o))th, resp. (j 1% stage of the buffer for K, resp. V.
16: Copy Smext to Scur-
17: end for
18: Wait for Vr.-1 to be loaded shared memory.
19: Compute O, 0, Pinss VI.-I using WGMMA. Commit and wait.
20: Epilogue: Rescale O, based on m,- Compute L, based m, and &, Write O, and L, to HBM as the i-th block of O and L.
Algorithm 2 fiuncionn as a replacement for the consumer path of Algorithm to comprise the complete FLASHATTENTION-3 algorithm for FP16 precision. At high-level, for asynchronous GEMM. Within the mainloop (lines 8 to 16), the second WGMMA operation of iteration (line IT is overlapped with softmax operations from iteration j+1 (line 13)
WGMMA metonym
While the pipelined structure illustrated above offers theoretical performance gains, there are several practical aspects to consider:
Note that the number of stages of the overlapping scheme is bounded by, but need not equal, the number of stages in the circular SMEM buffer.

---

TO fdo, d T1 |do, d1)
T2 |do, d1}
T3 [do, d1}
TO [d4, d5)
T1 [d4, d5)
T2 [d4, d5
T3 (d4, d5|
TO [d2, T1 [d2, d3)
T2 [d2, d3)
T3 {d2, )
TO (d6,d7)
T1 [d6, d7) T2 [d6, d7]
T3 d7)
Figure 3: FP32 accumulator regisstt WGMMA layout rows 0 and 8, threads 0-3, entries 0-7.
TO (a0, ]
TO [a2, a3)
T1 (a0, )
T1 [a2,a3)
T2 (a0, a1)
T2 [a2,
T3 [a0, al)
T3 (a2, a3)
TO [a4, a5) TO (a6, a7)
T1 a5)
T1 (a6, a7]
T2 (a4,
T2 (a6, a7)
T3 (a4,a5)
T3 (a6, a7)
Figure 4: FP8 operand A register WGMMA layout O and 8, threads 0-3, entries 0-7.
Compiler reordering The pseudocode represents an idcalized execution order but the compiler (NVCC) often rearranges instructions for optimization. This can disrupt the carefully crafted WGMMA and non-WGMMM operation pipelining sequence, potentially leading to unexpected behavior or diminished performance gains. An analysis of the SASS code shows that the compiler generates overlapped code as expected (Section B.22..
Register pressure To maintain optimal performance, register spilling should be minimized. However, the 2-stage pipeline requires additional registers to store intermediate results and maintain context between stages. Specifically, an extra Spext must be kept in registers, leading to extra register usage of size B, x xBe Xsizeof float) per threadblock. This increased register demand may conflict with using larger block sizes (another common optimization), which is also eegister-hungry In practice, trade-offs should be made based on profiling results.
3-stage pipelining Extending the 2-stage algorithm described above, we propose a 3 3-stage variant that would further overlap the second WGMMA with softmax. While this approach offers the potential for even higher Tensor Core utilization, it requires even more registers due to additional stage in the pipeline, making the trade-off between tile size and pipeline depth more difficult to balance. A detailed description of the 3-stage algorithm and its evaluation results can be found in Appendix B.3
3.3 Low-precision with FP8
Effficieny: FP8 layout for accumulator and operand. Computing the forward pass of FLASHAITTENTION-3 in FP8 precision poses two additional challenges not encountered for FP16 in terms of conformmnce The relates to the datatype conversion the first WGMMA's FP32 ac- cumulator to the second WGMMA's lower-precision (FP16 or FP8) operand, which was left implicit in Algorithm [ Specifically, after downccating to FP8, we need to transform the register ownership pattern from that depicted in Fig. into Fig. repeated per every four threads of the consumer warpgroups.
In Appendix B.7 we give solution for this in code using shuffle instructions.
Secondly, the k-major constraint on FP8 WGMMA explained in $2.2 entails clashing assumptions on how Q, K. and V are laid out in global memory, since the TMA load cannot change the contiguous dimension. Namely, Q and K should be contiguous in the head dimension, whereas should be contiguous in the sequence length dimension. We perform in-kernel transposition of the V, tiles in SMEM prior invoking the second FP8 WGMMA, since practice. typically assumed to be contiguous in the head dimension. In Appendix B.8 we describe in details how to perform the transpose as part of the attention kernel itself
Accuracy: block quantization and incoherent processing. With FP8 (c4m3) format, one only uses 3 bits to store the mantissa and 4 bits for the exponent. This results in higher numerical error than FP16/BF16. Moreover, large models typically have outlier values (3] that are much larger in magnitude than most other values, making quantization difficult. One typically use per-tensor scaling (37 by keeping one scalar per tensor (e.g., one for Q, for K, and for V). To reduce the numerical error of attention in FP8, we employ two techniques:
1. Block quantization: we keep one scalar per block, so that for each of Q, K, we split the tensor into blocks of size B, xd or B, xd and quantize them separately. This quantization can be fused with an operation right before attention (e.g., rotary embedding) with no additional slow down (since rotary embedding is memory- memory-bandwidit bound). As the FL.ASH ATTSSHATTENTOONSS algorotn naturally operates on blocks, we can scale each block of to account for this block quantization at no computation cost.
7

---

2. Incoherent processing: to even out outliers, we multiply Q and K with random orthogonal matrix M before quantizing to FP8. Since M is orthogonal, MM = and (QM)(MM) QK", i.e. multiplying both Q and K with M does not change the attention output. This serves to "spread out" the outliers since each entry of QM KM random sum of entries of Q or K, thus reducing quanti- zation error. In practice, we follow Chee et al. B and Tseng et al. [57] and choose M to he the product of random diagonal matrices of 1 and Hadamard matrix, which can be multiplied O (dloggd) instead of (dd)), and can also be fused with the rotary embedding at no extra computation cost.
We validate that these two techniques reduces numerical error by up 2.6x in $4.3
4 Empirical Validation
We use the primitives from CUTLASS [56] such as WGMMA and TMA abstractions to implement FLASHITTENNIIN--3 and evaluate its efficiency and accuracy.
Benchmarking attention. We measure the runtime of FLASIAATEN across different sequence lengths and compare it to standard implementation in PyTorch, FLASIIATTENTION-2 PLASHATTENTION-2 Triton (which uses H100-specific instructions), as well as vendor's implementation of AALASHATTENTION that FLASHAATENTIOON- 3 is up 2.0x faster than FLASHATTENTION- and 1.5x faster than FLASHATTENTION-2 in Triton. PLASHATTENTION- reaches up to 840 TFLOPs/s, 85% of the theoretical maximum TFLOPs/s on H100 GPUs.
optimized for H100 GPUs from cuDNN. We confirm
Ablation study. We confirm that our algorithmic improvements with warp-specialization and GEMM-softmax pipelining contribute to the speedup of FLASHATTENTION-33
Accuracy of FP8 attention. We validate that block quantization and incoherent processing reduces the numerical error of FP8 FLASHATTENTION-3 by 2.6x.
4.1
Benchmarking Attention
We measure the runtime of different attention methods on an H100 80GB SXM5 GPU for different settings (without / with causal mask, head dimension 64 or 128) for BF16 imputs. We report the results in Fig- 5 and Fig. [ showing that FLASHATTENTION-3 is around 1.5-2.0x faster than FLASHATTENTION in the forward pass and 1.5-1 .75x faster in the backward pass. Compared to a standard attention implementation, FLASHATTENTION-3 can be up to 3-16x faster. For medium and long sequences (1k and above), PLASHATTENTION. even surpasses the speed of vendor's library (cuDNN closed source) that has been optimized for H100 GPUs.
Benchmark settings: We vary the sequence length as 512, 1k, ..., 16k, and set batch size so that the total number of tokens is 16k. We set the hidden dimension to 2048, and head dimension to be either 64, 128, or 256 (i.e., 32 heads, 16 heads, 8 heads). To calculate the FLOPs of the forward pass, use:
4 4-seqlen head dimension number of heads.
With causal masking, 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, we 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).
We also measure the runtime for FP8 for the forward pass under similar settings. We report the results for headdim 256 in Fig. and give the full results in Appendix C.2
4.2 Ablation Study: 2-Stage Pipelining Experiments
We ablate both the 2-stage WGMMMA-softmal pipelining and warp-specializzaton for non-causal FP16 FLASHATTENTION-33 with fixed parameters Seatch,seqlen, nheads, hdim) = (4,8448, 16, 128). The result in 2 that our algorithmic improvements (asynchrony with warp-specializatioo and overlapping between GEMM and softmax) lead to significant speedup, from 570 to 661 TFLOPs.
4.3 Numerical Error Validation
As there has been interest in the numerical error 21 of FLASHATTENTION, we compare FLASHATTENTION-2 FLASHATTENTION-3 and standard implementation of attention against
8

---

speed,
speed, (H100
Standard attention
Standard attention
FladtMttantiin222
FlasMMtteniinn22
Speed (TFLOPS/S)
CUDNN
(TFLORS
CUDNN
FlastVittentio-3
FlasnAttention-3
400
400
Speed
200
200
Sequence length
Sequence length
(a) Forward, without causal mask, head dim
(b) Forward, with causal mask, head dim 64 Attention speed, 128 (H100 80GB SXM5)
Attention speed. 128
Standard ataeniion
Standard attention
2
Flashwttertiin222
Speed (TFLOPs)s)
Speed (TFLOPs/s)
600
3last-Ateentton-3
600
FlashAttenkion
#0
200
200
512
1kK
4k
8k
16k
512
1k
2k
4x
8k
16k
Sequence length
Sequence length
(c) Forward, without causal mask, head dim 128
(d) Forward, with causal mask, head dim 128
Attention forward speed, head dim 256 (H100 80GB SXM5)
Attention forward speed, head dim 256 (H100 80GB SXMS)
Flast-Attention22
FlashAterttinn22
cJDMN FlasshAtrentinnn3
cLDKIN
Speed (TFLOPs1S)
Speed (TFLOPs/s)
Flash-tttentom-3
600
400
4.00
Sequence length
Sequence length
(c) Forward, without causal mask, head dim 256
Forward, with causal mask, head dim 256
Figure 5: Attention forward speed (BF16) on H100 GPU
Table 2: Pipelining ablation measurements
Configuration FLASHATTENTION-3
Time
TFLOPs/s 661
3.538
No GEMM-Softmax Pipelining, Warp-Speccalizzation
4.021 ms
582
GEMM-Softmar Pipelining, No Warp-Specciaiization
4.105 ms
570
reference implementation in FP64. To simulate outlier features and activations in LLMs 20 we generate the entries Q.K.V V with the following distribution:
53
(O( (O,, Bernoulli(O.00))
That is, each entry is normally distributed with zero mean and standard deviation 1, but for 0.1% of entries we add an independent term that's normally distributed with standard deviation 10. We then measure the root mean squared error (RMSE) in Table In FP16, both PITSSAATTNNION- and PLASHATTENTION-S achieves 1. lower RMSE compared the standard implementation since inter- mediate results (softmax) are kept in FP32. The baseline attention in FP8 uses per-tensor scaling, with

---

backward speed, (H100 SXMS)
Attention backward speed, head dim 128 (H100 80GB SXMS)
Standard attention
Standard allenion
600
2
600
FlashAttntion2
Speed (TFLOP//5)
3
Speed (TF (TFLOPs/S)
FlashAltentiond
400
400
200
200
512
1k
2k
4
8k
14k
512
3k
2k.
8k
16k
Sequence length
Sequence length
(a) Backward, without causal mask, head dim 64
(b) Backward, without causal mask, head dim 128
Figure 6: Attention backward speed (BF16) on H100 GPU
Attention forward spood, head dim 256 (H100 80GB SXMS)
Attertion forward spoed, head dim 256 (H100 80GB SXMS)
1200
1200
FlashAttentoos11
Fashveteetions]1
Speed (TFLOPS/S)
Speed (TFLOPs/s)
4000
400
512
1k
26
Iic
1.kc
512
12
2k
4k
8kk
1ikc
Sequence length
Sequence length
(a) Forward, without causal mask, head dim 256
(b) Forward, with causal mask, head dim 256
Figure 7: Attention forward speed (FP8) H100 GPU
matmul accumulator in FP32 and intermediate softmax results kept in FP16. Thanks block quantiza- tion and incoherent processing, FLASHATTENTION-3 in FP8 2,6x more accurate than this baseline.
Table 3: Numerical error comparisons in FP16 and FP8 (c4m3).
Method
Baseline FP16
FLASHATTENTION-2 FP16 PLASHATTENTION-3 HP16
RMSE
3.2e-4
1.9e-4
1.9e-4
Method RMSE
Baseline FP8
FLASIIÁTTENTION-3 FP8
No block quant
No incoherent processing
2.4e-2
9.1e-3
9.3e-3
2.4e-2
5 Dicussion, Limitations, Conclusion
With PLASHATTENTION-33 we have demonstrated that new programming techniques and hardware features such as asynchrony and low-precision can have a dramatic impact on the efficiency and accuracy of attention. We are able to speed up attention by 1.5-2.0x times compared to FLASIIATTENTION-2. and reduce FP8 numerical error by 2.6x compared to standard per-tensor quantization. Some limitations of our work that we hope address in the future include: optimizing LLM inference, and understanding the effects of low-precision attention in large-scale training. Though we have focused on Hopper GPUs in this work, we expect that the techniques developed here will apply to other hardware accelerators. We hope that a faster and more accurate primitive such as attention will unlock new applications in long- cong-context tasks.
10

---

References
[1] Ahmad Abdelfattah, Azzam Haidar, Stanimire Tomov, and Jack Dongarra. Performance, and autotuning of batched gemm for gpus. pages 21-38, 06 2016. ISBN 978-3-339994411220- doi: 1111007978-3311---1000
[2] A121. jaabduduingjaabbar Ai21's groundbreaking ssm -sm-transformer model. A121 blog. 2024.
[3] Joshua Ainslie, James Lee-Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprini arXiv:2305. 13245, 2023.
[4] Michael Bauer, Henry Cook, and Brucek Khailany. CudaDMA: Optimizing GPU Memory Bandwidth via Warp Specialization. In Proceedings of 2011 International Conference for High Performance Computing, Networking, Storage and Analysis, SC "1,,,, New York, NY, USA, 2011. Association for Computing Machinery. ISBN 9781450307710. doi: 10. 00.1142206384420644000 2063400
https:t/:://di. org/ 10.
[5] Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp. Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. arXiv preprint arXiiv2400.04511, 2024.
[6] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05/50, 2020.
[7] William Brandon, Aniruddha Nrusimha, Kevin Qian, Zachary Ankner, Tian Jin, Zhiye Song, and Jonathan KelleyKelley, Striped attention: Faster ring attention for causal transformers. arXiv preprint 2rrXiv:2311004933 2023.
[8] Jerry Chee, Yaohui Cai, Volodymyr Kuleshov, and Christopher M De Sa. Quip: 2-bit quantization of large language models with guarantees. Advances in Neural Information Processing Systems, 36, 2024.
[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] Richard Chen, Chengkuan Chen, Yicong Li, Tiffany Y Chen, Andrew D Trister, Rahul G Krishnan, and Faisal Mahmood. Scaling vision transformers to gigapixel images via hierarchical self-supervised learning. In Proceedings of the IEFF/CVF Conference on Computer Vision and Puttern Recognition, pages 16144-16155, 2022.
[11] Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 001199440.0000 10509, 2019.
[12] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In The International Conference on Learning Representations (ICLR), 2021.
[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. In International Conference Learning Representations 2020. 2020.
[14] Tri Dao. FlashAttention-2:: Faster Attention with Better Parallelism and Work Partitioning, 2023, 08691
[15] Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms with structured state space duality. In International Conference on Machine Learning (ICML), 2024.
[16] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with 1O-awareness. In Advances in Neural Information Processing Systems, 2022.

---

[17] Tri Dao. Daniel Y Fu, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. In The International Conference on Learning Representations (ICLR), 2023.
[18] Tri Dao, Daniel Haziza, Francisco Massa, and Grigory Sizov. Flash- Flash-Decoding for long-context inference, 2023. Utt.jttg.e/pprrorc...o///////be/
[19] DeepSeek-AL. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. arXiv preprint
2024.
[20] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm. int8 O: 8-bit matrix multiplication for transformers at scale. CoRR abs/2208. 2022.
[21] Alicia Golden, Samuel Hsia, Fei Sun, Bilge Acun, Basil Hosmer, Yejin Lee, Zachary DeVito, Jeff Johnson, YeoYeon Wei, David Brooks, et al. Is flash attention stable? arXiv preprint arXiv:2405.02803, 2024.
[22] Albert Gu and Tri Dao. Mamba: Linear-ting sequence modeling with selective state spaces. 2023.
[23] Anmol Gulati, James Qin, Chung-Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang. Yonghui Wu, et al. Conformer: Convolution-aagmennte transformer for speech recognition. arXiv preprint Xiv:220050081000 2020.
[24] Mandy Guo, Joshua Ainslie, David Uthus, Santiago Ontanon, Jianmo Ni, Yun-Hsuan Sung, and Yinfei Yang. Longt5: Efficient text-to-text transformer for long sequences. arXiv preprint arXiv:2112.079166 2021.
[25] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. Advances in Neural Information Processing Systems, 35:8633- 8646, 2022.
[26] Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length Ilm inference with kv cache quantization. arXiv preprint wrXiv:2401.1801 2024.
[27] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. International Conference on Machine Learning, pages 51 5156-5100 5165. PMLR, 2020.
[28] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In The International Conference on Machine Learning (ICML), 2020.
[29] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with PagedAttentionn Principles, pages 61 1-626, 2023.
In Proceedings of the 29th Symposium on Operating Systems
[30] Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Mare Marone, Christopher Akiki, Jia Li, Jenny Chim, et al. Starcoder: may the source be with you! arXiv preprint arXiv.2305.06111. 2023.
[31] Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context. arXiv preprint arXiv:2310.0188)) 2023.
[32] Hao Liu, Wilson Yan, Matei Zaharia, and Pieter Abbeel. World model million-length video and language with ringattention. arXiv preprint arXiv:2402.08268, 2024.
[33] Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. Kivi: A tuning-free asymmetric 2bit quantization for kv cache. arXiv preprint arXiv:2402.027550 2024.
[34] Weile Luo, Ruibo Fan, Zeyu Li, Dayou Du, Qiang Wang, and Xiaowen Chu. Benchmarking and Dissecting the Nvidia Hopper GPU Architecture, 2024. URL
https://arxiv. org/ ab////////222 13499
12

---

[35] Xuezhe Ma. Chunting Zhou, Xiang Kong, Jumxian He, Liangke Gui, Graham Neubig. Jonathan May, and Luke Zettlemoyer. Mega: Moving average equipped gated attention. In The International Conference on Learning Representations' (ICI.), 2023.
[36] Xuezhe Ma, Xiaomeng Yang, Wenhan Xiong, Beidi Chen, Lili Yu, Hao Zhang, Jonathan May, Luke Zettlemoyer, Omer Levy, and Chunting Zhou. Megalodon: Efficient I1m pretraining and inference with unlimited context length. arXiv preprint arXiv:2404. 08801, 2024.
[37] Paulius Micikevicius, Dusan Stosic, Neil Burgess, Marius Cornea, Pradeep Dubey, Richard Grisenthwaite, Sangwon Ha, Alexander Heinecke, Patrick Judd, John Kamalu, et al. Fp8 formats for deep learning. XrXiv preprint 2rX1iv:209.005433 2022.
[38]
NVIDIA. CUDA Programming Guide Version 12.4, 2024. UtLhttps://docs www.i/iiaoonniiiiiaaaa.iaa nvidia. com/
cuda/ caaaa/cuaaaa c- progr www.aaaaaniinnnnnneeeo -ww.gudeeidd..hhaa
[39] Nvidia.
Accelerating transformers
with nvidia cudnn 9.
Nvidia
blog. 2024.
URL https://devell.: //developer .vviiii cwt..looccooooooeeo
accelerating accelerating-traasse teteerrtttnttttrasffmmrr----e- transformers-witee witth- Sevriiiia-cudn-99 sudnn-9/
[40]
NVIDIA. Parallel Thread Execution ISA Version 8.4, 2024.
URL
https
/docs. nvidia. httm//udappd/ppx....tm 1ttp:///txxisa_..ppppppp 8.4.
[41]
Muhammad Osama, Duane Merrill, Cris Cecka, Michael Garland, and John D. Owens. Stream-k: Work-centric parallel decomposition for dense matrix-matrix multiplication on the gpu. In Proceedings of the 28th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming, PPoPP '23, pages 429 4314, New York, NY, USA, 2023. Association doi: 10. 10.1145/3572888.3777799 URL
for Computing Machinery. ISBN 9798400700155 156. 1ttp://...11//3772888 3577479
[42] Bo Peng. Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao. Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, et al. RWKV: Reinventing RNNs for the Transformer era. arXiv preprint 2rXiv:2305. 13048, 2023.
[43] Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension large language models. arXiv preprint arXiv:2309.00m 00001, 2023.
[44] Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. In The International Conference on Learning Representations (ICLR), 2021.
[45] Markus N Rabe and Charles Staats. Self-attention does not need (n²) memory. arXiv preprint arXiv:2112.05668, 2021.
[46] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content- content-based sparse attention with routing Transformers. arXiv preprint arXiv:2003.05997 2020.
[47] Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Tal Remez, Jérémy Rapin, et al. Code Ilama: Open foundation models for code. arXiv preprint arXiv:2308, / /2950, 2023.
[48] Rya Sanovar, Srikant Bharadwaj, Renee St. Amant, Victor Rühle, and Saravan Rajmohan. Lean attention: Hardware-aware scalable attention mechanism for the secode-phase of transformers. 2024.
[49] Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, et al. Scrolls: Standardized comparison over long language sequences. arXiv preprint wtxXi:2201.0355.. 2022.
[50] Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint
1rXiv:/111.021500 2019.
[51] Benjamin Spector, Aaryan Singhal, Simran Arora, and Christopher Ré, 2024. URL
thub. http://www......eeoaaaaaaaaaaata
13

---

[52] Fei Sun, Jun Liu, Jian Wu, Changhua Pei, Xiao Lin, Wenwu Ou, and Peng Jiang. Bert4ree: Sequential recommendation with bidirectional encoder representations from transformer. In Proceedings of the 28th ACM international conference on information and knowledge management, pages 1441-1450, 2019.
[53] Mingjie Sun, Xinlei Chen, J Zico Kolter, and Zhuang Liu. Massive activations in large language models, arXiv preprini arXiv: 240. 2024.
[54] Yutao Sun, Li Dong, Shaohan Huang. Shuming Ma, Yuging Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXtv:2307.08621, 2023.
[55] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. arXiv preprint arXiv:2009.06733 2020.
[56] Vijay Thakkar, Pradeep Ramani, Cris Cecka, Aniket Shivam. Honghao Lu, Ethan Yan, Jack Kosaian, Mark Hoemmen, Haicheng Wu, Andrew Kerr, Matt Nicely, Duane Merrill, Dustyn Blasig, Fengqi Qiao. Piotr Majcher, Paul Springer, Markus Hohnerbach, Jin Wang, and Manish Gupta. CUTLASS, January 2023. BRRLhttps://githuu com/ won/WVVDIA//utllee....
[57] Albert Tseng, Jerry Chee, Qingyao Sun, Volodymyr Kuleshov, and Christopher De Sa. Quipit: Even better Ilm quantization with hadamard incoherence and lattice codebooks. arXiv preprint arXiv:2402. 04396, 2024.
[58] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[59] Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, et al. An empirical study of mamba-based language models. arXiv preprint arXiv:2406.07888880 07884400000772220 2024.
[60] Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A 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.
[61] Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao. React: Synergizing reasoning and acting in language models. arXiv preprint arXiv:2210.03622 2022.
[62] 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.
[63] Zyphra. Zyphra unveils zamba: compact 7b ssm hybrid model. Zyphra blog, 2024.
14

---

A Related Work
Attention variants and distributed attention Ever since attention became popular with the Trans- former architecture (58), there has been a large body of work approximating attention to scale it to longer sequences. These approximation methods can generally be categorized into two classes: sparse and low-rank. Sparse attention only computes some entries of the attention matrix (softmax(QK" and assumes that other entries are zero. Different methods have different ways of choosing which entries should be zero, either with fixed pattern [T]), with sliding window [6], with dynamic pattern through hashing [28] or routing [46]. The low-rank approach instead assumes that the attention matrix has low-rank structure, and apply pointwise nonlinearity to the query and key (27] with random projection 12 60]. One can also combine the sparse and low-rank approximation for better quality [ 62 However, these approximation methods typically do not offer the same model quality as standard attention [55] and so most large-scale models do not employ these techniques.
There are other variants of attention aimed at reducing the size of the KV cache to improve inference efficiency. Multi-query attention [50] and grouped query attention B] tie different heads of K and V, and multiple query heads interact with the same key and value head. Multi-head latent attention [9] parameterizes the K and V as low-rank projections of a shared matrix to further reduce the KV cache size. However, all of these approaches do not change the core computation softmax (QK") ) V
during training and simply change how Q.K, V are obtained. As result, any efficiency or accuracy improvement to the standard attention computation benefits these methods.
To extend even longer context, attention computation can be distributed across multiple GPUs. Methods such as Ring attention [2] and variants 77] can reach a context length of up to million. They use FLASHATTENTION (or PTTSSIATTNTION-2 as a primitive, and so the improvement from FLASUATTENTION-3 would benefit these distributed attention methods as well.
Alternative architectures Motivated by the limitations of attention, a variety of alternative architectures have been proposed. They build on the connection between linear attention (2] and recurrent neural networks (RNNs). RWKV (42) H3 [7]) MEGA [35], Retnet (54) enhance the expressivity of the simple cumulative sum in linear attention with more sophisticated recurrences. Mamba [22] and XLSTM [5] use learnable weighting for the recurrence and can match the quality of Transfooreess in language modeling at small or medium scale. These approaches can be connected to generalizations of linear attention through the lens of the structure of the token-mixing matrix (15). These models have started to some traction, seeing usage in some medium large-scale models such as Jamba 2 Zamba [33 Megalodon (36. and Mamba2-hybrid [59]. For the highest quality, these SSM- and RNN- --sed models still employ many layers of attention. We expect that techniques to speed up attention presented in this work will be useful to speedup these alternative architectures.
cow-precision attention Quantization is a promising approach to speed up attention, but they have mostly focused on reducing the space for KV cache for inference efficiency. QulP [ and (5IP#[57) use incoherent processing to reduce the quantization, and we adapted this technique for FP8 FLASHATTENTION Recent work suggests that for inference the KV cache is highly compressible down to 4-, 3-, or even 2-bits [26] 33). However, quantization during training is still challenging as higher precision is typically required for stable training.
Hardware-aware Algorithms Our work presented in this paper focuses on the micro-architecture specific tuning to leverage new instruction sets and adopt natively asynchronous programming model. There are other orthogonal axes for hardware-awand algorithm co-design being explored. A recent example of this is LeanAttention [48), which recognizes the poor GPU occupancy and high memory bandwidth requirements of the sequential token generation phase as primary bottlenecks for inference and optimizes it via a smarter load balancing strategy similar to Stream-K load balancing [4] to achieve nearly peak occupancy. There is large literature optimizing GEMM for specific hardware that employs many of the same techniques. As an example, Abdelfattah et al. [ presents a high performance batched GEMM kernel on K40c Graphics Processing Units (GPU) for both fixed and variable sizes, proposing specialized GEMM designs and a comprehensive autotuning process to deliver state-of-the-ar performance.
15

---

B Addition Details on Algorithms
B.1 Asynchrony Through Warp Specialization for the Backward Pass
Similar to the forward pass $3.1 we use warp specialization to handle asynchrony. Instead of just a simple producer-consumer pattern in the forward pass, we add one extra role of dQ writer, since need to accumulate the value of dQ produced by each thread block to the global value of dQ. This dQ accumulation introduces memory contention (many thread blocks writing to the same location) so having separate warp to handle this (along with asynchrony) will avoid blocking the rest of the warps in the thread block to perform the next computation (matmul).
We include the backward pass with warp specialization in Algorithm
Algorithm 3 FLASHATTEETIONN-3 backward pass with warp specialization
Require: Matrices Q.K. Nxd HBM, logsumexp vector LERN N in HBM, block sizes Be, Br-
1: In a preprocessing kernel, compute rowsum dO+ E Rd (pointwise multiply), write D to HBM and divide it into T, blocks Dr.....Dr, of size B, each.
2: Divide Q into T, A, blocks QI.... Qr, of size B, x each, and divide K,V in T.
N B.
blocks K,,...,,, Ky, and Vi..... V......r. of size X. each.
3: Divide dO into T, blocks dO,,..,dOr, of size B,Xd each, and divide L into T, blocks Li,....Lr of size B, each.
4: Initialize pipeline object to manage barrier synchronization with s-stage circular SMEM buffer.
5: in producer warpgroup then
6: Deallocate predetermined number of registers.
7: Issue load K, and V, from HBM to shared memory.
8: Upon completion, commit notify consumer of the load of Ky and Vy-
9:
for <isT, do
Wait for the (i%s) stage the buffer be consumed.
11: Issue loads of Q,,O, from HBM to shared memory at the (i%ss) stage of the buffer.
Upon completion, commit notify consumers of the loads of Q,,O,,
13:
end for
else if in consumer warpgroups then
15: Reallocate predetermined number of registers as function of number of consumer warps.
On-chip, Initialize d (0)B, B..dd. .d.ddddVaaa.... (0)
17:
Wait for Kj and V, to be loaded in shared memory.
for <i do
19:
Wait for Q; to be loaded in shared memory. Load L, from HBM to on-chip SRAM.
20:
21:
On chip, compute S = E
B,xB, (SS-GEMM). Commit.
22:
Wait for dO, to be loaded shared memory.
23:
On chip, compute = =do"=do,V)
B B,XB, (SS-GEMM), Commit.
On chip, wait for (()) then compute p(") =xxp(S;; L.) B.
24:
On chip, wait for ( then compute ds (i) -p!/) (dp! -D,) E BRc
25:
26:
On chip, compute dV, - dV (V TdO,ERR.xt Buxd (RS-GEMM). Commit.
27:
On chip, compute dk, - dK +ds dV, and dk,.
Q e. (RS-GEMM). Commit and wait for both
On chip, compute do, (local) =ds (1) Notify the dQ-writer. end for
K, Ky B.xd (SS-GEMM), and write dQ (local) to smem.
28:
29:
30:
else if in dQ-writer warp then
31:
for 1 T do
Wait for do, (local) to be ready in smem.
32:
Using semaphore, atomically add dQ, (Incal) to dQ, in global memory.
33:
34:
end for
35:
end
16

---

B.2 2-Stage Pipelining SASS Analysis
We give simplified SASS code for the inside of the consumer warpgroup mainloop.
1/ Compute row_mas FMIMX. FTZ RO, R24, R6, !PT SHFL.BFLY PT, R185, R2, 0x2, Ox1f FMNMX and SHFL.BFU
// Apply exp2 and row_aum. Rescale 0.
FMUL. FTZ R2, R4, UR9
MUFU EX2 R185, R184
FFM R24, R24, UR9, reuse .FZ R24, R211, R24
FMUL, FFMA, FMUL, MUFU EX2,
// FP32 -> FP16 conversion are interleaved with exp2, som_aum and D rescaling. F2FP.F16. F32. PACK_AE AB R231, R25, R231 F2FP, FMUL, NUFU, FFMA, FADD
// Start the first NGMMA. Broken down into 8 HGMMAв. // The first 7 HGMMAs are packed together. WARPGROUP, ARRIVE
HGMMA. 64w.1/x66x122.6...mmmm F32 R24, [deac) RZ, !UPT
HGMMA x
// FP32->FP16, exp2, row_sum, 0 rescaling are interleaved with HGMMA. F2FP_F16 F16.FFF2.2222. 933.PACK. R214, R214, R187
MUFU.EX, EX2 R234, R5 FADD.FTZ R237, R187, R2 F2FP, MUFU, FADD
// The last HGMMA is issued here. No need to wait. HGMMA. 64x192x16. F32 R24, gdesc[URARe R24, gab0
// Start the second WGMMA. Broken down into 12 HGMMAs. // All 12 HGMMAs are packed together. Not interleaved with other instructions.
WARPGROUP ARRIVE
HGMMA. 64t.18x66x122.16...mm HGMMA x 64x128x16. HGMMA.64x128:
__226_FF2 R120, R228, gdeac [URS] tnspB, R120
HGMMA.
R120, R184, gdeac [UR8] tnspB, R120, gab0
// wgmma. wait_group the end.
WARPGROUP DDEERAR.EEEEEREREE LE gabo, 0x0
We make the following observations:
1. Softmax is reordered the very beginning, before the first WGMMA.
2. The first WGMMA is interleaved with softmax and FP32 FP16 datatype conversion of
S. This indicates that WGMMA and non WGMMAs are executed in parallel.
3. exp2, Betr_oowSsssm O rescaling and FP32 FP16 conversions are interleaved together.
4. The second WGMMA not overlapped with other instructions, as expected.
Overall, SASS shows that the 2-stage pipelining idea works as expected.
B.3 3-Stage Pipelining Algorithm
We experiment with a 3-stage pipelining algorithm to parallelize the first WGMMA from iteration j+2 softmax from iteration 1+1, and the second WGMMA from iteration j. We describe this algorithm in Al- 4 This algorithm behaves worse than the 2-stage pipelining algorithm due to the reasons below:
17

---

WGMMAD
2
N-1
Softmax
3
N-2
WGMMA1
time
Figure 3-Stage Pipelining
Algorithm 4 FLASHATTENTION 3-stage pipelining consumer warpgroup forward pass
Require: Matrices Q.K.VER.MA in HBM, block sizes B., B, Each warpgroup reads block Qi of size B, x,, T B, blocks K Ky...,,y, Ky. and V...., Vr, of size B. xd. Each warpgroup writes output block O, of size B, xd, and logsumexp block L; of size B,.
1: Initialization. Load Q, from HBM to on-chip SRAM. Initialize m,,f,m,,ssaaee
2: Wait for the producer warpgroup loading Ko from HBM to on-chip SRAM.
3: Compute S ? using WGMMA. Commit and wait.
4: Compute m, Pi. 4, scale based on S.
5: Wait for the producer warpgroup loading K₁ from HBM to on-chip SRAM.
6: Compute K using WGMMA. Commit and wait.
7: for 2 <T do
8: Wait for the producer warpgroup loading Kj from HBM to on-chip SRAM.
9: Compute S =__nex=== using WGMMA. Commit but do not wait.
10:
Wait for the producer warpgroup loading j--2 from HBM to on-chip SRAM. Rescale O; based on scale_o.
12:
Compute O, 0 +P, J-2 using WGMMA. Commit but do not wait.
13:
Compute m; P. ne_next, l, scale_o based on
14:
Wait for all previous WGMMAs.
15: Copy S next to S.
16: Copy P,_next to P/.
17: end for
18: Wait for the producer warpgroup loading VT.-2 from HBM to on-chip SRAM.
19: Rescale O; based scale_o.
20: Compute O,=O,+P.Vr.- using WGMMA, Commit and wait.
21: Compute m, P,. 4, scale_o based on S.
22: Wait for the producer warpgroup loading Vz.-1 from HBM to on-chip SRAM.
23: Rescale O, based on scale_o.
24: Compute O, P using WGMMA. Commit and wait.
25: Epilogue. Rescale O based on G. Compute L, based on &, and m, Write O and L, to HBM as the i-th block of O and L.
Overlapping. We expected that softmax can be overlapped with (the first WGMMA + the second WGMMA). However, the compiler doesn't cooperate this way. SASS code shows that only the first WGMMA is overlapped with softmax, while the second WGMMA is not. It's not clear why the compiler chooses reorder instructions in this way.
Register pressure. This algorithm requires more registers compared to the 2-stage pipelining algorithm. In theory, it needs to store an extra P, and scale_o, which is size B, x Bg x sizeof((inutt data data type)+ xsizcof float). As result, smaller block size needs to be chosen.
B.4 Variable Sequence Length
Some optimizations mentioned above cannot be directly used for variable sequence lengths and require special handling.
TMA To enable TMA to handle variable sequence lengths directly, additional steps are required. These include modifying tensormap using the PTX instruction mansormap.replacee and store the
18

---

tensormap in global memory, which adds overhead and complexity. To address this, during the forward pass when loading Q, we make TMA consistently loads tile size rows of data. For out-of-bound access, TMA sets zeros for rows beyond the original tensor, while S tensor masking masks out unused rows a tile. When writing O, we leverage memory-coalesee and does not require asynchrony. In the backward pass, preprocess kernel pads each sequence in dQ, dPSum, and LSE tensors with an additional 128 (tile sile_size) elements, allowing us utilize TMA store for efficient data transfer.
writes directly, as this the final step
Threadblock cluster and TMA multi-cast We utilize TMA multi-cast with cluster size of 2 fixed sequence length data loads, allowing every 2 threadblocks processing the same sequence to collaboratively read KV tensors. However, this approach is not well-suited for variable sequence lengths or cases like causal masking and window attention, where some threadblocks may exit earlier and cannot collaborate with other threadblocks in the same cluster. Not utilizing clustering for variable sequence lengths results in a performance drop of around 2% compared to fixed sequence lengths.
B.5 Masks: causal, local attention, variable sequence length
We apply masks to the S tensor to handle causal and local attention, as well as out-of-bound access for variable sequence lengths. Since masking is expensive, we apply it only when necessary. For instance, in the forward pass, the minimum and maximum KV block indices are calculated and iterated over in the main loop. For causal variable sequence lengths, masking applied only the maximum K block index. For local attention, masking is applied only to the first and last few K block indices based on local attention configurations. Masking is skipped for other K block indices.
B.6 Persistent Kernel
During the execution of the attention kernel, there a prologue (loading Q) and epilogue (writing output) where the Tensor Cores are not running. To maximize efficiency, we implement persistent kernel that can overlap the epilogue of one iteration with the prologue of the next iteration reduce this overhead and keep the Tensor Cores busy. In particular, we launch as many thread blocks as there are streaming multiprocessors (e.g.. 132 on the H100 SXM5) and implement scheduler that assigns tiles to each of the thread block. Each thread block might perform attention for more than one tile.
B.7
Register data exchange required for second WGMMA in FPS FLASHATTENTION-
TO
TO
TO
TO
T1
T1
TI
T1
T2
T2
T2
T2
T3
T3
T3
T3
TO
TO
TO
TO
T1
TI
TI
T1
T2
T2
T2
T2
T3
T3
T3
T3
TO
TO
T1
T1
T3
T3
T2
T2
T1
T1
TO
TO
T2
T2
T3
T3
TO
T1
TO
T1
T2
T3
T2
T3
TO
T1
TO
T1
T2
T3
T2
T3
Figure 9: Register data movement to satisfy layout conformance requirements of FP8 WGMMA.
In code, we can effect the register-to-regise data exchange that transforms the register ownership pattern of Fig. iinto Fig. 4 through invoking a combination of the following two CUDA intrinsics:
byte pprtep Given two 32-bit unsigned integers and y and selector s, the byte permute instruction returns bytes from the 8 input bytes as specified by s.
shfl_sync: shuffle instruction exchanges register data from source lane index j into its own destination register.
Our method is illustrated in Fig. First, we can swap the order of data held within thread's 's registers by using byte permute as follows. Referring the top row Fig. for given thread let upper be the first bytes (those in light and dark blue) and let lower be the last bytes (those in light and dark yellow). Then for the data held by threads and 2, we do the swap by calling byte_perm with the indicated selectors:
auto upper_mis auto Lower_mid = __bytepperm (upper, lower, 0x3210);
b_bbyt__eeeeeeeeeeeeeeeeee__epe_ perm (upper, lower, 0x7654)
19

---

Now between the second and third rows, we exchange data among threads by using shuffle instructions. Observe that the upper and lower blocks of 4 bytes should be each exchanged among themselves. Moreover, the shuffling of the upper blocks differs from that of the lower blocks, and both shuffles depend on the thread index (mod 4). We account for this using two pre-defined arrays to call __shfl_sync with the correct srcLane parameter as follows:
int upper_map [4] (0,3,1,2);
int lower_map map [4] (1,2,0,3); upper_mic __shf_ssymc (uint32_t(-1), upper_mid, poper_map(threeatt xx/4] 4) lower _mid = __shfl_syne (uint32_t(-1)), lower_mid, lower_mar [threadIdx. xX4]] 4)
Finally, between the third and fourth rows, repeat the technique with byte_perm. but now for all four threads and with the selector depending on the thread index (mod 4). For threads 0 and 3, we have:
upper_ _llastttt b__byt__pp_____ee_e_eeeeeeee____ _errte_eerr_e_e (upper_mise) lower lower_mid, 0x5410); lower _last __byte_perrs (upper_mil ,id, lower_mide 0x7632);
whereas for threads and 2, we have:
upper_ upper_last last perte__ermmupper _acippprrrridpp-pp lower lower_mid, 0x1054) Jower_lase beteeyterermpppppppep (upper _mid lower_mide 0x3276)
B.8 In-kernel transposition of V for FP8 FLASHATTENTION-3
We describe how fuse the memory transpose of V needed for the second FP8 WGMMA into FLASIATTENNTON-33. This is handled as an out-of-place SMEM to RMEM to SMEM transfer that executed in the producer warpgroup.
Specifically, within the producer mainloop, after issuing the TMA load a tile of V, the producer warpgroup waits for the load to complete. Then, producer warps effect the transpose by issuing LDSM (1dmatrix) and STSM (stmatrix) instructions, which involve a warp of threads collectively loading SMEM to RMEM and storing RMEM to SMEM at granularity of 128 bytes. Finally, we have addi- tional pipeline object to manage synchronization between the producer warpgroup and consumers, since the producer pipeline for the TMA load of V now instead has the producer warpgroup its consumer.
We choose LDSM/STSM instructions as they are both register efficient, allowing us to execute them in the producer warpgroup even after register deallocation, and capable of transposing layouts when doing memory copy. Note that as SMEM requirements are first reduced by the smaller memory footprint of the FP8 datatype, we find that we have enough SMEM for the separate buffer used to store the transpose.
There is technical obstacle to using LDSM and STSM in the context of FP8 datatype that is worth mentioning. Note that in the PTX documentation, LDSM/STSM are described as copying x matrices with 16-bit entries [40 $9.7.13.4.15.- LDSM/STSM in the context of FP8 precision. However, the transpose versions LDSM/STSM cannot split packed 8-bit entries, which necessitates certain register movements in between LDSM and STSM to actually perform a tile-wise transpose. The use of byte permute to split and reorder packed 8-bit entries between LDSM and STSM is depicted in the following code snippet:
but can pack 8-bit entries two time to
cute: : copy ctoedettile__opyy_11l _ldsm, tXsX, tXrX); auto data= tXxX. #pragma unroll
for (int
size(tXrX); 8)
uint32_t t #data_32bit reinterpret_castcuuie auto upper
*> (&data [n]);
auto lower data [2]]_]]bi((]];
data_32bit [0] pyre_peermuupper lower, 0x6420);
data_ [1] _byte_permmuppeer lower, 0x7531); }
cute: : (tiledd_cpy_stteeee tXrX, tXsX_outh
Since this permutes the eventual rows of the transposed V tile, we also need to modify the register movements on the consumer side that transform accumulator to operand P. We the mathematical fact that
20

---

P. P.V=colperm' " rowwoerm (V)
for o' a permutation of the common inner dimension P and V. Moreover, for the modified register exchange, we can eliminate the use of warp shuffles, but not byte permute, as each thread will already own all the entries it needs for WGMMA.
B.9 FLASHATTENTION: for inference
For decoding inference, the query sequence length is much shorter than the key/value sequence length, typically the order of one few tokens compared to the thousands stored in the KV cache. In this situation, attention becomes a memory-bound workload, and the relevant metric is not tensor core utilization measured by FLOPs/s, but loading the KV cache fast possible measured by memory bandwidth. Furthermore, since the TLASSATTENTION alo-3algoone algorithm described in $3.1 parallelizes the query sequence length, it can suffer from lack of parallelism for decoding.
We make two modifications FLASHATTENTION: to introduce more parallelism for decoding:
1. Split KV (or Flash-Decoding): We split the attention kernel along the key/value sequence length, with the number of splits determined by heuristic at launch, and combine the resulting outputs using a separate post-processing reduction kernel. "Splitting" according that n threadblocks load the same tile of Q and n different segments of the KV cache, computing different output tiles O 1,...,,,, O, and Ise vectors Ise:....,Ise,, which then the reduction kernel. We also allow for early exit of threadblocks whose given segment of the KV cache doesn't I contribute the final output, in which case the threadblock writes
parameter means
compute O in
as its lse.
This amounts to essentially the same implementation as described in [18].
2. GQA packing: For multi-query attention grouped-query attention, we can restructure the attention mainloop in order to pack multiple query heads per KV head, where each threadblock now loads its Q tile across different query heads. When query length is short, this achieves additional parallelism "for free" thanks to the large width of the first operand WGMMA tile, given as 64 per warpgroup. For example, we could have model architecture with 16 query heads per KV head and a query sequence length of 8, in which case threadblock can pack all 16 query heads into its Q tile without any change Algorithm 2 In practice, this yields up to Nx speedup over implementation that doesn't GQA packing, where N is the GQA ratio.
FLASHATTENTION-3 for inference also features an implementation of Attennitonnionn 29 that was contributed by Kai Londenberg. Recall that PagedAttention is memory optimization technique for efficiently storing the KV cache terms of fixed-size pages. This entails separating the logical position KV blocks from their physical addresses, with block table defining the address translation [29] $4.2].
Now, prior implementations of TMA load in CUTLASS construct the tensor map object such that TMA tensor coordinates are determined using the physical GMEM tensor. To use block table with TMA, Londenberg defines a SM90 TMA_ TMA_LOAD _LAD_PAGED
PAGED __P class and tensor map constructor that instead determines TMA tensor coordinates in terms of the virtual shape. The block table is then passed into the TMA copy method additional argument.
C Addition Details on Experiments and Benchmarking
C.1 System and libraries
We benchmark the speed on an H100 80GB SXM5 (700W). We generally use the latest versions of the libraries, at the time of writing (October 2024). Specifically, we use:
CUDA 12.3
cuDNN 9.5.0.50
CUTLASS 3.6
FLASHATTENTION 2.6.3
Triton 3.1
PyTorc 2.5.0
21

---

To reduce variability, we fix the GPU clock speed to 1830MHZ (clock speed used to calculate the 989 TFLOPS FP16 theoretical max throughput). We repeat the benchmarks 10 times and take the average timing.
C.2 FP8 Attention Full Results
We use following sequence lengths: 512, 1024, 2048, 4096, 8192, 16384.
Attention forward speed, head dim 64 (H100 80GB SXMS)
Attention forward speed, head dim 64 (H100 80GB SXMS5
HeslWileniion
Speed (TFLOPS/S)
Speed LTFL)))
800
800
400
400
Sequence length
Sequence length
(a) Forward, without causal mask, head dim 64
(b) Forward, with causal mask, head dim 64
Attention forward head dim (H100 BOGB SXMS)
Attertion forward head dim (H100 80GB SXMS)
Tr
Tritnn
cu DNN
CuDNN
3
Fashietenniin
Speed (TFLOPS/S)
Speed (TFLOPs/s)
800
800
400
400
Sequence length
Sequence length
(c) Forward, without causal mask, head dim 128
(d) Forward, with causal mask, head dim 128
Attention forward spoed, head dim 256 (H100 80GB SXMS)
Attention forward spoed, head dim 256 (H100 80GB SXMS)
Triton
F
FashiAttenioon3
Speed (TFLOPSiS)
Speed (TFLOPs(s)
800
800
4:00
400
512
1k
15kc
512
1k
2k
15k
Sequence length
Sequence length
(c) Forward, without causal mask, head dim 256
(f) Forward, with causal mask, head dim 256
Figure 10: Attention forward speed (FP8) on H100 GPU
22

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Abstract and intro reflects the paper's 's contribution.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Discussed in
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. speech-to -oottexttttttttte system might not used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification: The paper does not include theoretical results.
23

---

Guidelines:
(c) If the contribution is new model (e.g., large language model), then there should either be way to access this model for reproducing the results or way to reproduce The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof' can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proof' provided in appendix supplemental material.
reliee Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question:
Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that affects the main claims conclusions the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Detailed information in Appendix C.I
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
the model (e.g., with an open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-soue -ouuree models, it may be that access to the model is limited some way (e.g., g., registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification:
The code will be released with permissive license in the near future.
24

---

Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https
for more details.
While we encourage the release of code and data, we understand that this might not be pos- sible, "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines Sendides/CooeSumiiii
for more details. The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the and and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperpa- rameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [NA]
Justification: The paper does not include training models.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Not necessary for speed benchmarks since we already take average of large number (30) of trials.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
25

---

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative enror rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics httpppp//eeuurppsscppplllll/////
Answer: [Yes]
Justification: Yes
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification:
The paper focuses on foundational research and not tied to a particular
application.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does address societal impact.
Examples negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
26

---

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
No release of high-risk data or models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Yes
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should included for each asset.
For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, |aaeesswithcoee thcode. com/ dataests has curated licenses for some datasets. Their licensing guide help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No assets.
Guidelines:
27

---

The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the datasct/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsourcing or human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based on the requirements of your country institution) were obtained?
Answer: [NA]
Justification: No crowdsourcing or human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
28


Reviews:
[{"review_id": "4VPNGuiok3", "reviewer": "Reviewer_WEiA", "summary": "This paper presents FlashAttention-3, which speedup the commonly-used attention operator on Hopper GPUs. The paper proposes to leverage the asynchronous execution of the Tensor Cores and Tensor Memory Acceleator to better utilize the GPU hardware. Specifically, the paper proposes three techniques: (1) It overlaps the data movement and GEMM computation by warp specialization; (2) It proposes to overlap the GEMM and softmax operations by interleaving these operations; and (3) It leverages FP8 tensor core by block quantization. The experimental results show that FlashAttention-3 can achieve superior speedup by 1.5-2.0x with FP16 and can outperform vendor libraries.", "questions": "Thanks for submitting the excellent paper to NeurIPS. The paper is in general well-written and the ideas are novel and insightful. However, I would like to make the following comments for further polishing the paper:\n\nThe paper has mentioned a few hyperparameters. However, how do you tune these hyperparameters to achieve the best performance? Do you use grid search or simply heuristics? This is important because the inputs to language models often have varying sequence lengths, and tuning each of them is not practical. It would be great to see a systematic way to find these hyperparameters.\n\nThe paper achieves superior speedup on the single flash attention operator. However, given the existence of FlashAttention-2, how much percentage does the attention operation takes in the context of whole LLM inference or training pipeline? This could help better understanding the positioning of this paper.", "limitations": "* The paper is only specifically optimized for the Hopper architecture and may not applies to other or future GPUs.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "* It optimze attention operator on Hopper architecture, which is a timely problem and is of great importance;\n* It provides new insights for optimizing for Hopper architecture;\n* It achieves superior speedup compared to prior work and comparable speedup to vendor libraries;", "weaknesses": "* The method introduces a few hyper-parameters such as number of pipeline stages, block sizes $B_r$ and $B_c$. However, the paper has never discussed how to tune or select this parameters in practice.\n* The paper does not evaluate with the end-to-end LLM inference task, which makes it obscure to see its real performance gain in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the enthusiastic support from the reviewer.\n\n1. Setting of hyperparameters: We select tile size and stage count hyperparameters as a function of the head dimension (64, 128, 256) and datatype (16 or 8 bit) based on available register and smem budget. This is similar to matrix multiply (e.g. from cuBLAS or Triton), where the tile sizes and number of stages are tuned either with heuristics (cuBLAS) or by an autotuner (Triton). For FA3, we found 2 stages to be consistently best, and we choose tile sizes (B_r and B_c) to simply maximize the registers and shared memory that we can use. There are not that many choices for these tile sizes, (e.g. 128 x 128, 128 x 256, 256 x 128) since some dimensions have to be divisible by 64 or 128 due to hardware constraint, so we tune them by hand. These hyperparameters are not functions of seqlen, so different sequence lengths use the same tile sizes.\n2. For LLM inference, FA3 would be most helpful for the prefill stage. For the decode stage, the bottlenecks are different (loading KV cache as fast as possible), and other techniques are more relevant (e.g Flash-Decoding, KV cache quantization). \n\nThanks for the suggestion on end-to-end evaluation. We have mentioned the impact of FA3 on LLM inference in a production setting in the shared response. We include more detailed experiment here, measuring the savings from switching FA2 -> FA3 on Llama-3.1 405B with different context length (percentage of time saved vs original time).\n\n| Model size / Seq length | 8K | 16K | 32K | 128K | \n|-------------------------|-----|------|------|-------| \n| 405B | 4.6%| 7.3% | 5.4% | 26.2% |\n\nWe see that even for some of the largest models (where MLPs and QKV projections are very large), we can save up to 26.2% time, i.e. 1.35x speedup. Roughly speaking, for 405B at 128K sequence length, attention takes 50% of the time and MLP takes 50% of the time, so speeding up attention by 2x would bring around 1.3x end-to-end speedup. \nAt smaller model sizes, the speedup is larger at the same sequence length, since attention would take proportionally more time.", "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": "bRqIeGdhfg", "reviewer": "Reviewer_qxbQ", "summary": "This is a system paper that introduces FlashAttention-3, an optimized version of FlashAttention for NVIDIA's SM90+ GPUs. The key contributions (summarized by the authors) are:\n\n- Taking advantage of warp specialization in NVIDIA SM90+ and designing a producer-consumer computation paradigm to allow better intra-CTA overlapping between producers and consumers;\n- Designing 2-stage software pipeline to allow better overlapping within consumers;\n- Designing on-the-fly layout transformation to support FP8 WGMMA and using incoherent processing to reduce quantization error by 3x.", "questions": "Please address my comments in the Weaknesses section. Additional questions related to ablation studies:\n\n- Will it be also possible to open source the benchmarking code in Table 2? This will be super helpful for system researchers to understand the impact of warp specialization and 2-stage computation pipeline on the end-to-end throughput in different workloads.\n- Please comment why your implementation achieves performance advantage over OpenAI Triton ones. What is the most important design space extension that leads to this improvement?\n- Please comment on how the register re-allocation feature provided by NVIDIA SM90+ impacts the performance. Are there any important design spaces enabled by this feature which would otherwise lead to register spilling to local memory in SM8x?\n- Please comment on whether GEMM-SoftMax 2-stage pipelining bring about performance improvement on SM8x (it's related to the question above. If there are a lot of register spills then I think this is also a specialized optimization for SM90+).", "limitations": "The authors have adequately addressed the limitations of this paper. There is no potential negative societal impact.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The empirical results of this paper are very strong, and are expected to generate immediate impact to the community after code release (promised by the authors).\n- The presentation of this paper is clear. It is easy for people who are familiar with NVIDIA SM8x but not experts in SM90+ to follow the writing and essence of the optimizations. I learned a lot from the descriptions related to the warp specialization design.\n- This is the first (to-be) open-sourced system paper that discusses the challenges and solutions related to FP8 attention, which has been implemented and close-sourced by NVIDIA for a long time. The insights related to FP8 such as layout transformation and reducing quantization error are very helpful for the readers.", "weaknesses": "[Results, Benchmarking]\n\n- There has already been a paper from Colfax (arXiv 2312.11918) available online more than half a year ago describing how to implement efficient FlashAttention on Hopper GPUs. It will be better if the authors could cite this paper and compare the performance of FlashAttention-3 and this paper.\n- NVIDIA's TensorRT-LLM has a close-sourced implementation of FP8 attention. The result comparison with this library is missing.\n- While the authors mention RMSE reduction from incoherence processing using simulated data, there is no evaluation on how this step improves real LLM inference workloads (such as Wikitext perplexity, MMLU accuracy, etc.) \n- How many end-to-end LLM inference / training speedup can we obtain by switching to this FP8 attention? For inference, I am interested in a TensorRT-LLM-like implementation with FP8 GEMM + FP16 attention (FlashAttention-3) VS FP8 GEMM + FP16 attention (FlashAttention-2) VS FP8 GEMM + FP8 FlashAttention-3.\n\n[Methodology, Novelty]\n\n- While all the methodologies mentioned in the paper make sense to me, it is important to know that some of them occurred in previous literature / open source implementation on the same workload. For example, FlashInfer from University of Washington implements within-consumer software pipelining on SM80+ devices. Incoherent processing is also not new, and has been adopted in related work such as QUIP#, QuaRot, etc. \n- The authors mentioned that a lot of operations can be fused into the preceding RoPE operation **with no additional overhead**. For example, transposing V, applying Hadamard transformation, block quantization, etc. While I agree that intuitively computation operations could be fused with memory bound Ops, it is also important to notice that the CUDA cores of H100 are not that strong. I want to see real measured numbers in end-to-end workloads that justify the authors' claims on \"no additional overhead\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thorough review and helpful suggestions.\n\n1. Existing literature: The Colfax arXiv paper discusses use of WGMMA and TMA for FA2, but not more sophisticated techniques with asynchrony. This is similar to the Triton implementation that we compared to. We will cite and include this in the discussion. FA3, with asynchrony (warp specialization, GEMM-softmax overlapping) can reach 800 TFLOPS for headdim 256 FP16, while the Colfax implementation can reach 550 TFLOPS for the same setting. This highlights the importance of the new algorithm.\n\nWe have benchmarked the closed-source TensorRT-LLM implementation and found its speed similar to the cuDNN implementation. We have included the comparison with cuDNN in the paper. In the newest version, FA3 is about 10-25% faster than cuDNN for FP16/BF16, even though cuDNN was already optimized for Hopper GPUs. We are hopeful that by open-sourcing FA3, other researchers can build on it to find other hardware-aware algorithms such as approximate (sparse, low-rank) attention.\n\nWe plan to also open source the benchmarking code, thanks for this suggestion.\n\n2. End-to-end evaluation: we include this in the common response, in an LLM inference production setting. \nWe include a more detailed experiment here, measuring the savings from switching FA2 -> FA3 on Llama-3.1 405B with different context length (percentage of time saved vs original time). Note that GEMMs are done in FP8 and attention is done in BF16.\n\n| Model size / Seq length | 8K | 16K | 32K | 128K | \n|-------------------------|-----|------|------|-------| \n| 405B | 4.6%| 7.3% | 5.4% | 26.2% |\n\nWe see that even for some of the largest models (where MLPs and QKV projections are very large), we can save up to 26.2% time, i.e. 1.35x speedup. Roughly speaking, for 405B at 128K sequence length, attention takes 50% of the time and MLP takes 50% of the time, so speeding up attention by 2x would bring around 1.3x end-to-end speedup. \n\n3. Existing methodologies: we leverage existing techniques and build on a rich literature. We will add these to the discussion. We find it very encouraging that these conceptually simple techniques can speed up attention, one of the core operations in ML and has been studied intensely by the community, by up to 2x on modern hardware.\n\n4. Fusing operations: we show here the time taken by RoPE, Hadamard transform, and their fusion.\nWe measure on an H100 80GB SXM5 (3.35TB/s memory bandwidth), with Q and K of size 2 x 8192 x 32 x 128 each (BF16). \n| Operation | Timing | \n|-----------|-----| \n| Memcopy | 110us|\n| RoPE | 115us | \n| Hadamard transform | 107us |\n| RoPE + Hadamard transform | 120us |\nWe see that all of the combinations reach about 70-80% of memory bandwidth.\n5. Comparison with Triton: The most significant improvement of FA3 over it is fine-grained exploitation of asynchrony / scheduling techniques, Though Triton already uses Hopper-optimized instructions (WGMMA and TMA), there is no automatic warp-specialization or overlapping between GEMMs and other operations like softmax yet. We are collaborating with Triton developers to implement some of these techniques in the compiler.\n6. Register re-allocation is important for warp-specialization, since the warps doing copying (TMA) needs very few registers while the warps doing matmul (WGMMA) need a lot of registers to hold the accumulator, and more registers means we can use larger WGMMA instructions and more opportunity to overlap with other operations. The other techniques (e.g. overlapping GEMM and Softmax) can still apply without warp-specialization / register re-allocation.\n7. Ampere / Ada GPUs (SM8x). One can still overlap GEMM and softmax in SM8x, but warp specialization is difficult and generally not done (e.g. for matmul). We note that FA2 is already close to optimal on Ampere / Ada GPUs, reaching up to 70% of theoretical max FLOPS (while matmul can reach around 80%). For newer hardware, the tensor cores are simply much faster, so there is a greater need for asynchrony / overlapping. We expect this trend to hold for future accelerators (please see the common response for more detailed discussion). We are also exploring FP8 on Ada, even without warp specialization we believe that FP8 can still substantially improve the attention throughput on Ada cards.", "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": "t26VhnPBVv", "reviewer": "Reviewer_KCED", "summary": "The paper builds on the existing work on FlashAttention, and introduces an advanced method to accelerate the attention mechanism in Transformer model. It specifically targets the newer GPU architectures, like NVIDIA H100, by exploiting asynchrony in Tensor Cores and Tensor Memory Accelerators. There are three techniques that are proposed: 1) producer-consumer asynchrony across warps; 2) compute asynchrony within consumer warps; 3) hardware accelerated low precision GEMM. The extensive empirical results show impressive performance, with significant speedups compared to state-of-the-art methods, and with reduced numerical errors.", "questions": "The backward pass algorithm is only included in the appendix. It would be useful if it could be included in the main paper.\nCan you comment if FlashAttention-3 could be used on older GPU architectures that lack some of the advanced features of the Hopper GPUs?", "limitations": "Yes.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The authors present novel techniques that exploit hardware features like asynchrony and low-precision computation.\nThe performance gains are impressive, in speed and accuracy.\nThe empirical validation includes comprehensive benchmarks and comparisons with previous methods.\nIt is nice to see the open source commitment.", "weaknesses": "The proposed methods are highly complex, and may be difficult to implement and understand for practitioners.\nThe techniques are tailored to NVIDIA Hopper GPUs, and may have limited applicability to other hardware achitectures.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are very happy that the reviewer appreciates the impact of the paper.\n1. Generality: we are excited about generalizing our techniques for other hardwares (AMD GPUs, Google TPUs, and Nvidia Ada GPUs). Please refer to the common response for more details.\nFor Ampere cards (e.g. A100, 3090) and Ada cards (e.g. 4090), with BF16/FP16, FA2 is already close to optimal (reaching 70% of theoretical max FLOPS, while matmul can reach around 70-80% of theoretical max FLOPS). For Ada cards that support FP8 (e.g. 4090), there is an opportunity to potentially double the attention throughput. Though they lack asynchronous Tensor Memory Accelerator (TMA) and the matmul instructions are synchronous (mma.sync instead of wgmma.mma_async), there’s still some form of asynchronous copy (cp.async). Though it might be harder to overlap GEMM and softmax, we are still exploring FP8 on Ada cards, since they offer the highest FLOPS per dollar.\n\n2. Backward pass: it is more involved than the forward pass, and we would include it in the main text if space permits.", "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": "FRZvhUddcR", "reviewer": "Reviewer_TcT4", "summary": "FLASH ATTENTION sped up attention on GPUs but achieved only 35% utilization on H100 GPUs. To address this, FLASH ATTENTION-3 introduces three techniques for Hopper GPUs: warp-specialization to overlap computation and data movement by assigning warps to producer and consumer roles; interleaving operations to combine block-wise matrix multiplication (matmul) and softmax operations for asynchronous execution; and support for FP8 precision using block quantization and incoherent processing to leverage hardware capabilities. These techniques result in a 1.5-2.0× speedup with FP16, up to 1.2 PFLOPs/s with FP8, and 2.6× lower numerical error than baseline FP8 attention, improving computational efficiency and hardware utilization on modern GPUs.", "questions": "A few question and notes are listed below:\n\nHow would the methods are applicable to other GPU devices?\n\nIs speed is slower for flashattention3 in small sequence len only seen in H100?\n\nIn 4.1: '4·seqlen2 ·head dimension·number of heads'. would be better to format as equation", "limitations": "The authors adequately addressed the limitations and, if applicable, potential negative societal impact of their work", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The strengths of the paper are that it is well presented and shows good results. The paper addresses the important topic of accelerating Transformer modules by presenting an incremental improvement over previous work, specifically FLASHATTENTION-2, and applies these advancements to the new H100 GPU. The paper demonstrates how the enhanced methods leverage the capabilities of the H100 GPU to achieve significant performance gains, showcasing advancements in computational speed and efficiency. This progression highlights the continuous evolution in optimizing Transformer architectures for cutting-edge hardware, emphasizing the relevance and impact of the research in the field of deep learning and hardware acceleration.", "weaknesses": "The work primarily focuses on the H100 GPU, demonstrating the performance enhancements of FLASHATTENTION-3 for this hardware. However, evaluating how FLASHATTENTION-3 performs across other GPU families would increase the work's impact, showing its versatility and potential for wider adoption in different computing environments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewers for the support and appreciate the thoughtful questions.\n1. Other hardware: please refer to the common response. We are collaborating with other researchers and engineers on FA3 for AMD cards, Google TPUs, and Nvidia Ada cards (e.g. 4090). As suggested by the reviewer, this would enable wider adoption of FA3.\n2. Small seqlens: FA3 in May was slower than cuDNN for small seqlens (e.g. 512, 1024). We have since developed a persistent scheduler that loads the Q, K, V for the next subproblem while writing out the output of the current subproblem, and thus matching / exceeding cuDNN for small seqlens (FA3 was already faster than cuDNN for large seqlens). This overlapping is important since the prologue / epilogue (loading initial Q, K, V and writing out O) takes non-negligible time (e.g. up to 10-15% of the overall computation) when the sequence is short.", "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:
We thank the reviewers for their enthusiastic support, their careful read of the paper, and their thoughtful questions and suggestions. We are very happy that the reviewers find the paper “has immediate impact to the community”, the ideas “novel and insightful”, and the writing “clear and easy to understand”.

Since the submission in May, we have made FA3 even faster, more accurate, easier to use, and more general:
- New optimizations: (1) A persistent scheduler to start loading the Q, K, V for the next subproblem while writing out the output O of the current subproblem, thus speeding up the cases of small seqlen (2) Optimized for the case where sequences in the batch have variable lengths (e.g. in Llama 3 training) (3) Reduced the numerical error for GQA/MQA by accumulating the gradient dK and dV of different heads in FP32, before converting the sum to BF16. This brings 2-3x smaller numerical error compared to the standard method of computing dK and dV for different heads in BF16, then summing the gradients.
- Speed: FA3 version in August is 6-25% faster, reaching up to 700 TFLOPS for headdim 128 FP16 (vs 650 in May), 800 TFLOPS for headdim 256 FP16 (vs 640 in May), and 1300 TFLOPS for headdim 256 FP8 (vs 1230 in May). This is thanks to a technique called pingpong scheduling that better overlaps the GEMM of one warpgroup with the softmax of another warpgroup. For FP16/BF16, attention speed from FA3 is now on par with matmul speed from cuBLAS (arguably the most optimized operation) for similar sizes (700-770 TFLOPS), suggesting that we are using the hardware as efficiently as possible. This is another validation that asynchrony can play a major role in optimizing for modern hardware.
- Integration with other libraries: we are working with PyTorch developers to integrate FA3 to PyTorch to benefit the largest number of researchers and engineers. We are also working with inference and training libraries (Hugging Face, vLLM, Megatron-LM, TransformerEngine, DeepSpeed) to help with their integration effort to speed up transformer training and inference with FA3.
- Generality: We are working with Triton and Jax/XLA developers to implement some of the techniques in FA3 in the Triton & XLA compilers. We are also collaborating with other researchers and engineers on hardware other than H100 (elaborated below).

We now respond to some common questions from the reviewers.
1. **Generality**: The techniques developed in FA3 are not limited to Nvidia Hopper GPUs. As mentioned in the intro, asynchrony and low-precision are the general trend of AI accelerators, due to Amdahl's law. Amdahl’s law states that the overall performance improvement gained by optimizing a single part of a system is limited by the fraction of time that the improved part is actually used. As tensor cores in hardware accelerate matmuls exponentially with each new generation of hardware, the serial sections of the code (softmax and causal masking in attention) can dominate the runtime. The only way to overcome these latencies, other than adding expensive hardware units that compute them, is to overlap their execution with other concurrent work. We chose to validate our algorithms and ideas on Hopper GPUs due to good hardware and software support (async barriers, CUTLASS), but these ideas also apply to other modern accelerators. In fact, we are collaborating with AMD engineers to implement the overlapping of GEMM and Softmax on AMD GPUs such as MI300X, as well as Google researchers on a version of FA3 on TPUs. We are also exploring FP8 support for Ada GPUs (e.g. 4090) to potentially double the attention throughput on these consumer cards.
2. **Impact**: With FA3, we have significantly sped up attention, one of the two main layers in Transformers. Attention speedup can have a major impact when the sequences are long, which is increasingly common (e.g. Llama 3.1 has 128k context length). In a LLM inference production setting (with an already very optimized inference engine), we have measured that switching from FA2 to FA3 yields up to 30-40% speedup in time to first token for Llama-3 405B. This is a major improvement for inference services running on hundreds or thousands of GPUs. For 405B at 128K sequence length, attention takes about 50% of the time and MLP takes 50% of the time, and making attention 2x faster will yield about 1.3x speedup. One can combine FA3 with other techniques to speed up the MLP (e.g. sparsity, FP6/FP4) to get even higher speedup. 
3. **Interoperability**: FA3 is a drop-in replacement for FA2, save for the FP8 case where the version of FA3 in May requires V to have layout BHDS (sequence length dimension must be contiguous), due to the constraint of the FP8 tensor cores. We have since developed a variant of the FP8 kernel that supports V input in the standard BSHD layout (with the head dimension as the contiguous mode). This is important to integrate FP8 attention into standard libraries, such as PyTorch and vLLM, and also circumvents a technical challenge with variable sequence length in terms of loading memory addresses not aligned to 16 bytes. 
Recall from the paper that the FP8 WGMMA instruction only supports k-major operand B, so this variant necessitates doing an "in-kernel" transpose on tiles of V after they are loaded from GMEM to SMEM. Moreover, in order to overlap this transpose with other operations to minimize its impact on speed, we change the kernel scheduling. In short, we take advantage of special LDSM-STSM instructions that both minimize register usage and support transposition for SMEM <=> RMEM copy in order to place the in-kernel transpose of V in the producer, thereby further leveraging warp specialization with warpgroup register reallocation.
Our implementation achieves comparable performance to cuDNN, exceeding it for headdim 64 and matching for headdim 128/256.

Overall we are excited about generalizing our techniques to other hardwares to unlock new use cases for long context models.

Meta-review: The work introduces FlashAttention 3, an optimized version of FlashAttention for modern GPUs, which is a key building block for efficient transformer models. While a main focus is on NVIDIA's Hopper GPUs which currently lack such efficient kernels, the innovations of this paper are also promising and valuable for GPUs from other vendors and various other compute backends. Compatibility with new lower-precision floating point formats such as FP8 is another strength of this work, which is very valuable for the wider community.
Decision: Accept (spotlight)