Part 2: The Attention Sink Detector
Part 2 of a 4-part series on system-level LLM inference internals. Part 1 built the entropy tracker; this one looks one step earlier in the pipeline — at prefill, before a single token is generated.
Where This Fits
Part 1 tracked entropy during decode: how confident the model is, one generated token at a time. This post goes further upstream. Before decode even starts, the model runs a single forward pass over the entire prompt, and that pass is where attention gets distributed across every token in the context.
The unifying thesis from Part 1 said context quality shapes attention distribution during prefill, and attention distribution shapes generation confidence during decode. This post is where that first link gets measured directly.
| Part | What We Build |
|---|---|
| 1 | Per-token entropy tracker, visualized in real time |
| 2 — this post | Attention sink detector, context health scoring |
| 3 | Entropy-guided adaptive speculative decoding |
| 4 | Empirical study: correlation plots proving the causal chain |
The Phenomenon
Across a wide range of transformer LLMs, a small number of tokens, very often just the first token in the sequence — receive a hugely disproportionate share of attention from almost every later query, in almost every head, almost regardless of semantic relevance. StreamingLLM (Xiao et al., 2023) is the paper that named this: attention sinks, and the tokens absorbing the attention are sink tokens.
The mechanism comes straight out of the softmax constraint. Attention weights over keys are:
a_i = exp(q · k_i) / Σ_j exp(q · k_j)
This always sums to 1. A query can't express "nothing here is relevant" by producing low weight everywhere — the weights must still add up to 1 across all keys. So the model needs a release valve: somewhere to dump residual probability mass when nothing in the context is a strong match. Position 0 is the natural candidate, since it's visible to every query in a causal model and it's the one position guaranteed to exist as a valid attention target in every training sequence, at every length, regardless of anything else.
Why this matters practically, per the StreamingLLM paper: if you're managing a KV cache for long-context or streaming generation and evict the sink token to save memory, perplexity spikes — even though position 0 seems completely irrelevant by the time you're thousands of tokens deep. Their fix is to pin the first few tokens' KV pairs permanently, alongside a sliding window of recent tokens.
Phase 0: Getting the Numbers to Exist at All
The first obstacle isn't conceptual, it's mechanical. SDPA and FlashAttention kernels — the fast, default attention implementations — never materialize the full seq_len × seq_len attention matrix. That's the entire point of them: they compute attention output through tiled, online-softmax kernels, discarding each block's partial scores as soon as it's consumed. Great for speed and memory at long context. Useless if you want to look at the numbers.
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen2.5-3B-Instruct",
attn_implementation="eager", # non-negotiable for this phase
torch_dtype=torch.float16,
).to("mps")
attn_implementation="eager" is the only path where output_attentions=True returns real tensors instead of None. It's slower and considerably more memory-hungry — for a seq_len × seq_len matrix per head per layer, that adds up fast — but for a single prefill pass on a 3B model on Apple Silicon, it's a completely reasonable trade.
Phase 1: The Prefill Analyzer
Aggregating Attention
A single forward pass with output_attentions=True returns one tensor per layer, each shaped (batch, num_heads, seq_len, seq_len). Averaging across layers and heads gives one matrix: attn[i, j] = how much query token i attends to key token j.
def aggregate_attention(attentions):
"""Mean attention matrix across layers and heads.
Returns (seq_len, seq_len); attn[i, j] = query i's attention to key j."""
stacked = torch.stack([layer_attn[0].float() for layer_attn in attentions], dim=0)
return stacked.mean(dim=(0, 1))
Causal-Normalized Attention Received
A raw column sum over this matrix favors early tokens for a boring reason: token j is only visible to queries i >= j, so early positions have simply had more chances to be attended to. Dividing by the number of queries that could actually see each position gives a fair, comparable-across-positions score:
def mean_attention_received(attn_matrix):
seq_len = attn_matrix.shape[0]
col_sums = attn_matrix.sum(dim=0)
valid_queries = torch.arange(seq_len, 0, -1, dtype=torch.float32)
return col_sums / valid_queries
Structural Tokens Aren't Content
Every ChatML-templated prompt carries fixed scaffolding — <|im_start|>, <|im_end|>, role-name tokens, and the newlines the template inserts around them. These tokens are content-independent: their identity is fixed by the template, not by what was actually said. They need to be excluded from the baseline used to judge what "normal" attention looks like, or they distort the comparison for every real content token in the prompt.
def find_structural_positions(input_ids, tokenizer):
"""Positions fixed by the ChatML template rather than by content."""
ids = input_ids[0].tolist()
special_ids = set(tokenizer.all_special_ids)
im_start_id = tokenizer.convert_tokens_to_ids("<|im_start|>")
im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
structural = {0}
for i, tid in enumerate(ids):
if tid in special_ids:
structural.add(i)
if i > 0 and ids[i - 1] in (im_start_id, im_end_id):
structural.add(i) # role name (after im_start) or newline (after im_end)
if ids[i - 1] == im_start_id and i + 1 < len(ids):
structural.add(i + 1) # newline after the role name specifically
return structural
Phase 2: Getting the Statistics Right (Two Times)
The first version of the sink-scoring logic looked reasonable and produced numbers that were, on inspection, badly wrong. Getting this right took two separate corrections.
Correction 1 — z-score, not skew-blind mean/std
Attention weight is strictly non-negative and heavily right-skewed: most tokens sit near zero, a handful sit higher. A plain mean/std z-score assumes something closer to a normal distribution, and against a skewed population that assumption quietly inflates what counts as "anomalous." A first pass at a threshold of 2.0 reliably flagged over a dozen tokens per prompt as sinks — and on inspection they were just ordinary salient content words ("regarding," "potential," "biases"), not sinks in any mechanistic sense.
The fix: log-transform attention values before computing spread, and use median and median absolute deviation (MAD) instead of mean/std — MAD doesn't get dragged upward by the same skew it's measuring.
def detect_sinks(mean_attn, threshold=6.0, exclude_positions=(0,)):
mask = torch.ones_like(mean_attn, dtype=torch.bool)
for p in exclude_positions:
if 0 <= p < len(mask):
mask[p] = False
population = mean_attn[mask]
log_population = torch.log(population.clamp_min(1e-8))
log_all = torch.log(mean_attn.clamp_min(1e-8))
median = log_population.median()
mad = (log_population - median).abs().median()
scaled_mad = mad * 1.4826 # comparable scale to std under normality
z = torch.zeros_like(mean_attn) if scaled_mad.item() < 1e-8 else (log_all - median) / scaled_mad
return z > threshold, z
Correction 2 — structural doesn't mean sink
An earlier version force-set every structural position to "is a sink," on the reasoning that BOS-as-sink is well-established enough to just assume. That reasoning breaks the moment it's generalized to every scaffolding token, most of the non-BOS structural positions turned out to have z-scores near zero or negative once actually measured. The fix was to stop overwriting measured evidence: a position is only a sink if its z-score says so. "Structural" and "sink" are now two independent flags, not one collapsing into the other.
What It Looks Like End to End
Here's an actual run, on a prompt about training-data provenance and self-supervised bias:
╭─ Context health ─╮
│ 6.7/10 │
╰──────────────────╯
Structural positions (template-fixed, excluded from baseline)
┌──────────┬────────────────┬────────────┬─────────┬───────┐
│ Position │ Token │ Mean attn. │ Z-score │ Sink? │
├──────────┼────────────────┼────────────┼─────────┼───────┤
│ 0 │ '<|im_start|>' │ 0.3716 │ 7.23 │ yes │
│ 1 │ 'system' │ 0.0023 │ -1.06 │ no │
│ 2 │ '\n' │ 0.0359 │ 3.42 │ yes │
│ 19 │ '<|im_end|>' │ 0.0027 │ -0.82 │ no │
│ 20 │ '\n' │ 0.0038 │ -0.26 │ no │
│ 21 │ '<|im_start|>' │ 0.0018 │ -1.44 │ no │
│ 22 │ 'user' │ 0.0018 │ -1.50 │ no │
│ 23 │ '\n' │ 0.0062 │ 0.55 │ no │
│ 274 │ '<|im_end|>' │ 0.0648 │ 4.38 │ yes │
│ 275 │ '\n' │ 0.0684 │ 4.47 │ yes │
│ 276 │ '<|im_start|>' │ 0.0726 │ 4.57 │ yes │
│ 277 │ 'assistant' │ 0.0734 │ 4.59 │ yes │
│ 278 │ '\n' │ 0.1230 │ 5.43 │ yes │
└──────────┴────────────────┴────────────┴─────────┴───────┘
Anomalous sink tokens
┌──────────┬──────────────┬────────────┬─────────┐
│ Position │ Token │ Mean attn. │ Z-score │
├──────────┼──────────────┼────────────┼─────────┤
│ 73 │ ' variety' │ 0.0208 │ 2.53 │
│ 246 │ ' ###' │ 0.0187 │ 2.36 │
│ 247 │ ' Self' │ 0.0173 │ 2.23 │
│ 253 │ ' Bias' │ 0.0217 │ 2.59 │
│ 254 │ '.' │ 0.0182 │ 2.31 │
│ 255 │ ' When' │ 0.0213 │ 2.57 │
│ 256 │ ' an' │ 0.0152 │ 2.02 │
│ 259 │ ' trains' │ 0.0217 │ 2.60 │
│ 260 │ ' on' │ 0.0186 │ 2.35 │
│ 264 │ ',' │ 0.0304 │ 3.15 │
│ 265 │ ' several' │ 0.0249 │ 2.83 │
│ 266 │ ' factors' │ 0.0288 │ 3.06 │
│ 267 │ ' come' │ 0.0165 │ 2.15 │
│ 269 │ ' play' │ 0.0209 │ 2.53 │
│ 270 │ ' regarding' │ 0.0247 │ 2.81 │
│ 271 │ ' potential' │ 0.0202 │ 2.48 │
│ 272 │ ' biases' │ 0.0215 │ 2.58 │
│ 273 │ ':' │ 0.0319 │ 3.23 │
└──────────┴──────────────┴────────────┴─────────┘
Reading this correctly, in three parts:
Position 0 is a real, strong sink. z = 7.23 on a prompt where the log-MAD baseline is built entirely from the remaining ~270 positions. This is the textbook case — BOS absorbing residual softmax mass regardless of content, exactly as StreamingLLM describes.
The five-token run right at the generation boundary (<|im_end|> \n <|im_start|> assistant \n, positions 274–278) is a second, milder sink cluster — z climbing from 4.38 to 5.43 as you approach the point where generation actually starts. This wasn't something I expected going in; the working theory is that these tokens function as a second "nothing decided yet" boundary, structurally similar to BOS in that every later query can see them and their identity is fixed by the template rather than by the conversation. Whether this is a general pattern or specific to this template/model is exactly the kind of question the 50-prompt run should answer.
The "anomalous" table is still showing false positives, and that's on me, not the model. Every row here — "regarding," "biases," "trains," even bare punctuation like , and : — is an ordinary content word from a real sentence about training data and bias, not a token with no semantic relationship to its context. The log-MAD correction from Phase 2 fixed how the spread is measured, but this run still uses threshold=2.0, which is a leftover from when the metric was mean/std-based and needs recalibrating against the new scale. With BOS at z=7.23 and the boundary run at z=4.4–5.4, a threshold somewhere in the 5–6 range would separate real structural/positional sink behavior from what's currently just the natural right tail of content-word attention. I'm leaving this table in rather than cropping it out, because the gap between "the statistic is now sound" and "the threshold is correctly tuned" is itself worth showing — they're different bugs, and fixing the first doesn't automatically fix the second.
How This Connects to Parts 1 and 3
Back to entropy (Part 1): The thesis was that degraded attention during prefill should show up as elevated generation entropy during decode. Now that sink detection is trustworthy rather than noisy, this phase produces a single per-prompt "context health" number that Part 4 can actually correlate against Part 1's mean-entropy metric, instead of correlating against a score that was partly measuring content-word variance by accident.
Forward to speculative decoding (Part 3): KV cache eviction under speculative decoding runs into exactly the sink-eviction problem StreamingLLM identified — evict the wrong position and acceptance rates degrade. Knowing precisely which positions in a given prompt are genuine sinks, rather than assuming "it's probably just BOS," is a direct input into cache retention policy once Part 3 starts trading off draft length against verification cost.
Links
| Resource | Link |
|---|---|
| GitHub repo | https://github.com/cyprus09/llm-inference-lab |
| StreamingLLM | https://arxiv.org/abs/2309.17453 |
| Speculative Decoding | https://arxiv.org/abs/2211.17192 |
| Lost in the Middle | https://arxiv.org/abs/2307.03172 |
| The Illustrated Transformer | https://jalammar.github.io/illustrated-transformer/ |
Series:
- Part 1 — Entropy Tracker
- Part 2 — Attention Sink Detector (you are here)
- Part 3 — Speculative Decoding with Entropy-Guided Draft Length (coming)
- Part 4 — The Empirical Study (coming)
Top comments (0)