TL;DR: FlashAttention computes dense attention in tiles while retaining a running maximum, denominator and weighted output. It avoids writing the full N×N scores and probabilities to HBM. When a later tile raises the maximum, it rescales both accumulated sums. This saves intermediate storage and traffic, but dense attention still has quadratic arithmetic and K/V tiles can be read repeatedly.
How to approach it
Write scaled attention and identify the large intermediate. Follow one query through two key/value tiles, including the rescale. Then separate avoided intermediate traffic from unavoidable input/output traffic and actual measured runtime. Finish with masking, numerical checks and backward recomputation.
A strong answer
IO-aware means the algorithm is designed around data movement to and from HBM, the GPU's large but comparatively slow main memory. Doubling context length makes an explicit attention score tensor four times larger. For one head at sequence length 8192, that score matrix has 67,108,864 elements. At two bytes each it occupies 128 MiB, even though an individual score row is only 16 KiB. The matrix, not one row, is the large allocation to eliminate.
Keep only the state needed for the final weighted average
For query q and keys K, scores are q @ K.T / sqrt(d) plus the intended mask or bias. Softmax turns those scores into weights; multiplying by V forms the output. FlashAttention evaluates blocks of these operations together. A score tile can be consumed and discarded without storing the full matrix.
The running state has three parts: the largest score seen so far, m; the sum of exponentials relative to that maximum, l; and the weighted value sum on the same scale, o. A new tile may contain a larger score, so first choose m_new = max(m, max(s)). The earlier sums are now expressed on the wrong scale. Multiply both by alpha = exp(m - m_new) before adding the new tile's contributions: l = alpha*l + sum(exp(s-m_new)) and o = alpha*o + exp(s-m_new) @ V_tile. After the final tile, divide o by l to obtain the output.
The toy scores in the diagram are already scaled. The first tile gives m=0, l=1, o=10. The second raises m to ln 2, so alpha=1/2, l=1.5 and o=25. The result is 16⅔, matching the full weights 1 and 2 applied to values 10 and 20.
Executable recurrence, with an empty-state policy
This CPU reference accepts precomputed scaled/masked scores for one query and finite value vectors. It illustrates the recurrence, not the GPU kernel's tiling or performance. An all-masked or empty row returns zeros by explicit policy. It skips all-masked chunks so the first update never evaluates negative infinity minus negative infinity.
import numpy as np
def online_row(scores, values, tile=2):
scores = np.asarray(scores, dtype=np.float64)
values = np.asarray(values, dtype=np.float64)
if scores.ndim != 1 or values.ndim != 2 or len(scores) != len(values):
raise ValueError("one score per value row required")
if type(tile) is not int or tile <= 0:
raise ValueError("tile must be a positive integer")
if np.any(np.isnan(scores) | np.isposinf(scores)) or not np.all(np.isfinite(values)):
raise ValueError("finite values; scores finite or negative infinity")
maximum, denominator = -np.inf, 0.0
numerator = np.zeros(values.shape[1])
for start in range(0, len(scores), tile):
s = scores[start:start+tile]
chunk_max = np.max(s)
if np.isneginf(chunk_max):
continue
new_max = max(maximum, chunk_max)
alpha = 0.0 if np.isneginf(maximum) else np.exp(maximum-new_max)
weights = np.exp(s-new_max)
numerator = alpha*numerator + weights @ values[start:start+tile]
denominator = alpha*denominator + weights.sum()
maximum = new_max
return numerator/denominator if denominator else numerator
Tests compare it with direct stable softmax over random finite scores, extreme finite offsets, masks and several tile sizes. The GPU implementation still needs its own forward/backward numerical tests, race checks and target-hardware timing. For causal attention, the mask uses absolute query and key positions across tile boundaries; each tile's local indices are insufficient.
What the byte arithmetic does and does not say
Take one non-causal head with N=8192, d=128 and two-byte Q, K, V and output. Each of those four tensors is 2 MiB, so their one-read/write minimum totals 8 MiB. A deliberately materializing baseline that writes scores, reads them for softmax, writes probabilities and reads probabilities for the value product adds four 128 MiB transfers. Its logical total is 520 MiB, or 545,259,520 bytes, before other passes and scratch.
One head: N = 8,192; d = 128; tensor element size = 2 B
Four score/probability transfers = 4 × N² × 2 = 536,870,912 B
Minimum Q/K/V/output traffic = 4 × N × d × 2 = 8,388,608 B
Materializing total = 545,259,520 B = 520 MiB
Illustrative bandwidth floor = total / (3.35 × 10^12 B/s)
= 0.000163 s = 0.163 ms
Check: 520 MiB is 65 times the 8 MiB input/output minimum.
This is a lower-bound calculation, not a measurement of either implementation or a 65× speedup. Its two matrix products require about 4N²d = 34.36 GFLOPs; whether memory or compute wins depends on achieved resource rates and the actual implementation.
FlashAttention removes the full-matrix stores and loads. It does not guarantee that each K/V element is fetched from HBM only once: different query tiles may revisit it, and cache behavior affects where those reads land. Treat 8 MiB as an input/output lower bound, not a FlashAttention traffic prediction. Memory-Bound vs Compute-Bound Kernels explains why a faster implementation can expose a different bottleneck without automatically becoming compute-bound for every shape.
Exact attention, finite precision and backward
“Exact” here means the dense attention formula is preserved rather than approximated by dropping key interactions. Floating-point operation order and intermediate precision still change numerical results. Compare against a higher-precision reference with an explicit tolerance, including masks and awkward dimensions. Kernel Fusion and Numerics: FP32, BF16, FP8 and FP4 explain the storage and accuracy tradeoffs.
Backward can recompute score/probability tiles using saved normalization information rather than retaining N×N probabilities from forward. That saves memory at the cost of recomputation. The total training-memory budget still includes model weights, gradients, optimizer state, other activations and any required RNG state; removing this intermediate does not prove that a proposed context length fits.
Choose the fused backend when its mask, dtype, head dimension, layout and differentiation support match the workload and numerical checks pass. Verify the backend that actually ran; an API call alone does not prove FlashAttention was selected. Reverse the choice for unsupported semantics or a shape where the fallback is faster. Use Nsight Systems and allocator statistics to examine peak allocated memory, kernel launches and forward-plus-backward time separately from a serving KV-cache budget.
What interviewers probe next
- Why rescale the numerator? Its previous weighted values were expressed relative to the old maximum too.
- Is attention now linear time? No. Avoiding N×N storage does not remove dense N² interactions.
- Can a fully masked row be ignored? Only after defining the output policy and implementing it safely.
Common mistakes
- Confusing a 16 KiB row with the 128 MiB score matrix in this example.
- Calling minimum input/output bytes the measured FlashAttention HBM traffic.
- Forgetting the scale, global causal positions, rescale of o, or backward memory obligations.
Key takeaways
- Running state replaces the full score/probability matrices; both l and o rescale when m rises.
- The illustrative materializing baseline moves 520 MiB, versus an 8 MiB input/output lower bound, not a guaranteed 65× speedup.
- Verify numerical behavior and the backend actually executed before making performance claims.
The original FlashAttention paper establishes the IO-aware algorithm; FlashAttention-2 describes the unnormalized-output recurrence used to explain the calculation here.