Attention Memory Calculator
Inputs
| Sequence length | 4,096 |
|---|---|
| Attention heads | 32 |
| Batch size | 1 |
| Precision | FP16 / BF16 (2 bytes) |
Attention Memory Calculator
Estimate the memory of the materialized attention score matrix in standard transformer attention, from sequence length, head count, batch size, and bytes per element — the quadratic term FlashAttention removes.
Inputs
Workload
Results
Enter a value to see results.
Attention Memory
Standard transformer attention compares every token with every other token, building a score matrix before the softmax. That matrix is square in the sequence length, so its memory grows with the square of the context — the famous quadratic cost of attention. This calculator estimates the size of that materialized matrix from the sequence length, the number of attention heads, the batch size, and the bytes used per score. It is precisely the memory that FlashAttention avoids storing.
The quadratic term
For a sequence of tokens, attention forms an matrix of scores: row , column holds how much token attends to token . Storing it takes memory proportional to . This is unlike the model weights or the key-value cache, which grow linearly with context — the score matrix grows with its square, so at long context it becomes the dominant intermediate buffer and the reason naive attention runs out of memory.
The formula
Each head builds its own score matrix and each sequence in the batch carries its own copy, so the memory in bytes is
Abytes=B⋅H⋅s2⋅ewhere is the batch size, the number of attention heads, the sequence length, and the bytes per stored score. Dividing by gives gigabytes. The square on is what makes long context costly: the head count and batch only multiply linearly.
What this memory is
The figure measures the materialized score matrix that standard attention writes to memory between the query-key product and the softmax-weighted sum over values. It is the activation buffer that a naive implementation must hold. Whether it persists depends on the setting: in a plain forward pass it is transient and can be reused across layers, while training frameworks may keep a per-layer copy for the backward pass unless checkpointing or a fused kernel intervenes.
FlashAttention
FlashAttention computes the identical output without ever writing the full matrix. It walks the keys and values in small blocks, maintaining running softmax statistics so it only needs one block in memory at a time, which turns the storage from quadratic to linear in sequence length. The matrix this calculator sizes is exactly what FlashAttention declines to store — so the result here is a good proxy for the memory a fused attention kernel saves at a given context length.
Worked example
Take a single sequence of 4,096 tokens with 32 attention heads in 16-bit precision:
Abytes=1×32×40962×2=1,073,741,824about 1.07 GB for one layer's score matrix. Double the context to 8,192 tokens and the square term takes over: the same expression gives roughly 4.29 GB, four times as much for twice the length. The key-value side of the same attention budget is covered in the KV Cache Size Calculator, and the model-weight portion in the LLM Inference VRAM Calculator.
Frequently Asked Questions (FAQ)
Why is attention memory quadratic in sequence length?
Attention compares every token with every other token, producing a score matrix whose dimensions are both the sequence length. Storing that full matrix therefore takes memory proportional to the sequence length squared. Doubling the context quadruples the score matrix, which is why standard attention becomes memory-bound at long context — the quadratic term outpaces the linear cost of the weights and activations.
How does FlashAttention reduce this memory?
FlashAttention computes the same result without ever materializing the full score matrix. It streams the attention over small blocks of keys and values, keeping running softmax statistics so it only holds one block at a time.
The quadratic matrix this calculator measures is exactly the memory FlashAttention avoids storing, which is what lets it scale to much longer sequences. The arithmetic is unchanged; only the intermediate storage shrinks from quadratic to linear.
Is this memory per layer or for the whole model?
It depends on the implementation. The figure here is the size of one layer’s score matrix. In a naive forward pass that buffer can be freed and reused between layers, so it is transient rather than multiplied by layer count. During training, however, frameworks may retain per-layer attention tensors for the backward pass unless gradient checkpointing or FlashAttention is used, in which case the total can scale with the number of layers.
Disclaimer
This estimate covers only the materialized score matrix at the stated precision; it excludes the query, key, and value tensors, the output projection, and other activations. It describes standard attention — kernels such as FlashAttention do not store this matrix at all.