KV-Surgeon improves offline key-value cache compaction accuracy and speed
The KV-Surgeon: Compensation-Aware Offline KV-Cache Compaction
Machine LearningArtificial Intelligence
Summary
When computers use key-value caches to remember information, they sometimes need to shrink these caches to save space or speed. The difficulty is picking which parts to keep or remove because removing one part affects how the rest works. The paper's authors created KV-Surgeon, a method that looks at how the parts interact and chooses the best ones to keep, making the cache smaller without losing as much accuracy. Their tests show KV-Surgeon works better and faster than previous methods on several tasks.
What this means in practice
- •For language model engineers: Reduce memory and compute costs by optimizing which cache entries to keep for long context processing in large language models.
- •For software developers building ai assistants: Improve the efficiency and accuracy of AI assistants that rely on repeated queries over fixed contexts by using compacted offline KV-caches.
Authors
Moritz Wagner, Christophe Roux, Max Zimmer, Sebastian Pokutta
Abstract
Offline compaction of the key--value cache (KV-cache) reduces the cost of repeatedly querying a fixed context by constructing a smaller cache once and reusing it subsequently. The central challenge of compaction is deciding which cache entries to delete and which to retain. Existing methods score each candidate for deletion on its own, with the rest of the cache being held fixed. However, softmax renormalization of the attention operation couples these decisions in a nonlinear manner: removing one entry changes the attention weights of every remaining one. In this work, we introduce KV-Surgeon, which uses calibration data to take interactions into account in a tractable way. We derive a quadratic surrogate that admits closed-form expressions for jointly making the optimal greedy removal decision as well as compensation within the remaining entries. On QuALITY, LongHealth, and Long Code Arena, and at every evaluated compaction ratio, KV-Surgeon substantially advances the cost--accuracy Pareto frontier of dense and mixture-of-experts models from 4B to 30B parameters. At $50\times$ compaction, KV-Surgeon improves question-answering accuracy by up to four points over attention matching (Zweiger et al., 2026), the state of the art, at comparable cost. Lastly, we introduce an even faster variant, which also exceeds the state of the art while requiring only two fifths of its cost.