The mechanism
Standard attention can create a score matrix with one entry for each query–key pair. Materializing that matrix in HBM becomes expensive as sequence length grows. FlashAttention reorganizes the calculation into tiles, keeping useful intermediates in faster local memory and maintaining the normalization information needed to combine blocks. The algorithm still computes exact attention in its stated numerical setting; it is different from deleting attention connections or approximating the model. Backward computation can recompute some intermediates rather than storing them all. The benefit depends on the memory hierarchy, tile shape, sequence length and implementation, not simply on changing the name of the attention kernel.
- Partition query, key and value matrices into tiles that fit the chosen local-memory budget.
- Form one score block and update stable running softmax statistics for the relevant query rows.
- Accumulate weighted values while rescaling partial results when the running normalization changes.
- Discard the temporary score block rather than writing the full score matrix to HBM.
- Use the retained statistics and permitted recomputation for backward execution, then verify numerical outputs and measured traffic.
Science and assumptions
Score storage
A materialized dense score matrix contains query positions times key positions times bytes per score.
Symbols and units
- M_score
- Dense score-matrix payload for one head and request
bytes
- T_query
- Query-position count
positions
- T_key
- Key-position count
positions
- s
- Stored bytes per score
bytes/score
Tquery and Tkey count positions and s is bytes/score; Mscore is bytes for one dense score matrix.
Assumptions One head and one request, no padding; causal storage and implementation fusion change actual allocation.
Engineering example
Synthetic one-head score matrix with 4,096 queries, 4,096 keys and two-byte scores has 33,554,432 bytes, or 32 MiB. Avoiding its materialization removes one large temporary, while keys, values and outputs still require storage.
Distinctions and tradeoffs
IO-aware exact attention reduces materialized intermediates; it does not by itself reduce the quadratic number of dense attention pair interactions.
Less outer-memory traffic can justify additional recomputation and local bookkeeping; oversized tiles can exhaust scratchpad or reduce execution residency.
Long contexts amplify intermediate-memory costs, so attention performance depends on IO organization as well as matrix arithmetic.
Specifications and interfaces
- Score shape
- Tquery × Tkey
Element count grows quadratically when both lengths scale together.
- Exactness boundary
- Same mathematical attention with finite-precision effects
Order and rounding can differ from another implementation.
Interfaces
- Query, key and value tiles Input
- Blocks staged from accelerator memory
- Attention output Output
- Normalized weighted values without a full stored score matrix
Check your understanding
What is the main avoided object in tiled IO-aware attention?
- The full attention score/probability matrix in outer memory
- All model weights
- The causal relationship between tokens
Answer and reasoning
The full attention score/probability matrix in outer memory
Tiles and running normalization allow the result to be accumulated without retaining every pairwise score in HBM.
Evidence and scope
The example is a storage calculation, not a promised speedup. Exactness refers to the algorithm’s attention definition; finite-precision evaluation still needs testing.
No current calculator directly models this lesson’s main mechanism. Use the library for related bounds while retaining their stated omissions.