Why FlashAttention Can Avoid a Square Attention Array Without Making Attention Linear

The requested output is smaller than the square intermediate

Consider one self-attention head with N query positions, N key/value positions and fixed head dimension d. Q and K hold the query and key vectors; V holds the values to combine. A scaled query-key dot product supplies each score. Row-wise softmax converts scores into weights, and the weighted combination of V produces the output.

Each output row has d values, while a fully materialized score array has N values for each of N queries. The conceptual square relationship does not require the implementation to retain every score simultaneously. A mathematical intermediate and a live device-memory allocation are different things.

This explanation concerns the dense algorithm in arXiv:2205.14135v2, revised June 23, 2022, with a June 24 cover date. Its separate block-sparse extension and later kernel designs are outside this mechanism comparison.

References: [AT1][FA1]

A two-byte square grows from 512 MiB to 2 GiB

Stipulate one full N-by-N score array for one head and one sequence, with exactly two bytes per element. At N = 16,384, the count is 16,384 x 16,384 = 268,435,456 elements. Multiplying by two gives 536,870,912 bytes, or 512 MiB. Here 1 MiB = 2²⁰ bytes.

At N = 32,768, the same calculation gives 1,073,741,824 elements and 2,147,483,648 bytes, or 2 GiB, where 1 GiB = 2³⁰ bytes. Doubling both dimensions quadruples this particular array. That arithmetic says nothing about a fourfold change in total model memory.

These are invented storage assumptions for a complete square, not observations of a kernel allocation. They omit any second probability array, other tensors, allocator overhead and precision differences. Do not multiply this figure by all model layers without establishing which allocations coexist.

A causal mask excludes future positions from the attention calculation. It does not, by itself, turn a stipulated full square allocation into packed triangular storage. Retain the full-square byte count unless the actual representation establishes something else.

References: [AT1]

Tiles need a shared softmax normalization

The tiling idea is to bring blocks of Q, K and V from GPU high-bandwidth memory, or HBM, into much smaller on-chip SRAM. A score tile is consumed while updating a compact output state. The implementation can then reuse its working space instead of writing a whole square score or probability array to HBM. Tri Dao's later FlashAttention Recap describes this original mechanism; its newer parallelism changes are not assumed here.

The difficulty is softmax: the denominator for one query depends on all participating keys. Computing an independent softmax in every tile and adding the resulting outputs would give each tile its own normalization. It would not recover the full-row result.

For numerical stability, subtract a maximum m before exponentiation. Retain m and a running sum l of exp(score - m). When another block arrives, choose the larger maximum and rescale the old and new sums to that common reference before adding them. The 2018 online-normalizer paper establishes this rescaling principle. In attention, the value-weighted contribution must be rescaled consistently as well.

References: [FA2][SM1]

A two-block example shows why rescaling matters

Use an original scalar illustration with four already-computed scores: 0, 0, ln(2), ln(2), where ln is the natural logarithm. Assign corresponding scalar values 0, 0, 10, 10. Ignore masking and dropout. These numbers illustrate algebra, not a model workload or a floating-point test.

For the first block, the maximum is 0, the shifted exponential sum is 2, and the value-weighted numerator is 0. For the second block, the maximum is ln(2), its shifted exponential sum is also 2, and its numerator is 20.

To combine them, use ln(2) as the common maximum. The first block's sum is rescaled by exp(0 - ln(2)) = 1/2, becoming 1. The second block's sum remains 2. The combined numerator is 0 + 20, and the denominator is 1 + 2, giving output 20/3, approximately 6.667.

An equal average of the two independently normalized block outputs would instead give (0 + 10)/2 = 5. The blocks contain the same number of elements, yet carry different softmax mass. Keeping the normalization state preserves that difference without retaining the complete row of probabilities. The example expresses the combining identity; it does not specify a kernel's internal accumulator layout.

References: [SM1][AT1]

Backward can recreate tiles instead of retaining the square

Training also needs gradients. The original FlashAttention algorithm retains Q, K, V, the output and row normalization statistics so that score and probability tiles can be reconstructed on chip during backward. It avoids retaining full square intermediates from forward for that purpose. With dropout enabled, the full algorithm also retains pseudorandom-generator state to regenerate the same dropout mask during backward, rather than storing a square mask.

Recomputation adds arithmetic while avoiding the corresponding storage and HBM round trips. Whether that tradeoff reduces elapsed time on a particular system requires evidence for that implementation. An inference-only comparison has no training backward pass to count.

References: [FA1]

Linear extra storage and quadratic work can coexist

The 2022 paper states O(N²d) arithmetic and O(N) additional memory beyond inputs and output for its dense forward algorithm. With d fixed, the arithmetic remains quadratic in N. Q, K, V and the output each have N-by-d shape, so their storage also grows linearly under these assumptions.

A tile makes the live working set smaller; it does not remove the dense query-key relationships. A causal calculation still has N(N + 1)/2 participating positions, which remains quadratic. Avoiding a square allocation is therefore compatible with substantial growth in arithmetic as context grows.

Memory traffic asks how many bytes cross a boundary over time. Repeated reads count repeatedly, even when they reuse the same allocation. Smaller peak storage does not establish linear HBM traffic. The original algorithm can revisit input and output blocks. With head dimension and on-chip capacity fixed, the paper's dense IO bound remains quadratic in sequence length, even while reducing HBM accesses relative to its materialized baseline.

Use the related roofline guide to interpret work per transferred byte. For this decision, keep three separate questions: which tensors coexist, how much data moves through HBM, and how much dense arithmetic is performed.

References: [FA1][FA2]

The model still has weights, state and other working memory

Removing square attention intermediates leaves other memory populations. NVIDIA's TensorRT-LLM memory reference distinguishes model weights, internal activation tensors and input/output tensors, including KV cache. It also describes runtime buffers, attention-plugin workspace and memory reuse between layers. Those categories are useful accounting examples, rather than a universal allocation formula.

In autoregressive serving, retained key/value state supports subsequent generation. Eliminating a score workspace does not eliminate that state. The related weight-quantization and prefill/decode guides address these separate populations and serving phases.

The square example concerns equal query and key lengths. PyTorch 2.14 documents target and source lengths separately: the attention-weight shape contains L x S. A single new query attending to N retained keys has a 1-by-N score relationship for that step, not an N-by-N square. Do not apply the example's allocation arithmetic to every decode step.

Ask which phase a memory claim describes and what remains live at its peak. Training, prompt processing and generation need their own scope. A saving in one intermediate can enable useful headroom without determining the memory allocation required for the complete workload.

References: [N1][PT1]

A PyTorch API label does not prove execution or numerical identity

PyTorch 2.14's scaled_dot_product_attention documentation describes automatic implementation selection based on inputs and specific input limitations for fused kernels. It names FlashAttention-2 among its implementations. That is a versioned API reference, not evidence that a given run executed the classic 2022 kernel or any particular optimized path.

The same page warns that results can vary with the selected backend because floating-point operations are fused differently. Exact attention means preserving the mathematical attention operation; it does not promise bitwise equality across implementations. Numerical acceptance therefore needs an agreed comparison and tolerance appropriate to the application.

Request the implementation and version actually used, the supported input conditions and evidence of the executed path alongside memory and timing results. Neither enabling an API nor attaching a family name supplies those results.

References: [PT1]

Decide which part of the long-context claim is supported

Consider two fictional proposals for the same long-document workload. Proposal A cites the two-byte square calculation and says FlashAttention makes twice the context fit in the same total memory. Proposal B identifies the model revision, hardware, runtime and executed attention backend, and offers a comparison report for the proposed context and concurrency. No measured result is stipulated for either proposal.

A explains why one conventional intermediate becomes expensive, but leaves total fit unresolved. B offers the relevant kind of evidence; its report still needs to show the complete peak, metric and units, elapsed time and acceptable outputs under matching conditions. Keep batch size, head dimensions, precision, mask and dropout conditions, prompt/output lengths and concurrent requests visible. Identify whether timing covers attention alone or the complete workload.

Accept an allocation or speed claim only to the extent that this evidence supports the requested use. If only attention storage has been established, retain that narrower conclusion. The related benchmark guide provides the broader system-comparison framework.

If the selected equipment needs physical installation, Leviathan's rack-and-stack service covers agreed equipment mounting, IT cabling, labels and installation records. Use the related project inquiry with the equipment list, site and released rack/cable package. Keep software evaluation and workload-performance acceptance with the appointed platform team.

References: [L1]