{"id":"attention-io-tiling","level":"software","name":"IO-aware attention and online softmax","shortName":"IO-aware attention and online softmax","role":"Computes attention while reducing large intermediate transfers to outer memory.","description":"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.","specs":[{"label":"Score shape","value":"Tquery × Tkey","context":"Element count grows quadratically when both lengths scale together."},{"label":"Exactness boundary","value":"Same mathematical attention with finite-precision effects","context":"Order and rounding can differ from another implementation."}],"interfaces":[{"name":"Query, key and value tiles","direction":"in","signal":"Blocks staged from accelerator memory"},{"name":"Attention output","direction":"out","signal":"Normalized weighted values without a full stored score matrix"}],"tradeoff":"Less outer-memory traffic can justify additional recomputation and local bookkeeping; oversized tiles can exhaust scratchpad or reduce execution residency.","whyAI":"Long contexts amplify intermediate-memory costs, so attention performance depends on IO organization as well as matrix arithmetic.","related":["attention","kernel-locality-fusion","gpu-occupancy-pipeline"],"sources":[{"title":"FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness","url":"https://arxiv.org/abs/2205.14135"}],"scale":"One model step to accelerator memory hierarchy","kind":"workload","branch":"Model–machine co-design","subbranch":"Training, prefill and decode","mechanism":["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."],"misconception":"IO-aware exact attention reduces materialized intermediates; it does not by itself reduce the quadratic number of dense attention pair interactions.","check":{"question":"What is the main avoided object in tiled IO-aware attention?","options":["The full attention score/probability matrix in outer memory","All model weights","The causal relationship between tokens"],"answer":0,"explanation":"Tiles and running normalization allow the result to be accumulated without retaining every pairwise score in HBM."},"science":[{"title":"Score storage","equation":"Mscore = Tquery Tkey s","explanation":"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."}],"prerequisites":["attention","kernel-locality-fusion"],"learningObjective":"Explain how tiling and online normalization reduce attention IO without confusing exact and sparse methods.","workstream":"W11","visualFamily":"model-workload","engineeringExample":"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.","evidenceNotes":"The example is a storage calculation, not a promised speedup. Exactness refers to the algorithm’s attention definition; finite-precision evaluation still needs testing.","labId":"library","scope":"No current calculator directly models this lesson’s main mechanism. Use the library for related bounds while retaining their stated omissions.","investigationId":"execution","labScope":"No current calculator directly models this lesson’s main mechanism. Use the library for related bounds while retaining their stated omissions."}