wikillm

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

  1. 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.
  2. The producer waits on a buffer slot only after the first s iterations have filled the ring §3.1.
  3. Registers are reallocated per role (dealloc on the producer, alloc on the consumers) with setmaxnreg §2.2 §3.1.
  4. 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.
  5. 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.
Pseudocode
# 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, lse

The 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)

  1. 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.
  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.
  3. 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.
Pseudocode
# 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, lse

The 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

  1. 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.
  2. 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.
  3. 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.
  4. 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).
Pseudocode
# 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_scales

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

NameValueWhereNotes
Benchmark GPUH100 80 GB SXM5, 700 Wapp. C.1
Clockfixed at 1830 MHzapp. C.1matches the 989 TFLOPS FP16 theoretical maximum
Repeats100, averagedapp. C.1
Library versionsCUDA 12.3, cuDNN 9.1.1.17, CUTLASS 3.5, FlashAttention 2.5.8, Triton 3.0.0 nightly, PyTorch 2.3.0app. C.1
Sequence lengths (FP16)512, 1k, ..., 16k§4.1
Tokens per batch16k (batch size set to fit)§4.1
Hidden size2048§4.132, 16 or 8 heads for head dim 64, 128, 256
Head dimension64, 128, 256§4.1
Causal maskingbenchmarked with and without§4.1causal FLOPs are halved
FP8 sequence lengths512, 1024, 2048, 4224, 8448, 16896app. C.2lengths at or above 4k divisible by 132 SMs
Ablation configbatch 4, seqlen 8448, 16 heads, head dim 128, non-causal FP16§4.2
FP8 formate4m3§3.3
Block sizes B_r, B_cnot stated§3.1code: see Paper vs code
Pipeline stages snot stated§3.1code: 2 for SM90 (see Paper vs code)
Overlap depth2-stage§3.23-stage was slower app. B.3
Random mixing matrixrandom +1/-1 diagonal times Hadamard§3.3
Numerical test distributionN(0,1) + N(0,100) x Bernoulli(0.001)§4.30.1% of entries are outliers
Warp register splitnot stated§3.1code: 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.

SettingMetricBaselineThis paperSource
FP16 forwardspeedupFlashAttention-2 (CUDA)1.5-2.0x§4.1 fig. 5
FP16 forwardthroughputH100 theoretical FP16 peak, 989 TFLOPSup to 740 TFLOPs/s (75% of peak)§4.1 §1 app. C.1
FP16 forwardspeedupFlashAttention-2 in Triton1.5x§4.1
FP16 forwardspeedupPyTorch standard attention3-16x§4.1
FP16 backwardspeedupFlashAttention-21.5-1.75x§4.1 fig. 6
FP16 forward, 1k+ sequencesvs cuDNNcuDNN (closed source)faster (stated in text)§4.1
FP8 forwardthroughputnot stated in textclose to 1.2 PFLOPs/s§1 fig. 7
FP8 forwardvs cuDNN FP8cuDNNahead at head dim 64; on par at 128 and 256 without causal mask; behind with causal mask§1 footnote
Pipelining ablationTFLOPs/sno overlap, with warp-specialisation661 full vs 582tab. 2
Pipelining ablationTFLOPs/soverlap, no warp-specialisation570tab. 2

Ablation tab. 2:

ConfigurationTimeTFLOPs/s
FlashAttention-33.538 ms661
No GEMM-softmax pipelining, warp-specialisation4.021 ms582
GEMM-softmax pipelining, no warp-specialisation4.105 ms570

Numerical error, RMSE against an FP64 reference, synthetic inputs with 0.1% outliers §4.3 tab. 3:

MethodBaselineFlashAttention-2 / FA3Source
FP163.2e-41.9e-4 (both)tab. 3
MethodBaseline FP8FA3 FP8FA3 FP8, no block quantFA3 FP8, no incoherent processingSource
RMSE2.4e-29.1e-39.3e-32.4e-2tab. 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_V is 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 - 1 never 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.