On this page
A Transformer's compute and memory grow with the number of tokens n; computing in small tiles, caching past keys and values, sharing them across heads and letting each token read only a recent window reduce that cost. The cost is paid in every attention layer, over all n tokens at once in training and on the prompt, then once per new token during generation.
Each token is d numbers wide (the model width); after attention, a feed-forward network (FFN) widens each token to dff and back. A decoder-only model like GPT generates one token per step and feeds it back in.
Each layer has two costs, attention and the FFN; attention's grows with the square of the sequence length, the FFN's only linearly. Costs are in FLOPs, one per multiply or add; projections, softmax and norms are left out.
Memory follows the same square: at 100,000 tokens, one head's score table holds 10 billion entries.
Figure 2. Below n = 4d the FFN dominates; above it the square term takes over.
So long-context models restrict attention: windowed (each token reads only a recent window), sparse (each token reads a chosen subset) or linear (cost growing in a straight line with length).
Standard attention writes the n by n score table to the GPU's large, slow memory (HBM) and reads it back; Dao et al. (2022) found this traffic, more than the arithmetic, sets the time.
FlashAttention computes exactly the same output in tiles in the GPU's small fast on-chip memory (SRAM) of size M; each tile holds Bc rows of K and V, and the full table never lands in HBM. Softmax subtracts each row's maximum before exp, so exp never overflows; a tile sees only part of a row, so it keeps a running maximum and sum, and when a later tile has a bigger maximum, the earlier sum is rescaled by exp(old max − new max). The backward pass keeps only those per-row statistics and recomputes the tiles, so it does more arithmetic, not less.
The math behind the running sum, on one row of 8 scores split into two tiles.
import numpy as np
TILE = 4
s = np.random.randn(2 * TILE) # one row of 8 scores
tile1, tile2 = s[:TILE], s[TILE:]
m1, m2 = tile1.max(), tile2.max()
sum1 = np.exp(tile1 - m1).sum()
sum2 = np.exp(tile2 - m2).sum()
m = max(m1, m2)
total = sum1 * np.exp(m1 - m) + sum2 * np.exp(m2 - m)
full = np.exp(s - m).sum()
print(np.allclose(total, full)) # TrueThe FlashAttention paper timed attention in GPT-2 medium (1024 tokens, 16 heads of size 64, batch 64, forward and backward passes, one A100).
A naive generation loop reruns the whole sequence at every step. Under the causal mask (a token never sees later tokens) past keys and values never change, so real systems store them in a KV cache. Each step computes only the new token's key and value, appends them, and scores the new query against everything stored; prefill runs the prompt through at once to fill the cache first.
In one decode step with the cache, one new token goes in and one output row comes out.
import torch
B, h, T, d_k = 1, 8, 10, 64
K = torch.randn(B, h, T, d_k) # cache so far
V = torch.randn(B, h, T, d_k)
q, k, v = torch.randn(3, B, h, 1, d_k) # the new token
K = torch.cat([K, k], dim=2)
V = torch.cat([V, v], dim=2)
S = q @ K.transpose(-2, -1) / d_k ** 0.5 # (B, h, 1, T + 1)
out = S.softmax(dim=-1) @ V # (B, h, 1, d_k)Cost per generated token, with and without the KV cache.
Figure 5. Each block is one stored key and value; the bars compare the newest token's cost with and without the cache.
The cache holds a key and a value for every token, in every head of every layer:
The cache size for a 12-layer model, read off the card.
L, h_kv, d_k = 12, 12, 64 # layers, KV heads, head size
T, b = 4096, 2 # tokens, bytes per number
BATCH = 32
per_seq = 2 * L * h_kv * d_k * T * b # 2 = key + value
per_batch = BATCH * per_seq
print(per_seq, per_batch) # 150994944 4831838208, about 151 MB and 4.8 GBShazeer (2019) notes that each decode step reloads the large key and value tensors, so memory bandwidth, not arithmetic, sets the speed.
Multi-query attention (MQA) keeps all h query heads but shares one key head and one value head. Grouped-query attention (GQA, Ainslie et al. 2023) splits the query heads into g groups, each sharing one key head and one value head.
Grouped-query attention in PyTorch, with 8 query heads sharing 2 key-value heads.
import torch
import torch.nn.functional as F
sdpa = F.scaled_dot_product_attention
B, T, d_k = 1, 16, 64
h, g = 8, 2 # query heads, KV heads
Q = torch.randn(B, h, T, d_k)
K = torch.randn(B, g, T, d_k) # the cache stores g heads
V = torch.randn(B, g, T, d_k)
per_group = h // g
K_rep = K.repeat_interleave(per_group, dim=1) # temporary
V_rep = V.repeat_interleave(per_group, dim=1)
out = sdpa(Q, K_rep, V_rep, is_causal=True)
print(out.shape) # torch.Size([1, 8, 16, 64])Sharing trades a little quality for much faster decoding. The MQA paper timed an 8-head English to German translator with 211M weights on one TPUv2, a Google chip built for neural networks. The GQA paper took T5 XXL, Google's pretrained 11-billion-weight text model, averaged each group's key and value heads, trained 5% more steps and scored summarization, translation and question answering.
In online speech recognition (often trained with CTC loss, which needs no frame-by-frame labels) or live translation the model can't see the whole sequence: the future hasn't been spoken yet. Windowed attention lets each query look back only a fixed number of frames (time slices of the input, such as 10 ms of audio) and masks everything else to negative infinity. The cost then grows in a straight line with length.
The windowed mask for n frames and a window of w.
import torch
n, w = 10, 3 # frames, window
NEG_INF = float("-inf")
mask = torch.full((n, n), NEG_INF) # added to the scores
for i in range(n):
start = max(0, i - w + 1) # oldest frame query i reads
mask[i, start:i + 1] = 0
print(mask)Size the window from the smallest unit you must capture (a phoneme, one distinct sound of a language; a gesture; a word), converted to frames.
Each layer's tokens already carry their window's context, so reach grows with depth, like a receptive field.
Chunked streaming emits output in fixed-size chunks, and every token may see its whole chunk, so with chunk size 4, counting from 0, token 9 can see tokens 10 and 11.
Chunked streaming, token by token.
Figure 9. A chunk's rows light up only when its last token arrives, so output lags input by up to one chunk.
Conformer puts convolution back beside attention for speech; Quantization and Pruning make a trained Transformer cheaper to serve, and Adapters fine-tune one without retraining all of it.
| cause | symptom | fix |
|---|---|---|
| Post-LN without warmup | loss diverges early | warm up, or use Pre‑LN |
| a fully masked row | NaN spreads through training | mask padding keys only |
| inputs past the trained length | quality falls off | ALiBi, or train longer |
| a KV cache too big | out of memory | GQA or MQA |
Generation runs once over the prompt, then once per token against the cache.
import torch
import torch.nn.functional as F
sdpa = F.scaled_dot_product_attention
B, h, h_kv, T, d_k = 1, 8, 2, 16, 64
q = torch.randn(B, h, T, d_k)
K = torch.randn(B, h_kv, T, d_k) # cache after the prompt
V = torch.randn(B, h_kv, T, d_k)
prompt_out = sdpa(q, K, V, is_causal=True, enable_gqa=True)
q_new = torch.randn(B, h, 1, d_k)
k_new, v_new = torch.randn(2, B, h_kv, 1, d_k)
K = torch.cat([K, k_new], dim=2) # cache grows by one
V = torch.cat([V, v_new], dim=2)
step_out = sdpa(q_new, K, V, enable_gqa=True)
print(step_out.shape) # torch.Size([1, 8, 1, 64])| FlashAttention | Every attention call, forward and backward; same output, and it keeps only each row's softmax maximum and sum (two numbers per row, so this memory grows only linearly with n) for the backward pass. |
|---|---|
| KV cache | Lives through generation; 2·L·hkv·dk·T·b bytes. |
| MQA, GQA | Fixed in the architecture; the cache shrinks by h/g. |
| Streaming | Each chunk of C frames as it arrives; earlier chunks stay readable. |
| Kernel | PyTorch picks FlashAttention or another kernel; FlashAttention needs fp16 or bf16, the two 16-bit number formats. |
| Cache | Hugging Face's default DynamicCache grows per token; cache_implementation="static" preallocates. |
| Cache length T | Grows the bytes card linearly, one key and value per layer per token. Longer remembers more and costs more; capped by the trained context, 2048 for GPT-3. |
| Groups g | KV heads; the cache keeps g/h of full size. More approach multi-head quality, fewer MQA speed; 1 is MQA, h is multi-head. GQA chose 8. |
| Block size | K and V rows per tile, Bc = M/4dk for on-chip memory M. Larger tiles cut HBM traffic and run faster up to about 256 rows. Past the on-chip memory they no longer fit. |
| Chunk size C | Sets which keys a query may read, the chunk card in Section 5. Larger lets each token see more of its own chunk and adds up to one chunk of delay; smaller cuts the delay. At 1 it is the plain causal mask; at the full length it is ordinary attention with no streaming. No library default; set it from the delay the application allows. |
| Window w | Frames each query looks back over, in the windowed mask of Section 5. Wider sees more context per layer at w scores per query; narrower is cheaper and needs more layers to reach far. At 1 a token sees only itself; at the full length it is plain causal attention. Mistral 7B (Jiang et al. 2023) uses 4096 tokens. |
| Code | F.scaled_dot_product_attention (PyTorch docs), flash-attention (Dao-AILab), Cache strategies (Hugging Face). |
is_causal=True on the one-query step; it would show only the first key.enable_gqa needs query heads divisible by KV heads.