MagicPIG: don't pick the tallest

112 GiB for every new token

An LLM generates one token at a time: a word, a piece of a word, or punctuation. In each layer, attention heads produce a query, a vector describing what to look for. They compare it with the cached keys of earlier tokens, turn those scores into weights with a softmax, and average the corresponding values. Keys decide how much to listen; values carry what is said. That is the soft nearest-neighbour lookup of 7.3 The same math in other places, with the KV cache as the database.

Do the napkin arithmetic for Llama-3.1-8B holding a 96k-token document (here 96 × 1,024 = 98,304 tokens): 32 layers × 8 key–value heads × 128 numbers × 2 (key and value) × 2 bytes is 128 KiB per token. That is 12 GiB, read in full, for every new token it writes. At a hypothetical sustained 1 TB/s, those bytes alone take about 13 ms to read. This is a bandwidth estimate for a single sequence, not a measured decoding time.

The obvious idea, if you have read Part 2: when most of the softmax weight lands on a handful of keys, find those with an index and skip the rest. This is TopK attention, and it is what most fast-attention methods approximate.

MagicPIG’s first move is to ask whether TopK is even the right target.

Attention does not want the nearest keys. It wants an average.

A search engine returns a list, and you judge it by recall. Attention returns a weighted mean of the values, and a mean has its own rules.

2Picking the tallest

Here is a picture that makes the difference visible. Pretend each value is a single number, and put each token on a number line at its value. Draw it as a stick whose height is its attention weight. The attention output is then the balance point of all the sticks. In the formula, $i$ runs over cached tokens, $d$ is the number of coordinates in a head, and $q\cdot k_i$ is its query–key dot product. Softmax exponentiates these scaled scores and divides by their sum, so the weights $w_i$ are positive and add to 1:

TopK keeps the tallest sticks and finds their balance point. Sampling draws tokens independently with replacement, each with probability equal to its weight, and averages their values. A token drawn twice counts twice, even though its value only needs to be loaded once. The height analogy motivates sampling; the stick picture shows the actual weighted problem.

That last point is the whole argument, and one line of algebra makes it exact.

Sampling has the opposite character. It is unbiased, and its typical miss shrinks like $1/\sqrt{B}$ with the number of draws $B$. It does not care how the weight is spread, only how much the values disagree:

Try the sink scenario before moving on. One token takes half the weight and contributes almost nothing; everyone else is flat. TopK grabs the sink and a few neighbours, and lands near 0.2 when the truth is 0.35. Remember this one: real models do it, and we will meet it again.

3The catch: you cannot sample without looking

There is an obvious problem. To sample keys in proportion to their weights you need the weights, and the weights need every score $\kq{q}\cdot\kbook{k_i}$. Computing every score still requires reading all the keys. It saves the value reads, but leaves much of attention’s work intact. The paper calls this the oracle, and it saves almost nothing.

So we need to sample from something cheap that is only roughly right, and then repair the “roughly”. Statisticians have done this for seventy years. It is importance sampling: draw from a convenient guess $u$, and weigh each draw by how much the guess under- or over-sampled it.

There is a price for cancelling that unknown denominator. This is self-normalized importance sampling: a ratio of two random sums. Unlike the direct sampling average above, it is generally biased at finite sample sizes. With independent draws and a positive proposal probability wherever $w_i>0$, it converges to the right answer as the number of draws grows.

Look at what the correction needs: for every key we drew, we must know its sampling probability $u_i$. It must also be positive for every key that contributes. Ordinary deterministic HNSW or IVF retrieval leaves some keys with zero chance, so it cannot supply this correction for the omitted keys. A randomized index could in principle work; SimHash makes the probabilities particularly easy to compute.

So the question becomes: is there a cheap way to pick keys that knows its own odds? Before answering, look at what the keys are like.

4Where the keys live

Part 1 taught us to look at the cloud before choosing a tool (1.6 The shape of the cloud). The paper does exactly that for keys and queries, and finds something that shapes everything after. We repeated the measurement on a different model, a different text and a different tokenizer.

That first-token key is the attention sink. The softmax must hand out a total weight of 1, even when nothing in the context is relevant. One interpretation is that training gives the model a place to park weight on a token whose short value changes the output relatively little. It acts like “none of the above”, and it is the sink scenario from the toy above.

Now think of hashing these keys with random hyperplanes. It goes badly.

5A hash that knows its odds

5.1 Hash functions that like their neighbours built SimHash from one fact: a random hyperplane separates two vectors with probability $\theta/\pi$, where $\theta$ is the angle between them. Take $K$ hyperplanes, write down which side of each a vector falls on, and you have a $K$-bit bucket. Vectors pointing the same way tend to share buckets.

Here is the picture in two dimensions, with keys shaped like the real ones.

With raw keys, nothing works. The keys sit in one narrow cone and the query points elsewhere, so the query’s bucket is nearly empty. The paper’s ablation found that fewer than 0.1% of keys ever collided with the query, and retrieval accuracy fell to almost zero.

Now press centred keys. The cone slides onto the origin, the lines slice through it, and the keys that land in the query’s wedge point its way. This can favour high scores, although key length matters too. And attention does not notice, because every score shifts by the same amount:

and the softmax ignores anything added equally to every score. This is the same one-line fix that every compression method in Part 3 applies first. Here it is free, and it decides whether the method works at all.

The toy drew repeated independent samples from a distribution whose probabilities sum to 1. MagicPIG instead selects a set $S$: each key is included once or left out. Now $u_i$ means the probability that key $i$ is included, not the probability of a single draw; the $u_i$ need not sum to 1. The same ratio uses one $\tilde w_i/u_i$ contribution per selected key. Keys’ inclusion events may be correlated because they share hyperplanes.

Here is the payoff. Because SimHash collides with probability $1 - \theta/\pi$ per bit, the chance that a key is picked is a formula of its angle to the query, and nothing else. MagicPIG builds $L$ tables of $K$ bits and samples a key if it shares the query’s bucket in at least two of them:

Assume the bits and tables use independent random hyperplanes. Then $p_i^K$ is a collision in one table, and the formula subtracts the probabilities of zero hits and exactly one hit from 1. The first two terms are the S-curve from 5.2 Turning a gentle slope into a cliff: none of the $L$ tables collide, versus at least one. The third removes keys that collide in exactly one, which makes the curve steeper. Here it is over one real query’s 16,384 keys:

The core estimator, for one query, is five lines. $S$ excludes the exact set (sink and recent tokens), so no token is counted twice. s_exact contains their ordinary attention scores. This is a mathematical sketch; production code also needs stable probability arithmetic and handling of zero-length vectors:

s = (K_S @ q) / np.sqrt(d)                     # scores of the sampled keys only
p = 1 - np.arccos(cos(q, K_S - mu)) / np.pi    # SimHash: chance one bit agrees
u = 1 - (1 - p**K)**L - L * p**K * (1 - p**K)**(L - 1)
w = softmax(np.concatenate([s - np.log(u), s_exact]))   # sink + recent tokens: exact
o_hat = w @ np.concatenate([V_S, V_exact])

6The bill

LSH has a well-known price, and 5.2 Turning a gentle slope into a cliff warned about it: every table stores a pointer to every key.

At the paper’s default (10, 150), the tables are bigger than the KV cache they index: 600 bytes of 32-bit token indices beside a 512-byte key–value pair. Across 98,304 tokens, 32 layers and 8 KV heads, that is 14.1 GiB of table entries plus 12 GiB of KV data, before bucket metadata and other overhead. Table 8 reports roughly 14 GB of tables. The figure also shows packed hash bits as a separate, optional storage cost; they are not another mandatory copy of the tables.

Finding candidates is cheap in bytes, under the simplifying assumption of evenly filled buckets. A query scans about $L n/2^K = 14{,}400$ entries, or 57,600 bytes of indices. That is 0.23% of the 25,165,824 key bytes for one KV head in one layer. It excludes reading the selected keys and values, computing hashes, and counting duplicate hits. Small byte traffic does not mean zero elapsed time.

The system’s aim is to save GPU memory. Most cached keys and values, the tables, and sampled attention computation move to the CPU; the GPU runs the model, computes hashes, and handles the exact attention subset. This lets larger contexts or batches fit in VRAM, at the cost of more total memory and a dependency on CPU bandwidth.

7Did it work? We checked

We re-implemented the sampling estimator in Algorithm 1 from the paper’s description and ran it on every head of every layer of Qwen2.5-1.5B, at the last 64 positions of the same 16k-token text: 21,504 attention rows. We compared it with three oracles that may see every score: TopK, sampling from the true weights, and a hybrid that reads the top half of its budget exactly and samples the rest.

Error here means $\lVert \hat o-o\rVert_2/\lVert o\rVert_2$, averaged over attention rows. Budgets count distinct keys read, including the exact subset; the oracle sampling curves count repeated draws only once. This checks attention outputs on one fixed input, not generation quality or system speed. We used one hash seed and did not reproduce the CPU–GPU serving system.

8Raised eyebrows

The core idea is right, and it is a good one. These are the questions we would put to the authors.

9What it suggests