LaProx: Redefining KV Cache Eviction Strategy in Long-Context LLM Inference

LaProx proposes a new output-aware KV cache eviction framework. By explicitly modeling the multiplicative interaction between attention maps and projected value states, it achieves a globally unified token importance assessment. This strategy maintains model performance even when only 5% of the cache is retained, providing an efficient solution to the memory bottleneck problem in long-context LLM inference.

With the widespread application of LLMs in scenarios such as document analysis and code understanding, long-context inference has become an essential need. However, KV cache memory usage increases linearly, which can easily cause GPU memory overflow. Traditional strategies adopt head-level weighted averaging, relying on local attention weights while ignoring value vector representations, the impact of output projection matrices, and cross-head dependencies—leading to a sharp performance decline when compression rates are high.

LaProx redefines the KV cache eviction problem as an output-aware hierarchical matrix multiplication approximation problem. Its core lies in considering the complete computation chain of the attention mechanism (interaction between Query, Key, Value, and projection results) rather than isolated attention weights. By explicitly modeling the multiplicative interaction between attention maps and projected value states, it accurately quantifies the actual contribution of each token to the final output.

LaProx proposes the first globally unified token eviction strategy, breaking the limitation of traditional intra-head local decision-making. It assigns comparable importance scores to all tokens, enabling unified eviction decisions at the model level. In extreme compression scenarios, it can identify core token sets and avoid redundant retention.

In LongBench and Needle-In-A-Haystack benchmark tests (19 datasets), LaProx maintains the original model performance even when only 5% of the cache is retained, stably outperforming existing baselines. In extreme compression scenarios (2-3% cache), the accuracy loss is reduced by up to 2x compared to state-of-the-art methods, with minimal computational overhead that barely affects inference latency.

LaProx marks the shift of KV cache management from heuristic compression to principle-driven optimization, laying the foundation for subsequent theoretical analysis and potentially inspiring research on attention mechanism structures. For engineering practitioners, it provides a plug-and-play solution that requires no model architecture modifications or retraining, and will become a key component of long-context inference infrastructure.