FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision

Jay ShahGanesh BikshandiYing ZhangVijay ThakkarPradeep RamaniTri Dao

article2024NeurIPS456 citations

Introduces FlashAttention-3, a method that speeds up Transformer attention by up to 2x on Hopper GPUs by overlapping computation with data movement via hardware asynchrony and executing FP8 low-precision operations with 2.6x lower numerical error.

Listen

The attention mechanism is the core computational foundation of modern artificial intelligence models, including large language models. However, it represents a primary processing bottleneck because its computational demands scale quadratically with context length. While earlier solutions like FlashAttention-2 reduced data transfer between memory layers, they achieved only 35% utilization on modern GPU architectures such as NVIDIA Hopper H100 because they followed a synchronous processing model and did not fully exploit specialized asynchronous hardware features or low-precision numerical formats.

The article demonstrates an optimized attention algorithm, FlashAttention-3, designed to maximize throughput and numerical accuracy on modern GPUs. It evaluates how asynchronous hardware execution, software pipelining, and 8-bit floating-point (FP8) computations can accelerate both the training and execution of transformer models.

To achieve these gains, the researchers developed three main techniques: separating data movement and computation into specialized worker groups (warp-specialization), interleaving lower-throughput softmax calculations with asynchronous matrix multiplications across a two-stage software pipeline, and incorporating hardware-accelerated FP8 operations with block quantization and incoherent processing to prevent quantization errors caused by outlier values. The approach was empirically validated on NVIDIA H100 SXM5 GPUs across various sequence lengths (512 to 16,384 tokens), head dimensions, and masking configurations, averaging execution metrics across 100 benchmark iterations against standard PyTorch, FlashAttention-2, Triton, and proprietary cuDNN implementations.

The findings show that for 16-bit floating-point (FP16) inputs, FlashAttention-3 achieves a 1.5× to 2.0× speedup in the forward pass over FlashAttention-2, reaching up to 740 teraflops per second (75% GPU utilization) and a 1.5× to 1.75× speedup in the backward pass. For low-precision FP8 operations, performance approaches 1.2 petaflops per second, roughly doubling FP16 throughput. Furthermore, for long sequence lengths, FlashAttention-3 matches or outperforms closed-source, vendor-optimized libraries. In terms of numerical accuracy, FP8 FlashAttention-3 reduces error by 2.6× compared to standard baseline FP8 attention when processing outlier features.

These performance improvements have immediate practical implications: they substantially reduce computational training costs, increase inference speeds, and make expanding AI context windows to handle multi-document analysis, long video streams, and extensive codebases commercially feasible. The algorithm provides a drop-in replacement that delivers higher hardware utilization without requiring model architecture compromises.

Organizations developing or deploying large-scale language models should adopt FlashAttention-3 via open-source integrations (such as PyTorch and Hugging Face) to reduce infrastructure overhead. Teams should leverage FP8 low-precision pipelines where supported by hardware, utilizing the recommended block quantization to protect output quality against numerical degradation.

While confidence in the benchmarked hardware performance is high, several limitations remain. FlashAttention-3 was benchmarked specifically on Hopper architectures; deeper 3-stage pipelining showed diminished returns due to register pressure and compiler reordering; and the FP8 implementation currently lacks persistent kernel load-balancing optimizations, which slightly reduces efficiency on small sequence lengths with causal masking. Additional empirical validation is recommended when extending FP8 attention to full-scale foundation model pretraining.

Colfax Research Georgia Institute of Technology Meta NVIDIA Princeton University Together AI encoder-decoder gpus language model FlashAttention-3 FP8 attention producer-consumer asynchrony

arXiv: 2407.08608 Dao-AILab/flash-attention

Paper: FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, Tri Dao (2023). FlashAttention-3 directly builds on FlashAttention-2's work partitioning and tiling techniques to push GPU throughput further using Hopper-specific asynchronous features.
Paper: FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, Tri Dao et al. (2022). This foundational paper introduces the exact, IO-aware tiling and online softmax recomputation algorithm that forms the core baseline for all subsequent FlashAttention architectures.

No sufficiently relevant recommendations were found.

Cover for 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 $\times$ with FP16 reaching up to 740 TFLOPs/s (75% utilization), and with FP8 reaching close to 1.2 PFLOPs/s. We validate that FP8 FlashAttention-3 achieves 2.6 $\times$ lower numerical error than a baseline FP8 attention.

Topics

encoder-decoder gpus language model large language models quantization interleavers numerical solution gpgpu

Keywords

FlashAttention-3 FP8 attention producer-consumer asynchrony asynchronous execution FP8 Tensor Cores WGMMA instruction multi-head attention NVIDIA Hopper

1 Introduction
2 Background: Multi-Head Attention and GPU Characteristics
2.1 Multi-Head Attention
2.2 GPU hardware characteristics and execution model
2.3 Standard Attention and Flash Attention
3 FlashAttention-3: Algorithm
3.1 Producer-Consumer asynchrony through warp-specialization and pingpong scheduling
3.2 Intra-warpgroup overlapping GEMMs and softmax
3.3 Low-precision with FP8
4 Empirical Validation
4.1 Benchmarking Attention
4.2 Ablation Study: 2-Stage Pipelining Experiments
4.3 Numerical Error Validation
5 Dicussion, Limitations, Conclusion
References
A Related Work
B Addition Details on Algorithms
B.1 Asynchrony Through Warp Specialization for the Backward Pass
B.2 2-Stage Pipelining SASS Analysis
B.3 3-Stage Pipelining Algorithm
C Addition Details on Experiments and Benchmarking
C.1 System and libraries
C.2 FP8 Attention Full Results

Knowls

Knowl 1 — FlashAttention-3 Intra-Warpgroup GEMM-Softmax Two-Stage Pipelining
algorithm
In attention computation on NVIDIA Hopper GPUs, non-matrix operations (such as floating-point exponential and row reductions in softmax) execute on the multi-function unit at much lower throughput than matrix multiplications on Tensor Cores. Within a single consumer warpgroup, sequential data dependencies serialize execution because softmax requires the output of the first GEMM ( $S = Q K^T$ ), while the second GEMM ( $O = P V$ ) consumes the output of softmax. FlashAttention-3 breaks this dependency by pipelining execution across loop iterations using additional register space ( $B_r \times B_c \times \text{sizeof}(\text{float})$ per threadblock to buffer $S_{\text{next}}$ ). While the asynchronous warpgroup matrix-multiply-accumulate (WGMMA) instruction computes the matrix product $S_{\text{next}} = Q_i K_j^T$ and $O_i = O_i + \tilde{P}_{\text{cur}} V_{j-1}$ in the background, the CUDA cores execute the softmax reductions and scaling for the next iteration.
Input: Query block $Q_i \in \mathbb{R}^{B_r \times d}$ , key matrix $K \in \mathbb{R}^{N \times d}$ , value matrix $V \in \mathbb{R}^{N \times d}$ loaded in shared memory, block size $B_c$ , total column blocks $T_c = \lceil N / B_c \rceil$ .
Output: Output tile $O_i \in \mathbb{R}^{B_r \times d}$ and log-sum-exp vector $L_i \in \mathbb{R}^{B_r}$ written to global memory (HBM).
Initialize $O_i = (0) \in \mathbb{R}^{B_r \times d}$ , $\ell_i = (0) \in \mathbb{R}^{B_r}$ , and $m_i = (-\infty) \in \mathbb{R}^{B_r}$ in registers.
Wait for $Q_i$ and $K_0$ in shared memory.
Compute $S_{\text{cur}} = Q_i K_0^T$ using WGMMA; commit and wait.
Release buffer stage for $K_0$ .
Compute $m_i = \text{rowmax}(S_{\text{cur}})$ , $\tilde{P}_{\text{cur}} = \exp(S_{\text{cur}} - m_i)$ , $\ell_i = \text{rowsum}(\tilde{P}_{\text{cur}})$ , and rescale $O_i$ .
for $j = 1$ to $T_c - 2$ do
    Wait for $K_j$ in shared memory.
    Compute $S_{\text{next}} = Q_i K_j^T$ using WGMMA (commit without waiting).
    Wait for $V_{j-1}$ in shared memory.
    Compute $O_i = O_i + \tilde{P}_{\text{cur}} V_{j-1}$ using WGMMA (commit without waiting).
    Wait for WGMMA of $Q_i K_j^T$ to finish.
    Compute $m_{\text{next}} = \max(m_i, \text{rowmax}(S_{\text{next}}))$ , $\tilde{P}_{\text{next}} = \exp(S_{\text{next}} - m_{\text{next}})$ , and update $\ell_i$ .
    Wait for WGMMA of $\tilde{P}_{\text{cur}} V_{j-1}$ to finish, then rescale $O_i = \text{diag}(\exp(m_i - m_{\text{next}}))^{-1} O_i$ .
    Release shared memory buffer stages for $K_j$ and $V_{j-1}$ .
     $S_{\text{cur}} \leftarrow S_{\text{next}}$ , $\tilde{P}_{\text{cur}} \leftarrow \tilde{P}_{\text{next}}$ , $m_i \leftarrow m_{\text{next}}$ .
end for
Wait for $V_{T_c-1}$ in shared memory.
Compute $O_i = O_i + \tilde{P}_{\text{last}} V_{T_c-1}$ using WGMMA; commit and wait.
Rescale $O_i = \text{diag}(\ell_i)^{-1} O_i$ and compute $L_i = m_i + \log(\ell_i)$ .
Write $O_i$ and $L_i$ to global memory.
Knowl 2 — FlashAttention-3 Forward Pass with Producer-Consumer Warp Specialization
algorithm
FlashAttention-3 partitions warps within each cooperative thread array (CTA) into producer and consumer roles to exploit hardware asynchrony. Producer warpgroups manage data movement via the Tensor Memory Accelerator (TMA) to transfer tiles of $Q, K, V$ from global memory (HBM) into an $s$ -stage circular shared memory (SMEM) buffer without stalling. Consumer warpgroups execute Tensor Core matrix multiplications (WGMMA) and softmax reductions. Using setmaxnreg, the GPU dynamically allocates the bulk of registers to consumer warpgroups, while producer warps operate with minimal register footprints.
Input: Query block $Q_i \in \mathbb{R}^{B_r \times d}$ , key/value sequences $K, V \in \mathbb{R}^{N \times d}$ in HBM, block sizes $B_r, B_c$ , number of buffer stages $s$ , $T_c = \lceil N / B_c \rceil$ .
Output: Output block $O_i \in \mathbb{R}^{B_r \times d}$ and log-sum-exp block $L_i \in \mathbb{R}^{B_r}$ written to HBM.
Initialize pipeline object for barrier synchronization over an $s$ -stage circular SMEM buffer.
if in producer warpgroup then
    Deallocate unused registers via setmaxnreg.
    Issue TMA load of $Q_i$ from HBM to SMEM; commit and notify consumer.
    for $j = 0$ to $T_c - 1$ do
        Wait for the $(j \bmod s)$ -th buffer stage to be released by consumer.
        Issue TMA loads of $K_j$ and $V_j$ from HBM to the $(j \bmod s)$ -th stage in SMEM.
        Commit loads and notify consumer.
    end for
else
    Reallocate registers to consumer warps via setmaxnreg.
    Initialize on-chip accumulators $O_i = (0) \in \mathbb{R}^{B_r \times d}$ , $\ell_i = (0) \in \mathbb{R}^{B_r}$ , $m_i = (-\infty) \in \mathbb{R}^{B_r}$ .
    Wait for $Q_i$ in SMEM.
    for $j = 0$ to $T_c - 1$ do
        Wait for $K_j$ in SMEM.
        Compute $S_i^{(j)} = Q_i K_j^T$ via SS-WGMMA; commit and wait.
         $m_i^{\text{old}} \leftarrow m_i$ ; $m_i \leftarrow \max(m_i^{\text{old}}, \text{rowmax}(S_i^{(j)}))$ .
        Compute $\tilde{P}_i^{(j)} = \exp(S_i^{(j)} - m_i)$ and update $\ell_i = \exp(m_i^{\text{old}} - m_i)\ell_i + \text{rowsum}(\tilde{P}_i^{(j)})$ .
        Wait for $V_j$ in SMEM.
        Compute $O_i = \text{diag}(\exp(m_i^{\text{old}} - m_i))^{-1} O_i + \tilde{P}_i^{(j)} V_j$ via RS-WGMMA; commit and wait.
        Release the $(j \bmod s)$ -th buffer stage for the producer.
    end for
    Normalize $O_i = \text{diag}(\ell_i)^{-1} O_i$ and compute $L_i = m_i + \log(\ell_i)$ .
    Write $O_i$ and $L_i$ to HBM.
end if
Knowl 3 — FP8 Layout Transformation via In-Kernel Transposition and Register Swizzling
model/method
Executing back-to-back GEMMs in FP8 on NVIDIA Hopper architectures imposes strict layout constraints that do not exist in FP16:
1. In-kernel transpose of $V$ : FP8 WGMMA requires both matrix operands to be in $k$ -major layout (contiguous in the contraction dimension). For the second GEMM ( $P V$ ), $V$ must be contiguous along its sequence-length dimension rather than its default head dimension. To avoid external global memory transposition kernels, FlashAttention-3 executes an in-kernel transpose of $V$ tiles in shared memory using ldmatrix (LDSM) and stmatrix (STSM) instructions within the producer warpgroup, scheduled in the background shadow of prior GEMMs.
2. Accumulator register swizzling: The FP32 accumulator layout resulting from the first FP8 WGMMA ( $S = Q K^T$ ) does not match the register layout expected for operand $A$ in the subsequent FP8 WGMMA ( $P V$ ). FlashAttention-3 rearranges the thread register entries from sequence: $\{d_0, d_1, d_2, d_3, d_4, d_5, d_6, d_7\}$ into: $\{d_0, d_1, d_4, d_5, d_2, d_3, d_6, d_7\}$ using hardware byte permute instructions replicated every 8 bytes. This register permutation corresponds to permuting the columns of tile $P$ , which is matched with a corresponding row permutation applied during the in-kernel shared memory transpose of $V$ , completely avoiding cross-thread shuffle instructions.
Knowl 4 — FP8 Attention Quantization via Block-Wise Scaling and Incoherent Processing
model/method
FP8 (e4m3 format, using 4 exponent bits and 3 mantissa bits) exhibits limited dynamic range and suffers severe accuracy loss in the presence of activation outliers in large language models. FlashAttention-3 reduces quantization error through two complementary techniques:
1. Block Quantization: Instead of maintaining a single scale factor per tensor, FlashAttention-3 computes independent scale factors for individual blocks of $Q, K, V$ of size $B_r \times d$ or $B_c \times d$ . Because attention computation natively processes blocks, the block-wise scaling is integrated into the score scaling of $S$ at zero additional compute cost.
2. Incoherent Processing: To disperse outlier values across feature dimensions, $Q$ and $K$ are multiplied prior to FP8 quantization by a randomized orthogonal matrix $M \in \mathbb{R}^{d \times d}$ such that $M M^T = I$ . By orthogonality, $(Q M)(K M)^T = Q M M^T K^T = Q K^T$ , leaving the mathematical output invariant. $M$ is structured as the product of a random diagonal sign matrix $D \in \{-1, 1\}^{d \times d}$ and a normalized Walsh-Hadamard matrix $H$ , enabling multiplication in $\mathcal{O}(d \log d)$ time instead of $\mathcal{O}(d^2)$ . This step is fused directly into the preceding rotary positional embedding operation.
Knowl 5 — FlashAttention-3 Backward Pass with Warp-Specialized dQ Accumulation
algorithm
The backward pass of FlashAttention-3 computes gradients $dQ, dK, dV$ given output gradient $dO \in \mathbb{R}^{N \times d}$ and forward outputs $O, L$ . To eliminate contention caused by multiple threadblocks writing concurrently to the global memory of $dQ$ , the algorithm introduces a dedicated $dQ$ -writer warp alongside producer and consumer warpgroups.
Input: Matrices $Q, K, V, O, dO \in \mathbb{R}^{N \times d}$ in HBM, logsumexp $L \in \mathbb{R}^N$ , block sizes $B_c, B_r$ , $T_r = \lceil N / B_r \rceil$ , $T_c = \lceil N / B_c \rceil$ .
Output: Gradients $dQ, dK, dV \in \mathbb{R}^{N \times d}$ in HBM.
Compute $D = \text{rowsum}(dO \circ O) \in \mathbb{R}^N$ in a preprocessing kernel and partition into blocks $D_1, \dots, D_{T_r}$ .
Partition $Q, dO, L, D$ into $T_r$ blocks of row dimension $B_r$ , and $K, V$ into $T_c$ blocks of row dimension $B_c$ .
Initialize pipeline object with an $s$ -stage circular SMEM buffer.
if in producer warpgroup then
    Deallocate registers; issue TMA loads of $K_j, V_j$ into SMEM and commit.
    for $i = 1$ to $T_r$ do
        Wait for buffer stage $(i \bmod s)$ to be consumed.
        Issue TMA loads for $Q_i, dO_i$ into SMEM stage $(i \bmod s)$ ; commit and notify consumers.
    end for
else if in consumer warpgroups then
    Initialize on-chip accumulators $dK_j = (0) \in \mathbb{R}^{B_c \times d}$ , $dV_j = (0) \in \mathbb{R}^{B_c \times d}$ .
    Wait for $K_j, V_j$ in SMEM.
    for $i = 1$ to $T_r$ do
        Wait for $Q_i, dO_i$ in SMEM; load $L_i, D_i$ to SRAM.
        Compute $S_i^{(j)} = Q_i K_j^T \in \mathbb{R}^{B_r \times B_c}$ and $dP_i^{(j)} = dO_i V_j^T \in \mathbb{R}^{B_r \times B_c}$ via SS-WGMMA.
        Compute $P_i^{(j)} = \exp(S_i^{(j)} - L_i)$ and $dS_i^{(j)} = P_i^{(j)} \circ (dP_i^{(j)} - D_i)$ .
        Accumulate $dV_j \leftarrow dV_j + (P_i^{(j)})^T dO_i$ and $dK_j \leftarrow dK_j + (dS_i^{(j)})^T Q_i$ via RS-WGMMA.
        Compute $dQ_i^{(\text{local})} = dS_i^{(j)} K_j \in \mathbb{R}^{B_r \times d}$ via SS-WGMMA; write to SMEM and notify dQ-writer.
    end for
    Write $dK_j, dV_j$ to HBM.
else if in dQ-writer warp then
    for $i = 1$ to $T_r$ do
        Wait for $dQ_i^{(\text{local})}$ in SMEM.
        Atomically add $dQ_i^{(\text{local})}$ to $dQ_i$ in HBM using a semaphore.
    end for
end if
Knowl 6 — Pingpong Scheduling for Inter-Warpgroup GEMM and Softmax Overlap
model/method
On the NVIDIA H100 SXM5 GPU, FP16 Tensor Core throughput reaches 989 TFLOPS, whereas the multi-function units executing special functions (such as $\exp$ for softmax) provide only 3.9 TFLOPS. In a standard single-warpgroup flow, the low-throughput softmax calculation creates an execution bottleneck.

Pingpong scheduling coordinates two consumer warpgroups within a CTA using barrier synchronization (bar.sync). The execution is scheduled such that while warpgroup 1 is computing its softmax reduction on the multi-function unit, warpgroup 2 executes its GEMMs ( $P V$ and the subsequent iteration's $Q K^T$ ) on the Tensor Cores. Once completed, the roles reverse: warpgroup 2 executes its softmax while warpgroup 1 performs GEMMs. On an H100 GPU running FP16 forward attention with head dimension 128 and sequence length 8192, pingpong scheduling increases throughput from 570 TFLOPS to 620–640 TFLOPS.

Knowl 7 — Numerical Error of FlashAttention-3 under Activation Outliers

data/table

To evaluate numerical precision under activation outliers characteristic of large language models, entries of input matrices $Q, K, V \in \mathbb{R}^{N \times d}$ are generated according to the mixture distribution: $X \sim \mathcal{N}(0, 1) + \mathcal{N}(0, 100) \cdot \text{Bernoulli}(0.001)$ Root mean squared error (RMSE) is measured against a 64-bit floating-point (FP64) reference implementation.

Precision	Method	Quantization / Rescaling	Incoherent Processing	RMSE
FP16	Baseline Attention	FP16 intermediate softmax	N/A	3.2e-4
FP16	FlashAttention-2	FP32 intermediate softmax	N/A	1.9e-4
FP16	FlashAttention-3	FP32 intermediate softmax	N/A	1.9e-4
FP8 (e4m3)	Baseline Attention	Per-tensor scaling, FP16 softmax	No	2.4e-2
FP8 (e4m3)	FlashAttention-3	Block quantization	Yes	9.1e-3
FP8 (e4m3)	FlashAttention-3 (ablation)	Per-tensor scaling	Yes	9.3e-3
FP8 (e4m3)	FlashAttention-3 (ablation)	Block quantization	No	2.4e-2

In FP16, FlashAttention-3 matches FlashAttention-2 and achieves $1.7\times$ lower RMSE than baseline attention by retaining softmax accumulations in FP32. In FP8, combining block quantization with incoherent processing reduces RMSE by $2.6\times$ compared to standard per-tensor quantized FP8 attention.

Knowl 8 — Speedup from Warp Specialization and Two-Stage Pipelining in FlashAttention-3

data/table

An ablation study isolating the effects of producer-consumer warp specialization and two-stage GEMM-softmax pipelining was performed on an NVIDIA H100 SXM5 GPU for non-causal FP16 forward attention with parameters: batch size $= 4$ , sequence length $= 8448$ , number of heads $= 16$ , and head dimension $= 128$ .

Warp specialization alone improves performance by hiding memory latencies via decoupled TMA producers, and two-stage pipelining provides an additional $13.6\%$ throughput improvement by hiding non-GEMM softmax operations behind asynchronous WGMMA executions.

Knowl 9 — Throughput Performance of FlashAttention-3 on Hopper GPUs
empirical result
Benchmarked on an NVIDIA H100 80GB SXM5 GPU (fixed 1830 MHz clock, total batch tokens $= 16\text{k}$ , hidden dimension $= 2048$ , head dimensions $d \in \{64, 128, 256\}$ , sequence lengths 512 to $16\text{k}$ ):
1. FP16 Forward Pass: FlashAttention-3 achieves 1.5–2.0 $\times$ speedup over FlashAttention-2 and up to 1.5 $\times$ over FlashAttention-2 in Triton. For $d=256$ , throughput reaches 756 TFLOPs/s (without causal mask) and 642 TFLOPs/s (with causal mask), representing up to $75\%$ of the theoretical hardware peak (989 TFLOPs/s).
2. FP16 Backward Pass: FlashAttention-3 achieves 1.5–1.75 $\times$ speedup over FlashAttention-2, reaching up to 474 TFLOPs/s ( $d=64$ ) and 561 TFLOPs/s ( $d=128$ ).
3. FP8 Forward Pass: FlashAttention-3 achieves nearly 1.2 PFLOPs/s (reaching 1171 TFLOPs/s for non-causal $d=256$ , and 1008 TFLOPs/s for non-causal $d=128$ ), approximately doubling the throughput of FP16 execution.
Knowl 10 — Overhead and Compiler Constraints in Three-Stage Attention Pipelining
limitation
Attempting to extend intra-warpgroup pipelining to 3 stages (simultaneously overlapping iteration $j+2$ first GEMM, iteration $j+1$ softmax, and iteration $j$ second GEMM) performs worse than 2-stage pipelining due to two bottlenecks:
1. Compiler Instruction Reordering: SASS analysis reveals that the NVCC compiler reorders instructions such that only the first WGMMA is overlapped with softmax, while the second WGMMA remains serialized.
2. Increased Register Pressure: Buffering additional intermediate states requires holding an extra $\tilde{P}_i$ and scaling parameters in registers, adding $B_r \times B_c \times \text{sizeof}(\text{input\_data\_type}) + B_r \times \text{sizeof}(\text{float})$ of register memory per threadblock. This forces kernels to use smaller tile block sizes ( $B_r, B_c$ ), degrading overall Tensor Core efficiency.
Knowl 11 — Current Hardware and Architectural Limitations of FlashAttention-3
limitation
FlashAttention-3 exhibits three main limitations:
1. Missing Persistent Scheduling in FP8: The FP8 implementation does not incorporate a persistent CTA kernel with dynamic load balancing (which is implemented for FP16). Consequently, FP8 FlashAttention-3 exhibits lower utilization on short sequence lengths and causal masking relative to vendor-tuned closed-source kernels (such as cuDNN).
2. Focus on Prefill/Training: The kernel optimizations target the compute-heavy prefill and training regimes; they are not optimized for the memory-bandwidth-bound single-token generation (decode) phase in LLM inference.
3. Unverified Large-Scale Training Stability: While block quantization and incoherent processing suppress outlier-induced RMSE in isolated attention tests, the stability and convergence characteristics of FP8 attention within end-to-end large-scale model pretraining remain unverified.

Coverage note — None was omitted; all primary architectural, algorithmic, low-precision, empirical, and ablation contributions of FlashAttention-3 are covered.

References

1.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.
2.AI21. Introducing jamba: Ai21’s groundbreaking ssm-transformer model. AI21 blog, 2024.
3.Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245, 2023.
4.Michael Bauer, Henry Cook, and Brucek Khailany. CudaDMA: Optimizing GPU Memory Bandwidth via Warp Specialization. In Proceedings of 2011 International Conference for High Performance Computing, Networking, Storage and Analysis, SC ’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.
5.Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. arXiv preprint arXiv:2405.04517, 2024.
6.Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.
7.Ganesh Bikshandi and Jay Shah. Delivering 1 PFLOP/s of Performance with FP8 FlashAttention-2, 2024. URL https://research.colfax-intl.com/adding-fp8-to-flashattention/.
8.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.
9.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.
10.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.
11.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.
12.Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.
13.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.
14.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.
15.Tri Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023. URL https://arxiv.org/abs/2307.08691.
16.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.
17.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.
18.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.
19.DeepSeek-AI. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.04434, 2024.
20.Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm. int8 (): 8-bit matrix multiplication for transformers at scale. CoRR abs/2208.07339, 2022.
21.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.
22.Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. 2023.
23.Anmol Gulati, James Qin, Chung-Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, et al. Conformer: Convolution-augmented transformer for speech recognition. arXiv preprint arXiv:2005.08100, 2020.
24.Mandy Guo, Joshua Ainslie, David Uthus, Santiago Ontanon, Jianmo Ni, Yun-Hsuan Sung, and Yinfei Yang. Longt5: Efficient text-to-text transformer for long sequences. arXiv preprint arXiv:2112.07916, 2021.
25.Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. Advances in Neural Information Processing Systems, 35:8633–8646, 2022.
26.Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length llm inference with kv cache quantization. arXiv preprint arXiv:2401.18079, 2024.
27.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.
28.Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In The International Conference on Machine Learning (ICML), 2020.
29.Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, pages 611–626, 2023.
30.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.
31.Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context. arXiv preprint arXiv:2310.01889, 2023.
32.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.
33.Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. Kivi: A tuning-free asymmetric 2bit quantization for kv cache. arXiv preprint arXiv:2402.02750, 2024.
34.Weile Luo, Ruibo Fan, Zeyu Li, Dayou Du, Qiang Wang, and Xiaowen Chu. Benchmarking and Dissecting the Nvidia Hopper GPU Architecture, 2024. URL https://arxiv.org/abs/2402.13499.
35.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.
36.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.
37.Paulius Micikevicius, Dusan Stosic, Neil Burgess, Marius Cornea, Pradeep Dubey, Richard Grisenthwaite, Sangwon Ha, Alexander Heinecke, Patrick Judd, John Kamalu, et al. Fp8 formats for deep learning. arXiv preprint arXiv:2209.05433, 2022.
38.NVIDIA. CUDA Programming Guide Version 12.4, 2024. URL https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html.
39.Nvidia. Accelerating transformers with nvidia cudnn 9. Nvidia blog, 2024. URL https://developer.nvidia.com/blog/accelerating-transformers-with-nvidia-cudnn-9/.
40.NVIDIA. Parallel Thread Execution ISA Version 8.4, 2024. URL https://docs.nvidia.com/cuda/pdf/ptx_isa_8.4.pdf.
41.Muhammad Osama, Duane Merrill, Cris Cecka, Michael Garland, and John D. Owens. Stream-k: Work-centric parallel decomposition for dense matrix-matrix multiplication on the gpu. In Proceedings of the 28th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming, PPoPP ’23, pages 429–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.
42.Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, et al. RWKV: Reinventing RNNs for the Transformer era. arXiv preprint arXiv:2305.13048, 2023.
43.Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.00071, 2023.
44.Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. In The International Conference on Learning Representations (ICLR), 2021.
45.Markus N Rabe and Charles Staats. Self-attention does not need $\mathcal{O}(n^2)$ memory. arXiv preprint arXiv:2112.05682, 2021.
46.Colfax Research. Tutorial: Matrix Transpose in CUTLASS, 2024. URL https://research.colfax-intl.com/tutorial-matrix-transpose-in-cutlass/.
47.Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing Transformers. arXiv preprint arXiv:2003.05997, 2020.
48.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.
49.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.
50.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.
51.Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150, 2019.
52.Benjamin Spector, Aaryan Singhal, Simran Arora, and Christopher Ré, 2024. URL https://github.com/HazyResearch/ThunderKittens.
53.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.
54.Mingjie Sun, Xinlei Chen, J Zico Kolter, and Zhuang Liu. Massive activations in large language models. arXiv preprint arXiv:2402.17762, 2024.
55.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.
56.Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. arXiv preprint arXiv:2009.06732, 2020.
57.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.
58.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.
59.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.
60.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.
61.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.
62.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.
63.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.
64.Zyphra. Zyphra unveils zamba: A compact 7b ssm hybrid model. Zyphra blog, 2024.

Citation

MLA

Shah, J., et al. “FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision”. arXiv, 2024, http://arxiv.org/abs/2407.08608v2.

APA

Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., & Dao, T. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv. http://arxiv.org/abs/2407.08608v2

Chicago

Shah, J., G. Bikshandi, Y. Zhang, V. Thakkar, P. Ramani, and T. Dao. 2024. “FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision”. arXiv. http://arxiv.org/abs/2407.08608v2.

Harvard

Shah, J. et al. (2024) “FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision”, arXiv [Preprint]. Available at: http://arxiv.org/abs/2407.08608v2.

Vancouver

1. Shah J, Bikshandi G, Zhang Y, Thakkar V, Ramani P, Dao T (2024) FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv

BibTeX

@article{shah2024flashattention,
  title = {FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision},
  author = {Shah, Jay and Bikshandi, Ganesh and Zhang, Ying and Thakkar, Vijay and Ramani, Pradeep and Dao, Tri},
  year = {2024},
  journal = {arXiv},
  url = {http://arxiv.org/abs/2407.08608v2},
  eprint = {2407.08608}
}

Metadata:arXiv

Source Code

This paper has an official code repository available. Click below to access the source code.

View Repository

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

Configuration	Execution Time (ms)	Throughput (TFLOPs/s)
FlashAttention-3 (Full)	3.538	661
No GEMM-Softmax Pipelining, with Warp-Specialization	4.021	582
GEMM-Softmax Pipelining, without Warp-Specialization	4.105	570

FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — FlashAttention-3 Intra-Warpgroup GEMM-Softmax Two-Stage Pipelining

Knowl 2 — FlashAttention-3 Forward Pass with Producer-Consumer Warp Specialization

Knowl 3 — FP8 Layout Transformation via In-Kernel Transposition and Register Swizzling

Knowl 4 — FP8 Attention Quantization via Block-Wise Scaling and Incoherent Processing

Knowl 5 — FlashAttention-3 Backward Pass with Warp-Specialized dQ Accumulation

Knowl 6 — Pingpong Scheduling for Inter-Warpgroup GEMM and Softmax Overlap