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:

flowchart LR
    QK["Q and K"] --> GEMM1["GEMM: QK^T"]
    GEMM1 --> S["Write N x N scores to HBM"]
    S --> SM["Read scores and apply softmax"]
    SM --> P["Write N x N probabilities to HBM"]
    P --> GEMM2["Read probabilities and multiply by V"]
    V["V"] --> GEMM2
    GEMM2 --> O["Output O"]

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.

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:

flowchart LR
    Q["Q tile"] --> SCORE["Score tile Q_i K_j^T"]
    KV["K/V tile"] --> SCORE
    SCORE --> ONLINE["Online-softmax update"]
    ONLINE --> OUT["Running output O_i"]
    OUT --> SCORE

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.

5. Concrete GPU Implementation

5.1 Kernel structure

A simplified forward kernel has the following structure:

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. Excessive per-block resource use reduces occupancy, leaving fewer independent thread blocks available to hide latency. A production kernel therefore selects tile dimensions based on head dimension, data type, causal masking, and target architecture.

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.

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.

The algorithm has subsequently been rescheduled for newer GPU architectures without changing this principle. Those implementation changes are discussed separately in The Evolution of FlashAttention: From Ampere to Blackwell.

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.

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.