LAVa: Layer-wise KV Cache Eviction with Dynamic Budget Allocation

2025-09-11

AI Summary

To address the excessive KV cache memory overhead in long-context LLM inference and the limitations of existing compression methods—namely, their reliance on hand-crafted heuristics and lack of dynamic budget allocation—this paper proposes the first training-free hierarchical KV cache compression framework. Our method models information flow via Transformer residual streams to theoretically derive an attention output loss metric, enabling cross-layer and cross-head information fidelity analysis. We introduce, for the first time, a dual-level (layer-wise and head-wise) dynamic budget allocation scheme coupled with a unified dynamic eviction mechanism. Experiments reveal a task-dependent preference: generative tasks benefit more from dynamic layer-level budgeting, whereas extraction tasks rely critically on dynamic head-level budgeting. Our framework achieves significant improvements over state-of-the-art cache compression methods on LongBench, Needle-In-A-Haystack, RULER, and InfiniteBench.

Technology Category

• Data Mining & Knowledge Management: Data Compression
• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

KV Cache is commonly used to accelerate LLM inference with long contexts, yet its high memory demand drives the need for cache compression. Existing compression methods, however, are largely heuristic and lack dynamic budget allocation. To address this limitation, we introduce a unified framework for cache compression by minimizing information loss in Transformer residual streams. Building on it, we analyze the layer attention output loss and derive a new metric to compare cache entries across heads, enabling layer-wise compression with dynamic head budgets. Additionally, by contrasting cross-layer information, we also achieve dynamic layer budgets. LAVa is the first unified strategy for cache eviction and dynamic budget allocation that, unlike prior methods, does not rely on training or the combination of multiple strategies. Experiments with benchmarks (LongBench, Needle-In-A-Haystack, Ruler, and InfiniteBench) demonstrate its superiority. Moreover, our experiments reveal a new insight: dynamic layer budgets are crucial for generation tasks (e.g., code completion), while dynamic head budgets play a key role in extraction tasks (e.g., extractive QA). As a fully dynamic compression method, LAVa consistently maintains top performance across task types. Our code is available at https://github.com/MGDDestiny/Lava.

Problem

• Reducing KV cache memory usage in long-context LLM inference
• Enabling dynamic budget allocation across layers and heads
• Minimizing information loss during cache compression without training

Innovation

• Minimizes information loss in Transformer residual streams
• Uses layer-wise compression with dynamic head budgets
• Achieves dynamic layer budgets by contrasting cross-layer information