policybook / kv-cache
kv-cache
Policies for deciding which tokens to forget while a language model generates. A transformer's KV cache grows by one entry per token, per layer, per head. At a long context that is gigabytes, and the cost is linear in the sequence while the value of any individual token is not.
Run these side by side → Same trace, same seed, stepped in lockstep.
Choosing one
| If you need | Use | Because |
|---|---|---|
A simple default that needs no attention weights | A ring buffer plus four pinned positions. 0.7862 retained mass at budget 512, from a policy with no arithmetic in it. | |
The best retained attention, and you can read weights | 0.8240 at budget 512, the highest here. Sizes its own recency protection from the data, so it has nothing to misconfigure. | |
To find old tokens that still matter | Best heavy-hitter recall at the wider budgets, 0.9137 at 512, because a cumulative score remembers what a current-step score cannot. | |
Importance to decay rather than accumulate | Counts how often a token mattered, so a stale early spike stops defending its slot. A third less memory than H2O. | |
To keep phrases intact, not just their peaks | The only policy here whose scoring is neighbour-aware. Costs the most memory in the domain. | |
To divide a fixed budget across layers | The only policy that answers “how many”, not “which”. Inert on this single-layer trace. See its README. | |
The baseline, to measure against | Keep the most recent and forget the rest. Included to be measured, not recommended. |
Benchmarks
Retained mass on each trace, measured by this repository's benchmarks.
| Policy | decode-4096@256 | decode-4096@512 | decode-4096@1024 |
|---|---|---|---|
| TOVA | 0.7781 | 0.8240 | 0.8859 |
| PyramidKV | 0.7834 | 0.8212 | 0.8880 |
| SnapKV | 0.7834 | 0.8212 | 0.8880 |
| StreamingLLM | 0.7367 | 0.7862 | 0.8662 |
| H2O | 0.6439 | 0.7149 | 0.7844 |
| Scissorhands | 0.6413 | 0.7114 | 0.7847 |
| Sliding Window | 0.5966 | 0.6555 | 0.7541 |
Every policy
TOVA
Drop whichever token the model just stopped looking at. Only the latest step's attention counts, so nothing a position did earlier defends it: the most forgetful policy here, and the most responsive.
PyramidKV
Spend more cache on early layers than late ones. Attention widens near the input and narrows with depth, so a uniform per-layer budget overfeeds the layers that need it least.
SnapKV
Score on the last few steps only, then max-pool across neighbouring positions so an important token drags its neighbours in with it. A forgetting window plus a defence against keeping half a phrase.
StreamingLLM recommended
A sliding window that also pins the first few tokens. Those "attention sinks" absorb attention no matter what they contain, and keeping four of them is most of the quality a cheap policy can recover.
H2O recommended
Keep the tokens that have received the most attention so far, plus a recent window. The first policy here that reads attention, and so the first that finds important old tokens, not merely early ones.
Scissorhands
Count how many steps a token mattered for, not how much. Persistence of importance: a position that clears its share of the attention again and again beats one that spiked once and went quiet.
Sliding Window
Keep the most recent tokens and forget the rest. The honest baseline: most attention is recent, so this is far better than it sounds and far worse than it needs to be.