IO-aware exact attention
Trace one exact tiled-attention iteration and account for the HBM tensors that it avoids
Before you begin
Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.
- Online softmax
- Tiled GEMM
- Memory hierarchy
A conventional attention implementation for batch 1, 32 heads, and sequence length 8192 can create one score value for each of roughly 2.1 billion query-key pairs. An FP16 score matrix of that size would occupy about 4 GiB before counting probabilities, gradients, or other tensors. A framework memory report shows the allocation, while Nsight Systems shows separate matrix multiplication, softmax, and value-combination launches with high-bandwidth-memory transfers between them.
FlashAttention computes the same exact attention result without storing the complete score or probability matrix. A CUDA thread block retains a tile of query rows and their running softmax state in on-chip storage, streams key and value tiles through that state, and discards each temporary score tile after use. Peak-memory measurements and profiler byte counters provide the empirical signature: the quadratic intermediate allocation disappears and transfers between high-bandwidth memory and on-chip memory decrease.
Standard attention is written as softmax(QKᵀ)V. A literal operator sequence writes the score matrix, reads it for softmax, writes probabilities, and reads them for the value product. At long sequence length these S×S intermediates dominate memory traffic and capacity even though the final output is only S×D. FlashAttention reorganizes the exact computation around the memory hierarchy.
The main benefit comes from lower input/output complexity. The tiled algorithm performs the same attention arithmetic, with possible recomputation. The algorithm keeps query blocks and online-softmax state on chip while streaming key and value blocks. The kernel consumes scores without storing a complete tensor in high-bandwidth memory.
The term exact refers to the attention equation, not to bitwise equality with every implementation. Tiling and floating-point reduction order can change rounding. FlashAttention differs from approximate linear-attention methods because it still computes the complete selected set of query-key interactions and applies the ordinary softmax semantics.
Traffic in the materialized implementation
Let Q, K, and V each have sequence length N and head dimension d. The score matrix S = QKᵀ has N² elements. A separated implementation writes S to HBM, reads it for softmax, writes the probability matrix P, and reads P for the product PV. At long sequence length, these N² transfers and allocations can dominate the N×d inputs and output.
Kernel fusion alone is not enough if the fused kernel still requires the complete score matrix on chip, because the matrix is much larger than registers or shared memory. The algorithm must partition the calculation and retain only information needed to combine partitions. Online softmax supplies that information for each query row.
One tiled forward-pass iteration
Select a block of query rows Qi and initialize its row maxima, normalization sums, and output numerators. Load a K block Kj and V block Vj. Compute the local score block QiKjᵀ, apply the scale and mask, and find each row's local maximum and sum. Merge those statistics with the running state and multiply the local exponential weights by Vj.
The local score block is no longer needed after it updates the state, so it can be discarded. Continue with the next K/V block. When all relevant key positions have been processed, divide each output numerator by its normalization sum and store the output block. At no point must the complete N×N score or probability matrix exist in HBM.
Tile order and on-chip capacity
The tile order determines which values are reused and which values are loaded again. Keeping a larger query block resident increases output accumulator and normalization state. Keeping larger K and V blocks resident consumes shared memory but can reduce the number of passes over Q or output state. The FlashAttention paper analyzes this tradeoff in terms of HBM accesses as a function of available SRAM.
Causal attention visits only key positions at or before each query position. Entire blocks above the causal diagonal can be skipped. Blocks intersecting the diagonal require element-level masking. Variable-length or packed sequences add bounds and indexing. Masking and sequence bounds affect useful work and control flow, so they belong in both the correctness tests and the performance shape definition.
Backward recomputation
A conventional autograd path might save P because gradients use the softmax probabilities. Saving P requires O(N²) storage. The memory-efficient backward instead saves compact per-row normalization information and the ordinary inputs and output. The backward kernel reconstructs score and probability blocks as needed while accumulating dQ, dK, and dV.
The reconstruction method performs additional matrix and exponential arithmetic, but avoids writing and reading a much larger saved tensor. The trade can reduce both peak memory and wall-clock time on hardware where the removed high-bandwidth-memory traffic costs more than the recomputation. The result is a measured trade, not a general rule that additional arithmetic is faster.
Why implementations form a family
Head dimension changes the register and matrix-instruction shape. Sequence length changes the number of tiles and the balance between setup and steady state. Datatype changes storage, throughput, and error. Dropout, bias, causal masks, grouped-query structure, and variable-length packing each modify the calculation or data access.
Hardware generations also change shared-memory capacity, asynchronous movement, matrix instructions, and warp-group organization. Preserve the algorithmic invariants—exact selected interactions, stable online normalization, and bounded on-chip score tiles—while selecting an implementation for the declared inputs and target architecture.
Worked example: Comparing intermediate traffic
Consider one attention head with sequence length 8192 in BF16. The conceptual score matrix contains about 67 million elements.
- Materializing scores alone requires roughly 134 MB; writing and rereading probabilities adds another similar-scale round trip, before accounting for Q, K, V, and output.
- The score tensor must also be read by softmax, so score materialization contributes at least one 134 MB write and one 134 MB read. Perform the same accounting for the probability tensor.
- A tiled exact algorithm holds only score blocks and running statistics on chip. HBM traffic scales with Q, K, V, output, and repeated tile reads rather than full stored S×S intermediates.
- Build a byte model for a candidate tile order, including how often K and V blocks are reread for different query tiles. Compare with the naive materialization floor.
- Measure sequence-length scaling and peak memory. Verify output and gradients against a reference over causal, noncausal, and ragged cases.
- Repeat the comparison at short sequence lengths. The tiled implementation can lose when setup, masking, or a poorly fitted tile costs more than the eliminated intermediate traffic.
Conclusion: The algorithm is faster because it changes the movement schedule of exact attention. Calling it an approximation or merely a fused softmax misses the essential mechanism.
Common errors
- Claiming linear attention complexity. FlashAttention remains exact quadratic arithmetic in sequence length while reducing memory IO and materialization.
- Counting only eliminated score bytes. Tile order may reread K and V; the complete IO model determines the gain.
- Using one schedule for all inputs. Decode, small heads, unusual masks, and new architectures can require different schedules.
- Equating exact attention with bitwise equality. A different floating-point reduction order can change low-order bits while preserving the specified numerical accuracy.
Reference summary
Conventional attention can materialize an O(sequence²) score or probability matrix in high-bandwidth memory. FlashAttention tiles queries, keys, and values so score blocks are formed on chip, incorporated into online softmax state, used to update output accumulators, and discarded. FlashAttention computes exact attention while reducing transfers between high-bandwidth memory and on-chip memory. Recomputation in backward can be cheaper than saving the large intermediate.
- Outer tiling chooses which Q rows keep output and normalization state resident.
- K/V tiles stream through and update the resident state.
- Causal and other masks constrain which tile pairs are valid.
- Backward reconstructs needed probabilities from compact saved statistics and inputs.
- Sequence length, head dimension, datatype, mask, and architecture influence the best schedule.
Knowledge check
How can performing extra arithmetic in attention backward reduce total wall-clock time?
Show the answer guide
- A conventional attention backward pass can read an N-by-N probability or score matrix that the forward pass stored in HBM. The matrix grows quadratically with sequence length and can dominate memory traffic and capacity.
- A tiled backward pass can retain only compact row statistics from forward and recompute score and probability tiles from Q and K when each tile is needed. The recomputation adds matrix arithmetic and exponentials but avoids writing and later reading the quadratic intermediate.
- Modern GPUs can execute the added tensor-core arithmetic faster than they can transfer the eliminated bytes when the original pass is bandwidth-bound. Smaller saved state can also prevent out-of-memory failures or allow a larger batch.
- Recomputation reduces wall-clock time only when the saved HBM and allocation costs exceed the added arithmetic, synchronization, and any loss of kernel efficiency. The profiler must confirm that trade for the target shape.
Measure stored-state and recomputed attention backward
Run exact scaled dot-product attention backward for batch 2, 16 heads, head dimension 64, and sequence lengths 512, 1024, 2048, and 4096. Compare a path that materializes the score or probability matrix with a memory-efficient path that recomputes tiles. Use identical inputs and loss gradients, and verify dQ, dK, and dV against a reference.
Evidence to keep: Submit the benchmark program, gradient-correctness tests with tolerances, a table of backward latency and peak allocated memory by sequence length, and profiler reports that compare HBM traffic and arithmetic work for one long sequence. State the measured sequence length at which recomputation pays for itself.
The same attention equation appears under distinct training and serving conditions. Training, prefill, and decode change tile shapes, cache access, and which dimension offers reuse.
Sources and further reading
The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.