Paper · 2024-07-11 · cs.LG, cs.AI
FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
Exact attention kernel for Hopper (H100) GPUs that overlaps TMA loads, WGMMA matrix multiplies and softmax across warpgroups, with an FP8 path using block quantization. Author-reported: FP16 forward 1.5-2.0x faster than FlashAttention-2, up to 740 TFLOPs/s; FP8 near 1.2 PFLOPs/s; 2.6x lower FP8 error.
Self-attentionFused attention kernelLow-precision arithmeticWarp specializationattentiongpu-kernelslow-precisionhopperlong-context
arXivPDFCodeMarkdown15 min read
AI-compiled by claude-haiku-5-5 on 2026-10-18, not yet reviewed by a human. Every claim cites its source; check it when a detail matters.
In plain words
Attention is the part of a language model that lets every token look at every other token. Doing this naively means building a huge table of scores, which is slow to write to and read back from the graphics card's main memory. An earlier method, FlashAttention, avoids that table by working through the input in small pieces that fit in the chip's fast local memory.
This paper is about making that trick fit the newest NVIDIA chips, the Hopper generation. Those chips can copy data, multiply matrices and do the softmax (the step that turns scores into weights) at the same time, but only if the program is arranged so that these jobs run side by side. The authors split the work between groups of threads so one group loads data while another multiplies and a third does the softmax. They also rearrange the softmax and the matrix multiplies so they overlap.
They then add a low-precision mode that uses 8-bit numbers, which the chip runs about twice as fast. Such numbers are less accurate, so the authors use two tricks to keep the error down: scaling each small block of numbers separately, and mixing the values with a random rotation so that a few very large values do not dominate the scale.
TL;DR
- Exact attention kernel for NVIDIA Hopper (H100): a warp-specialised producer/consumer design plus an overlap of the softmax with the matrix multiplies §3.1 §3.2.
- FP8 variant: a transpose of V inside the kernel, an accumulator-to-operand register permutation, per-block quantization and an orthogonal (Hadamard-based) mixing step §3.3.
- Author-reported FP16 forward speedup over FlashAttention-2 is 1.5-2.0x, reaching 740 TFLOPs/s; FP8 reaches close to 1.2 PFLOPs/s §1 §4.1.
- Relevant if you write or tune attention kernels on Hopper; the FP8 accuracy technique is not in the released code at the pinned commit (see Paper vs code).
Problem
Standard attention materialises the N x N score and probability matrices in high-bandwidth memory, so its cost grows quadratically with sequence length §2.1 §2.3. FlashAttention (Dao et al., 2022, cited in the paper) fuses the operations into one kernel so these matrices never leave on-chip memory §2.3. FlashAttention-2 parallelises the work further but reaches only about 35% utilisation on the H100, against 80-90% for optimised matrix-multiply kernels §1. The paper attributes this partly to FlashAttention-2 not using Hopper-specific instructions, and more fundamentally to a synchronous design that makes no use of asynchronous data movement or low-precision units §1.
Two hardware features are targeted: asynchrony (the Tensor Cores and the Tensor Memory Accelerator, TMA, run separately from the general-purpose cores) and FP8 tensor cores, which double matrix-multiply throughput over FP16/BF16 §2.2. The FP8 path has its own difficulties: the (e4m3) format has few mantissa bits, and outlier values in LLM activations make per-tensor scaling inaccurate §3.3.
Method
Overview
The paper gives the algorithm for the forward pass in detail; the backward pass uses the same warp-specialisation pattern with an extra role that accumulates the query gradient into global memory §B.1 alg. 3. The forward pass is parallel over batch, heads and query blocks, so the description is given per CTA for one query block Q_i §3.1.
Notation: B_r is the query block size, B_c the key/value block size, T_c = ceil(N / B_c) the number of key blocks, d the head dimension §3.1. The paper does not state values for B_r, B_c or the pipeline depth s; see Hyperparameters.
3.1 Producer-consumer asynchrony and pingpong scheduling
- Warps are split into a producer warpgroup, which only issues TMA loads of Q_i and of the K_j, V_j blocks into a circular shared-memory buffer with s stages, and consumer warpgroups, which only compute §3.1 alg. 1.
- The producer waits on a buffer slot only after the first s iterations have filled the ring §3.1.
- Registers are reallocated per role (dealloc on the producer, alloc on the
consumers) with
setmaxnreg§2.2 §3.1. - Consumers run two asynchronous GEMMs per key block: S = Q K_j^T with both operands in shared memory (SS), and O += P V_j with P taken from registers (RS) alg. 1 §3.1.
- Pingpong scheduling: with two consumer warpgroups, barriers force the GEMMs of one warpgroup to be issued while the other runs its softmax. The paper reports this moves FP16 forward with head dimension 128 and sequence length 8192 from 570 TFLOPs/s to 620-640 TFLOPs/s, and notes the real schedule is less clean than the figure §3.1 fig. 1.
# FlashAttention-3 CTA, warp-specialised, no GEMM-softmax overlap [alg. 1]
# Input: query block (B_r, d); keys, values (N, d) in HBM
# Output: output block (B_r, d); logsumexp (B_r,)
def producer_warpgroup(query_block, keys, values):
# Free registers, then load the query block once [alg. 1]
release_registers()
load_tma(query_block)
for j in range(T_c):
# Wait until ring slot j % s has been consumed [alg. 1]
wait_slot_free(j % s)
# Stream key and value block j into that slot [alg. 1]
load_tma(keys[j], values[j], slot=j % s)
def consumer_warpgroup(query_block):
# Take registers back and zero the running state [alg. 1]
acquire_registers()
output = zeros(B_r, d)
row_max = full(B_r, -inf)
row_sum = zeros(B_r)
for j in range(T_c):
# GEMM 1 (SS): scores for key block j [alg. 1]
scores = query_block @ keys[j].T # (B_r, B_c)
# Update the running max and rescale factor [alg. 1]
old_max = row_max
row_max = maximum(old_max, rowmax(scores))
scale = exp(old_max - row_max) # (B_r,)
probs = exp(scores - row_max) # (B_r, B_c)
# Update the running sum [alg. 1]
row_sum = scale * row_sum + rowsum(probs)
# GEMM 2 (RS): add probs times values; rescale the old part [alg. 1]
output = scale[:, None] * output + probs @ values[j]
# Free ring slot j % s for the producer [alg. 1]
release_slot(j % s)
# Normalise and emit output and logsumexp [alg. 1]
output = output / row_sum[:, None]
lse = row_max + log(row_sum)
return output, lseThe scale line is written here as the factor exp(old - new). The paper's
listing prints the output update with the inverse of that factor; see Pitfalls.
3.2 Overlapping GEMMs and softmax inside a warpgroup (2-stage)
- The dependency chain in the loop above (GEMM 1, then softmax, then GEMM 2) serialises the units. The fix is to keep one extra score block in registers, so the softmax of block j runs while the GEMM for block j+1 is in flight §3.2 alg. 2 fig. 2.
- The extra block costs B_r x B_c x 4 bytes of registers per thread block; the paper flags this as a conflict with large tiles §3.2.
- A 3-stage variant (overlapping a third GEMM) was implemented and performed worse, because the compiler did not overlap the second GEMM and register pressure forced smaller tiles app. B.3 alg. 4 fig. 8.
# Consumer warpgroup, FP16, 2-stage GEMM-softmax pipeline [alg. 2]
# Input: query block (B_r, d); keys, values (N, d) in shared memory
# Output: output block (B_r, d); logsumexp (B_r,)
def softmax_block(scores, row_max, row_sum):
# Max, rescale factor, probabilities and sum for one block [alg. 1]
new_max = maximum(row_max, rowmax(scores))
scale = exp(row_max - new_max) # (B_r,)
probs = exp(scores - new_max) # (B_r, B_c)
new_sum = scale * row_sum + rowsum(probs)
return new_max, new_sum, probs, scale
def consumer_forward(query_block, keys, values):
# Zero output, max at -inf, sum at zero [alg. 2]
output = zeros(B_r, d)
row_max = full(B_r, -inf)
row_sum = zeros(B_r)
# Prologue: score block 0, nothing to overlap yet [alg. 2]
scores_cur = query_block @ keys[0].T
row_max, row_sum, probs_cur, _ = softmax_block(
scores_cur, row_max, row_sum)
# Steady state: score GEMM j overlaps PV GEMM j-1 and softmax [alg. 2]
for j in range(1, T_c):
# Issue score GEMM for block j, asynchronously [alg. 2]
scores_next = query_block @ keys[j].T
# Issue PV GEMM for the previous block, asynchronously [alg. 2]
output = output + probs_cur @ values[j - 1]
# Softmax of block j while both GEMMs run [alg. 2]
row_max, row_sum, probs_next, scale = softmax_block(
scores_next, row_max, row_sum)
# Wait for the PV GEMM, then rescale it to the new max [alg. 2]
output = scale[:, None] * output
scores_cur = scores_next
probs_cur = probs_next
# Epilogue: last PV GEMM, then normalise and emit logsumexp [alg. 2]
output = output + probs_cur @ values[T_c - 1]
output = output / row_sum[:, None]
lse = row_max + log(row_sum)
return output, lseThe loop bound here is range(1, T_c), which runs the score GEMM for every
key block after the prologue. The paper's listing writes the bound as
1 <= j < T_c - 1, which would leave the last key block unscored; see Pitfalls.
3.3 FP8 forward pass
- Layout of V. FP8 WGMMA accepts only K-major operands for both inputs of the second GEMM, so V tiles in shared memory must be contiguous along the sequence dimension, while Q, K and V arrive contiguous in the head dimension §3.3. Options: transpose V in global memory beforehand (rejected: too costly when memory-bound) or transpose each V tile in shared memory inside the kernel. The paper chooses the in-kernel transpose, using ldmatrix and stmatrix in the producer warpgroup, arranged to overlap with the two WGMMAs of the previous step §3.3 fig. 3 fig. 4.
- Register permutation. The FP32 accumulator layout of an FP8 WGMMA differs from the layout expected for operand A. The paper permutes the register order of the first GEMM's accumulator (to d0 d1 d4 d5 d2 d3 d6 d7, repeated every 8 bytes) with byte-permute instructions, and writes the matching row permutation into the transposed V tile §3.3 fig. 3 fig. 4.
- Block quantization. One scale per block: Q and K are split into B_r x d and B_c x d blocks, V likewise, and each block is quantised with its own scale. The scale can be fused with the preceding rotary embedding at no extra cost, since that step is memory-bound, and each block of S is rescaled to account for the quantization at no extra computation cost §3.3.
- Incoherent processing. Q and K are multiplied by an orthogonal matrix M before quantization. Because M M^T = I, (Q M)(K M)^T = Q K^T and the output is unchanged, while outliers are spread across entries. M is chosen as a product of a random diagonal matrix with entries +1 or -1 and a Hadamard matrix, which costs O(d log d) and can be fused with the rotary embedding §3.3 (following QuIP and QuIP#, cited there).
# FP8 preprocessing of queries and keys [§3.3]
# Input: Q, K (N, d) in FP16/BF16
# Output: quantised blocks and per-block scales
def prepare_fp8_qk(queries, keys, signs, hadamard):
# Mix the head dimension with a random orthogonal matrix [§3.3]
mixed_q = queries @ (diag(signs) @ hadamard) # (N, d)
mixed_k = keys @ (diag(signs) @ hadamard) # (N, d)
# One scale per block, then cast to e4m3 [§3.3]
q_blocks, q_scales = quantize_blocks(mixed_q, B_r)
k_blocks, k_scales = quantize_blocks(mixed_k, B_c)
return q_blocks, q_scales, k_blocks, k_scalessigns is a random vector of +1/-1 values and hadamard a Hadamard matrix
of size d; both are the paper's choice §3.3. quantize_blocks is not
specified by the paper beyond "one scalar per block" §3.3.
Hyperparameters
Memory levels of the H100 SXM5, as given in the paper's hardware table tab. 1:
- Off-chip global memory (GMEM, HBM): 80 GiB, about 3.35 TB/s; the grid runs here.
- L2 cache, shared by a GPC of SMs: 50 MiB, about 12 TB/s; threadblock clusters run here.
- Shared memory (SMEM) of each SM: 228 KiB per SM; the threadblock (CTA) runs here.
- Registers (RMEM) of each SM: 256 KiB per SM; individual threads hold their own registers.
Settings and tuning values:
| Name | Value | Where | Notes |
|---|---|---|---|
| Benchmark GPU | H100 80 GB SXM5, 700 W | app. C.1 | |
| Clock | fixed at 1830 MHz | app. C.1 | matches the 989 TFLOPS FP16 theoretical maximum |
| Repeats | 100, averaged | app. C.1 | |
| Library versions | CUDA 12.3, cuDNN 9.1.1.17, CUTLASS 3.5, FlashAttention 2.5.8, Triton 3.0.0 nightly, PyTorch 2.3.0 | app. C.1 | |
| Sequence lengths (FP16) | 512, 1k, ..., 16k | §4.1 | |
| Tokens per batch | 16k (batch size set to fit) | §4.1 | |
| Hidden size | 2048 | §4.1 | 32, 16 or 8 heads for head dim 64, 128, 256 |
| Head dimension | 64, 128, 256 | §4.1 | |
| Causal masking | benchmarked with and without | §4.1 | causal FLOPs are halved |
| FP8 sequence lengths | 512, 1024, 2048, 4224, 8448, 16896 | app. C.2 | lengths at or above 4k divisible by 132 SMs |
| Ablation config | batch 4, seqlen 8448, 16 heads, head dim 128, non-causal FP16 | §4.2 | |
| FP8 format | e4m3 | §3.3 | |
| Block sizes B_r, B_c | not stated | §3.1 | code: see Paper vs code |
| Pipeline stages s | not stated | §3.1 | code: 2 for SM90 (see Paper vs code) |
| Overlap depth | 2-stage | §3.2 | 3-stage was slower app. B.3 |
| Random mixing matrix | random +1/-1 diagonal times Hadamard | §3.3 | |
| Numerical test distribution | N(0,1) + N(0,100) x Bernoulli(0.001) | §4.3 | 0.1% of entries are outliers |
| Warp register split | not stated | §3.1 | code: see Paper vs code |
FLOP accounting used for all throughput numbers: forward is 4 x seqlen^2 x d x heads, halved for causal masking; backward is 2.5 times forward §4.1.
Results
Author-reported. Setup: kernel benchmarks on one H100 SXM5 with the settings above, FP16 or FP8 inputs, against FlashAttention-2 (CUDA and Triton), PyTorch standard attention and cuDNN §4.1 app. C.1.
| Setting | Metric | Baseline | This paper | Source |
|---|---|---|---|---|
| FP16 forward | speedup | FlashAttention-2 (CUDA) | 1.5-2.0x | §4.1 fig. 5 |
| FP16 forward | throughput | H100 theoretical FP16 peak, 989 TFLOPS | up to 740 TFLOPs/s (75% of peak) | §4.1 §1 app. C.1 |
| FP16 forward | speedup | FlashAttention-2 in Triton | 1.5x | §4.1 |
| FP16 forward | speedup | PyTorch standard attention | 3-16x | §4.1 |
| FP16 backward | speedup | FlashAttention-2 | 1.5-1.75x | §4.1 fig. 6 |
| FP16 forward, 1k+ sequences | vs cuDNN | cuDNN (closed source) | faster (stated in text) | §4.1 |
| FP8 forward | throughput | not stated in text | close to 1.2 PFLOPs/s | §1 fig. 7 |
| FP8 forward | vs cuDNN FP8 | cuDNN | ahead at head dim 64; on par at 128 and 256 without causal mask; behind with causal mask | §1 footnote |
| Pipelining ablation | TFLOPs/s | no overlap, with warp-specialisation | 661 full vs 582 | tab. 2 |
| Pipelining ablation | TFLOPs/s | overlap, no warp-specialisation | 570 | tab. 2 |
Ablation tab. 2:
| Configuration | Time | TFLOPs/s |
|---|---|---|
| FlashAttention-3 | 3.538 ms | 661 |
| No GEMM-softmax pipelining, warp-specialisation | 4.021 ms | 582 |
| GEMM-softmax pipelining, no warp-specialisation | 4.105 ms | 570 |
Numerical error, RMSE against an FP64 reference, synthetic inputs with 0.1% outliers §4.3 tab. 3:
| Method | Baseline | FlashAttention-2 / FA3 | Source |
|---|---|---|---|
| FP16 | 3.2e-4 | 1.9e-4 (both) | tab. 3 |
| Method | Baseline FP8 | FA3 FP8 | FA3 FP8, no block quant | FA3 FP8, no incoherent processing | Source |
|---|---|---|---|---|---|
| RMSE | 2.4e-2 | 9.1e-3 | 9.3e-3 | 2.4e-2 | tab. 3 |
The FP8 baseline uses per-tensor scaling, FP32 matmul accumulation and FP16 softmax intermediates §4.3. The paper's headline FP8 claim of 2.6x lower error is the ratio of the baseline and FA3 FP8 entries §4.3 §1.
Paper vs code
Code: https://github.com/Dao-AILab/flash-attention, BSD-3-Clause, commit
965c9212e6be33c821ee9f2b461564c5c91a2503 (committed 2026-10-16). The paper's
footnote links this repository §1. The Hopper kernels are in hopper/. The
commit is newer than the paper, so the code may differ from the version the
paper evaluated; every difference below is about the pinned commit.
- Incoherent processing is not in the code. The Hadamard and random-sign mixing of Q and K §3.3 does not appear anywhere in the repository at this commit (no occurrence of "hadamard" or "incoherent" in any file). The FP8 accuracy claim therefore cannot be reproduced from the code.
- Scale granularity differs. The paper uses one scale per B_r x d or B_c x d block §3.3. The mainloop reads descale values indexed by batch and KV head only mainloop_fwd_sm90_tma_gmma_ws.hpp and applies the V descale in the finalize step mainloop_fwd_sm90_tma_gmma_ws.hpp. I found no per-block scale in the mainloop.
- Block sizes are not in the paper but are in the code. For FP16 with head dim 128 the code uses a 128-row query block and a key block of 176 (non-causal) or 128 (causal or local) tile_size.h. This is a hand-written table keyed by head dim, dtype and causal/local flags.
- Overlap flag. The 2-stage overlap (IntraWGOverlap) is set to true in the tile-size table for almost all FP16 head dims, with exceptions for some head value dimensions tile_size.h. For head dim 128 FP16 it is set tile_size.h.
- Pingpong. The paper's barrier schedule corresponds to the scheduler barriers in the mainloop mainloop_fwd_sm90_tma_gmma_ws.hpp, enabled by a rule that depends on the overlap flag, the warpgroup count and the head dimension (and differs for FP8, which uses it at head dim 128 and above) mainloop_fwd_sm90_tma_gmma_ws.hpp. The code comment says this only affects speed, not correctness.
- Pipeline depth. The paper's listing uses an s-stage ring buffer with s unstated. The launch template sets 2 stages for SM90 flash_fwd_launch_template.h.
- Register split. The paper says only that registers are reallocated. The kernel uses different counts by number of consumer warpgroups: 24 or 40 for the producer and 240 or 232 for the consumers with two warpgroups, 56 and 256 with one, 32 and 160 with three flash_fwd_kernel_sm90.h.
- Transpose of V. The paper's FP8 path transposes V inside the kernel §3.3.
The code does this by default (
Transpose_Vis true for FP8 unless V is already column-major) mainloop_fwd_sm90_tma_gmma_ws.hpp. The FP8 benchmark script transposes V on the host before the call benchmark_flash_attention_fp8.py. - Benchmark scales. The FP8 benchmark passes descale values of 1.0 benchmark_flash_attention_fp8.py. Its timings are valid, but it says nothing about accuracy.
- FLOP formula. The benchmark FLOP count matches the paper's formula and the backward factor of 2.5 benchmark_attn.py benchmark_attn.py.
No discrepancies found in the FP16 forward structure (warp-specialised producer, SS and RS GEMMs, 2-stage overlap).
Pitfalls
- Output rescale sign in Algorithm 1. The listing updates the output as the inverse of exp(m_old - m_new) while the sum uses exp(m_old - m_new) directly. These only agree if the output is kept normalised, but the final step divides by the sum. The code multiplies the output by exp(old max - new max) softmax.h softmax.h. This page uses that form.
- Loop bound in Algorithm 2. The listing's loop
1 <= j < T_c - 1never scores the last key block, and its epilogue uses an undefined P_last. The page uses one more iteration (range(1, T_c)), which is the only reading that makes the epilogue consistent. Check against the code before implementing. - Listing is idealised. The paper says NVCC reorders instructions. In the 2-stage SASS the first GEMM interleaves with the softmax, but the second is not overlapped app. B.2. The overlap is a schedule, not a guarantee.
- Register cost of overlap. The extra score block must stay in registers. Larger tiles or deeper pipelines compete for the same registers, and the paper decides tile size by profiling §3.2.
- First FP8 tile is not hidden. The paper hides the transpose of the next V tile behind two WGMMAs only from the second iteration on §3.3. The first tile's transpose is not covered by that schedule.
- Accuracy claims use synthetic data. The error test uses N(0,1) plus rare large outliers, not real LLM activations §4.3. The 2.6x figure depends on that distribution.
- FP8 is weaker in some settings. The paper says FP8 is behind cuDNN with causal masks and loses to cuDNN on small sequences, and that FP8 has no persistent kernel in the paper's build (footnote in the Discussion) §1 §5.
- Speed numbers are from figures for FP8. The FP8 TFLOPs values beyond the headline 1.2 PFLOPs/s are only in figures fig. 7 fig. 9, so the text does not give exact per-configuration numbers.
- Pingpong numbers are for a different configuration. The 570 to 620-640 TFLOPs/s improvement is for head dim 128 and sequence length 8192 §3.1. The ablation table uses sequence length 8448 with 16 heads and batch 4 tab. 2 §4.2. Do not compare the two.
- FLOP conventions. Throughput counts the causal half and the backward factor of 2.5 §4.1. Compare only numbers computed with the same convention.
- Lowest-level behaviour is not specified. The paper gives no numbers for stage count, block sizes or register split (see Hyperparameters). Take them from the code, and say so when you do.
Lineage
- Builds on: FlashAttention (Dao et al., 2022; no arXiv id given in the paper's bibliography) §1 §2.3; FlashAttention-2 (arXiv:2307.08691) §1, which it restructures for asynchrony and FP8.
- Uses ideas from: incoherent processing from QuIP and QuIP# (cited in §3.3); the idea of the in-kernel FP8 transpose is credited to the cuDNN team in the acknowledgments; CUTLASS primitives are used §4 [acknowledgments].
- Compared against: FlashAttention-2 (CUDA and Triton), PyTorch standard attention, and cuDNN attention (closed source) §4.1.
- Related but different: multi-query and grouped-query attention, which change K and V sharing and not the attention computation itself app. A.
- Superseded by: not known from content already in this wiki.
Source and license
Adapted from “FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision” by Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao (arXiv:2407.08608), used under CC BY 4.0. Changes: condensed and restructured for implementation, with tables, equations and algorithms taken from the paper and its code, each cited to its source. Not affiliated with or endorsed by the authors. Code citations refer to https://github.com/Dao-AILab/flash-attention at commit 965c9212e6be (BSD-3-Clause). Any quoted code is shown under that license.