Fast Attention and Long Context
Attention sinks and streaming
Models dump huge amounts of attention onto their first few tokens because softmax forces every query to spend all its weight somewhere, and throwing those tokens away breaks the model.
- 14 min read
- 3 reading levels
- Updated
Read these first
On this page 10
One lesson, three depths. Pick the one that fits you today — you can switch any time.
Beginner — No maths. Plain English.
The short answer
A model must spread all its attention across the words it can see, even when none are useful. It dumps the leftovers on the first few words.
The plate at a wedding
You are handed a plate and pushed into the food line. You are not hungry, and nothing on the counter appeals to you.
You cannot walk out with an empty plate. So you take a spoon of plain rice from the start of the counter and move on.
Nobody chose the rice. The rice was chosen because the plate could not be empty. The first words in a document are that plain rice.
Why the plate can never be empty
The step that turns attention scores into percentages is built to always total one hundred percent. There is no option for "none of these matter to me".
So a head that finds nothing worth looking at still has to place its full weight somewhere. It needs a harmless place to put it.
The first token of the text is perfect for the job. Every query can see it, no matter how far along the text has gone. It becomes the model's dustbin.
The measurement that surprised people
Look inside a trained model and the numbers are startling. In GPT-2, roughly forty percent of all attention weight lands on the very first word of the text.
Not because that word is important. In the test further down this page it is the word "The".
Why this breaks long conversations
Sliding-window attention keeps only the most recent words and drops old ones. Sensible, and it caps memory.
Then the window slides past the beginning of the text, and the first tokens are thrown away.
The dustbin is gone. Every head that was quietly parking weight there must now dump it on real words instead. Those words get wrongly amplified, and the output falls apart.
works breaks
----- ------
[t0 t1 t2] ... [recent 500] [ ..... recent 500 ..... ]
^ kept forever first tokens evicted
the dustbin survives nowhere to park the leftoversThe fix, which is almost embarrassing
Keep the first four tokens. Forever. Then slide the window over everything else.
That is the whole repair. A model that produced nonsense after a few thousand words now runs over millions of words without drifting.
The better fix, built into new models
Newer models add a dustbin on purpose. Each head gets its own extra slot that holds no word at all.
The head can put weight into that slot when nothing deserves it. Because the slot holds nothing, that weight disappears. The remaining words then get honest shares, adding up to less than one hundred percent.
This is a two-line change to the architecture. It has to be there during training, so it cannot be bolted onto an old model.
Where you have already seen this
- A chat assistant that stays coherent through a very long session.
- Live transcription and translation that runs for hours without restarting.
- Local model servers with a "keep first N tokens" setting you never understood.
Remember this
- Attention must spend all its weight. The first tokens absorb whatever is left over.
- Evicting those tokens from a rolling window breaks the model badly.
- Keeping four of them fixes it. New models add a dedicated empty slot instead.
What to learn next
- Sparse and block attention patterns — richer skipping schemes, all of which must preserve sinks.
- Sliding-window attention — the technique whose failure mode this lesson explains.
- Context window — what a model can hold, and why streaming is not the same as remembering.
Developer — Code and libraries.
Setup
pip install numpy torch transformersWritten against NumPy 1.26.4, PyTorch 2.5.1 and transformers 5.6.2, Python 3.10. Both scripts run on CPU. The second one downloads GPT-2, about 500 MB, once.
Why the plate cannot be empty
import numpy as np
# 1. Softmax has no way to say "none of these".
scores = np.array([-8.0, -7.5, -9.0, -8.2]) # every key looks irrelevant
w = np.exp(scores - scores.max()); w /= w.sum()
print("scores:", scores)
print("weights:", np.round(w, 3), " sum =", round(float(w.sum()), 6))
strong = np.array([-8.0, 6.0, -9.0, -8.2]) # one key is a great match
w2 = np.exp(strong - strong.max()); w2 /= w2.sum()
print("weights when one key fits:", np.round(w2, 3), " sum =", round(float(w2.sum()), 6))
# 2. A learned sink: one extra logit per head, dropped after the softmax.
def attn_with_sink(scores, sink):
combined = np.concatenate([scores, [sink]])
p = np.exp(combined - combined.max()); p /= p.sum()
return p[:-1] # the sink column is discarded
for sink in (-10.0, 0.0, 5.0):
out = attn_with_sink(scores, sink)
print(f"sink logit {sink:>6.1f} -> weights {np.round(out, 3)} sum = {out.sum():.4f}")scores: [-8. -7.5 -9. -8.2] weights: [0.261 0.43 0.096 0.213] sum = 1.0 weights when one key fits: [0. 1. 0. 0.] sum = 1.0 sink logit -10.0 -> weights [0.252 0.415 0.093 0.206] sum = 0.9659 sink logit 0.0 -> weights [0. 0.001 0. 0. ] sum = 0.0013 sink logit 5.0 -> weights [0. 0. 0. 0.] sum = 0.0000
The first two lines are the whole problem. Scores of -8 and scores of +6 produce weights that sum to exactly 1.0 either way. Softmax is scale-invariant in a way that destroys the information "nothing here is relevant". A head with no good match still emits a confident-looking distribution.
The sink rows show the repair. With a very negative sink logit the sink absorbs 3.4% and the real weights barely move. Push the sink logit to 0 and it swallows 99.9% of the weight, so the head contributes almost nothing to the output.
That is the point: the head can now choose to contribute nothing. The sink logit is a learned parameter, one per head, and the model learns per head how eager it is to abstain.
This is exactly what gpt-oss does. In transformers, GptOssAttention holds self.sinks = nn.Parameter(torch.empty(config.num_attention_heads)), concatenates it to the attention logits, softmaxes, then drops the last column with the comment # we drop the sink here.
Measuring a real sink
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
torch.manual_seed(0)
name = "gpt2" # 124M parameters, downloads once (~500 MB)
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(name, attn_implementation="eager").eval()
text = ("The train left the station at dawn. Fields of sugarcane ran past the window "
"for an hour. A vendor walked the aisle selling tea in paper cups. By noon the "
"line climbed into the hills and the air turned cool. Passengers pulled down "
"their bags and began to gather at the doors.") * 3
ids = tok(text, return_tensors="pt").input_ids
print("tokens:", ids.shape[1], "| model:", name, "| torch:", torch.__version__)
with torch.no_grad():
att = model(ids, output_attentions=True).attentions # tuple: one tensor per layer
# For queries at position 32 onward, how much weight lands on the very first token?
print(f"\n{'layer':>6}{'weight on token 0':>20}{'weight on last 8 tokens':>26}")
for L, a in enumerate(att):
a = a[0][:, 32:, :] # heads x queries x keys
on_first = a[:, :, 0].mean().item()
on_recent = sum(a[:, q, max(0, q + 32 - 7):q + 33].sum() for q in range(a.shape[1]))
on_recent = (on_recent / (a.shape[0] * a.shape[1])).item()
print(f"{L:>6}{on_first:>20.3f}{on_recent:>26.3f}")
allmass = torch.stack([a[0][:, 32:, 0].mean() for a in att])
print(f"\naverage weight on token 0 across all {len(att)} layers: {allmass.mean():.3f}")
print("token 0 is:", repr(tok.decode(ids[0, :1])))tokens: 180 | model: gpt2 | torch: 2.5.1+cu121
layer weight on token 0 weight on last 8 tokens
0 0.012 0.455
1 0.070 0.350
2 0.165 0.640
3 0.299 0.544
4 0.363 0.448
5 0.480 0.177
6 0.445 0.273
7 0.513 0.127
8 0.568 0.143
9 0.612 0.085
10 0.640 0.044
11 0.589 0.070
average weight on token 0 across all 12 layers: 0.396
token 0 is: 'The'output_attentions=True requires attn_implementation="eager". Fused kernels never materialise the attention matrix, so they cannot return it. Model download aside, this runs in a few seconds on CPU.
Reading that output
Layer 0 puts 1.2% on token 0. Layer 10 puts 64%. The sink is not present at the input. It is built up through the network, and by the upper half of the model it is the single largest destination for attention weight.
The two columns trade off almost perfectly. As weight on token 0 climbs from 0.012 to 0.640, weight on the eight most recent tokens falls from 0.455 to 0.044. Upper layers are not doing local mixing. They are, for most heads and most queries, abstaining.
Token 0 is the word "The". It carries no information anyone needs at position 150. Its value is structural: it is visible to every query, so it is a reliable place to park weight.
39.6% average. If a rolling cache evicts this token, 39.6% of the attention weight in the model has to be redistributed onto tokens that did not earn it. That is why naive window eviction fails, and it is why the fix is so cheap.
Putting it to work
The rule from StreamingLLM: keep the first few tokens pinned in the cache forever, then roll the window over everything after them.
# conceptual cache policy - keep sinks, roll the rest
def evict(cache_len, n_sink=4, window=2048):
if cache_len <= n_sink + window:
return list(range(cache_len)) # nothing to drop yet
return list(range(n_sink)) + list(range(cache_len - window, cache_len))One detail that trips everyone up: positions are assigned within the cache, not in the original text. If the cache holds tokens 0-3 and 9000-11047, the rotary positions used are 0, 1, 2, 3, 4, 5, … not 0, 1, 2, 3, 9000, 9001, … The model was never trained on a gap of nine thousand positions, and using the true indices reintroduces the failure you were fixing.
Common mistakes
Keeping zero sink tokens. The most common bug in hand-rolled streaming. Output stays grammatical for a while and then degrades into fluent nonsense, which makes it hard to spot in casual testing.
Keeping too many. Four is enough in the published experiments. Keeping 512 wastes cache without further benefit.
Using original positions for the retained window. Covered above. Symptom: perplexity that is bad immediately rather than degrading gradually.
Expecting streaming to mean long-term memory. StreamingLLM keeps a model fluent over four million tokens. It cannot recall what happened a million tokens ago, because that content was evicted. Fluency and recall are different properties.
Trying to add learned sinks to an existing checkpoint. The sink parameter must be present during training. Retrofitting it to trained weights changes every attention distribution at once.
Try it yourself
Change the text so it starts with a comma or a newline, and re-measure. The sink usually stays on position 0 regardless of what that token is, which is the cleanest evidence that its role is structural rather than semantic. Then compare layer 0 and layer 11 side by side to see the sink form.
What to learn next
- Sparse and block attention patterns — richer skipping schemes, all of which must preserve sinks.
- Sliding-window attention — the technique whose failure mode this lesson explains.
- Context window — what a model can hold, and why streaming is not the same as remembering.
Researcher — Mathematics and papers.
The phenomenon
Xiao, Tian, Chen, Han and Lewis (2024), Efficient Streaming Language Models with Attention Sinks (arXiv:2309.17453, ICLR 2024) named the effect. Their observation: window attention collapses when the first tokens leave the cache, and "keeping the KV of initial tokens will largely recover the performance of window attention".
Their StreamingLLM keeps a small number of initial tokens — four in the reported experiments — alongside a rolling window, and demonstrates stable generation on Llama-2, MPT, Falcon and Pythia "up to 4 million tokens and more", with up to 22.2x speedup over a sliding-window-with-recomputation baseline.
Why softmax causes it
For query $q$ and keys $k_1 \dots k_n$:
$$ \alpha_i = \frac{\exp(q^{\top} k_i / \sqrt{d})}{\sum_{j=1}^{n} \exp(q^{\top} k_j / \sqrt{d})}, \qquad \sum_i \alpha_i = 1 $$
The constraint $\sum_i \alpha_i = 1$ is unconditional. There is no assignment of scores that yields a small total. A head whose useful pattern is absent from the current context must still emit a full unit of weight, and whatever it lands on is added to the residual stream.
The learned escape route is a token that is (a) visible to every query, since causal masking guarantees position 0 always is, and (b) trained to carry a value vector near zero, so that receiving weight is harmless. Initial tokens satisfy both.
The architectural fix
Append a per-head learned logit $s_h$ to the score vector before normalising, and discard its probability afterwards:
$$ \tilde{\alpha}_i = \frac{\exp(q^{\top} k_i/\sqrt{d})}{\exp(s_h) + \sum_{j} \exp(q^{\top} k_j/\sqrt{d})}, \qquad \sum_i \tilde{\alpha}_i = \frac{1}{1 + e^{s_h - \text{lse}}} < 1 $$
where $\text{lse}$ is the log-sum-exp of the real scores. This is the $\text{softmax}_1$ construction proposed in Evan Miller's 2023 note Attention Is Off By One, where the extra term is a fixed 1 (equivalently $s_h = 0$). Making $s_h$ learned and per-head is what gpt-oss ships.
The transformers implementation is worth reading directly, in models/gpt_oss/modeling_gpt_oss.py: the sink is concatenated to attn_weights, softmax runs over the concatenation, and scores = probs[..., :-1] discards the sink column. Two lines of substance.
Related findings
Darcet, Oquab, Mairal and Bojanowski (2023), Vision Transformers Need Registers (arXiv:2309.16588) find the same pathology in vision transformers: high-norm artefact tokens appear in low-information patches and corrupt dense feature maps. Their fix, extra learnable "register" tokens with no image content, is the sink mechanism arrived at independently from a different symptom. The convergence is strong evidence that this is a property of softmax attention rather than of language.
Related work on massive activations in the residual stream ties sinks to a small number of extreme-magnitude feature dimensions that act as implicit bias terms. Read that literature before assuming the sink is only an attention-level effect; the evidence points at a whole-network mechanism.
Practical consequences
- Cache policies must pin sinks. Any eviction strategy — window, H2O-style heavy-hitter selection, block eviction — needs the first tokens exempt. This is now standard in vLLM, SGLang and llama.cpp.
- Position remapping is required. Assign positions by cache slot, not by original index.
- Quantising the sink is dangerous. Those key and value vectors have unusual magnitudes; naive per-tensor int8 quantisation of the KV cache degrades them badly. Keep the sink entries in higher precision.
- Interpretability tooling must account for it. Attention-weight visualisations that include position 0 are dominated by it. Most published attention maps renormalise after removing the sink, often without saying so.
The open question
Sinks are a workaround for a normalisation constraint that arguably should never have been imposed. Attention weights are asked to be a probability distribution, but nothing in the operation requires that; the requirement is inherited from the original formulation rather than derived.
Alternatives exist — sigmoid attention, $\text{softmax}_1$, learned per-head temperature — and none has yet displaced softmax at scale. Whether that is because softmax is genuinely better, or because sinks work well enough that nobody has had to find out, is not settled.
What to learn next
- Sparse and block attention patterns — richer skipping schemes, all of which must preserve sinks.
- Sliding-window attention — the technique whose failure mode this lesson explains.
- Context window — what a model can hold, and why streaming is not the same as remembering.