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

Researchers propose OBCache, a novel KV cache pruning framework that optimizes memory efficiency for long-context LLM inference by measuring token importance based on actual impact to attention outputs rather than heuristic attention weights. The method, grounded in Optimal Brain Damage theory, demonstrates consistent accuracy improvements over existing eviction strategies on LLaMA and Qwen models.

OBCache addresses a critical bottleneck in modern LLM deployment: the quadratic memory scaling of key-value caches with context length and batch size. As applications demand longer context windows, memory constraints become the primary limiting factor for inference efficiency. Existing solutions rely on attention sparsity patterns, but they use accumulated attention weights as proxies for token importance—a heuristic that doesn't necessarily reflect true output impact.

The research builds on established pruning theory by reformulating cache eviction as a structured pruning problem with closed-form mathematical scores. This principled approach quantifies token saliency by measuring perturbation in attention outputs when tokens are removed, accounting for both attention weights and value state information. The methodology provides separate scores for isolated keys, isolated values, and joint key-value pairs, offering flexibility across different pruning scenarios.

For the AI infrastructure industry, this work has practical implications for deployment costs and scalability. Reducing memory overhead directly translates to higher throughput, lower latency, and reduced computational requirements for serving long-context applications. This efficiency gain becomes increasingly valuable as enterprises deploy models for document analysis, code generation, and retrieval-augmented generation tasks requiring extended context windows.

The experimental validation on LLaMA and Qwen models demonstrates that output-aware signals consistently outperform heuristic baselines. Future developments might explore adaptive pruning strategies that balance memory savings with accuracy across different token types, or integration with other optimization techniques like quantization. The open-source availability suggests rapid adoption potential within the AI infrastructure community.

• →OBCache formulates KV cache pruning as structured pruning using Optimal Brain Damage theory with closed-form mathematical scores.
• →Output-aware token saliency measurement outperforms heuristic attention-weight-based approaches in long-context accuracy.
• →Method accounts for attention weights, value states, and attention outputs to enhance existing eviction strategies.
• →Tested on LLaMA and Qwen models with consistent improvements in long-context inference efficiency.
• →Open-source code availability enables rapid adoption in LLM inference optimization pipelines.