Flash Attention: Why Memory Traffic, Not Compute, Is the Bottleneck

Attention is often introduced as one short equation. On a long sequence, the equation creates a very large score matrix. An 8,192-token sequence has 8,192 rows and 8,192 columns of scores for each head. Writing that matrix out to GPU memory, then reading it back for softmax and values, can cost more time than the arithmetic itself.

Each query position can have a score against every key position.

The Problem Is Moving Intermediate Data

GPU memory, often called HBM, holds large tensors but is farther from the compute units. Fast on-chip memory is much smaller but can be read and written far more quickly. Standard attention tends to create an intermediate score matrix in HBM, then revisit it in later steps.

The cost is often traffic between memory and computation, not a lack of multiplication units.

For attention, the usual path is QK^T, then softmax, then multiply by V. The full matrix between those steps has shape T x T. It grows with the square of sequence length.

A Square Matrix Grows Fast

At 1,024 tokens, one attention matrix has about one million entries. At 8,192 tokens, it has about 67 million entries. The exact memory depends on data type and implementation, but the growth pattern is the important part: doubling length makes four times as many score entries.

Long contexts make the intermediate matrix much larger than the input token list.

This does not mean attention is wrong. It means the order of operations matters on real hardware.

For 8,192 positions, the score grid contains exactly $8192^2=67,108,864$ entries per head. If one stored entry uses two bytes, one such matrix occupies 128 MB. Thirty-two heads would require 4 GB for that intermediate alone if all were materialized this way.

That count is an illustrative layout, not the memory consumption of every framework call. Implementations can fuse operations, use different intermediate precision, or avoid storing the full grid. The point is to identify which intermediate is expensive: the temporary scores or probabilities, not the final width-$d_v$ output vector for each query.

Doubling sequence length doubles the number of query rows and the number of key columns, hence four times as many cells. The final output instead has $T{d}_v$ entries. Avoiding the large square intermediate is therefore different from avoiding the output.

Flash Attention Uses Tiles

Flash Attention divides queries, keys, and values into small blocks called tiles. A tile of queries and a tile of keys are brought into fast on-chip memory. Their partial scores are used immediately, then the temporary tile can disappear instead of being written as part of one giant score matrix.

Small blocks fit in fast memory while the full score and probability matrices are not written out to large GPU memory.

The algorithm visits enough key tiles for each query tile to include the same information as standard attention. It changes data movement, not which tokens can attend to which tokens.

Softmax Looks Like an Obstacle

Softmax normally needs all scores in a row so it can subtract the row maximum and divide by the sum of exponentials. At first glance, this seems to force us to save every score before we can finish the row.

Flash Attention uses an online softmax update. For each query row, it keeps a running maximum, a running normalizing sum, and a running output. When a new key tile arrives, it updates those running values carefully.

A row can be completed tile by tile without storing every earlier score.

This is exact attention in its mathematical definition, not a sparse or approximate replacement. Floating-point implementations can still differ slightly because the operation order changes. The running maximum rescales earlier contributions when a later tile contains a larger score.

For each query row, keep three running quantities: $m$, the largest score seen; $\ell$, the sum of exponentials relative to that maximum; and $u$, the value-weighted exponential sum. The normalized output is $u/\ell$.

When a new tile has a larger maximum $m^{\prime }$, old contributions must be multiplied by $e^{m-m^{\prime }}$. Their old exponentials used $m$ as a reference, while the new ones use $m^{\prime }$. Rescaling puts both groups on the same reference scale.

We must update the denominator and the weighted-value numerator together. Correcting only the sum of weights while leaving the output accumulator untouched would give the wrong answer. This running calculation is the central bridge between tiling and the softmax formula.

A Small Row Example

Imagine a query sees scores from two key tiles. The first tile has scores [2, 1]; the second has [3, 0]. After the second tile arrives, the row maximum changes from 2 to 3. The earlier exponential contributions are rescaled to use the new maximum before the final division.

Rescaling makes tile-by-tile softmax agree with processing the entire row together.

The value output is updated with the same corrected weights. We keep the final output row, not the entire score history.

Carry the Actual Values Through Both Tiles

Use already-scaled scores $[2,1,3,0]$ and one-number teaching values $[10,20,30,40]$. After tile 1, $m=2$, $\ell =1+e^{-1}\approx 1.367879$, and $u=10+20e^{-1}\approx 17.357589$.

Tile 2 changes the maximum to 3. Rescale the old sum by $e^{-1}$ and add the new contributions:

$ \ell' = e^{-1}(1+e^{-1})+1+e^{-3} $

$ u'=e^{-1}(10+20e^{-1})+30+40e^{-3} $

Dividing $u^{\prime }$ by ${\ell }^{\prime }$ produces the same mathematical weighted value as softmax over all four scores. We can test that equality in a few lines without implementing a GPU kernel.

import math

We first calculate the tiled result using the two accumulators. The shared scale factor is used in both updates.

normalizer = math.exp(-1) * (1 + math.exp(-1)) + 1 + math.exp(-3)
weighted_sum = math.exp(-1) * (10 + 20 * math.exp(-1)) + 30 + 40 * math.exp(-3)
tiled_output = weighted_sum / normalizer

Then compute the full-row reference. Its shifted exponentials use the final maximum, 3, from the beginning.

scores = [2., 1., 3., 0.]
values = [10., 20., 30., 40.]
exp_scores = [math.exp(score - max(scores)) for score in scores]
reference = sum(w * v for w, v in zip(exp_scores, values)) / sum(exp_scores)
assert math.isclose(tiled_output, reference, rel_tol=1e-12)

The assertion tests the online-softmax arithmetic. It does not test GPU performance, memory use, or a FlashAttention library installation.

Causal Masks Still Apply

Decoder attention must still hide future positions. Flash Attention applies the causal rule inside the relevant tiles. A query at position 3 can use columns 0 through 3 and must ignore columns 4 and later.

Tiles near the diagonal contain allowed and blocked regions just like the full matrix.

The algorithm is therefore compatible with the causal attention we studied in Day 16. It is a faster way to organize the same rule.

Why This Helps Prefill Most Clearly

Prefill processes many prompt positions and can need a large attention matrix. Avoiding HBM writes for that matrix often makes long-prompt processing faster and uses less temporary memory. Decode has a one-token query at each step, so its shape is different, though Flash Attention kernels can still be useful in several settings.

The biggest square-matrix pressure appears when many known prompt positions are processed together.

Flash Attention does not remove the KV cache or the sequential nature of output generation. It solves a different part of inference and training: unnecessary large intermediate storage.

FlashAttention does not make dense attention linear-time. Query-key comparisons still require work that grows roughly as $T^{2}d$ for a full sequence, where $d$ is the feature width involved. The saving is in how intermediate data is stored and moved. During training, implementations can also recompute some intermediates in the backward pass instead of saving a large matrix.

A bottleneck is the resource that limits completion time for a particular workload. Reducing memory traffic helps when that traffic is expensive enough to matter. It is not proof that compute never limits attention, or that every short sequence becomes faster. The FlashAttention paper studies this input/output-aware organization of attention.

We Usually Call It Through an API

In PyTorch, we rarely write tiled GPU code by hand. We give the framework queries, keys, values, and a causal setting. It can select an appropriate backend, such as a Flash Attention implementation, when hardware and tensor shapes allow it.

Recap

This is how Flash Attention reduces attention memory traffic. Standard attention can write a huge T x T score matrix to GPU memory. Flash Attention works in small query-key tiles, keeps softmax statistics while it goes, and avoids materializing the full intermediate matrix.

It gives the same attention result while moving less temporary data. Day 25 shows the PyTorch SDPA interface that lets our code request attention without tying it to one kernel.