OBCache: Optimal Brain KV Cache Pruning for Efficient Long-Context LLM Inference

Large language models (LLMs) with extended context windows enable powerful applications but impose significant memory overhead, as caching all key-value (KV) states scales linearly with sequence length and batch size. Existing cache eviction methods address this by exploiting attention sparsity, yet they typically rank tokens heuristically using accumulated attention weights without considering their true impact on attention outputs. We propose Optimal Brain Cache (OBCache), a principled framework that formulates cache eviction as a layer-wise structured pruning problem. Building upon the Optimal Brain Damage (OBD) theory, OBCache quantifies token saliency by measuring the perturbation in attention outputs induced by pruning tokens, with closed-form scores derived for isolated keys, isolated values, and joint key-value pairs. Our scores account not only for attention weights but also for information from value states and attention outputs, thereby enhancing existing eviction strategies with output-aware signals. Experiments on LLaMA and Qwen models demonstrate that replacing the heuristic scores in existing works, which estimate token saliency across different query positions, with OBCache's output-aware scores consistently improves long-context accuracy. Code is available at https://github.com/DreamSoul-AI/OBCache.

Large language models increasingly need to read very long inputs, such as books, codebases, or long conversations. To do this efficiently, they store a temporary memory called a key-value cache, but this memory grows with input length and can quickly exceed GPU limits. A common solution is to delete parts of the cache, but existing methods often rely on simple attention scores that may not reflect how much a token affects the model's output. We introduce OBCache, a method that asks: if this cached token were removed, how much would the model's attention output change? By estimating this change efficiently, OBCache assigns each cached token an importance score using both attention weights and stored cache information. These scores can be plugged into existing attention-based cache-eviction methods. Experiments with LLaMA and Qwen models show that OBCache consistently improves long-context accuracy while keeping memory use low, making long-context language-model applications more practical and affordable.