LAV a: Layer-wise KV Cache Eviction with Dynamic Budget Allocation (2025)

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 four benchmarks (LongBench, Needle-In-A-Haystack, Ruler, and InfiniteBench) demonstrate its superiority over strong baselines. 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.

## introduction large language models (llms) have shown remarkable capability in handling long-text scenarios, enabling advancements in tasks such as question answering @xcite , code generation @xcite , and multi-turn dialogues @xcite . to further enhance external knowledge integration, state-of-the-art models like claude 3.5 @xcite , @xcite , and qwen2.5 max (qwen and * corresponding author et al., 2025) have extended their context lengths beyond 128k tokens. however, supporting such long contexts comes with increased computational challenges. one common approach to accelerating llm inference is caching key and value vectors (kv cache), but its high memory demand necessitates efficient cache compression techniques. while existing compression methods have shown promise, they are largely heuristic, relying on statistical measures such as accumulated attention scores @xcite . these metrics are derived from empirical observations rather than a theoretical foundation. additionally, although dynamic head allocation @xcite and dynamic layer allocation @xcite have been explored, no method, to our knowledge, fully adapts head and layer budgets. to address this gap, we propose a unified framework for cache compression, which is formulated through the lens of minimizing information loss in transformer residual streams (see figure 1 , and sec. 3). many existing methods can be formulated within our framework. specifically, context compression methods @xcite aim to minimize global information loss at the logits layer. kv cache compression methods @xcite @xcite primarily focus on local information loss at the head or layer levels. our framework provides a principled approach to designing new algorithms. this paper introduces a novel method based on layer attention output loss, which measures the impact of compression on the information retained in each layer after multi-head attention. the layer-wise loss function provides a balanced perspective on both local information within layers and global information flow across layers. within each layer, the loss function guides the design of a scoring mechanism to assess token importance across heads, allowing for simultaneous head budget allocation and cache eviction. across layers, it enables dynamic layer budget allocation by comparing information between layers. our method is theoretically grounded, and significantly simpler than cake, the only training-free method with dynamic layer budgets. extensive experiments were conducted using various llm series on the longbench and needle in a haystack benchmarks. the results consistently demonstrate lava's strong ability to preserve the model's long-text comprehension under various memory constraints. additionally, compared to a full cache implementation of flashattention-2, lava significantly reduces memory consumption while simultaneously reducing latency (9× faster decoding for 128k-token sequences). our empirical findings highlight that dynamic layer budgets are essential for generation tasks, while dynamic head budgets are crucial for text extraction tasks. achieving dynamic budget allocation at both the head and layer levels is key to optimizing performance across different tasks. our contributions: 1) we introduce a principled framework for kv cache eviction by analyzing the information flow through transformer residual streams, accounting for information loss at various points during decoding. 2) building on this framework and the notion of information loss at the layer-wise attention output, we propose lava-a unified method that simultaneously performs kv cache eviction and dynamic budget allocation. to the best of our knowledge, lava is the first training-free method to achieve dynamic budget allocation without relying on multiple combined metrics, making it simple for practical purposes. 3) evaluations on longbench, needle in a haystack, ruler and infinitebench demonstrate that our simple method outperforms strong baselines. 4) experiments reveal new insights into the role of dynamic budget allocation across different tasks, offering guidance for the adaptive selection of strategies. ## the information flow of llm decoding process with kv cache kv cache is initialized at prefilling stage, which basically computes the key and value for tokens in the initial prompts in the standard way @xcite . in the following, we assume that there exists a kv cache of (n -1) previous tokens and demonstrate how decoding is performed at step-n . notations the llm has l layers, each has h heads. the model and head dimensions are d and d @xmath0; k l , v l are the kv cache for the l-th layer up to the current time step (the n -th token), which are of [h, (n -1), d h ] sizes. the full notation decoding process according to @xcite , llm decoding can be viewed as operating on the current (n -th) residual stream, as illustrated in figure 1 . specifically, suppose that x n l is the current input for layer l, we first calculate 13673 the corresponding q n l , k n l , v n l as follows: ## a principled framework for kv cache eviction based on information loss given the kv cache, compression can be seen as masking entries in the kv tensors so that the attention heads cannot copy masked information to the later residual streams. formally, one can define the attention mask i l,h for layer-l and head-h: the goal is to find a kv cache eviction policy so that to minimize the information loss for the logits at the last layer (p n ) for all subsequent residual streams (from n to n e ; see figure 1 ). let p denote this logit loss, and b be the memory constraint. the unified problem for budget allocation and cache eviction can be defined as follows: here, b l,h represents the budget for layer-l and head-h, b l denotes the total budget for layer-l. the final constraint ensures that the most recent tokens within a window of size w are retained for all heads, aligning with the common practice in the literature. as computing the loss over future, unseen tokens is impractical. to address this, we approximate the loss by considering only residual streams up to the current step n . considering the current step-n , one can define p as the cross-entropy loss between p n and pn , which is the logit obtained with the attention mask @xcite . additionaly, since the search space for the mask matrix is combinatorial, we instead search for a scoring function s, where s l,h [i] assigns an importance score to token i at layer l and head h. this scoring function allows us to greedily choose the least important entries to be masked @xmath1, b). all in all, we have the following (surrogate) optimization problem: the aforementioned approach to minimizing global logit loss can be impractical for online inference when the scoring function is computationally expensive. a more feasible alternative is to focus on local information and apply localized kv cache eviction. for instance, head attention loss can be used for head-wise eviction, a strategy adopted by most existing methods @xcite @xcite . in this case, the scoring functions are lightweight, relying on simple statistical features, like head-wise attention weights. ## related work recently, various kv cache compression methods have been proposed, leveraging different policies such as recency @xcite , accumulated attention scores @xcite , last-token attention scores @xcite , and recent attention scores @xcite . while most approaches assume a uniform budget, recent efforts have been made for dynamic budget allocation across layers @xcite and heads @xcite . some methods aim at layer-dependent budgets but fix the patterns across all samples @xcite . in general, kv cache eviction and budget allocation are typically treated as separate problems, requiring a combination of independent strategies. in contrast, we develop a principled framework based on information loss in the residual stream and propose a unified method for both cache compression and dynamic budget allocation. closely related to lava is @xcite , which aims at minimizing the layer output perturbation. however, this study only applies the derived metric locally for head budget allocation. in contrast, we propose a metric for layer-wise cache eviction with dynamic layer budgets. ## conclusion this paper provided a comprehensive of current kv cache compression into a unified framework, grounded in the principle of minimizing information loss in transformer residual streams. by analyzing the layer attention output loss, we proposed lava, a novel layer-wise compression method that enables fully dynamic head and layer budget allocation. our experiments demonstrate that dynamic layer budgets are crucial for generation tasks, whereas dynamic head budgets are important for extraction tasks. as a fully dynamic compression method, lava consistently maintains top performance across task types and llm architectures, while achieving the same speedup of 9× with 128k context length compared to full cache. future directions include exploring new compression algorithms based on our framework, as well as extending our framework for model compression. by advancing efficient methods for llms, our work contributes to making llm more accessible and scalable for diverse applications.