Attention sinks: the four tokens that stabilize infinite context on a 16 GB Mac

Jun 3, 2026

A 2023 paper from MIT and Microsoft — StreamingLLM — ran a simple experiment: slide a fixed KV cache window across a long document and evict the oldest tokens as the window fills up. Standard stuff. The result: perplexity on Llama-2 jumped from 5.4 to 5,158 the instant the initial tokens fell out of the window.

Not a gradual degradation. A cliff.

The fix: keep four tokens pinned permanently and rotate everything else. Perplexity drops back to 5.40 and stays stable across 4 million tokens. mlx-lm ships this as RotatingKVCache(keep=4).

This post explains why it works, maps it to the exact code, and shows what the memory numbers actually look like on a 16 GB M3.

Why context length is a memory problem

Every token a model generates requires attending to every previous token. The attention mechanism stores those previous tokens in a KV cache — matrices of keys and values, one entry per token per layer.

For Llama-3.1-8B:

memory per token = num_layers × 2 × n_kv_heads × head_dim × dtype_size
                 = 32 × 2 × 8 × 128 × 2 bytes
                 = 131,072 bytes ≈ 131 KB

At 4,096 tokens that’s 0.54 GB of KV cache. The model weights themselves are ~4.5 GB at 4-bit. So a 16 GB Mac with normal working apps starts feeling pressure around 8,000–10,000 tokens, and hits the wall somewhere around 84,000 tokens.

The chart below shows real measurements taken on M3 with mx.get_peak_memory() — the formula holds to four decimal places:

The naive approach to fitting longer context: cap the KV cache at a fixed size and slide a window over the most recent tokens. When the window moves past the oldest tokens, throw them out.

This works in theory. In practice, it causes catastrophic failure.

The sliding window trap

The StreamingLLM paper (Xiao et al., 2023) quantified what happens when a naive sliding window drops the initial tokens. They measured perplexity on Llama-2 across 65,000 tokens:

That’s not a gradual degradation. It’s a cliff. The moment the initial tokens fall out of the window, the model breaks — effectively producing random output. Then add four tokens back, and it recovers completely.

Why?

What softmax does to attention

The attention mechanism computes a score for each previous token, then converts those scores to weights using softmax. Softmax has one constraint: the weights must sum to 1.

This creates a structural problem. When a model is decoding and there’s nothing relevant to attend to in the recent window — a common situation during generation — the weights still need to sum to 1. The model can’t zero out all of them.

During training, models learn to route this “surplus” attention somewhere safe: tokens that are always present in the context. Initial tokens are the natural choice. By the time a model generates position n, the initial tokens are the only positions that every earlier position could have attended to. Repeated exposure trains them to absorb excess attention weight without distorting the computation.

These are called attention sinks. They hold excess weight, not because they’re semantically important, but because softmax demands it.

When a naive sliding window evicts the initial tokens, the attention sinks disappear. The model tries to route surplus attention elsewhere and fails. Perplexity spikes.

Why four tokens specifically

The StreamingLLM paper tested different sink sizes. The relationship is not smooth:

Four is the threshold where quality fully recovers. Fewer sinks and you’re still paying for the sink structure without getting the full benefit. More and there’s diminishing return.

The specific value is empirical, not derived. Four tokens, across model families (Llama-2, MPT, Falcon, Pythia), is where the curve flattens.

How mlx-lm implements this

mlx-lm’s RotatingKVCache (in mlx_lm/models/cache.py) implements exactly this design:

class RotatingKVCache:
    def __init__(self, max_size, keep=0):
        self.keep = keep     # number of sink tokens to preserve
        self.max_size = max_size
        ...

    def _update_in_place(self, keys, values):
        ...
        # When the buffer is full, rotate — but skip the first `keep` slots
        if self._idx == self.max_size:
            self._idx = self.keep   # new tokens overwrite from position `keep` onward

Positions 0 through keep - 1 are never overwritten. They stay in the cache permanently. New tokens rotate through positions keep through max_size - 1.

When you call make_prompt_cache with a cache size limit:

from mlx_lm.models.cache import make_prompt_cache

cache = make_prompt_cache(model, max_kv_size=4096)
# → RotatingKVCache(max_size=4096, keep=4)

You get a 4,096-token window with the first four tokens pinned. Memory is capped at 0.54 GB for Llama-3.1-8B, regardless of how long the conversation runs.

Or from the CLI:

mlx_lm.generate \
  --model mlx-community/Meta-Llama-3.1-8B-Instruct-4bit \
  --max-kv-size 4096 \
  --prompt "..."

Measured memory: what the cache actually costs

I measured actual KV cache memory at increasing context lengths using mx.get_peak_memory() on a 16 GB M3 MacBook Air with mlx-community/Meta-Llama-3.1-8B-Instruct-4bit:

The formula holds to four decimal places. At 4,096 tokens with a capped window, the cache stays at 0.54 GB forever — adding more tokens just rotates out older ones.

The practical memory ceiling without any cap: peak RAM = ~4.9 GB (model) + (131 KB × context length). On 16 GB with normal apps using ~2–3 GB, you run out around 60,000–70,000 tokens.

With --max-kv-size 4096, you can run an arbitrarily long conversation at a fixed 6.9 GB peak. Quality stays near the full-context baseline, because the four sink tokens are always there.

One thing to know

If a model defines its own make_cache() method — which includes Llama, Gemma, Qwen, Mistral, and most other major families — max_kv_size is not applied to those caches. This is current, documented behavior: the model’s own cache structure takes precedence, so --max-kv-size has no effect for them.

That matters more than it first looks for interleaved-attention models. On Gemma 3, for example, most layers already slide a window, but every few layers use global attention over the full context. Those global layers are exactly the ones a blanket cap would convert to a sliding window — trading memory for long-range recall. So “just cap everything” isn’t free; it depends on the architecture.

There’s an open proposal to extend max_kv_size to these models (mlx-lm #1343), with that memory-vs-recall tradeoff as the open design question.

The actual takeaway

The KV cache grows linearly with context. On a 16 GB Mac, that matters. The attention sink mechanism — keeping four specific tokens pinned while rotating the rest — is what makes bounded-memory long-context generation viable without meaningful quality loss.

The implementation in mlx-lm is RotatingKVCache(keep=4). The flag is --max-kv-size. The memory cost is fixed once you set the window.

I work on making LLMs run well on-device. The benchmark tool used here is ondevice-bench, open source.