Illustrative tensor reads and writes for one attention head.
Attention memory path
The canvas tracks one query tile against key/value tiles. Colored blocks are resident in fast on-chip memory; dim cells remain in HBM.
Ready: global score matrix is not resident.
Memory needed to materialize all pairwise attention scores.
Query-key tile interactions for the selected block size.
Standard softmax versus streaming online-softmax reference.
Exact attention, reorganized around the memory hierarchy.
The quadratic matrix is a data-movement problem.
Standard attention can write the N by N score matrix to HBM, read it for softmax, and read it again to multiply by V. Arithmetic is not the only cost: repeatedly moving that matrix can dominate runtime.
S = QK^T / sqrt(d), P = softmax(S), O = PVTiles stay close to the compute units.
FlashAttention loads blocks of Q, K, and V into SRAM, updates partial outputs, and writes only the final output back. Tile sizes are constrained by fast-memory capacity.
HBM -> SRAM -> tensor cores -> HBMOnline softmax preserves exactness.
A running maximum and normalization sum let each row absorb another key block without storing every score. Rescaling keeps earlier contributions consistent when a larger maximum arrives.
Smaller tiles fit more hardware.
Reducing tile size lowers SRAM demand but increases the number of block interactions. Larger tiles reduce scheduling overhead until they exceed the selected budget.
Longer context raises the stakes.
The materialized score matrix grows with N squared. Q, K, V, and O grow linearly in N, so avoiding the intermediate matrix matters increasingly at long sequence lengths.
Read each kernel phase.
Load Q
A query tile remains resident while key and value tiles stream past it.
q_block = Q[i : i + Bq]Score a K tile
Compute a local score block. Causal masks can be applied before the softmax update.
scores = q_block @ k_block.T / sqrt(d)Update online softmax
Merge the block maximum and normalization with prior blocks, rescaling the partial output exactly.
m_new = max(m_old, rowmax(scores))Accumulate V
Add the probability-weighted value tile, then write the completed output tile once.
o += exp(scores - m_new) @ v_block