Flash Attention: IO-Aware Exact Attention

Keep the attention, change its execution

Scaled dot-product attention computes a weighted sum of values. A straightforward implementation stores every query-key score, then every attention probability, before producing the output. Those intermediate matrices grow quadratically with sequence length.

FlashAttention computes the same attention function while avoiding those full matrices in GPU memory. It streams small tiles through fast on-chip storage and carries a few statistics per query row. Exact refers to the mathematics: different reduction orders and floating-point precision can produce small numerical differences.

Follow one query through both paths

Choose a query, then step through the same fixed keys and values in both lanes. The dense lane retains scores until it can normalize the whole row. The tiled lane updates a running maximum, denominator and weighted-value numerator after each tile. At the end, their outputs agree within numerical tolerance.

This is a deterministic full-attention illustration with one head, six query positions and no mask or dropout. Its numbers come from the displayed vectors, not model activations. The HBM disclosure counts the whole example under an explicit query-resident schedule; the main walkthrough shows only your selected row. Browser arrays do not demonstrate actual SRAM residence or GPU throughput.

Why intermediate matrices cost memory traffic

For one head, let Q and K have shape N × d and V have shape N × d_v:

A dense materialized path performs three operations:

1. Read Q and K, compute S, then write the N × N scores to HBM.
2. Read S, apply row-wise softmax, then write the N × N probabilities P.
3. Read P and V, compute O, then write the output.

Just writing and reading S and P transfers 4N² elements. Input and output transfers are additional, and tiled matrix multiplies may reload inputs. The explorer includes those reloads rather than assuming Q, K and V are each read exactly once.

GPU HBM holds large arrays. Registers and on-chip SRAM hold smaller working sets with lower access cost. Kernel fusion consumes intermediates before they leave the chip; tiling makes the working set fit. Whether a workload is limited by memory traffic or arithmetic depends on its shapes, kernel and hardware. The data movement analysis in transformers gives the broader context.

Online softmax makes tiling exact

Softmax couples all keys in a row. Averaging separately normalized tile outputs loses that coupling. Instead, carry three quantities after the keys seen so far:

Here m and ℓ are scalars for each query, and a is a d_v-dimensional unnormalized numerator. Start with m = −∞, ℓ = 0 and a = 0. For a new tile of scores t and values V_tile:

When a later tile raises the maximum, α changes the scale of both old sums. The explorer's Q0 raises the maximum in each tile, and its values differ across keys, so rescaling the numerator matters. Before any keys are seen, α is defined as zero. The partial ratio a/ℓ attends only to processed keys; it becomes the full output after every tile has contributed.

The original paper's Algorithm 1 carries a normalized O and rescales it at each update. Substituting a = ℓO gives the equivalent numerator above, with one final division. FlashAttention-2 uses an unnormalized output accumulator to reduce non-matmul work.

Loop order changes which data is reloaded

Original FlashAttention Algorithm 1 keeps a K/V block on-chip and visits all query blocks inside that loop. It reloads Q and the current output statistics for each K/V block, then writes updated row state. It does not read every input just once.

The explorer and CPU example below use the alternative Q-block outer, K/V-block inner forward order described in FlashAttention-2's parallelism section. Q and its accumulator remain resident while all K/V tiles pass through. Each new query block reloads K and V. This helps explain the recurrence; it is not the complete FlashAttention-2 GPU implementation or a reconstruction of the original paper's traffic count.

In the explorer, N = 6, d = d_v = 2 and each query block has two rows. There are three query blocks. Both declared paths read 12 Q elements, 36 K elements and 36 V elements, then write 12 output elements. Dense additionally writes and reads 36 S and 36 P elements: 240 transfers versus 96 for tiled. These analytical totals assume the working set fits on-chip and exclude caches, register spills, backward statistics and backward computation. Changing the number of keys per tile changes the walkthrough, but not these totals under the declared schedule.

Check the recurrence on CPU

This executable NumPy example uses the explorer's vectors. It compares the online output with an independently computed dense softmax for key tiles of sizes 1, 2, 3, 4 and 7, including nondividing and oversized tiles. It performs full unmasked forward attention in float64. It checks algebra, not a GPU kernel, timing, memory residence or traffic.

# CPU forward algebra for the fixed teaching example.
import numpy as np


def tiled_attention(q, k, v, query_block=2, key_block=2):
    """Compute full unmasked attention with query-resident row state."""
    output = np.empty((len(q), v.shape[1]), dtype=np.float64)
    for row in range(0, len(q), query_block):
        query = q[row:row + query_block]
        maximum = np.full(len(query), -np.inf)
        denominator = np.zeros(len(query))
        numerator = np.zeros((len(query), v.shape[1]))
        for column in range(0, len(k), key_block):
            scores = query @ k[column:column + key_block].T
            scores /= np.sqrt(q.shape[1])
            next_maximum = np.maximum(maximum, scores.max(axis=1))
            correction = np.exp(maximum - next_maximum)
            probabilities = np.exp(scores - next_maximum[:, None])
            denominator = correction * denominator + probabilities.sum(axis=1)
            numerator = correction[:, None] * numerator
            numerator += probabilities @ v[column:column + key_block]
            maximum = next_maximum
        output[row:row + len(query)] = numerator / denominator[:, None]
    return output


q = np.array([[1, 0], [0, 1], [1, 1], [-1, 0], [0, -1], [1, -1]], dtype=np.float64)
k = np.array([[-1, 0], [0, 1], [2, 0], [1, 1], [3, -1], [-1, 2]], dtype=np.float64)
v = np.array([[1, 0], [0, 2], [-1, 1], [2, -1], [0.5, 3], [-2, 0.5]], dtype=np.float64)
scores = q @ k.T / np.sqrt(q.shape[1])
weights = np.exp(scores - scores.max(axis=1, keepdims=True))
weights /= weights.sum(axis=1, keepdims=True)
dense = weights @ v
for key_block in (1, 2, 3, 4, 7):
    tiled = tiled_attention(q, k, v, query_block=2, key_block=key_block)
    np.testing.assert_allclose(tiled, dense, rtol=1e-12, atol=1e-12)
print("Q0 output:", dense[0])

For the original K/V-outer traversal, tail blocks, causal masks and backward-related row statistics, see the annotated FlashAttention paper implementation and its CPU reference.

Storage, arithmetic and IO are separate quantities

For fixed head dimensions, a materialized dense implementation uses quadratic global intermediate storage. FlashAttention needs the inputs and output, linear per-row statistics, and a bounded on-chip tile working set:

The second expression excludes the SRAM working set. Both methods still evaluate all query-key pairs for full attention; the forward matmul work is O(N^2(d+d_v)). A mask that skips blocks changes the attention pattern; it is a separate choice from computing full attention more efficiently.

For the original paper's square setup with d_v = d, Theorem 2 gives:

M is SRAM capacity in elements, and the theorem assumes d \leq M \leq Nd. Linear global storage does not mean linear total traffic: the original algorithm revisits query data as K/V blocks change. The explorer's finite transfer ledger illustrates its own declared schedule, not this asymptotic theorem.

During training, FlashAttention also avoids saving the full attention matrix for backward. It saves output and normalization statistics, then recomputes score/probability tiles when needed. This trades extra arithmetic for less memory traffic. The CPU code above covers forward only.

What changes in FlashAttention-2

FlashAttention-2 keeps exact attention and online normalization. It reduces non-matmul operations, parallelizes independent query blocks, and partitions work between warps to reduce shared-memory communication. It does not turn full attention's pairwise work into O(N log N), nor provide one speedup factor for every workload.

When selecting an implementation, compare equivalent masks, scaling and dropout, then benchmark the shapes and device used by the application. Avoid a universal sequence-length cutoff: kernel support and performance depend on the implementation. For decode-time history storage, KV caching addresses a different cost; sliding-window attention changes which historical keys can participate.

References

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, Dao et al. (2022). Algorithm 1 and Theorem 2 ground the recurrence and IO analysis.
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, Dao (2023). Algorithm, query-block parallelism and warp work partitioning.

Related concepts