Title: FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
Venue: NeurIPS 2024
Accepted: yes
Keywords: attention

hardware-aware algorithms

H100
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-3: Fast and Accurate Attention 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. FLASHATTENTION 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 GLYPH<2> 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 GLYPH<2> lower numerical error than a baseline FP8 attention.

1 Introduction

For the Transformer architecture [58], the attention mechanism constitutes the primary computational bottleneck, since computing the self-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, 47]), new 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, 27, 55] and software optimization ([16, 29, 45]), or even alternative architectures [22, 42, 54].

In this work, we build on the work of Dao et al. [16] on developing exact-attention algorithms that integrate knowledge of the GPU's execution model and hardware characteristics into their high-level design. In [16], Dao et al. introduced FLASHATTENTION, a novel tiling strategy for parallelizing attention that eliminates intermediate reads/writes to slow global memory through fusing all of the attention operations into a single GPU kernel. Dao [14] restructured the algorithm as FLASHATTENTION-2 to 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 FLASHATTENTION-2 nonetheless achieves poor utilization on newer GPUsrelative to optimized matrix-multiplication (GEMM) kernels, such as 35% vs. 80-85% on the Hopper H100 GPU. Partially, this may be attributed to implementation-level differences, such as not using Hopper-specific instructions in place of Ampere ones when targeting the Tensor Cores. Several

work such as ThunkerKitten [51] and cuDNN 9 [39] has shown that with Hopper-specific instructions and tile-based abstractions, one can speedup attention computation and simplify the implementation.

More fundamentally, FLASHATTENTION-2's algorithm adheres to a simplified synchronous model and makes no explicit use of asynchrony and low-precision in its design. Asynchrony is a result of hardware specialization to accelerate the most important operations in a ML workload: specific hardware units performing matrix multiplication (Tensor Cores) or memory loading (Tensor Memory Accelerator - TMA), separate from the rest of the CUDA cores performing logic, integer, and floating point computation. Low precision such as 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.2. The technical challenge is to redesign FLASHATTENTION-2 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-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 on newer GPU architectures: 2

Producer-Consumer asynchrony: Wedefine 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's ability to hide memory and instruction issue latencies.

Hiding softmax under asynchronous block-wise GEMMs: Weoverlap the comparatively lowthroughput non-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 algorithmtocircumventcertain sequential dependencies between softmax andtheGEMMs. Forexample,inthe2-stageversionofouralgorithm, whilesoftmaxexecutesonone block of the scores matrix, WGMMA executes in the asynchronous proxy to compute the next block.

Hardware-accelerated 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 to FP8 precision.

To validate our method empirically, we benchmark FLASHATTENTION-3 on the H100 SXM5 GPU over a range of parameters and show that (1) BF16 achieves 1.5-2.0 GLYPH<2> speedup over FLASHATTENTION2 in the forward pass (reaching up to 840 TFLOPs/s) and 1.5-1.75 GLYPH<2> in the backward pass, (2) FP8 achieves 1.3 PFLOPs/s, and (3) for large sequence length, BF16 outperforms and FP8 is on par compared to the state-of-the-art implementation of attention from NVIDIA's cuDNN library. We also validate that FP16 FLASHATTENTION-3 yields the same numerical error as FLASHATTENTION-2 and is better than the standard attention implementation as intermediate results (e.g., softmax rescaling) are kept in FP32. Moreover, FP8 FLASHATTENTION-3 with block quantization and incoherent processing is 2.6 GLYPH<2> more accurate than standard attention with per-tensor quantization in cases with outlier features.

Weopen-source FLASHATTENTION-3 with a permissive license 3 and plan to integrate it with PyTorch to benefit the largest number of researchers and developers.

2 Background: Multi-Head Attention and GPU Characteristics

2.1 Multi-Head Attention

Let Q GLYPH<150> K GLYPH<150> V 2 R 𝑁 GLYPH<2> 𝑑 be the query, key and value input sequences associated to a single head, where 𝑁 is the sequence length and 𝑑 is the head dimension. Then the attention output O is computed as:

where softmax is applied row-wise and one typically sets 𝛼 = 1 GLYPH<157> 𝑑 as the scaling factor. In practice, we subtract rowmax ' S ' from S to prevent numerical instability with the exponential function. For

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 to produce the full output tensor.

Now let 𝜙 be a scalar loss function and let d 'GLYPH<0>' = 𝜕𝜙 GLYPH<157> 𝜕 'GLYPH<0>' be notation for the gradient. Given the output gradient dO 2 R 𝑁 GLYPH<2> 𝑑 , we compute dQ , dK , and dV according to the chain rule as follows:

Here, we have that d 𝑠 = ' diag ' 𝑝 'GLYPH<0> 𝑝𝑝 > ' d 𝑝 for 𝑝 = softmax ' 𝑠 ' as a function of a vector 𝑠 , and we 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 GPUhardware characteristics and execution model

Wedescribe the aspects of the GPU's execution model relevant for FLASHATTENTION-3, with a focus on the NVIDIA Hopper architecture as a concrete instantiation of this model.

Memoryhierarchy: The GPU's memories are organized as a hierarchy of data locales, with capacity inversely related to bandwidth (Table 1) 4 . Global memory (GMEM), also known as HBM, is the offchip 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-managed 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 a CTA, whereas each thread has at most 256 registers (RMEM) private to itself.

Asynchrony and warp-specialization: GPUs are throughput processors that rely on concurrency and asynchrony to hide memory and execution latencies. For async memory copy between GMEM and SMEM, Hopper has the Tensor Memory Accelerator (TMA) as 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, §9.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 ability to generate optimal instruction schedules [4]. In addition, Hopper supports the dynamic reallocation of registers between warpgroups via setmaxnreg [40, §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).

Low-precision number formats: 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.

However, correctly invoking FP8 WGMMA entails understanding the layout constraints on its operands. Given a GEMM call to multiply 𝐴 GLYPH<2> 𝐵 > for an 𝑀 GLYPH<2> 𝐾 -matrix 𝐴 and an 𝑁 GLYPH<2> 𝐾 -matrix 𝐵 , we say that the 𝐴 or 𝐵 operand is mn-major if it is contiguous in the outer 𝑀 or 𝑁 dimension, and k-major if is instead contiguous in the inner 𝐾 -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, in situations such as attention where one wants to fuse back-to-back 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 FLASHATTENTION was to leverage a local version of the softmax reduction to avoid these expensive intermediate reads/writes and fuse attention into a single kernel. Local softmax corresponds to lines 18-19 of the consumer mainloop in Algorithm 1 together with the rescalings of blocks of O . The simple derivation that this procedure indeed computes O can be found in [14, §2.3.1].

3 FlashAttention-3: Algorithm

In this section, we describe the FLASHATTENTION-3 algorithm. For simplicity, we focus on the forward pass, with the backward pass algorithm described in Appendix B.1. 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-softmax 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 warp-specialization and pingpong scheduling

Warp-specialization As with FLASHATTENTION-2, the forward pass of FLASHATTENTION-3 is embarrassingly parallel in the batch size, number of heads, and query sequence length. Thus, it will suffice to give a CTA-level view of the algorithm, which operates on a tile Q 𝑖 of the query matrix to compute the corresponding tile O 𝑖 of the output. To simplify the description, we first give the warp-specialization scheme with a circular SMEM buffer that does not have in addition the GEMM-softmax overlapping. Let 𝑑 be the head dimension, 𝑁 the sequence length, and fix a query block size 𝐵 𝑟 to divide Q into 𝑇 𝑟 = d 𝑁 𝐵 𝑟 e blocks Q 1 GLYPH<150> GLYPH<149> GLYPH<149> GLYPH<150> Q 𝑇 𝑟 .

For our implementation of Algorithm 1 on Hopper, we use setmaxnreg for (de)allocations, TMA for loads of Q 𝑖 and f K 𝑗 GLYPH<150> V 𝑗 g 0 GLYPH<20> 𝑗GLYPH<159>𝑇 𝑐 , 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 1, 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 warpspecialization, 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 an example, the H100 SXM5 GPU has 989 TFLOPS of FP16 matmul but only 3.9 TFLOPS of special functions such as exponential 5 (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 to matmul. The situation is even worse with FP8, where the matmul throughput doubles but the exponential throughput stays the same.

Algorithm 1 FLASHATTENTION-3 forward pass without intra-consumer overlapping - CTA view

Require: Matrices Q 𝑖 2 R 𝐵 𝑟 GLYPH<2> 𝑑 and K GLYPH<150> V 2 R 𝑁 GLYPH<2> 𝑑 in HBM, key block size 𝐵 𝑐 with 𝑇 𝑐 = d 𝑁 𝐵 𝑐 e .

1: Initialize pipeline object to manage barrier synchronization with 𝑠 -stage circular SMEM buffer.

2: if in producer warpgroup then

3: Deallocate predetermined number of registers.

4: Issue load Q 𝑖 from HBM to shared memory.

5: Upon completion, commit to notify consumer of the load of Q 𝑖 .

6: for 0 GLYPH<20> 𝑗 GLYPH<159>𝑇 𝑐 do

7: Wait for the ' 𝑗 % 𝑠 ' th stage of the buffer to be consumed.

8: Issue loads of K 𝑗 GLYPH<150> V 𝑗 from HBM to shared memory at the ' 𝑗 % 𝑠 ' th stage of the buffer.

9: Upon completion, commit to notify consumers of the loads of K 𝑗 GLYPH<150> V 𝑗 .

10: end for

11: else

12: Reallocate predetermined number of registers as function of number of consumer warps.

13: On-chip, initialize O 𝑖 = ' 0 ' 2 R 𝐵 𝑟 GLYPH<2> 𝑑 and ℓ 𝑖 GLYPH<150> 𝑚 𝑖 = ' 0 ' GLYPH<150> 'GLYPH<0>1' 2 R 𝐵 𝑟 .

14: Wait for Q 𝑖 to be loaded in shared memory.

15: for 0 GLYPH<20> 𝑗 GLYPH<159>𝑇 𝑐 do

16: Wait for K 𝑗 to be loaded in shared memory.

17: Compute S ' 𝑗 ' 𝑖 = Q 𝑖 K 𝑇 𝑗 (SS-GEMM). Commit and wait.

18: Store 𝑚 old 𝑖 = 𝑚 𝑖 and compute 𝑚 𝑖 = max ' 𝑚 old 𝑖 GLYPH<150> rowmax ' S ' 𝑗 ' 𝑖 '' .

19: Compute e P ' 𝑗 ' 𝑖 = exp ' S ' 𝑗 ' 𝑖 GLYPH<0> 𝑚 𝑖 ' and ℓ 𝑖 = exp ' 𝑚 old 𝑖 GLYPH<0> 𝑚 𝑖 ' ℓ 𝑖 , rowsum ' e P ' 𝑗 ' 𝑖 ' .

20: Wait for V 𝑗 to be loaded in shared memory.

21: Compute O 𝑖 = diag ' exp ' 𝑚 old 𝑖 GLYPH<0> 𝑚 𝑖 '' O 𝑖 , e P ' 𝑗 ' 𝑖 V 𝑗 (RS-GEMM). Commit and wait.

22: Release the ' 𝑗 % 𝑠 ' th stage of the buffer for the producer.

23: end for

24: Compute O 𝑖 = diag ' ℓ 𝑖 ' GLYPH<0> 1 O 𝑖 and 𝐿 𝑖 = 𝑚 𝑖 , log ' ℓ 𝑖 ' .

25: Write O 𝑖 and 𝐿 𝑖 to HBM as the 𝑖 th block of O and 𝐿 .

26: end if

Since the exponential is performed by a separate hardware unit (the multi-function unit), ideally we'd want the exponential calculation to be scheduled when the Tensor Cores are performing the matmul. To do so, we use synchronization barriers ( bar.sync instructions) to force the GEMMs (GEMM1 PV of one iteration, and GEMM0 QK > of the next iteration) of warpgroup 1 to be scheduled before the GEMMsofwarpgroup 2. As a result, the softmax of warpgroup 1 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. 1. Though in practice the pingpong scheduling is not as clean as depicted in the figure, we generally find this to improve performance (e.g., from 570 TFLOPS to 620-640 TFLOPS for FP16 forward with head dimension 128 and sequence length 8192).

Warpgroup1

GEMMO

Softmax

GEMM1

GEMMO

Softmax

GEMM1

GEMMO

Softmax

1

1

I

1

1

1

1

I

1

1

Warpgroup2

GEMMO

Softmax

GEMM1

GEMMO

Softmax

GEMM1

GEMMO

Softmax

1

1

1

1

1

time

Attention variants For multi-query attention [50] and grouped query attention [3], we follow the approach in FLASHATTENTION-2 and adjust the tensor indexing to avoid duplicating K and V in HBM.

3.2 Intra-warpgroup overlapping GEMMs and softmax

Even within one warpgroup, we can overlap some instructions in the softmax with some 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 S ' 𝑗 ' 𝑖 of the first GEMM, while the second GEMM takes its result e P ' 𝑗 ' 𝑖 as an operand. Indeed, the wait statements in lines 17 and 21 of Algorithm 1 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 6 GEMM-softmax pipelining algorithm:

WGMMAO

0

2

N -1

Softmax

0

2

N-1

WGMMA1

0

1

N-2

N-1

time

Algorithm 2 FLASHATTENTION-3 consumer warpgroup forward pass

Require: Matrices Q 𝑖 2 R 𝐵 𝑟 GLYPH<2> 𝑑 and K GLYPH<150> V 2 R 𝑁 GLYPH<2> 𝑑 in HBM, key block size 𝐵 𝑐 with 𝑇 𝑐 = d 𝑁 𝐵 𝑐 e .

1: Reallocate predetermined number of registers as function of number of consumer warps.

2: On-chip, initialize O 𝑖 = ' 0 ' 2 R 𝐵 𝑟 GLYPH<2> 𝑑 and ℓ 𝑖 GLYPH<150> 𝑚 𝑖 = ' 0 ' GLYPH<150> 'GLYPH<0>1' 2 R 𝐵 𝑟 .

3: Wait for Q 𝑖 and K 0 to be loaded in shared memory.

4: Compute S cur = Q 𝑖 K 𝑇 0 using WGMMA. Commit and wait.

6: Compute 𝑚 𝑖 , ˜ P cur and ℓ 𝑖 based on S cur , and rescale O 𝑖 .

5: Release the 0 th stage of the buffer for K .

7: for 1 GLYPH<20> 𝑗 GLYPH<159>𝑇 𝑐 GLYPH<0> 1 do

8: Wait for K 𝑗 to be loaded in shared memory.

9: Compute S next = Q 𝑖 K 𝑇 𝑗 using WGMMA. Commit but do not wait.

10: Wait for V 𝑗 GLYPH<0> 1 to be loaded in shared memory.

11: Compute O 𝑖 = O 𝑖 , ˜ P cur V 𝑗 GLYPH<0> 1 using WGMMA. Commit but do not wait.

12: Wait for the WGMMA Q 𝑖 K 𝑇 𝑗 .

13: Compute 𝑚 𝑖 , ˜ P next and ℓ 𝑖 based on S next .

14: Wait for the WGMMA ˜ P cur V 𝑗 GLYPH<0> 1 and then rescale O 𝑖

15: Release the ' 𝑗 % 𝑠 ' th, resp. ' 𝑗 GLYPH<0> 1% 𝑠 ' th stage of the buffer for K , resp. V .

16: Copy S next to S cur .

17: end for

18: Wait for V 𝑇 𝑐 GLYPH<0> 1 to be loaded in shared memory.

19: Compute O 𝑖 = O 𝑖 , ˜ P last V 𝑇 𝑐 GLYPH<0> 1 using WGMMA. Commit and wait.

20: Epilogue: Rescale O 𝑖 based on 𝑚 𝑖 . Compute 𝐿 𝑖 based on 𝑚 𝑖 and ℓ 𝑖 . Write O 𝑖 and 𝐿 𝑖 to HBM as the 𝑖 -th block of O and 𝐿 .

Algorithm 2 functions as a replacement for the consumer path of Algorithm 1 to comprise the complete FLASHATTENTION-3 algorithm for FP16 precision. At a high-level, we use WGMMA as a metonym for asynchronous GEMM. Within the mainloop (lines 8 to 16), the second WGMMA operation of iteration 𝑗 (line 11) is overlapped with softmax operations from iteration 𝑗 , 1 (line 13).

While the pipelined structure illustrated above offers theoretical performance gains, there are several practical aspects to consider:

T0{d0, d1}

T1{d0, d1}

T0{d4, d5}

T1{d4, d5}

T2{d0, d1}

T3{d0, d1}

T2{d4, d5}

T3{d4, d5}

T0{d2, d3}

T1{d2, d3}

T0{d6, d7}

T1{d6, d7}

T2{d2, d3}

T3{d2, d3}

T2{d6, d7}

T3{d6, d7}

T0{a0, a1}

T0{a2, a3}

T1{a0, a1}

T1{a2, a3}

T2{a0, a1}

T2{a2, a3}

T3{a0, a1}

T3{a2, a3}

T0{a4, a5}

T0{a6, a7}

T1{a4, a5}

T1{a6, a7}

T2{a4, a5}

T2{a6, a7}

T3{a4, a5}

T3{a6, a7}

Compiler reordering The pseudocode represents an idealized execution order but the compiler (NVCC)often rearranges instructions for optimization. This can disrupt the carefully crafted WGMMA and non-WGMMA 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.2).

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 S next must be kept in registers, leading to extra register usage of size 𝐵 𝑟 GLYPH<2> 𝐵 𝑐 GLYPH<2> sizeof ' float ' per threadblock. This increased register demand may conflict with using larger block sizes (another common optimization), which is also register-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-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 an 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

Efficiency: FP8 layout for accumulator and operand. Computing the forward pass of FLASHATTENTION-3 in FP8 precision poses two additional challenges not encountered for FP16 in terms of layout conformance. The first relates to the datatype conversion of the first WGMMA's FP32 accumulator to the second WGMMA's lower-precision (FP16 or FP8) operand, which was left implicit in Algorithm 1. Specifically, after downcasting to FP8, we need to transform the register ownership pattern from that depicted in Fig. 3 into Fig. 4, repeated per every four threads of the consumer warpgroups.

In Appendix B.7, we give a 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 V should be contiguous in the sequence length dimension. We perform in-kernel transposition of the V 𝑗 tiles in SMEM prior to invoking the second FP8 WGMMA, since in practice, V is typically assumed to be contiguous in the head dimension. In Appendix B.8, we describe in details how to perform the V transpose as part of the attention kernel itself.

Accuracy: block quantization and incoherent processing. With FP8 (e4m3) 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 [20, 53] 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:

Block quantization : we keep one scalar per block, so that for each of Q , K , V we split the tensor into blocks of size 𝐵 𝑟 GLYPH<2> 𝑑 or 𝐵 𝑐 GLYPH<2> 𝑑 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-bandwidth bound). As the FLASHATTENTION-3 algorithm naturally operates on blocks, we can scale each block of S to account for this block quantization at no computation cost.

Incoherent processing : to even out outliers, we multiply Q and K with a random orthogonal matrix M before quantizing to FP8. Since M is orthogonal, MM > = 𝐼 and so ' QM '' KM ' > = 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 or KM is a random sum of entries of Q or K , thus reducing quantization error. In practice, we follow Chee et al. [8] and Tseng et al. [57] and choose M to be the product of random diagonal matrices of GLYPH<6> 1 and a Hadamard matrix, which can be multiplied in 𝑂 ' 𝑑 log 𝑑 ' instead of 𝑂 ' 𝑑 2 ' , and can also be fused with the rotary embedding at no extra computation cost.

Wevalidate that these two techniques reduces numerical error by up to 2.6 GLYPH<2> in §4.3.

4 Empirical Validation

Weuse the primitives from CUTLASS [56] such as WGMMA and TMA abstractions to implement FLASHATTENTION-3 and evaluate its efficiency and accuracy.

Benchmarking attention. We measure the runtime of FLASHATTENTION-3 across different sequence lengths and compare it to a standard implementation in PyTorch, FLASHATTENTION-2, FLASHATTENTION-2 in Triton (which uses H100-specific instructions), as well as a vendor's implementation of FLASHATTENTION-2 optimized for H100 GPUs from cuDNN. We confirm that FLASHATTENTION-3 is up to 2.0 GLYPH<2> faster than FLASHATTENTION-2 and 1.5 GLYPH<2> faster than FLASHATTENTION-2 in Triton. FLASHATTENTION-3 reaches up to 840 TFLOPs/s, 85% of the theoretical maximum TFLOPs/s on H100 GPUs.

Ablation study. We confirm that our algorithmic improvements with warp-specialization and GEMM-softmax pipelining contribute to the speedup of FLASHATTENTION-3.

Accuracy of FP8 attention. Wevalidate that block quantization and incoherent processing reduces the numerical error of FP8 FLASHATTENTION-3 by 2.6 GLYPH<2> .

4.1 Benchmarking Attention

Wemeasure 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 inputs. We report the results in Fig. 5 and Fig. 6, showing that FLASHATTENTION-3 is around 1.5-2.0 GLYPH<2> faster than FLASHATTENTION-2 in the forward pass and 1.5-1.75 GLYPH<2> faster in the backward pass. Compared to a standard attention implementation, FLASHATTENTION-3 can be up to 3-16 GLYPH<2> faster. For medium and long sequences (1k and above), FLASHATTENTION-3 even surpasses the speed of a vendor's library (cuDNN - closed source) that has been optimized for H100 GPUs.

Benchmark settings: Wevary 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, or 8 heads). To calculate the FLOPs of the forward pass, we use:

With causal masking, we divide this number by 2 to account for the fact that approximately only half of the entries are calculated. To get the FLOPs of the backward pass, 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).

Wealso measure the runtime for FP8 for the forward pass under similar settings. We report the results for headdim 256 in Fig. 7 and give the full results in Appendix C.2.

4.2 Ablation Study: 2-Stage Pipelining Experiments

Weablate both the 2-stage WGMMA-softmax pipelining and warp-specialization for non-causal FP16 FLASHATTENTION-3 with fixed parameters f batch GLYPH<150> seqlen GLYPH<150> nheads GLYPH<150> hdim g = f 4 GLYPH<150> 8448 GLYPH<150> 16 GLYPH<150> 128 g . The result in Table 2 confirms that our algorithmic improvements (asynchrony with warp-specialization 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 a standard implementation of attention against

512

1k

2k

4k

8k

16k

Sequence length

200

400

600

Speed (TFLOPs/s)

52

63

67

72

73

OOM

226

268

289

302

308

312

299

370

396

411

419

423

299

370

396

411

419

423

387

452

492

516

529

533

351

422

498

520

562

566

Attention forward speed, head dim 64 (H100 80GB SXM5)

Standard attention

FlashAttention-2

Triton

cuDNN

FlashAttention-3

(a) Forward, without causal mask, head dim 64

512

1k

2k

4k

8k

16k

Sequence length

200

400

600

Speed (TFLOPs/s)

74

100

119

133

139

OOM

232

289

324

346

358

364

314

373

412

436

449

451

314

373

412

436

449

451

514

599

646

666

678

681

482

607

678

717

750

760

Attention forward speed, head dim 128 (H100 80GB SXM5)

Standard attention

FlashAttention-2

Triton

cuDNN

FlashAttention-3

(c) Forward, without causal mask, head dim 128

Attention forward speed, head dim 256 (H100 80GB SXM5)

512

1k

2k

4k

8k

16k

Sequence length

200

400

600

Speed (TFLOPs/s)

228

276

305

322

332

337

223

260

279

295

302

305

223

260

279

295

302

305

501

619

690

728

749

758

512

661

749

795

831

842

FlashAttention-2

Triton

cuDNN

FlashAttention-3

(e) Forward, without causal mask, head dim 256

512

1k

2k

4k

8k

16k

Sequence length

200

400

600

Speed (TFLOPs/s)

16

18

18

18

18

OOM

146

201

241

271

289

298

185

249

296

331

353

370

185

249

296

331

353

370

215

306

390

450

487

514

213

278

391

444

497

520

Attention forward speed, head dim 64 (H100 80GB SXM5)

Standard attention

FlashAttention-2

Triton

cuDNN

FlashAttention-3

(b) Forward, with causal mask, head dim 64

512

1k

2k

4k

8k

16k

Sequence length

200

400

600

Speed (TFLOPs/s)

26

31

34

35

35

OOM

145

210

261

295

315

328

202

278

333

378

401

415

202

278

333

378

401

415

334

442

525

579

609

629

304

443

554

624

668

697

Attention forward speed, head dim 128 (H100 80GB SXM5)

Standard attention

FlashAttention-2

Triton

cuDNN

FlashAttention-3

(d) Forward, with causal mask, head dim 128

Attention forward speed, head dim 256 (H100 80GB SXM5)

512

1k

2k

4k

8k

16k

Sequence length

200

400

600

Speed (TFLOPs/s)

150

211

250

275

290

299

143

197

236

263

280

291

143

197

236

263

280

291

319

450

540

601

634

659

302

446

563

645

696

724

FlashAttention-2

Triton

cuDNN

FlashAttention-3

(f) Forward, with causal mask, head dim 256

a reference implementation in FP64. To simulate outlier features and activations in LLMs [20, 53], we generate the entries of Q GLYPH<150> K GLYPH<150> V with the following distribution:

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 3. In FP16, both FLASHATTENTION-2 and FLASHATTENTION-3 achieves 1.7 GLYPH<2> lower RMSEcomparedtothestandard implementation since intermediate results (softmax) are kept in FP32. The baseline attention in FP8 uses per-tensor scaling, with

512

1k

2k

4k

8k

16k

Sequence length

200

400

600

Speed (TFLOPs/s)

68

76

88

92

95

OOM

198

238

264

279

287

291

275

366

425

463

478

475

288

397

479

531

562

552

Attention backward speed, head dim 64 (H100 80GB SXM5)

Standard attention

FlashAttention-2

cuDNN

FlashAttention-3

(a) Backward, without causal mask, head dim 64

512

1k

2k

4k

8k

16k

Sequence length

200

400

600

Speed (TFLOPs/s)

104

131

159

174

181

OOM

214

260

291

310

318

322

304

426

509

562

587

576

317

440

531

593

616

615

Attention backward speed, head dim 128 (H100 80GB SXM5)

Standard attention

FlashAttention-2

cuDNN

FlashAttention-3

(b) Backward, without causal mask, head dim 128

512

1k

2k

4k

8k

16k

Sequence length

400

800

1200

Speed (TFLOPs/s)

529

664

766

854

897

903

529

664

766

854

897

903

686

878

1001

1087

1122

1139

761

1014

1172

1258

1302

1322

Attention forward speed, head dim 256 (H100 80GB SXM5)

Triton

cuDNN

FlashAttention-3

(a) Forward, without causal mask, head dim 256

512

1k

2k

4k

8k

16k

Sequence length

400

800

1200

Speed (TFLOPs/s)

299

425

520

591

628

663

299

425

520

591

628

663

304

449

768

1015

1056

1099

438

649

834

972

1070

1132

Attention forward speed, head dim 256 (H100 80GB SXM5)

Triton

cuDNN

FlashAttention-3

(b) Forward, with causal mask, head dim 256

matmul accumulator in FP32 and intermediate softmax results kept in FP16. Thanks to block quantization and incoherent processing, FLASHATTENTION-3 in FP8 is 2.6 GLYPH<2> more accurate than this baseline.

5 Dicussion, Limitations, Conclusion

With FLASHATTENTION-3, 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.0 GLYPH<2> times compared to FLASHATTENTION-2, and reduce FP8 numerical error by 2.6 GLYPH<2> compared to standard per-tensor quantization. Some limitations of our work that we hope to address in the future include: optimizing for 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-context tasks.

References

Ahmad Abdelfattah, Azzam Haidar, Stanimire Tomov, and Jack Dongarra. Performance, design, and autotuning of batched gemm for gpus. pages 21-38, 06 2016. ISBN 978-3-319-41320-4. doi: 10.1007/978-3-319-41321-1_2.

AI21. Introducing jamba: Ai21's groundbreaking ssm-transformer model. AI21 blog , 2024.

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 preprint arXiv:2305.13245 , 2023.

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 '11, New York, NY, USA, 2011. Association for Computing Machinery. ISBN 9781450307710. doi: 10.1145/2063384.2063400. URL https://doi.org/10.1145/2063384.2063400 .

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 arXiv:2405.04517 , 2024.

Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150 , 2020.

William Brandon, Aniruddha Nrusimha, Kevin Qian, Zachary Ankner, Tian Jin, Zhiye Song, and Jonathan Ragan-Kelley. Striped attention: Faster ring attention for causal transformers. arXiv preprint arXiv:2311.09431 , 2023.

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.

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.

Richard J 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 IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 16144-16155, 2022.

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 , 2019.

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.

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 on Learning Representations (ICLR) , 2020.

Tri Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023. URL https://arxiv.org/abs/2307.08691 .

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.

Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems , 2022.

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.

Tri Dao, Daniel Haziza, Francisco Massa, and Grigory Sizov. Flash-Decoding for long-context inference, 2023. URL https://pytorch.org/blog/flash-decoding/ .

DeepSeek-AI. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.04434 , 2024.

Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm. int8 (): 8-bit matrix multiplication for transformers at scale. CoRR abs/2208.07339 , 2022.

Alicia Golden, Samuel Hsia, Fei Sun, Bilge Acun, Basil Hosmer, Yejin Lee, Zachary DeVito, Jeff Johnson, Gu-Yeon Wei, David Brooks, et al. Is flash attention stable? arXiv preprint arXiv:2405.02803 , 2024.

Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. 2023.

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-augmented transformer for speech recognition. arXiv preprint arXiv:2005.08100 , 2020.

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.07916 , 2021.

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.

Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length llm inference with kv cache quantization. arXiv preprint arXiv:2401.18079 , 2024.

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning , pages 5156-5165. PMLR, 2020.

Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In The International Conference on Machine Learning (ICML) , 2020.

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 PagedAttention. In Proceedings of the 29th Symposium on Operating Systems Principles , pages 611-626, 2023.

Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Marc Marone, Christopher Akiki, Jia Li, Jenny Chim, et al. Starcoder: may the source be with you! arXiv preprint arXiv:2305.06161 , 2023.

Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context. arXiv preprint arXiv:2310.01889 , 2023.

Hao Liu, Wilson Yan, Matei Zaharia, and Pieter Abbeel. World model on million-length video and language with ringattention. arXiv preprint arXiv:2402.08268 , 2024.

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.02750 , 2024.

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/abs/2402.13499 .

Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: Moving average equipped gated attention. In The International Conference on Learning Representations (ICLR) , 2023.

Xuezhe Ma, Xiaomeng Yang, Wenhan Xiong, Beidi Chen, Lili Yu, Hao Zhang, Jonathan May, Luke Zettlemoyer, Omer Levy, and Chunting Zhou. Megalodon: Efficient llm pretraining and inference with unlimited context length. arXiv preprint arXiv:2404.08801 , 2024.

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. arXiv preprint arXiv:2209.05433 , 2022.

NVIDIA. CUDA Programming Guide Version 12.4, 2024. URL https://docs.nvidia.com/ cuda/cuda-c-programming-guide/index.html .

Nvidia. Accelerating transformers with nvidia cudnn 9. Nvidia blog , 2024. URL https://developer.nvidia.com/blog/ accelerating-transformers-with-nvidia-cudnn-9/ .

NVIDIA. Parallel Thread Execution ISA Version 8.4, 2024. URL https: //docs.nvidia.com/cuda/pdf/ptx_isa_8.4.pdf .

MuhammadOsama, 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-431, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400700156. doi: 10.1145/3572848.3577479. URL https://doi.org/10.1145/3572848.3577479 .

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 arXiv:2305.13048 , 2023.

Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.00071 , 2023.

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.

Markus N Rabe and Charles Staats. Self-attention does not need 𝑂 ' 𝑛 2 ' memory. arXiv preprint arXiv:2112.05682 , 2021.

Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing Transformers. arXiv preprint arXiv:2003.05997 , 2020.

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 llama: Open foundation models for code. arXiv preprint arXiv:2308.12950 , 2023.

Rya Sanovar, Srikant Bharadwaj, Renee St. Amant, Victor Rühle, and Saravan Rajmohan. Lean attention: Hardware-aware scalable attention mechanism for the decode-phase of transformers. 2024.

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 arXiv:2201.03533 , 2022.

Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150 , 2019.

Benjamin Spector, Aaryan Singhal, Simran Arora, and Christopher Ré, 2024. URL https://github.com/HazyResearch/ThunderKittens .

Fei Sun, Jun Liu, Jian Wu, Changhua Pei, Xiao Lin, Wenwu Ou, and Peng Jiang. Bert4rec: 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.

Mingjie Sun, Xinlei Chen, J Zico Kolter, and Zhuang Liu. Massive activations in large language models. arXiv preprint arXiv:2402.17762 , 2024.

Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08621 , 2023.

Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. arXiv preprint arXiv:2009.06732 , 2020.

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. URL https://github.com/NVIDIA/cutlass .

Albert Tseng, Jerry Chee, Qingyao Sun, Volodymyr Kuleshov, and Christopher De Sa. Quip#: Even better llm quantization with hadamard incoherence and lattice codebooks. arXiv preprint arXiv:2402.04396 , 2024.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

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.07887 , 2024.

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.

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.03629 , 2022.

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.

Zyphra. Zyphra unveils zamba: A compact 7b ssm hybrid model. Zyphra blog , 2024.

A Related Work

Attention variants and distributed attention Ever since attention became popular with the Transformer architecture [58], there has been a large body of work on 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 a fixed pattern [11], with a sliding window [6], or with a dynamic pattern through hashing [28] or routing [46]. The low-rank approach instead assumes that the attention matrix has a low-rank structure, and apply a pointwise nonlinearity to the query and key [27] with random projection [12, 44, 60]. One can also combine the sparse and low-rank approximation for better quality [9, 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 [3] tie different heads of K and V , and multiple query heads interact with the same key and value head. Multi-head latent attention [19] 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 GLYPH<150> K GLYPH<150> V are obtained. As a result, any efficiency or accuracy improvement to the standard attention computation benefits these methods.

To extend to even longer context, attention computation can be distributed across multiple GPUs. Methods such as Ring attention [31, 32] and variants [7] can reach a context length of up to 1 million. They use FLASHATTENTION (or FLASHATTENTION-2) as a primitive, and so the improvement from FLASHATTENTION-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 [27] and recurrent neural networks (RNNs). RWKV [42], H3 [17], 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 Transformers 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 see some traction, seeing usage in some medium to large-scale models such as Jamba [2], Zamba [63], Megalodon [36], and Mamba2-hybrid [59]. For the highest quality, these SSM- and RNN-based 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.

Low-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. QuIP [8] and QuIP#[57] use incoherent processing to reduce the quantization, and we adapted this technique for FP8 FLASHATTENTION-3. 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 a natively asynchronous programming model. There are other orthogonal axes for hardware-aware 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 [41] to achieve nearly peak occupancy. There is a large literature on optimizing GEMM for specific hardware that employs many of the same techniques. As an example, Abdelfattah et al. [1] 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-art performance.

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 a dQ writer, since we 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 a 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).

Weinclude the backward pass with warp specialization in Algorithm 3.

Algorithm 3 FLASHATTENTION-3 backward pass with warp specialization

Require: Matrices Q GLYPH<150> K GLYPH<150> V GLYPH<150> O GLYPH<150> dO 2 R 𝑁 GLYPH<2> 𝑑 in HBM, logsumexp vector 𝐿 2 R 𝑁 in HBM, block sizes 𝐵 𝑐 , 𝐵 𝑟 .

1: In a preprocessing kernel, compute 𝐷 = rowsum ' dO GLYPH<14> O ' 2 R 𝑑 (pointwise multiply), write 𝐷 to HBManddivide it into 𝑇 𝑟 blocks 𝐷 1 GLYPH<150> GLYPH<149> GLYPH<149> GLYPH<149> GLYPH<150>𝐷 𝑇 𝑟 of size 𝐵 𝑟 each.

2: Divide Q into 𝑇 𝑟 = l 𝑁 𝐵 𝑟 m blocks Q 1 GLYPH<150> GLYPH<149> GLYPH<149> GLYPH<149> GLYPH<150> Q 𝑇 𝑟 of size 𝐵 𝑟 GLYPH<2> 𝑑 each, and divide K GLYPH<150> V in to 𝑇 𝑐 = l 𝑁 𝐵 𝑐 m blocks K 1 GLYPH<150> GLYPH<149> GLYPH<149> GLYPH<149> GLYPH<150> K 𝑇 𝑐 and V 1 GLYPH<150> GLYPH<149> GLYPH<149> GLYPH<149> GLYPH<150> V 𝑇 𝑐 , of size 𝐵 𝑐 GLYPH<2> 𝑑 each.

3: Divide dO into 𝑇 𝑟 blocks dO 𝑖 GLYPH<150> GLYPH<149> GLYPH<149> GLYPH<149> GLYPH<150> dO 𝑇 𝑟 of size 𝐵 𝑟 GLYPH<2> 𝑑 each, and divide 𝐿 into 𝑇 𝑟 blocks 𝐿 𝑖 GLYPH<150> GLYPH<149> GLYPH<149> GLYPH<149> GLYPH<150>𝐿 𝑇 𝑟 of size 𝐵 𝑟 each.

4: Initialize pipeline object to manage barrier synchronization with 𝑠 -stage circular SMEM buffer.

5: if 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 to notify consumer of the load of K 𝑗 and V 𝑗 .

9: for 1 GLYPH<20> 𝑖 GLYPH<20> 𝑇 𝑟 do

10:

11:

12:

13:

Wait for the ' 𝑖 % 𝑠 ' th stage of the buffer to be consumed.

Issue loads of Q 𝑖 GLYPH<150> dO 𝑖 from HBM to shared memory at the ' 𝑖 % 𝑠 ' th stage of the buffer.

Upon completion, commit to notify consumers of the loads of Q 𝑖 GLYPH<150> dO 𝑖 .

end for

14: else if in consumer warpgroups then

15: Reallocate predetermined number of registers as function of number of consumer warps.

16: On-chip, Initialize dK 𝑗 = ' 0 ' 𝐵 𝑐 GLYPH<2> 𝑑 GLYPH<150> dV 𝑗 = ' 0 ' 𝐵 𝑐 GLYPH<2> 𝑑 .

17: Wait for K 𝑗 and V 𝑗 to be loaded in shared memory.

18: for 1 GLYPH<20> 𝑖 GLYPH<20> 𝑇 𝑟 do

19: Wait for Q 𝑖 to be loaded in shared memory.

20: Load 𝐿 𝑖 GLYPH<150>𝐷 𝑖 from HBM to on-chip SRAM.

21: Onchip, compute S ' 𝑗 ' 𝑖 = Q 𝑖 K 𝑇 𝑗 2 R 𝐵 𝑟 GLYPH<2> 𝐵 𝑐 (SS-GEMM). Commit.

22: Wait for dO 𝑖 to be loaded in shared memory.

23: Onchip, compute dP ' 𝑗 ' 𝑖 = dO 𝑖 V > 𝑗 2 R 𝐵 𝑟 GLYPH<2> 𝐵 𝑐 (SS-GEMM). Commit.

24: Onchip, wait for S ' 𝑗 ' 𝑖 , then compute P ' 𝑗 ' 𝑖 = exp ' S 𝑖 𝑗 GLYPH<0> 𝐿 𝑖 ' 2 R 𝐵 𝑟 GLYPH<2> 𝐵 𝑐 .

25: Onchip, wait for dP ' 𝑗 ' 𝑖 , then compute dS ' 𝑗 ' 𝑖 = P ' 𝑗 ' 𝑖 GLYPH<14>' dP ' 𝑗 ' 𝑖 GLYPH<0> 𝐷 𝑖 ' 2 R 𝐵 𝑟 GLYPH<2> 𝐵 𝑐 .

26: Onchip, compute dV 𝑗 dV 𝑗 ,' P ' 𝑗 ' 𝑖 ' > dO 𝑖 2 R 𝐵 𝑐 GLYPH<2> 𝑑 (RS-GEMM). Commit.

27: Onchip, compute dK 𝑗 dK 𝑗 , dS ' 𝑗 ' 𝑖 > Q 𝑖 2 R 𝐵 𝑐 GLYPH<2> 𝑑 (RS-GEMM). Commit and wait for both dV 𝑗 and dK 𝑗 .

28: On chip, compute dQ ' local ' 𝑖 = dS ' 𝑗 ' 𝑖 K 𝑗 2 R 𝐵 𝑟 GLYPH<2> 𝑑 (SS-GEMM), and write dQ ' local ' 𝑖 to smem. Notify the dQ -writer.

29: end for

30: else if in dQ -writer warp then

31: for 1 GLYPH<20> 𝑖 GLYPH<20> 𝑇 𝑟 do

32: Wait for dQ ' local ' 𝑖 to be ready in smem.

33: Using a semaphore, atomically add dQ ' local ' 𝑖 to dQ 𝑖 in global memory.

34: end for

35: end if

B.2 2-Stage Pipelining SASS Analysis

Wegive simplified SASS code for the inside of the consumer warpgroup mainloop.

Wemakethe following observations:

Softmax is reordered to the very beginning, even before the first WGMMA.

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.

exp2 , row\_sum , O rescaling and FP32 ! FP16 conversions are interleaved together.

The second WGMMA is 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

Weexperiment with a 3-stage pipelining algorithm to parallelize the first WGMMA from iteration 𝑗 , 2 , softmax from iteration 𝑗 , 1 , and the second WGMMAfromiteration 𝑗 . We describe this algorithm in Algorithm 4. This algorithm behaves worse than the 2-stage pipelining algorithm due to the reasons below:

1

1

1

WGMMAO

0

1

2

3

4

N - 1

1

1

Softmax

0

1

2

3

N-2

N-1

WGMMA1

0

1

2

N-3

N - 2

N - 1

time

Algorithm 4 FLASHATTENTION 3-stage pipelining consumer warpgroup forward pass

𝑖

Require: Matrices Q GLYPH<150> K GLYPH<150> V 2 R 𝑁 GLYPH<2> 𝑑 in HBM, block sizes 𝐵 𝑐 , 𝐵 𝑟 . Each warpgroup reads 1 block Qi of size 𝐵 𝑟 GLYPH<2> 𝑑 , 𝑇 𝑐 = l 𝑁 𝐵 𝑐 m blocks K 1 GLYPH<150> GLYPH<149> GLYPH<149> GLYPH<149> GLYPH<150> K 𝑇 𝑐 and V 1 GLYPH<150> GLYPH<149> GLYPH<149> GLYPH<149> GLYPH<150> V 𝑇 𝑐 of size 𝐵 𝑐 GLYPH<2> 𝑑 . Each warpgroup writes 1 output block O of size 𝐵 GLYPH<2> 𝑑 , and 1 logsumexp block 𝐿 of size 𝐵 .

𝑟

𝑖

𝑟

1: Initialization. Load Q 𝑖 from HBM to on-chip SRAM. Initialize O 𝑖 GLYPH<150>ℓ 𝑖 GLYPH<150> 𝑚 𝑖 GLYPH<150>𝑠𝑐𝑎𝑙𝑒 _ 𝑜 .

2: Wait for the producer warpgroup loading K 0 from HBM to on-chip SRAM.

3: Compute S = Q 𝑖 K 𝑇 0 using WGMMA. Commit and wait.

4: Compute 𝑚 𝑖 , ˜ P 𝑖 , ℓ 𝑖 , 𝑠𝑐𝑎𝑙𝑒 _ 𝑜 based on S .

5: Wait for the producer warpgroup loading K 1 from HBM to on-chip SRAM.

6: Compute S = Q 𝑖 K 𝑇 1 using WGMMA. Commit and wait.

7: for 2 GLYPH<20> 𝑗 GLYPH<159>𝑇 𝑐 GLYPH<0> 2 do

8: Wait for the producer warpgroup loading K 𝑗 from HBM to on-chip SRAM.

9: Compute S _ 𝑛𝑒𝑥𝑡 = Q 𝑖 K 𝑇 𝑗 using WGMMA. Commit but do not wait.

10: Wait for the producer warpgroup loading V 𝑗 GLYPH<0> 2 from HBM to on-chip SRAM.

11: Rescale O 𝑖 based on 𝑠𝑐𝑎𝑙𝑒 _ 𝑜 .

12: Compute O 𝑖 = O 𝑖 , ˜ P 𝑖 V 𝑗 GLYPH<0> 2 using WGMMA. Commit but do not wait.

13: Compute 𝑚 𝑖 , ˜ P 𝑖 _ 𝑛𝑒𝑥𝑡 , ℓ 𝑖 , 𝑠𝑐𝑎𝑙𝑒 _ 𝑜 based on S .

14: Wait for all previous WGMMAs.

15: Copy S _ 𝑛𝑒𝑥𝑡 to S .

16: Copy ˜ P 𝑖 _ 𝑛𝑒𝑥𝑡 to ˜ P 𝑖 .

17: end for

18: Wait for the producer warpgroup loading V 𝑇 𝑐 GLYPH<0> 2 from HBM to on-chip SRAM.

19: Rescale O 𝑖 based on 𝑠𝑐𝑎𝑙𝑒 _ 𝑜 .

20: Compute O 𝑖 = O 𝑖 , ˜ P 𝑖 V 𝑇 𝑐 GLYPH<0> 2 using WGMMA. Commit and wait.

21: Compute 𝑚 𝑖 , ˜ P 𝑖 , ℓ 𝑖 , 𝑠𝑐𝑎𝑙𝑒 _ 𝑜 based on S .

22: Wait for the producer warpgroup loading V 𝑇 𝑐 GLYPH<0> 1 from HBM to on-chip SRAM.

23: Rescale O 𝑖 based on 𝑠𝑐𝑎𝑙𝑒 _ 𝑜 .

24: Compute O 𝑖 = O 𝑖 , ˜ P 𝑖 V 𝑇 𝑐 GLYPH<0> 1 using WGMMA. Commit and wait.

25: Epilogue. Rescale O 𝑖 based on ℓ 𝑖 . Compute 𝐿 𝑖 based on ℓ 𝑖 and 𝑚 𝑖 . Write O 𝑖 and 𝐿 𝑖 to HBM as the 𝑖 -th block of O and 𝐿 .

Overlapping. Weexpected that softmax can be overlapped with (the first WGMMA + the second WGMMA). However, the compiler doesn't cooperate in 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 to 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 𝑠𝑐𝑎𝑙𝑒 _ 𝑜 , which is of size 𝐵 𝑟 GLYPH<2> 𝐵 𝑐 GLYPH<2> sizeof ' input_data_type ', 𝐵 𝑟 GLYPH<2> sizeof ' float ' . As a result, a 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 a tensormap using the PTX instruction 'tensormap.replace' and store the

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 in a tile. When writing O, we leverage memory-coalesced writes directly, as this is the final step and does not require asynchrony. In the backward pass, a preprocess kernel pads each sequence in dQ, dPSum, and LSE tensors with an additional 128 (tile_size) elements, allowing us to utilize TMA store for efficient data transfer.

Threadblock cluster and TMA multi-cast We utilize TMA multi-cast with a cluster size of 2 for 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

Weapply 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 or variable sequence lengths, masking is applied only to the maximum Kblock 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 is a prologue (loading 𝑄 ) and epilogue (writing output) where the Tensor Cores are not running. To maximize efficiency, we implement a persistent kernel that can overlap the epilogue of one iteration with the prologue of the next iteration to 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 a 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 FP8 FLASHATTENTION-3

T0

T1

T2

T3

T0

T1

T2

T3

T0

T1

T2

T3

T0

T1

T2

T3

T0

T3

T1

T2

T0

T3

T1

T2

T0

T3

T1

T2

T0

T3

T1

T2

T0

T3

T1

T2

T0

T3

T1

T2

T1

T2

T0

T3

T1

T2

T0

T3

T0

T2

T0

T2

T0

T2

T0

T2

T1

T3

T1

T3

T1

T3

T1

T3

In code, we can effect the register-to-register data exchange that transforms the register ownership pattern of Fig. 3 into Fig. 4 through invoking a combination of the following two CUDA intrinsics:

byte_perm : Given two 32-bit unsigned integers x and y and selector s , the byte permute instruction returns 4 bytes from the 8 input bytes as specified by s .

shfl_sync : The shuffle instruction exchanges register data from a source lane index j into its own destination register.

Ourmethodis illustrated in Fig. 9. First, we can swap the order of data held within a thread's registers by using byte permute as follows. Referring to the top row of Fig. 9, for a given thread let upper bethefirst 4 bytes (those in light and dark blue) and let lower bethelast 4 bytes (those in light and dark yellow). Then for the data held by threads 1 and 2, we do the swap by calling byte_perm with the indicated selectors:

Nowbetween 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:

Finally, between the third and fourth rows, we 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:

B.8 In-kernel transposition of V for FP8 FLASHATTENTION-3

We describe how to fuse the memory transpose of V needed for the second FP8 WGMMA into FLASHATTENTION-3. This is handled as an out-of-place SMEM to RMEM to SMEM transfer that is executed in the producer warpgroup.

Specifically, within the producer mainloop, after issuing the TMA load of a tile of V , the producer warpgroup waits for the load to complete. Then, producer warps effect the transpose by issuing LDSM ( ldmatrix ) and STSM ( stmatrix ) instructions, which involve a warp of threads collectively loading SMEMtoRMEMandstoring RMEM to SMEM at a granularity of 128 bytes. Finally, we have an additional 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 as its consumer.

Wechoose 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 a 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 8 GLYPH<2> 8 matrices with 16-bit entries [40, §9.7.13.4.15-16], but we can pack 8-bit entries two at a time to use LDSM/STSM in the context of FP8 precision. However, the transpose versions of LDSM/STSM cannot split packed 8-bit entries, which necessitates certain register movements in between LDSM and STSMtoactually perform a tile-wise transpose. The use of byte permute to split and reorder packed 8-bit entries in between LDSM and STSM is depicted in the following code snippet:

Since this permutes the eventual rows of the transposed V tile, we also need to modify the register movementsontheconsumersidethat transform accumulator to operand P . Weexploit the mathematical fact that

for 𝜎 a permutation of the common inner dimension of 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-3 for inference

For decoding inference, the query sequence length is much shorter than the key/value sequence length, typically on the order of one or a 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 as measured by FLOPs/s, but loading the KV cache as fast as possible as measured by memory bandwidth. Furthermore, since the FLASHATTENTION-3 algorithm described in §3.1 parallelizes over the query sequence length, it can suffer from a lack of parallelism for decoding.

Wemaketwo modifications to FLASHATTENTION-3 to introduce more parallelism for decoding:

Split KV (or Flash-Decoding ): We split the attention kernel along the key/value sequence length, with the number of splits determined by a heuristic at launch, and combine the resulting outputs using a separate post-processing reduction kernel. 'Splitting' according to a parameter 𝑛 means that 𝑛 threadblocks load the same tile of Q and 𝑛 different segments of the KV cache, computing 𝑛 different output tiles O 1 GLYPH<150> GLYPH<149> GLYPH<149> GLYPH<149> GLYPH<150> O 𝑛 and lse vectors lse 1 GLYPH<150> GLYPH<149> GLYPH<149> GLYPH<149> GLYPH<150> lse 𝑛 , which we then use to compute O in the reduction kernel. We also allow for early exit of threadblocks whose given segment of the KV cache doesn't contribute to the final output, in which case the threadblock writes out GLYPH<0>1 as its lse . This amounts to essentially the same implementation as described in [18].

GQA packing : For multi-query attention or 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 a model architecture with 16 query heads per KV head and a query sequence length of 8, in which case a threadblock can pack all 16 query heads into its Q tile without any change to Algorithm 2. In practice, this yields up to 𝑁 x speedup over an implementation that doesn't do GQA packing, where 𝑁 is the GQA ratio.

FLASHATTENTION-3 for inference also features an implementation of PagedAttention [29] that was contributed by Kai Londenberg. Recall that PagedAttention is a memory optimization technique for efficiently storing the KV cache in terms of fixed-size pages. This entails separating the logical position of KV blocks from their physical addresses, with a block table defining the address translation [29, §4.2].

Now, prior implementations of TMA load in CUTLASS construct the tensor map object such that TMAtensor coordinates are determined using the physical GMEM tensor. To use a block table with TMA, Londenberg defines a new SM90_TMA_LOAD_PAGED_OP class and a 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 as an additional argument.

C Addition Details on Experiments and Benchmarking

C.1 System and libraries

Webenchmark 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:

CUDA12.3

cuDNN9.5.0.50

CUTLASS3.6

FLASHATTENTION 2.6.3

Triton 3.1

PyTorch 2.5.0

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

Weuse following sequence lengths: 512, 1024, 2048, 4096, 8192, 16384.

Attention forward speed, head dim 64 (H100 80GB SXM5)

512

1k

2k

4k

8k

16k

Sequence length

400

800

1200

Speed (TFLOPs/s)

392

444

473

499

506

511

392

444

473

499

506

511

344

398

447

413

431

438

347

455

586

621

667

681

Triton

cuDNN

FlashAttention-3

(a) Forward, without causal mask, head dim 64

Attention forward speed, head dim 128 (H100 80GB SXM5)

512

1k

2k

4k

8k

16k

Sequence length

400

800

1200

Speed (TFLOPs/s)

408

502

563

605

630

635

408

502

563

605

630

635

617

751

886

864

971

1001

525

720

822

906

970

999

Triton

cuDNN

FlashAttention-3

(c) Forward, without causal mask, head dim 128 Attention forward speed, head dim 256 (H100 80GB SXM5)

512

1k

2k

4k

8k

16k

Sequence length

400

800

1200

Speed (TFLOPs/s)

529

664

766

854

897

903

529

664

766

854

897

903

686

878

1001

1087

1122

1139

761

1014

1172

1258

1302

1322

Triton

cuDNN

FlashAttention-3

(e) Forward, without causal mask, head dim 256

512

1k

2k

4k

8k

16k

Sequence length

400

800

1200

Speed (TFLOPs/s)

234

325

393

440

459

481

234

325

393

440

459

481

164

244

369

475

533

572

185

286

428

513

595

637

Triton

cuDNN

FlashAttention-3

Attention forward speed, head dim 64 (H100 80GB SXM5)

512

1k

2k

4k

8k

16k

Sequence length

400

800

1200

Speed (TFLOPs/s)

241

340

413

464

483

510

241

340

413

464

483

510

253

384

528

719

883

922

287

447

601

749

846

918

Triton

cuDNN

FlashAttention-3

(d) Forward, with causal mask, head dim 128

Attention forward speed, head dim 256 (H100 80GB SXM5)

512

1k

2k

4k

8k

16k

Sequence length

400

800

1200

Speed (TFLOPs/s)

299

425

520

591

628

663

299

425

520

591

628

663

304

449

768

1015

1056

1099

438

649

834

972

1070

1132

Triton

cuDNN

FlashAttention-3

(f) Forward, with causal mask, head dim 256

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 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 NAanswer 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.

It is fine to include aspirational goals as motivation as long as it 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 §5

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a 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 a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be 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, a 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 to 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.

Guidelines:

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 proofs can either appear in the main paper or the supplemental material, but if 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 be complemented by formal proofs provided in appendix or supplemental material.

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 it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Detailed information in Appendix C.1.

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 be 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 a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it 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 to replicate the results, access to a hosted model (e.g., in the case of a 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) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) Werecognize 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-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to 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 a permissive license in the near future.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

Whileweencouragethereleaseofcodeanddata, weunderstandthatthismightnotbepossible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https://nips.cc/public/guides/CodeSubmissionPolicy ) 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 a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to 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, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to 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, in 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 a 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 are 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 run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or 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 a 96% CI, if the hypothesis of Normality of errors is not verified.

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 error rates).

If error bars are 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 to reproduce the experiments?

Answer: [Yes]

Justification: In Appendix C.1

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 as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 or 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 is 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 not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness 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 is 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.

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.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to 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 or models that have a 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 a 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.

Werecognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: 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?

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 or dataset.

Theauthors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be 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 are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. NewAssets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: No new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as 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 is 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 nor 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 an equivalent approval/review based on the requirements of your country or 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.

Werecognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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 as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nFLASH 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nA 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\n\nLIMITATIONS:\nThe authors adequately addressed the limitations and, if applicable, potential negative societal impact of their work", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

comments: 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.

STRENGTHS:
The authors present novel techniques that exploit hardware features like asynchrony and low-precision computation.
The performance gains are impressive, in speed and accuracy.
The empirical validation includes comprehensive benchmarks and comparisons with previous methods.
It 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.
The techniques are tailored to NVIDIA Hopper GPUs, and may have limited applicability to other hardware achitectures.

QUESTIONS:
The backward pass algorithm is only included in the appendix. It would be useful if it could be included in the main paper.
Can 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n[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\".\n\nQUESTIONS:\nPlease 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+).\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations of this paper. There is no potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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;\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\nThanks 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.\n\nLIMITATIONS:\n* The paper is only specifically optimized for the Hopper architecture and may not applies to other or future GPUs.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-13"}