FlashAttention: Exact Attention Without Materialization

Abstract

Attention is commonly expressed as two matrix multiplications separated by a softmax. On modern GPUs, however, this operator is often constrained not by arithmetic throughput but by transfers between high-bandwidth memory (HBM) and on-chip SRAM. A conventional implementation materializes two $N \times N$ intermediate matrices: the attention scores and their normalized probabilities. FlashAttention avoids materializing either matrix in HBM. It computes score tiles on chip, consumes them immediately through an online softmax, and discards them. This article explains that central design independently of any particular FlashAttention version or GPU generation.

1. Introduction

Transformer attention computes

where $Q$, $K$, and $V \in \mathbb{R}^{N \times d}$, $N$ is the sequence length, and $d$ is the head dimension. This formulation suggests a straightforward implementation: compute the score matrix $S=QK^\top$, apply row-wise softmax to obtain $P$, and compute $O=PV$.

The straightforward implementation is inefficient because both $S$ and $P$ contain $N^2$ elements. Each intermediate is written to HBM and read by a subsequent kernel. GPU tensor cores can compute the matrix products rapidly, but the intermediate traffic repeatedly crosses a much slower memory interface.

FlashAttention is based on one systems observation:

The score matrix is required by the computation, but it does not need to exist in HBM.

The kernel computes one score tile in on-chip SRAM, immediately incorporates it into a running softmax and output, and then discards it. This reordering computes the same attention function with substantially less memory traffic.

The main contributions of the design are:

1. An IO-aware tiling of the complete attention operator, rather than either matrix multiplication in isolation.
2. An online-softmax recurrence that enables exact block-wise normalization.
3. A fused GPU kernel that retains score and probability tiles on chip.
4. A backward pass that recomputes inexpensive intermediates instead of storing quadratic state.

2. Background

2.1 GPU memory hierarchy

GPUs expose a large HBM capacity and a much smaller on-chip storage hierarchy composed of registers and shared memory. HBM holds model state and activations but has much lower bandwidth relative to the arithmetic rate of tensor cores. Shared memory is fast enough to feed matrix-multiply units, but only a small tile of the working set fits at one time.

This asymmetry makes data movement a first-class algorithmic cost. An implementation with the minimum number of floating-point operations can still underperform if it repeatedly materializes large intermediates in HBM.

2.2 Conventional attention

A conventional implementation launches separate kernels for three stages:

The inputs and output require $\Theta(Nd)$ storage, whereas each intermediate requires $\Theta(N^2)$ storage. For long sequences, the quadratic traffic dominates.

For example, with $N=16{,}384$ and $d=128$, each of $Q$, $K$, and $V$ contains approximately 2.1 million elements. The score matrix contains approximately 268 million elements. The principal problem is not that these scores must be computed; it is that a conventional schedule writes and rereads them.

The consequence is visible without any timing at all. Running the same attention over the same tensors, once with the two intermediates materialized and once through a fused kernel, and recording peak allocated memory:

Two things about this measurement are worth noting. First, it is a memory result, not a speed result: it holds regardless of how well the fused kernel is tuned, because the intermediates simply do not exist. Second, the ratio is not a constant — it is a slope difference. Measuring only at short sequence lengths would show 1.2× and support the wrong conclusion.

2.3 Numerically stable softmax

For one score row $x$, stable softmax is

Subtracting the row maximum prevents exponential overflow. The apparent obstacle to tiling is that both the maximum and normalization sum depend on the complete row.

3. Motivation

Tiling either matrix multiplication alone does not solve the problem. A tuned GEMM can efficiently produce $S$, but if $S$ is subsequently written to HBM, consumed by a softmax kernel, and written again as $P$, the implementation retains the dominant traffic.

The optimization boundary must therefore include the full operator:

This produces two requirements:

• The score and probability tiles must remain on chip and never become global intermediates.
• Softmax statistics computed from separate key blocks must be merged exactly.

The second requirement is the key algorithmic problem solved by FlashAttention.

4. Design

4.1 Tiled execution

Partition the query matrix into row blocks $Q_i$ and the key and value matrices into blocks $K_j$ and $V_j$. For each query block, the kernel scans all key/value blocks:

1. Load $Q_i$, $K_j$, and $V_j$ into on-chip storage.
2. Compute the score tile $S_{ij}=Q_iK_j^\top$.
3. Merge $S_{ij}$ into the running softmax statistics for $Q_i$.
4. Accumulate the weighted values into the running output $O_i$.
5. Discard $S_{ij}$ and proceed to the next key/value block.

Only the final $O_i$ is written to HBM.

4.2 Online softmax

For one query row, maintain three quantities after processing some key blocks:

• $m$: the largest score observed so far;
• $\ell$: the sum of exponentials relative to $m$;
• $z$: the unnormalized weighted-value sum relative to $m$.

For a new score block $s$ with corresponding value block $V_j$, compute

The previous state was expressed relative to $m$, whereas the merged state must be expressed relative to $m'$. Rescale the old contribution and add the new block:

After all blocks have been processed,

This recurrence is exact in real arithmetic. Every old exponential is converted from the old reference maximum to the new one because

4.3 Worked example

Consider one score row split into two blocks:

For the first block, $m=2$. Its unnormalized weights are $[e^{-1},1]$. The second block introduces a new maximum $m'=3$. The old weights are multiplied by $e^{2-3}=e^{-1}$, becoming $[e^{-2},e^{-1}]$. The new block contributes $[e^{-3},1]$.

The final normalization is therefore

which is exactly the result obtained by applying ordinary softmax to all four scores simultaneously. The algorithm has changed the evaluation order, not the function.

4.4 IO complexity

Let $M$ denote the available on-chip storage. FlashAttention chooses tile dimensions so that the active Q, K, V, score, and output tiles fit within $M$. The original analysis derives an HBM-access complexity of

compared with the $\Theta(N^2)$ intermediate traffic of conventional attention. The original paper further shows that this bound is asymptotically optimal for a range of practical SRAM capacities. The practical consequence is more important than the notation: increasing useful on-chip tile capacity directly reduces repeated HBM transfers.

4.5 Backward pass

Saving the full probability matrix for backpropagation would reintroduce quadratic memory consumption. FlashAttention instead stores the output and one log-sum-exp statistic per row. During the backward pass, it reloads Q, K, and V tiles and recomputes the local score and probability tiles on chip.

This is a deliberate trade: perform additional arithmetic to avoid substantially more HBM traffic. On modern GPUs, recomputation is often cheaper than storing and retrieving the quadratic intermediates.

Measured at $N=8{,}192$ with $B=1$, $H=8$, $d=128$ in BF16, the trade resolves clearly in favor of recomputation. Peak allocated memory for forward-plus-backward, against a conventional implementation:

Recomputation roughly doubles the backward cost in time relative to forward — a ratio of 3.5 rather than 1.8 — while keeping the footprint flat. Against the conventional path the fused backward is 19× smaller in memory. The extra arithmetic is real and cheap; the storage it avoids is not.

One detail from that measurement matters for anyone reproducing it: on a small batch, the host-side launch path can cost more than the backward pass itself. At $N=8{,}192$ the fastest kernel spent 0.365 ms on the GPU but 0.686 ms enqueueing the work, so wall-clock timing reported 0.950 ms — more than twice its device time. A backward-pass benchmark that reports only wall clock is measuring the launcher.

5. Concrete GPU Implementation

5.1 Kernel structure

A simplified forward kernel has the following structure. This is the version-1/version-2 structure: one thread block owns a query tile, and the online-softmax state is scalar per row. Later versions change the scheduling around this loop — which warp does what, when copies are issued, where accumulators live — but not the loop itself. Those changes are the subject of FlashAttention in Practice.

parallel for each query tile Q_i:
    load Q_i
    m = -infinity
    l = 0
    z = 0

    for each key/value tile (K_j, V_j):
        load K_j and V_j into shared memory
        S_ij = Q_i @ transpose(K_j)

        new_m = max(m, rowmax(S_ij))
        old_scale = exp(m - new_m)
        P_ij = exp(S_ij - new_m)

        l = old_scale * l + rowsum(P_ij)
        z = old_scale * z + P_ij @ V_j
        m = new_m

    O_i = z / l
    store O_i

One thread block typically owns one or more query tiles. Warps cooperate on tensor-core matrix multiplications, row-wise maximum and sum reductions, and movement between HBM, shared memory, and registers. Causal attention skips score tiles that lie entirely above the diagonal and masks only tiles intersecting the diagonal.

5.2 Fusion boundary

The implementation must fuse score computation, masking, softmax, dropout when required, and value accumulation into one kernel. Writing S_ij or P_ij to global memory between these stages defeats the design. The fusion boundary, rather than a novel matrix-multiplication primitive, is what removes the quadratic intermediate traffic.

5.3 Resource tradeoffs

Larger tiles improve data reuse but consume more shared memory and registers, and excessive per-block resource use reduces occupancy. How a given implementation resolves that tradeoff depends on the target architecture, which is the subject of the next article.

One consequence of that tradeoff is measurable, and it is the most commonly over-generalised number in this literature: how much fusion helps depends on the head dimension. At a fixed sequence length of 8,192 and a fixed FLOP count per query, varying $d_h$ alone changes the speedup by an order of magnitude:

The pattern is not what one would guess. Fused kernels get slower as heads get deeper — SDPA flash degrades 4.7× from $d_h=32$ to $d_h=256$ — because a deeper head spends more of the tile budget on depth instead of on length, which reduces reuse. The vendor kernel is the exception: nearly flat, within 25% across the whole range.

This is worth carrying into any comparison. A model with $d_h=128$ and a model with $d_h=32$ are not the same benchmark, and a fused kernel that wins comfortably on one may barely win on the other.

6. Evidence From the Original Implementation

The NeurIPS 2022 paper reports up to 7.6x kernel-level speedup over a PyTorch attention implementation. At the model level, it reports 15% end-to-end wall-clock improvement for BERT-large, approximately 3x speedup for GPT-2 at sequence length 1K, and approximately 2.4x speedup on Long Range Arena workloads [1]. These measurements are workload-specific, but they establish the practical consequence of the design: avoiding quadratic HBM intermediates improves both memory capacity and execution time.

6.1 Reproducing the shape of the result on current hardware

The absolute numbers above belong to the hardware of 2022. What transfers is the shape: speedup that grows with sequence length, and a memory footprint that stops being quadratic. Both reproduce on a current GPU, and reporting the shape is more useful than quoting an absolute figure measured on hardware nobody has.

Measured on a B300 (SM 10.3), $B=2$, $H=8$, $d_h=128$, BF16, causal:

Two features of this table are the same features as the paper’s: the fused path is nearly flat in latency while the conventional path is quadratic, so the speedup is a function of sequence length rather than a constant. And below roughly 2,048 tokens every fused backend lands on a launch-overhead floor of about 0.03 ms, which is why the short-sequence ratios say more about kernel launch than about attention.

A note on what “fused” means here. On this GPU generation, the fastest fused kernel is the vendor library’s, not the flash-attn package: the latter ships kernels compiled for SM80 and runs them through forward compatibility, which works but is roughly 6× slower at these shapes. Any claim about “the FlashAttention kernel” on a new architecture has to start by establishing which kernel actually executed.

7. Discussion

FlashAttention demonstrates three broader systems principles.

First, operators should be optimized across their full dataflow. Optimizing two GEMMs independently misses the traffic created by the nonlinear operation between them.

Second, algebra can expose a better execution schedule. Online softmax transforms a global normalization into a mergeable state, enabling the surrounding system optimization.

Third, recomputation can be preferable to storage. The backward pass performs more arithmetic but moves less data, a profitable exchange on machines whose compute throughput grows faster than memory bandwidth.

FlashAttention does not reduce the asymptotic arithmetic complexity of dense attention: it still performs $\Theta(N^2d)$ work. Its contribution is to make that arithmetic execute closer to the capabilities of the hardware by reducing avoidable IO.

7.1 A kernel speedup is not a step speedup

Every number in this article so far measures attention in isolation. In a training step, attention is one term among several, and the step can only improve by its share. Measured on a small model (8 layers, $D=512$, $D_{ff}=2048$, batch 8, BF16):

The first row is the trap: at a short sequence length attention is 3.5% of the step, so a “faster” attention kernel is invisible, and the measured attention speedup is below one because the comparison sits on the launch floor. Nothing was made worse; the measurement simply has no signal.

The middle row is ordinary Amdahl behaviour: a 4.2× improvement to a 10.7% component yields about 1.6× for the step.

The last row is the interesting one. The step improves by 10× — 154,000 to 1,544,000 tokens per second — which is more than attention’s measured share would predict. The reason is that the conventional path’s quadratic intermediates also inflate allocator pressure and the backward pass, so its cost spills outside the attention kernels the profiler attributes. Kernel-level attribution therefore tends to understate the gain from removing those intermediates, just as step-level timing understates the kernel improvement when the step is not attention-bound. Both numbers are worth reporting, and neither alone is the answer.

7.2 The energy bill is not the same as the time bill

There is one more cost that a latency table hides, and on shared hardware it is a real one. Measured at a sequence length of 16,384, with idle draw subtracted:

The fastest kernel draws the most power — 823 W above idle against 511 W for the conventional path — and is still 23× more efficient per joule. That looks paradoxical until the reason is stated: the conventional path leaves the chip idle waiting on memory, so it burns power producing nothing. Being fast and being efficient coincide here, because both follow from not doing redundant memory traffic.

This matters for a practical reason beyond accounting. When several processes share accelerators, the fastest kernel is also the one that draws the most instantaneous power for the shortest interval, and the slowest is the one that holds elevated power for thirty times as long. Neither situation is obviously better for a co-tenant, which is why a latency-only comparison is incomplete on shared hardware.

This article has treated the algorithm as a single design. In practice it has been rescheduled four times for successive GPU generations, with the non-materializing core unchanged each time. That sequence — and the parallel decomposition used for single-query decoding — is the subject of FlashAttention in Practice.

8. Limitations

• Reported throughput values are best-case measurements for specific shapes and GPU configurations; end-to-end model gains are smaller and workload dependent.
• Algorithmic exactness does not imply bitwise identity because tiled reductions change floating-point operation order.
• Very short sequences may not amortize fusion and scheduling overheads.
• Dense FlashAttention does not remove the quadratic arithmetic cost required by very long contexts.
• Hardware-specific kernels require continued retuning as memory systems and matrix units evolve.

The second point is usually stated as a caveat and then left unquantified, which invites the assumption that the tiled path is the less accurate one. Measured against a float64 reference at a sequence length of 8,192, the opposite is true:

Every fused implementation tested reported the identical maximum error, because they accumulate the softmax and the value product in FP32 regardless of storage precision, and share the same reduction order. The conventional path accumulates in the storage dtype and loses precision doing so.

Error also does not accumulate with sequence length the way long-sum intuition suggests. At FP32, the maximum absolute error stays between 8e-07 and 1.7e-06 from a sequence length of 512 through 8,192 — flat, not linear in $N$ — because each tile is normalized against a running maximum that only moves forward. And every backend tested was bitwise reproducible across identical calls at every shape, so “exact” here means exact enough that repeated runs agree, not merely close on average.

References

1. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022.