1 Introduction↩︎

Large language models (LLMs) [1]–[4] have recently revolutionized a wide range of natural language processing tasks, including document summarization [5], question answering [6], code generation [7], and dialogue systems [8], [9]. Despite their impressive capabilities, many of these applications require processing long sequences, which poses substantial challenges for efficient deployment. A major bottleneck in LLM inference stems from their autoregressive nature, which necessitates caching all key-value (KV) states across the context window. Because the cache size scales linearly with both sequence length and batch size, it leads to substantial memory and latency overheads. For instance, running a LLaMA-3.1-8B model [10] with a 1M-token context window requires storing over 120GB of KV cache, which exceeds the memory capacity of most GPUs.

A promising line of research addresses this challenge via KV cache eviction, a training-free technique that reduces inference costs by selectively discarding unimportant KV tokens [11], [12]. These methods are motivated by the observation that only a small subset of tokens significantly influences model predictions. By evicting redundant tokens during inference, such approaches can substantially reduce memory and computational costs while incurring moderate performance degradation. Early methods such as H\(_2\)O [11] evict tokens based on accumulated attention weights, a simple yet effective heuristic for estimating token saliency. More recent techniques, including TOVA [13] and SnapKV [14], refine this strategy by introducing more sophisticated attention-based scoring mechanisms to better preserve accuracy. However, these methods rely primarily on attention weights and often overlook the contribution of value states in shaping the final model outputs. Heuristically accumulating attention scores fails to fully capture the true impact of token removal on attention outputs, which directly affects the hidden states and downstream predictions. Consequently, these methods may mistakenly retain tokens with negligible influence or discard ones whose importance is not evident from attention weights alone.

To tackle these limitations, we propose Optimal Brain Cache (OBCache), a principled framework that formulates KV cache eviction as a layer-wise structured pruning problem. The key insight behind OBCache is that the actual impact of removing KV pairs, namely, their influence on future model outputs, can be effectively approximated by analyzing local perturbations in historical attention outputs when the corresponding KV vectors are pruned. This approximation enables us to estimate the contribution of each token to the model output without requiring access to future states at inference time. These perturbation estimates therefore serve as token-wise saliency scores to inform output-aware cache eviction and retention decisions.

In parallel, some recent works have also explored token saliency through output perturbation analysis and observed that value-state information, such as value norms, can provide useful signals for improved cache eviction. However, these approaches either lack a formal theoretical framework [15] or base their analysis on objective bounding and relaxation techniques that require additional assumptions on attention distributions [16]. In contrast, our analysis is grounded in the Optimal Brain Damage (OBD) theory [17], originally proposed for pruning static model weights. By treating cached KV states as pruning variables, we define three types of pruning units: isolated value vectors, isolated key vectors, and joint key-value pairs at the same token position. In each case, we derive closed-form expressions for pruning-induced output perturbations via second-order Taylor approximation.

Under this OBD-based formulation, we show that prior value-aware eviction scores [15], [16] correspond to the isolated-value case within our framework. Moreover, our framework offers flexibility in the choice of pruning objectives and pruning units. This flexibility allows it to extend beyond value-only criteria by capturing the contributions of key states and key-value interactions, and to naturally support both static eviction in the prefill stage and dynamic eviction in the decoding stage.

Furthermore, we show that existing attention-based scoring methods, which accumulate attention weights across different query positions, emerge as special cases under our framework. Specifically, the pruning objective is simplified to preserving the attention matrix within different perturbation windows, and the pruning units are reduced to individual attention columns. As such, OBCache generalizes and complements existing attention-based eviction methods, and can be seamlessly integrated into any score-based cache eviction pipeline to improve token selection.

In summary, our key contributions are as follows:

• We introduce OBCache, a principled scoring framework for KV cache eviction that estimates token saliency via eviction-induced perturbations in attention outputs, which can be seamlessly integrated into existing eviction pipelines to improve token selection.

• We provide the first theoretical formulation of KV cache eviction as a layer-wise structured pruning problem based on the Optimal Brain Damage theory, deriving closed-form saliency scores for isolated values, isolated keys, and joint key-value pairs, and showing that prior attention-based heuristic scores arise as special cases of our more general formulation.

• We conduct extensive experiments showing that replacing heuristic attention scores in existing eviction methods with OBCache’s output-aware scores consistently improves long-context performance on LLaMA-3.1 and Qwen-2.5 models across long retrieval, language modeling, and LongBench benchmarks.

2 Related Work↩︎

2.1 KV Cache Compression↩︎

Motivated by the observation that only a sparse subset of key-value (KV) tokens contributes to model predictions, cache eviction methods directly reduce the memory footprint by eliminating redundant KV states. For example, StreamingLLM [12] discovers the attention sink phenomenon and retains only the initial and most recent tokens to enable infinite-context decoding. H\(_2\)O [11] proposes accumulating attention weights across all query positions to dynamically identify salient tokens. TOVA [13] simplifies this approach by considering only the attention distribution of the most recent query. A more refined strategy, SnapKV [14], aggregates attention scores within a small observation window and applies a pooling-based clustering mechanism, which is effective in retrieval-centric tasks. To refine the attention-only saliency in earlier methods, VATP [15] and CriticalKV [16] show that scaling accumulated attention weights by the norm of the corresponding value states leads to improved long-context performance.

Orthogonal to eviction-based methods, sparse attention approaches [18], [19] reduce inference cost by dynamically selecting a subset of KV pairs during attention computation while retaining the full KV cache. Cache merging methods [20], [21] seek to mitigate the information loss caused by irreversible eviction by merging evicted KV states into the retained cache. Other works exploit distinctive sparsity patterns across different layers [22] or attention heads [23] to design adaptive KV cache allocation strategies. For example, PyramidKV [24] observes higher redundancy in deeper layers and statically assigns larger cache budgets to them, while AdaKV [25] dynamically allocates head-wise budgets based on global attention statistics across all heads.

2.2 Model Pruning↩︎

Another line of research for reducing LLM inference costs focuses on pruning model parameters. Classical pruning frameworks [17], [26] quantify the saliency of each pruning unit by estimating the perturbation it induces in a task-specific loss, often approximated to second order by a Taylor series. However, computing second-order statistics globally is still inefficient for LLM-scale models [27]. To reduce complexity, recent approaches adopt local formulations that operate at the level of individual layers [28], [29], where the objective becomes minimizing the change in layer outputs. By varying the pruning units, such methods can structurally remove redundant components such as attention layers, heads, or hidden channels. In this work, we extend the layer-wise pruning paradigm from static weight pruning to the eviction of the dynamic KV cache. By treating KV states as pruning units, we introduce a theoretically grounded framework to more accurately quantify token saliency, thereby connecting KV cache compression with the model-pruning literature.

3 Notation and Preliminaries↩︎

We review the KV caching mechanism in transformer-based LLMs, focusing on the prefill, decoding, and cache eviction phases. Bold capital letters denote matrices, and bold lowercase letters with subscripts denote row vectors. For clarity, we omit the batch, head, and layer indices.

Prefill. Let \({\boldsymbol{X}}\in \mathbb{R}^{l \times d}\) be the prompt embeddings, and let \({\mathbf{Q}}, {\mathbf{K}}, {\mathbf{V}}\in \mathbb{R}^{l \times d}\) be the projected query, key, and value matrices, where \(l\) is the prompt length and \(d\) is the hidden size. The attention output \({\mathbf{O}}\in \mathbb{R}^{l \times d}\) is computed as follows: \[{\mathbf{O}}= {\mathbf{A}}{\mathbf{V}}, \quad {\mathbf{A}}= \sigma({{\mathbf{Z}}}), \quad {\mathbf{Z}}= \tfrac{{\mathbf{Q}}{\mathbf{K}}^\top}{\sqrt{d}} \in \mathbb{R}^{l \times l},\] where \(\sigma(\cdot)\) is the row-wise softmax function, \({\mathbf{Z}}\) denotes the pre-softmax attention logits, and \({\mathbf{A}}\in \mathbb{R}^{l \times l}\) is the attention weight matrix. In this phase, \({\mathbf{K}}\) and \({\mathbf{V}}\) are cached to avoid recomputation in subsequent decoding steps.

Decoding. At decoding step \(t\), let \({\boldsymbol{x}}_t \in \mathbb{R}^d\) be the newly generated token embedding, and let \({\mathbf{q}}_t, {\mathbf{k}}_t, {\mathbf{v}}_t \in \mathbb{R}^d\) be its projected query, key, and value vectors. After appending \({\mathbf{k}}_t, {\mathbf{v}}_t\) to the cache, the key and value matrices extend to shape \(s \times d\), where \(s = l + t\). The step-\(t\) attention output is \[{\mathbf{o}}_{s} = {\mathbf{a}}_{s} {\boldsymbol{V}}, \quad {\mathbf{a}}_{s} = \sigma({{\mathbf{z}}_{s}}), \quad {\mathbf{z}}_{s} = \tfrac{{\mathbf{q}}_t {\mathbf{K}}^\top}{\sqrt{d}} \in \mathbb{R}^{s}.\] For the remainder of this paper, we use \({\mathbf{Q}}, {\mathbf{K}}, {\mathbf{V}}, {\mathbf{A}}, {\mathbf{Z}}, {\mathbf{O}}\) to denote the full matrices of sequence length \(s\). Under this notation, the prefill phase corresponds to the special case \(t = 0\), while the decoding phase corresponds to \(t > 0\).

Cache Eviction. Eviction is triggered when the sequence length \(s\) of the KV cache exceeds a predefined budget \(N\). After each forward pass, an eviction algorithm selects \(s - N\) rows from \({\mathbf{K}}\) and \({\mathbf{V}}\) for permanent deletion. By enforcing a fixed cache budget, the model can decode arbitrarily long sequences while maintaining a bounded memory footprint.

4 Optimal Brain Cache (OBCache)↩︎

In this section, we present the details of OBCache. In Section 4.1, we first formulate cache eviction as a layer-wise structured pruning problem, with the objective of minimizing perturbations in attention outputs. Section 4.2 details our analytical solution to this perturbation minimization problem based on second-order Taylor approximations. In Section 4.3, we show that our framework recovers existing attention-based scoring methods as special cases. We then provide a qualitative example to highlight the effectiveness of our formulation in Section 4.4. An overview of the OBCache scoring mechanism is shown in Figure 1.

4.1 Cache Eviction via Perturbation Minimization↩︎

Cache eviction can be viewed as a form of layer-wise structured pruning, where the saliency of each pruning unit (i.e., a key or a value vector) is quantified by the error it induces in the layer output. However, cache eviction introduces a unique challenge not encountered in classical model pruning: due to the autoregressive nature of generation, the actual error caused by removing a key–value pair at step \(s\) affects only future attention outputs, \({\mathbf{o}}_{s+1}, {\mathbf{o}}_{s+2}, \dots\), which are inaccessible at eviction time. We refer to this unobservable quantity as the true eviction error.

Although the true eviction error is not directly available, we observe that it can be effectively approximated by measuring perturbations in recent historical attention outputs, specifically \({\mathbf{o}}_s, {\mathbf{o}}_{s-1}, \dots\), when the corresponding KV vectors are pruned. We refer to this measurable surrogate as the pruning-induced eviction error, and use it as a proxy objective to estimate token saliency. Building on this insight, we formulate cache eviction as a layer-wise structured pruning problem by treating \({\mathbf{V}}\) and \({\mathbf{K}}\) as pruning variables.

To illustrate the effectiveness of the pruning-induced error as a proxy, we present an empirical example in Section 4.3.

4.2 Cache Pruning Scores↩︎

Directly recomputing Equation 1 for every position \(p\) requires repeated attention operations and is computationally infeasible. Inspired by Optimal Brain Damage (OBD) [17], we approximate it via a second-order Taylor expansion around the unperturbed point \(({\mathbf{V}}, {\mathbf{K}})\).

In Equation 2 , the first-order terms vanish at the expansion point since \(\widehat{\mathbf{O}}- {\mathbf{O}}= \boldsymbol{0}\). The simplification to Equation 3 uses the fact that the off-diagonal blocks of \(\mathbf{H}^{\textit{vv}}, \mathbf{H}^{\textit{kk}}, \mathbf{H}^{\textit{vk}}\) do not contribute to \({\mathcal{L}}\). This mirrors the diagonal assumption adopted in OBD. The full proof is provided in Appendix 7.

By explicitly evaluating the Hessian sub-blocks \(\mathbf{H}^{\textit{vv}}_{pp}\), \(\mathbf{H}^{\textit{kk}}_{pp}\), and \(\mathbf{H}^{\textit{vk}}_{pp}\), we next derive closed-form eviction scores that are both interpretable and efficient to compute. Currently, Equation 3 captures the joint impact of perturbing both \({\mathbf{V}}\) and \({\mathbf{K}}\). We also consider simplified variants where either \({\mathbf{V}}\) or \({\mathbf{K}}\) is perturbed independently while the other remains fixed. As a result, OBCache features three output-aware saliency scores, as demonstrated in Propositions 2–4.

This score computes the squared \(\ell_2\)-norm of the \(p\)-th column of the attention weight matrix, scaled by the squared \(\ell_2\)-norm of the value vector \({\mathbf{v}}_p\). It corresponds to the value-aware score proposed in VATP [15] and CriticalKV [16]. Note that these methods use the \(\ell_1\)-norm of the attention weights and value states. Our choice of a smooth \(\ell_2\) objective is required to derive key-pruning scores in closed form. If we instead adopt an \(\ell_1\) objective for value pruning, the resulting score reduces to the same form. The full proof, including the derivation of the Hessian sub-block \(\mathbf{H}^{\textit{vv}}_{pp}\) is provided in Appendix 7.2.

\[\begin{align} {\boldsymbol{S}}_p^{\textrm{key}} = \frac{1}{2} {\mathbf{k}}_p^\top \mathbf{H}^{\textit{kk}}_{pp} {\mathbf{k}}_p = \textstyle\sum_{i} |{\mathbf{A}}_{i, p} {\mathbf{Z}}_{i, p}|^2 \| {\mathbf{v}}_p - {\mathbf{o}}_i \|^2. \label{eq:score-k} \end{align}\tag{4}\] This score captures the deviation between the value vector and attention output, weighted by both the attention weights and the pre-softmax logits. Key pruning generally incurs larger errors than value pruning, as it alters the entire attention distribution. The proof is provided in Appendix 7.3.

This score captures both the individual and interactive effects of pruning the key and value states, providing the most comprehensive estimate of the pruning-induced eviction error (Equation 1 ). The derivation of the cross-term is provided in in Appendix 7.4.

We note that OBCache scores can be applied for both prefill and decoding. In the prefill phase, the saliency score \({\boldsymbol{S}}_p\) can be used to greedily evict multiple tokens to achieve a desired sparsity in a one-shot manner. During decoding, by accumulating \({\boldsymbol{S}}_p\) over time, OBCache can support real-time updates to token-wise saliency, enabling dynamic KV cache eviction as generation progresses. For models using Grouped-Query Attention [30], we include an additional score derivation in Appendix 7.5.

4.3 Connection to Existing Methods↩︎

Our framework naturally recovers existing attention-based eviction strategies as special cases. To show this, we introduce an additional index \(w \in [1,s]\) and use \({\mathbf{A}}_{w:s}\) to denote the rows of the attention matrix \({\mathbf{A}}\) corresponding to the query positions from \(w\) to \(s\). Consider an alternative formulation that minimizes perturbations not in the attention outputs \({\mathbf{O}}\), but in the historical attention rows \({\mathbf{A}}_{w:s}\). In this case, the pruning-induced error reduces to: \[\begin{align} {\boldsymbol{S}}_p^{\textrm{attn}} := {\mathcal{L}}_{\widehat{\mathbf{A}}_{:, p} =\boldsymbol{0}} (\widehat{\mathbf{A}}) = \bigg\| \widehat{\mathbf{A}}_{w:s} \Big|_{\widehat{\mathbf{A}}_{:, p} =\boldsymbol{0}} - {\mathbf{A}}_{w:s} \bigg\|_{1, 1}, \label{eq:eviction95p95reduce} \end{align}\tag{5}\] where the pruning unit also simplifies to an attention matrix column. We refer to query positions from \(w\) to \(s\) as the perturbation window in order to connect and formalize the accumulation strategies used in prior methods. Equation 5 can be directly simplified to the form: \[\begin{align} {\boldsymbol{S}}_p^{\textrm{attn}} = \sum_{i=w}^s | {\mathbf{A}}_{i, p}|, \end{align}\] which corresponds to an \(\ell_1\)-norm variant of our value-pruning score \({\boldsymbol{S}}_p^{\text{value}}\), but without incorporating any value-state information. This formulation connects directly to several recent methods. Specifically, H\(_2\)O [11] sets \(w = 1\), accumulating attention weights over the entire sequence history. TOVA [13] sets \(w = s\), targeting only the most recent query position. SnapKV [14] adopts a short window (i.e., \(w \gg 1\)), emphasizing recent attentions. By varying the choice of \(w\), these methods effectively target different perturbation windows.

In OBCache, when we also relax the objective to the output error within the perturbation window: \[\begin{align} { {\boldsymbol{S}}_p \!=\! \bigg\| \sigma\big({ \tfrac{{\mathbf{Q}}_{w:s} {\widehat{\mathbf{K}}}^\top}{\sqrt{d}} }\big) \widehat{\mathbf{V}} \Big|_{\boldsymbol{e}_p^\top [\widehat{\mathbf{V}}~ \widehat{\mathbf{K}}]=\boldsymbol{0}} \!-\! \sigma\big({ \tfrac{{\mathbf{Q}}_{w:s} {\mathbf{K}}^\top}{\sqrt{d}} }\big) {\mathbf{V}} \bigg\|_F^2, } \label{eq:eviction95p95window} \end{align}\tag{6}\] the resulting scores also become localized to the same set of query positions \(w\) through \(s\). Therefore, by accumulating the pruning-induced eviction error within different perturbation windows, OBCache scores generalize these prior approaches by further introducing output-aware signals, enabling more informed eviction decisions that go beyond raw attention statistics.

4.4 Effectiveness of Pruning-Induced Eviction Error↩︎

To further support our pruning-based formulation for cache eviction, we present a qualitative analysis on a Needle-In-A-Haystack [31] passkey retrieval task, as shown in Figure 2. We begin by establishing an oracle baseline based on the true eviction error, which is measured as the perturbation in the first decoding-step output \({\mathbf{o}}_{l+1}\) caused by evicting each cached token during the prefill phase. The top-\(k\) tokens with the largest oracle errors are treated as ground-truth important positions.

To evaluate the effectiveness of the pruning-induced eviction error as a proxy, we compute Equation 1 exactly for each candidate position in the prefill phase and select the top-\(k\) accordingly. As shown in the figure, when the perturbation window is appropriately chosen, the proxy achieves up to 85% recall of the oracle top-\(k\) selections. Next, we demonstrate the superiority of OBCache. Although derived via second-order Taylor approximation, OBCache scores achieve nearly identical ranking performance to the exact proxy, while being significantly more efficient due to their closed-form expression. Compared to attention-based methods, our scores consistently yield higher oracle recall, showcasing the benefit of incorporating output-aware signals into saliency estimation.

We observe that recall degrades when the perturbation window becomes larger. This issue stems from attention causality: earlier tokens attend to more queries and thus accumulate disproportionately higher saliency scores, introducing a structural bias that favors retaining initial tokens. To mitigate this, H\(_2\)O reserves a portion of the cache budget for a recent window, ensuring that the most recent tokens are never evicted. This policy is also complementary to OBCache, and we find that recall further improves when a fixed 20-token window is reserved, as shown in Figure [2][b].

5 Experiments↩︎

In this section, we conduct comprehensive experiments to evaluate the effectiveness of OBCache in improving existing score-based KV cache eviction methods across diverse long-context benchmarks. We evaluate OBCache in both static cache eviction in the prefill stage (Section 5.2) and dynamic cache eviction in the decoding stage (Section 5.3). We additionally report efficiency results to demonstrate the negligible computational overhead introduced by our output-aware scoring mechanism compared to attention-based methods in both prefill and decoding (Appendix 10.1).

5.1 Experimental Setup↩︎

5.1.1 Datasets.↩︎

For prefill-stage cache eviction, we evaluate on two widely used long-context benchmarks: RULER [32] and LongBench [33]. RULER is a synthetic benchmark suite consisting of 13 tasks designed to evaluate long-context retrieval, multi-hop reasoning, aggregation, and positional robustness under controlled context lengths and compression settings. Following prior KV cache eviction works, we evaluate under 4K (denoted as RULER-4K) and 32K (denoted as RULER-32K) context lengths. LongBench evaluates long-context understanding through 16 real-world datasets spanning six task categories: single-document QA, multi-document QA, summarization, few-shot learning, synthetic reasoning, and code completion. The average input length across all datasets is 6,711 words, making KV cache optimization critical for efficient inference. For decoding-stage cache eviction, we follow the setup in StreamingLLM [12] and report language modeling perplexity on PG19 [34], a dataset of 100 books with an average length of 70K tokens.

5.1.2 Baselines.↩︎

We compare OBCache with four representative KV cache eviction methods: H\(_2\)O [11], TOVA [13], SnapKV [14], and AdaKV [25], all of which rely on attention statistics as token importance indicators but differ in their eviction and budget allocation strategies. H\(_2\)O accumulates historical attention weights while retaining a fixed recent window. TOVA selects tokens based on the most recent query attention and does not require a recent window. SnapKV, designed for prefill-stage eviction, scores tokens based on a recent attention window and applies pooling to smooth token importance. AdaKV further introduces adaptive head-wise budget allocation, globally redistributing cache budgets across attention heads based on their cumulative attention weights. We integrate OBCache into these baselines by replacing their original attention-based scores with our output-aware scores while preserving their original eviction strategies. We refer to our eviction methods using the value-pruning score as OBCache-V, the key-pruning score as OBCache-K, and the joint pruning score as OBCache-V&K.

To compare against existing value-aware scoring methods, we additionally include comparisons with VATP [15] and CriticalKV [16] in Section 5.2.4. VATP uses an \(\ell_1\)-norm value-based importance score, while CriticalKV combines an attention-based preselection stage with a secondary value-aware scoring stage.

5.1.3 Implementation Details.↩︎

We use LLaMA-3.1-8B-Instruct [10] and Qwen-2.5-7B-Instruct [3] as our backbone LLMs, both of which natively support 128K context windows and adopt Grouped-Query Attention [30] to reduce KV cache memory. All experiments are implemented using the Transformers library [35] based on the KVPress repository [36]. To improve efficiency, the prefill-stage attention operation is implemented via FlashAttention-2 [37], which reduces memory overhead for long-context inputs.

In static cache eviction, KV compression is performed only once before decoding begins. This reflects realistic deployment scenarios where prefill-stage KV cache dominates memory usage compared to decoding-stage cache. We evaluate all methods under cache budget ratios of 10%, 20%, 30%, and 40%. The cache budget is determined as a fixed percentage of the input prompt length. For decoding-stage eviction, we follow H\(_2\)O and maintain a fixed recent window while evicting tokens at every decoding step to evaluate the dynamic impact of different saliency scores.

Following KVPress and AdaKV, we evaluate under both query-aware and query-agnostic cache compression settings. In the query-aware setting, cache eviction is performed after observing the downstream query, following the standard setup used in prior methods such as SnapKV. In the more challenging query-agnostic setting [25], [38], the context KV cache must be compressed before future queries are observed, which better reflects real-world cache reuse scenarios such as prefix caching and multi-turn dialogue. Full implementation details are provided in Appendix 9.

5.2 Evaluation of Prefill Cache Eviction↩︎

We present RULER and LongBench performance curves under varying cache budgets in Figure 3. We report results for OBCache-K, since it consistently outperforms OBCache-V and mostly matches OBCache-V&K’s performance while requiring lower complexity. Additionally, Table 1 reports detailed average RULER accuracy for LLaMA-3.1-8B, comparing baselines with their three OBCache-enhanced counterparts. Full results across all 13 RULER tasks and 16 LongBench tasks are provided in Appendix 10.2.

5.2.1 Results on RULER↩︎

Across all baselines, compression settings, cache budgets, context lengths, and model families, integrating OBCache scores consistently improves task accuracy. As illustrated in Table 1, when applied to H\(_2\)O, OBCache yields large average accuracy gains across all compression ratios, achieving more than 10% accuracy increase on RULER-4K and 5% on RULER-32K. TOVA is also consistently improved by 2%\(\sim\)​5% in average accuracy. For SnapKV, the improvement is relatively smaller in query-aware settings. This is likely because its heuristic 1D-pooling step is tailored to noisy attention-only scores. Since OBCache produces output-aware saliency scores that are already more accurate, the marginal benefit of such smoothing may be reduced. Notably, when integrated with AdaKV, our strongest baseline, OBCache yields substantial additional gains. For example, we observe nearly 15% accuracy increase on the query-agnostic RULER-4K setting with a 30% cache budget. These results demonstrate the consistent effectiveness of OBCache scores across diverse settings. Overall, all four baselines benefit from integrating OBCache, highlighting the advantage of our output-aware scoring mechanism over attention-based heuristics.

5.2.2 Comparison of OBCache Scores↩︎

In addition to the comparison with baselines, Table 1 also provides an ablation study of the three OBCache score variants derived under different pruning unit assumptions. We observe that eviction scores derived from key pruning, i.e., OBCache-K and OBCache-V&K, consistently outperform the value-only variant OBCache-V. This aligns with our expectation that pruning keys has a larger impact due to their role in shaping attention distributions and ultimately the model predictions. Accounting for key sensitivity therefore yields more accurate saliency estimation and better performance. Although the OBCache-V variant performs slightly worse, it remains appealing for its simplicity: it requires only an additional scaling factor based on value-state norms, making it nearly as efficient as attention-only scores while still delivering substantial performance gains.

5.2.3 Results on LongBench↩︎

As shown in Figure 3, we similarly observe consistent LongBench performance improvements when integrating OBCache into all four baselines for both LLaMA-3.1 and Qwen-2.5 models. The gains tend to become more pronounced under higher compression ratios, where output-aware saliency estimation becomes increasingly important. Notably, for AdaKV, our strongest baseline, OBCache-K improves performance by +1.2 in the query-aware setting and +2.6 in the query-agnostic setting under a 10% cache budget on LLaMA. This provides strong evidence that OBCache scores remain effective on real-world tasks and further advance the state of the art. Overall, these results confirm that replacing attention-only scores with OBCache scores enables more effective preservation of task-critical tokens under diverse compression settings.

5.2.4 Comparison to Value-Aware Baselines↩︎

As discussed in Section 4.2, OBCache-V reduces to the scores proposed in VATP and CriticalKV under different norm choices. Since our OBD-based formulation flexibly defines pruning units and objectives, it additionally yields two key-pruning-oriented scores that existing value-aware baselines do not consider. To demonstrate the effectiveness of our general framework, we reproduce these value-aware baselines when integrated into AdaKV on RULER-4K. As shown in Figure 4, in the query-aware setting, OBCache-K and OBCache-V&K achieve the strongest performance among all methods. OBCache-V follows and outperforms both CriticalKV and VATP, especially at higher compression ratios. In the query-agnostic setting, CriticalKV is competitive at the lowest compression ratio (40% budget); however, it still underperforms OBCache-K and OBCache-V&K at higher compression ratios. Overall, VATP provides the smallest improvements over AdaKV, CriticalKV underperforms our OBCache scores in most settings, and OBCache-K and OBCache-V&K achieve the best performance-compression trade-offs. These results demonstrate that the gains of OBCache do not arise merely from rederiving a known value-aware score within a cleaner framework, but from the broader OBD-based formulation and, importantly, from the additional key-aware scoring terms that more effectively capture output sensitivity.

5.3 Evaluation of Decoding Cache Eviction↩︎

Beyond static prefill-stage cache eviction, OBCache is also applicable to decoding scenarios, where eviction decisions must be made dynamically. To evaluate this setting, we measure cumulative language modeling perplexity at varying sequence lengths on the PG19 dataset using a fixed KV cache budget of 1024 tokens. Following [12], we retain 4 fixed initial tokens for all methods, as these tokens act as attention sinks and are crucial for long-sequence generation quality. As shown in Figure 5, the Sink baseline, which statically allocates the cache to fixed initial and recent tokens, exhibits the highest perplexity. The H\(_2\)O baseline, which dynamically selects important tokens based on accumulated attention statistics, performs moderately better. In contrast, all OBCache variants, which instead accumulate output-aware saliency scores over time, consistently outperform H\(_2\)O across all sequence lengths. We also observe that OBCache-V&K does not outperform OBCache-K, suggesting that an improved combination of key-pruning and value-pruning scores may exist beyond the current plain additive formulation. SnapKV and AdaKV are not evaluated, since they are specifically designed for static prefill-stage eviction and do not inherently support dynamic decoding eviction. Overall, these results demonstrate that our output-aware saliency scores can improve memory utilization while better preserving long-range dependencies critical for continuous high-quality generation.

6 Conclusion↩︎

In this work, we introduce OBCache, a principled scoring framework for KV cache eviction in language model inference grounded in the Optimal Brain Damage theory. By casting cache eviction as a layer-wise structured pruning problem, we derive token saliency scores that aim to minimize the impact of removal on attention outputs. Experiments across both prefilling and decoding scenarios on long-context benchmarks show that OBCache consistently improves state-of-the-art baselines, achieving superior performance–compression trade-offs with negligible overhead. Beyond the proposed scores, OBCache provides a flexible foundation for extending to additional KV cache compression tasks: by modifying the pruning objective or pruning units, the framework can be naturally applied to settings such as channel-wise KV pruning or cache merging via relaxed diagonal approximations. These directions open promising paths toward more effective and theoretically grounded KV cache management for long-context LLMs.

Impact Statement↩︎

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.

7 Theoretical Analysis↩︎

In this section, we present the full proof and derivations of the OBCache scores in Equation ?? , Equation 4 and Equation ?? . We use the notation previously defined in Section 3.

7.1 OBCache Objective Function↩︎

As demonstrated in Section 4.1 and Section 4.3, we formulate the saliency score for a candidate key-value token in Equation 1 and Equation 6 , which is the pruning-induced eviction error within a perturbation window for query positions \(w\) through \(s\): \[\begin{align} {\mathcal{L}}(\widehat{\mathbf{V}},\widehat{\mathbf{K}}) \;:=\; \big\| \widehat{\mathbf{O}}_{w:s} - {\mathbf{O}}_{w:s}\big\|_F^2 \;=\; \big\| \mathrm{softmax}\big({ \tfrac{{\mathbf{Q}}_{w:s} {\widehat{\mathbf{K}}}^\top}{\sqrt{d}} }\big) \widehat{\mathbf{V}} - \mathrm{softmax}\big({ \tfrac{{\mathbf{Q}}_{w:s} {\mathbf{K}}^\top}{\sqrt{d}} }\big) {\mathbf{V}} \big\|_F^2. \label{eq:obj95raw95app} \end{align}\tag{7}\] In what follows, we decompose the objective \({\mathcal{L}}\) into element-wise form. We let \({\mathbf{O}}_{i,j} = {\mathbf{a}}_i {\mathbf{V}}_{:,j}\) denote the \(j\)-th feature element in the \(i\)-th row vector of the attention output matrix \({\mathbf{O}}\), where \({\mathbf{a}}_i\) is the \(i\)-th row vector of the attention weight matrix. The squared Frobenius norm objective in Equation 7 can be explicitly decomposed into a summation form across all element-wise squared errors: \[\begin{align} \begin{aligned} {\mathcal{L}}(\widehat{\mathbf{V}},\widehat{\mathbf{K}}) &= \sum_{i=w}^s \sum_{j=1}^d |\widehat{\mathbf{O}}_{i,j}-{\mathbf{O}}_{i,j}|^2 \\ &= \sum_{i=w}^s \sum_{j=1}^d | \mathrm{softmax}\big({ \tfrac{{\mathbf{q}}_i {\widehat{\mathbf{K}}}^\top}{\sqrt{d}} }\big) \widehat{\mathbf{V}}_{:,j} - \mathrm{softmax}\big({ \tfrac{{\mathbf{q}}_i {\mathbf{K}}^\top}{\sqrt{d}} }\big) {\mathbf{V}}_{:,j} |^2 \\ &\overset{\triangle}{=} \sum_{i=w}^s \sum_{j=1}^d {\boldsymbol{E}}_{i, j}. \label{app-eq:32decomposeL} \end{aligned} \end{align}\tag{8}\] For convenience in the following derivations, we define \({\boldsymbol{E}}_{i,j} := |\widehat{\mathbf{O}}_{i,j}-{\mathbf{O}}_{i,j}|^2\) as the squared difference of the output element \({\mathbf{O}}_{i,j}\) before and after applying a perturbation \(\delta\mathbf{V}\) and \(\delta\mathbf{K}\). Expanding explicitly, the element-wise output perturbation for \({\mathbf{O}}_{i, j}\) is: \[\begin{align} \boldsymbol{E}_{i,j} (\widehat{\mathbf{V}}_{:,j}, \widehat{\mathbf{K}}) = | \mathrm{softmax}\big({ \tfrac{{\mathbf{q}}_i {\widehat{\mathbf{K}}}^\top}{\sqrt{d}} }\big) \widehat{\mathbf{V}}_{:,j} - \mathrm{softmax}\big({ \tfrac{{\mathbf{q}}_i {\mathbf{K}}^\top}{\sqrt{d}} }\big) {\mathbf{V}}_{:,j} |^2 . \label{app-eq:32element-ptb} \end{align}\tag{9}\] Given this element-wise objective, we follow the Optimal Brain Damage (OBD) theory [17] to apply a second-order Taylor expansion of \({\boldsymbol{E}}_{i,j}\) around \(({\mathbf{V}}_{:,j}, {\mathbf{K}})\), and analytically approximate the element-wise output perturbation: \[\begin{align} \boldsymbol{E}_{i,j}(\widehat{\mathbf{V}}_{:,j}, \widehat{\mathbf{K}}) \;\overset{\text{second}}{\underset{\text{order}}{=}} \; \boldsymbol{E}_{i,j}({\mathbf{V}}_{:,j}, {\mathbf{K}}) \; &+ \;\delta\mathbf{V}^\top_{:,j} \dfrac{\partial {\boldsymbol{E}}_{i,j}({\mathbf{V}}_{:, j}, {\mathbf{K}})}{\partial \widehat{\mathbf{V}}_{:, j}} \;+ \;\dfrac{1}{2} \delta\mathbf{V}^\top_{:,j} \dfrac{\partial^2 \boldsymbol{E}_{i,j}({\mathbf{V}}_{:,j}, {\mathbf{K}}) }{\partial \widehat{\mathbf{V}}_{:,j}^2} \delta\mathbf{V}_{:,j} \\ &+ \;\delta\mathbf{K}^\top \dfrac{\partial {\boldsymbol{E}}_{i,j}({\mathbf{V}}_{:, j}, {\mathbf{K}})}{\partial \widehat{\mathbf{K}}} \;+ \;\dfrac{1}{2} \delta\mathbf{K}^\top \dfrac{\partial^2 \boldsymbol{E}_{i,j}({\mathbf{V}}_{:,j}, {\mathbf{K}})}{\partial \widehat{\mathbf{K}}^2} \delta\mathbf{K}\\ &+ \;\delta\mathbf{V}_{:,j}^\top \dfrac{\partial^2 \boldsymbol{E}_{i,j}({\mathbf{V}}_{:,j}, {\mathbf{K}})}{\partial \widehat{\mathbf{V}}_{:,j} \partial \widehat{\mathbf{K}}} \delta\mathbf{K} \;+ \;{\mathcal{O}}(\| (\delta\mathbf{V}_{:, j}, \delta\mathbf{K}) \|^3) . \end{align}\] Through substitution into Equation 9 , the constant term \(\boldsymbol{E}_{i,j}({\mathbf{V}}_{:,j}, {\mathbf{K}})\) vanishes because we have \(\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j} = 0\) at the expansion point \(({\mathbf{V}}_{:,j}, {\mathbf{K}})\). Therefore, the full matrix-wise objective, pruning-induced eviction error, approximated via a second-order Taylor series, becomes: \[\begin{align} {\mathcal{L}}(\widehat{\mathbf{V}},\widehat{\mathbf{K}}) \;\overset{\text{second}}{\underset{\text{order}}{=}} \; &\sum_{i=w}^s \sum_{j=1}^d \delta\mathbf{V}^\top_{:,j} \dfrac{\partial {\boldsymbol{E}}_{i,j}({\mathbf{V}}_{:, j}, {\mathbf{K}})}{\partial \widehat{\mathbf{V}}_{:, j}} + \dfrac{1}{2} \delta\mathbf{V}^\top_{:,j} \dfrac{\partial^2 \boldsymbol{E}_{i,j}({\mathbf{V}}_{:,j}, {\mathbf{K}}) }{\partial \widehat{\mathbf{V}}_{:,j}^2} \delta\mathbf{V}_{:,j} \tag{10}\\ +&\sum_{i=w}^s \sum_{j=1}^d \delta\mathbf{K}^\top \dfrac{\partial {\boldsymbol{E}}_{i,j}({\mathbf{V}}_{:, j}, {\mathbf{K}})}{\partial \widehat{\mathbf{K}}} + \dfrac{1}{2} \delta\mathbf{K}^\top \dfrac{\partial^2 \boldsymbol{E}_{i,j}({\mathbf{V}}_{:,j}, {\mathbf{K}})}{\partial \widehat{\mathbf{K}}^2} \tag{11}\\ +&\sum_{i=w}^s \sum_{j=1}^d \delta\mathbf{V}_{:,j}^\top \dfrac{\partial^2 \boldsymbol{E}_{i,j}({\mathbf{V}}_{:,j}, {\mathbf{K}})}{\partial \widehat{\mathbf{V}}_{:,j} \partial \widehat{\mathbf{K}}} \delta\mathbf{K}. \tag{12} \end{align}\]

7.2 Isolated Value-Pruning Score↩︎

In isolated value pruning, the key-cache matrix is not perturbed and is considered a constant not affecting \({\mathcal{L}}\). Therefore, minimizing the pruning-induced eviction error reduces to minimizing Equation 10 . To derive saliency scores in closed form, we begin by deriving expressions for the gradient and Hessian with respect to \(\widehat{\mathbf{V}}_{:,j}\). The first-order derivative of \({\boldsymbol{E}}_{i, j}\) with respect to \(\widehat{\mathbf{V}}_{:,j}\) is: \[\begin{align} \dfrac{\partial \boldsymbol{E}_{i, j}}{\partial \widehat{\mathbf{V}}_{:,j}} = \dfrac{\partial}{\partial \widehat{\mathbf{V}}_{:,j}}|\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j}|^2 = 2(\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j}) \dfrac{\partial}{\partial \widehat{\mathbf{V}}_{:,j}} \hat{{\mathbf{a}}}_i \widehat{\mathbf{V}}_{:, j} = 2(\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j}) \hat{{\mathbf{a}}}_i. \end{align}\] When evaluated at \(({\mathbf{V}}_{:,j}, {\mathbf{K}})\), we again have \(\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j} = 0\), so the gradient term vanishes to zero. The Hessian of \({\mathcal{L}}\) with respect to \(\widehat{\mathbf{V}}_{:,j}\) evaluated at \(({\mathbf{V}}_{:,j}, {\mathbf{K}})\) is: \[\begin{align} \dfrac{\partial^2 \boldsymbol{E}_{i, j} ({\mathbf{V}}_{:, j}, {\mathbf{K}})}{\partial \widehat{\mathbf{V}}_{:,j}^2} = 2 \dfrac{\partial}{\partial \widehat{\mathbf{O}}_{i,j}} \Big( (\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j}) \hat{{\mathbf{a}}}_i^\top \Big) \Big( \dfrac{\partial \widehat{\mathbf{O}}_{i,j}}{\partial \widehat{\mathbf{V}}_{:,j}} \Big)^\top \Bigg|_{({\mathbf{V}}_{:, j}, {\mathbf{K}})} = 2 {\mathbf{a}}_i^\top {\mathbf{a}}_i. \end{align}\] Substituting the gradient and Hessian back into Equation 10 , we get the pruning-induced eviction error when value states are considered as isolated pruning units: \[\begin{align} {\mathcal{L}}^{\text{value}} = \sum_{i=w}^s \sum_{j=1}^d \delta\mathbf{V}_{:,j}^\top ({\mathbf{a}}_i^\top {\mathbf{a}}_i) \delta\mathbf{V}_{:,j} = \sum_{i=w}^s \sum_{j=1}^d | {\mathbf{a}}_i \delta\mathbf{V}_{:,j} |^2. \end{align}\] Next, the row-wise pruning constraint \(\boldsymbol{e}_p^\top \widehat{\mathbf{V}}=\boldsymbol{0}\) implies that when pruning the \(p\)-th token position, the value cache perturbation \(\delta\mathbf{V}_{:,j}\) is explicitly in the form: \[\begin{align} \delta\mathbf{V}_{t,j} = \begin{cases} - {\mathbf{V}}_{p, j}, &\text{when }t=p \\ 0, &\text{when }t\neq p \end{cases}, \quad \forall \;j = 1,..., d. \end{align}\] Substituting this value perturbation into \({\mathcal{L}}^{\text{value}}\), we obtain the output perturbation when pruning the \(p\)-th value vector from the value cache \({\mathbf{V}}\): \[\begin{align} {\boldsymbol{S}}_p^{\text{value}} = \sum_{i=w}^s \sum_{j=1}^d |{\mathbf{A}}_{i,p} \cdot {\mathbf{V}}_{p,j} |^2 = \sum_{i=w}^s |{\mathbf{A}}_{i,p} |^2 \cdot \sum_{j=1}^d | {\mathbf{V}}_{p,j} |^2 = \boxed{\sum_{i=w}^s |{\mathbf{A}}_{i,p}|^2 \cdot \|{\mathbf{v}}_p\|^2}. \end{align}\] This is also the value-pruning saliency score for evicting the \(p\)-th value vector \({\mathbf{v}}_p\) from the value cache. It is consistent with the expression in Equation ?? in Section 4.2.

7.3 Isolated Key-Pruning Score↩︎

In isolated key pruning, the value-cache matrix is not perturbed and is assumed a constant. Therefore, minimizing the pruning-induced eviction error reduces to minimizing Equation 11 . As with the value-pruning score, we begin by deriving closed-form expressions for the gradient and Hessian with respect to \(\widehat{\mathbf{K}}\). The first-order derivative of \({\boldsymbol{E}}_{i, j}\) with respect to \(\widehat{\mathbf{K}}\) is: \[\begin{align} \dfrac{\partial \boldsymbol{E}_{i,j}}{\partial \widehat{\mathbf{K}}} = \dfrac{\partial}{\partial \widehat{\mathbf{K}}}|\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j}|^2 = 2(\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j}) \dfrac{\partial \widehat{\mathbf{O}}_{i,j}}{\partial \widehat{\mathbf{K}}}. \end{align}\] Before deriving the first-order term explicitly, note that when evaluated at the point \(({\mathbf{V}}_{:,j}, {\mathbf{K}})\), we again have \(\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j} = 0\). Consequently, the first-order term is eliminated, a result that is the same as in isolated value pruning. We now explicitly derive the term \(\widehat{{\boldsymbol{M}}}^{(i,j)} := \tfrac{\partial \widehat{\mathbf{O}}_{i,j}}{\partial \widehat{\mathbf{K}}}\) above: \[\begin{align} \widehat{{\boldsymbol{M}}}_{p, r}^{(i,j)} &= \Big( \dfrac{\partial \widehat{\mathbf{O}}_{i, j}}{\partial \widehat{\mathbf{K}}} \Big)_{p,r} = \dfrac{\partial}{\partial \widehat{\mathbf{K}}_{p,r}} \sum_{m=1}^s \widehat{\mathbf{A}}_{i, m} \cdot \widehat{\mathbf{V}}_{m, j} = \sum_{m=1}^s \dfrac{\partial \widehat{\mathbf{O}}_{i,j}}{\partial \widehat{\mathbf{A}}_{i,m}} \dfrac{\partial \widehat{\mathbf{A}}_{i,m}}{\partial \widehat{\mathbf{K}}_{p,r}} \notag\\ &= \sum_{m=1}^s \dfrac{\partial \widehat{\mathbf{O}}_{i,j}}{\partial \widehat{\mathbf{A}}_{i, m}} \dfrac{\partial}{\partial \widehat{\mathbf{K}}_{p,r}} (\dfrac{e^{\widehat{\mathbf{Z}}_{i, m}}}{\sum_{u=1}^s e^{\widehat{\mathbf{Z}}_{i, u}}}) \notag\\ &= \sum_{m=1}^s \dfrac{\partial \widehat{\mathbf{O}}_{i,j}}{\partial \widehat{\mathbf{A}}_{i,m}} \sum_{u=1}^s \dfrac{\widehat{\mathbf{A}}_{i,m}}{\widehat{\mathbf{Z}}_{i,u}} \dfrac{\widehat{\mathbf{Z}}_{i,u}}{\partial \widehat{\mathbf{K}}_{p,r}} \notag\\ &= \sum_{m=1}^s \widehat{\mathbf{V}}_{m,j} \sum_{u=1}^s \widehat{\mathbf{A}}_{i,m}(\delta_{mu}-\widehat{\mathbf{A}}_{i,u}) \dfrac{1}{\sqrt{d}} {\mathbf{Q}}_{i,r} \delta_{up} \notag \\ &= \sum_{m=1}^s \widehat{\mathbf{V}}_{m,j} \cdot \widehat{\mathbf{A}}_{i,m}(\delta_{mp}-\widehat{\mathbf{A}}_{i,p}) \dfrac{1}{\sqrt{d}} {\mathbf{Q}}_{i,r} \;\;\;\text{(\delta_{up}=0 when u \neq p)} \notag \\ &=\dfrac{1}{\sqrt{d}} {\mathbf{Q}}_{i,r} \cdot ( \sum_{m=1}^s \widehat{\mathbf{V}}_{m,j} \cdot \widehat{\mathbf{A}}_{i,m} \cdot \delta_{mp} - \sum_{m=1}^s \widehat{\mathbf{V}}_{m,j} \cdot \widehat{\mathbf{A}}_{i,m} \cdot \widehat{\mathbf{A}}_{i,p} ) \notag \\ &= \dfrac{1}{\sqrt{d}} {\mathbf{Q}}_{i,r} \cdot (\widehat{\mathbf{V}}_{p,j} \cdot \widehat{\mathbf{A}}_{i,p} - \widehat{\mathbf{A}}_{i,p} \cdot \sum_{m=1}^s \widehat{\mathbf{V}}_{m,j} \cdot \widehat{\mathbf{A}}_{i,m}) \;\;\;\text{(\delta_{mp}=0 when m \neq p)} \notag \\ &=\dfrac{1}{\sqrt{d}} {\mathbf{Q}}_{i,r} \cdot \widehat{\mathbf{A}}_{i,p} \cdot (\widehat{\mathbf{V}}_{p,j} - \widehat{\mathbf{O}}_{i,j}). \end{align}\] Therefore, the first-order derivative of \({\boldsymbol{E}}_{i, j}\) with respect to the perturbed key cache \(\widehat{\mathbf{K}}\) is: \[\begin{align} \dfrac{\partial \boldsymbol{E}_{i, j}}{\partial \widehat{\mathbf{K}}} = 2(\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j}) \widehat{{\boldsymbol{M}}}^{(i,j)}, \text{ where } \widehat{{\boldsymbol{M}}}_{p,r}^{(i,j)} = \dfrac{1}{\sqrt{d}} {\mathbf{Q}}_{i,r} \cdot \widehat{\mathbf{A}}_{i,p} \cdot (\widehat{\mathbf{V}}_{p,j} - \widehat{\mathbf{O}}_{i,j}). \end{align}\]

We now proceed to derive the Hessian, which is a 4-th order tensor of size \(s \times d \times s \times d\). The second-order derivative of the element-wise perturbation \(\boldsymbol{E}_{i,j}\) with respect to \(\widehat{\mathbf{K}}\) is: \[\begin{align} \dfrac{\partial^2 \boldsymbol{E}_{i,j}}{\partial \widehat{\mathbf{K}}^2} &= \widehat{{\boldsymbol{M}}}^{(i,j)} \otimes \left( \dfrac{\partial}{\partial \widehat{\mathbf{K}}} 2(\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j}) \right) + 2(\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j}) \dfrac{\partial \widehat{{\boldsymbol{M}}}^{(i,j)}}{\partial \widehat{\mathbf{K}}} \notag \\ &= 2 \widehat{{\boldsymbol{M}}}^{(i,j)} \otimes \widehat{{\boldsymbol{M}}}^{(i,j)} + 2(\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j}) \dfrac{\partial \widehat{{\boldsymbol{M}}}^{(i,j)}}{\partial \widehat{\mathbf{K}}}, \label{app-eq:32hessian-k} \end{align}\tag{13}\] where \(\otimes\) denotes the matrix-wise outer product¹. Since the Hessian will be evaluated at the point \(({\mathbf{V}}_{:,j}, {\mathbf{K}})\), where the perturbed attention output again matches the original, i.e., \(\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j} = 0\), the second term in Equation 13 vanishes. Therefore, the Hessian is: \[\begin{align} \dfrac{\partial^2 \boldsymbol{E}_{i,j}}{\partial \widehat{\mathbf{K}}^2} \Big|_{({\mathbf{V}}_{:,j}, {\mathbf{K}})} &= 2 {\boldsymbol{M}}^{(i,j)} \otimes {\boldsymbol{M}}^{(i,j)} = 2 \text{vec} ({\boldsymbol{M}}^{(i,j)})^\top \text{vec} ({\boldsymbol{M}}^{(i,j)}), \\ &\text{ where } {\boldsymbol{M}}_{p,r}^{(i,j)} = \dfrac{1}{\sqrt{d}} {\mathbf{Q}}_{i,r} \cdot {\mathbf{A}}_{i,p} \cdot ({\mathbf{V}}_{p,j} - {\mathbf{O}}_{i,j}). \notag \end{align}\] Here, we use \(\text{vec}(\cdot)\) to denote the vectorization of a matrix. Substituting the gradient and Hessian back into Equation 11 , we get the pruning-induced eviction error when key states are considered as isolated pruning units: \[\begin{align} {\mathcal{L}}^{\text{key}} &= \sum_{i=w}^s \sum_{j=1}^d \text{vec}(\delta\mathbf{K})^\top \text{vec} \big({\boldsymbol{M}}^{(i,j)}\big)^\top \text{vec} \big({\boldsymbol{M}}^{(i,j)}\big) \text{vec}(\delta\mathbf{K}) \notag\\ &= \sum_{i=w}^s \sum_{j=1}^d | \text{vec} \big({\boldsymbol{M}}^{(i,j)}\big) \text{vec}(\delta\mathbf{K}) |^2. \end{align}\] Similarly, the row-wise pruning constraint in \(\boldsymbol{e}_p^\top \widehat{\mathbf{K}}=\boldsymbol{0}\) implies that when pruning the \(p\)-th token position, the key cache perturbation \(\delta\mathbf{K}\) is explicitly in the form: \[\begin{align} \delta\mathbf{K}_{t,j} = \begin{cases} - {\mathbf{K}}_{p, j}, &\text{when }t=p \\ 0, &\text{when }t\neq p \end{cases}, \quad \forall \;j = 1,..., d. \end{align}\] Substituting this key perturbation into \({\mathcal{L}}^{\text{key}}\), we obtain the output perturbation when pruning the \(p\)-th key vector from the key cache \({\mathbf{K}}\) (we let \({\mathbf{m}}_p^{(i,j)}\) denote the \(p\)-th row vector of \({\boldsymbol{M}}^{(i,j)}\)): \[\begin{align} {\boldsymbol{S}}_p^{\text{key}} = \sum_{i=w}^s \sum_{j=1}^d |{\mathbf{m}}_p^{(i,j)} ~ {\mathbf{k}}_p|^2 &= \sum_{i=w}^s \sum_{j=1}^d |\sum_{r=1}^d \dfrac{1}{\sqrt{d}} {\mathbf{Q}}_{i,r} \cdot {\mathbf{A}}_{i, p}({\mathbf{V}}_{p, j} - {\mathbf{O}}_{i, j}) \cdot {\mathbf{K}}_{p,r}|^2 \notag\\ &= \dfrac{1}{d} \sum_{i=w}^s \sum_{j=1}^d |{\mathbf{A}}_{i, p}({\mathbf{V}}_{p, j} - {\mathbf{O}}_{i, j}) \cdot \sum_{r=1}^d {\mathbf{Q}}_{i, r} {\mathbf{K}}_{p, r}|^2 \notag\\ &= \dfrac{1}{d} \sum_{i=w}^s \sum_{j=1}^d |{\mathbf{A}}_{i, p}({\mathbf{V}}_{p, j} - {\mathbf{O}}_{i, j}) \cdot {\mathbf{q}}_i {\mathbf{k}}_p^\top|^2 \notag\\ &= \dfrac{1}{d} \sum_{i=w}^s \left(|{\mathbf{A}}_{i, p}|^2 \cdot |{\mathbf{q}}_i {\mathbf{k}}_p^\top|^2 \sum_{j=1}^d |{\mathbf{V}}_{p, j} - {\mathbf{O}}_{i, j}|^2 \right) \notag \\ &= \boxed{ \sum_{i=w}^s |{\mathbf{A}}_{i, p}|^2 \cdot |{\mathbf{Z}}_{i, p}|^2 \cdot \|{\mathbf{v}}_p - {\mathbf{o}}_i\|^2_2 }. \notag \end{align}\] This is also the key-pruning saliency score for evicting the \(p\)-th key vector \({\mathbf{k}}_p\) from the key cache. It is consistent with the expression in Equation 4 in Section 4.2.

7.4 Joint-Pruning Score↩︎

When both \({\mathbf{V}}_{:,j}\) and \({\mathbf{K}}\) are treated as variables affecting \({\mathcal{L}}\), we need to compute the cross-term in the error function, which corresponds to deriving a closed-form expression for Equation 12 . Given that the first-order derivative of \(\boldsymbol{E}_{i,j}\) with respect to \(\widehat{\mathbf{V}}_{:,j}\) is known, we can compute the cross-term as follows: \[\begin{align} \dfrac{\partial^2 \boldsymbol{E}_{i,j}}{\partial \widehat{\mathbf{V}}_{:,j} \partial \widehat{\mathbf{K}}} &= \dfrac{\partial}{\partial \widehat{\mathbf{K}}} 2 (\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j}) \hat{\mathbf{a}}_i \notag\\ &= \hat{\mathbf{a}}_i \otimes \left( \dfrac{\partial}{\partial \widehat{\mathbf{K}}} 2(\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j}) \right) + 2(\widehat{\mathbf{O}}_{i,j} - {\mathbf{O}}_{i,j}) \dfrac{\partial \hat{\mathbf{a}}_i}{\partial \widehat{\mathbf{K}}}. \end{align}\] As in the derivation of key-pruning scores, the second term above will be zero when evaluated at the point \(({\mathbf{V}}_{:,j}, {\mathbf{K}})\). Therefore, the cross second-order term is: \[\begin{align} \dfrac{\partial^2 \boldsymbol{E}_{i, j}}{\partial \widehat{\mathbf{V}}_{:,j} \partial \widehat{\mathbf{K}}} \Big|_{({\mathbf{V}}_{:,j}, {\mathbf{K}})} = 2 {\mathbf{a}}_i \otimes {\boldsymbol{M}}^{(i,j)}, \text{ where } {\boldsymbol{M}}_{p,r}^{(i,j)} = \dfrac{1}{\sqrt{d}} {\mathbf{Q}}_{i,r} \cdot {\mathbf{A}}_{i,p} \cdot ({\mathbf{V}}_{p,j} - {\mathbf{O}}_{i,j}). \end{align}\] Substituting this back into Equation 12 , we get the cross-term of the pruning-induced eviction error when the value and key states are considered as combined pruning units: \[\begin{align} {\mathcal{L}}^{\text{cross}} = \sum_{i=w}^s \sum_{j=1}^d \delta\mathbf{V}_{:,j}^\top \big( 2 {\mathbf{a}}_i \otimes {\boldsymbol{M}}^{(i,j)} \big) \delta\mathbf{K} = 2 \sum_{i=w}^s \sum_{j=1}^d (\delta\mathbf{V}_{:,j}^\top {\mathbf{a}}_i) \cdot \langle {\boldsymbol{M}}^{(i,j)}, \delta\mathbf{K}\rangle_F ,\label{app-eq:32E-cross} \end{align}\tag{14}\] where \(\langle {\boldsymbol{M}}^{(i,j)}, \delta\mathbf{K}\rangle_F = \sum_{p,r} {\boldsymbol{M}}^{(i,j)}_{p,r} \cdot \delta\mathbf{K}_{p,r} = \boldsymbol{tr}({{\boldsymbol{M}}^{(i,j)}}^\top \delta\mathbf{K})\) is the Frobenius inner product.

When the \(p\)-th key and value are pruned, the row-wise pruning constraint in \(\boldsymbol{e}_p^\top [\widehat{\mathbf{V}}~\widehat{\mathbf{K}}] =\boldsymbol{0}\) implies that the value cache perturbation \(\delta\mathbf{V}\) and the key cache perturbation \(\delta\mathbf{K}\) are explicitly in the form: \[\begin{align} \delta\mathbf{V}_{t,j} = \begin{cases} - {\mathbf{V}}_{p, j}, &\text{when }t=p \\ 0, &\text{when }t\neq p \end{cases}, \quad \delta\mathbf{K}_{t,j} = \begin{cases} - {\mathbf{K}}_{p, j}, &\text{when }t=p \\ 0, &\text{when }t\neq p \end{cases}, \quad \forall \;j = 1,..., d. \end{align}\] Thus, we have each term in \({\mathcal{L}}^{\text{cross}}\): \[\begin{align} \delta\mathbf{V}_{:,j}^\top {\mathbf{a}}_i = -{\mathbf{A}}_{i,p} {\mathbf{V}}_{p,j}, \quad \langle {\boldsymbol{M}}^{(i,j)}, \delta\mathbf{K}\rangle_F = - \dfrac{1}{\sqrt{d}} {\mathbf{A}}_{i,p} \cdot ({\mathbf{V}}_{p,j} - {\mathbf{O}}_{i,j}) \cdot ({\mathbf{q}}_i {\mathbf{k}}_p^\top). \end{align}\] Substituting these two expressions back into Equation 14 , we obtain the cross term in closed form: \[\begin{align} {\mathcal{L}}^{\text{cross}} &= \dfrac{2}{\sqrt{d}} \sum_{i=w}^s \sum_{j=1}^d {\mathbf{A}}_{i,p}^2 \cdot {\mathbf{V}}_{p,j} \cdot ({\mathbf{V}}_{p,j} - {\mathbf{O}}_{i,j}) \cdot ({\mathbf{q}}_i {\mathbf{k}}_p^\top) \notag\\ &= \dfrac{2}{\sqrt{d}} \sum_{i=w}^s {\mathbf{A}}_{i,p}^2 \cdot ({\mathbf{q}}_i {\mathbf{k}}_p^\top) \sum_{j=1}^d {\mathbf{V}}_{p,j} ({\mathbf{V}}_{p,j} - {\mathbf{O}}_{i,j}) \notag \\ &= \boxed{2 \sum_{i=w}^s {\mathbf{A}}_{i,p}^2 \cdot {\mathbf{Z}}_{i,p} \cdot (\| {\mathbf{v}}_p \|^2_2 - {\mathbf{v}}_p^\top {\mathbf{o}}_i)}. \end{align}\] Finally, the saliency score of the token at position \(p\) when both value and key states are treated as pruning units is the summation of \({\mathcal{L}}^{\text{value}}\), \({\mathcal{L}}^{\text{key}}\) and \({\mathcal{L}}^{\text{cross}}\): \[\begin{align} \boxed{ {\boldsymbol{S}}_p^{\text{joint}} = 2 \sum_{i=w}^s {\mathbf{A}}_{i,p}^2 \cdot {\mathbf{Z}}_{i,p} \cdot (\| {\mathbf{v}}_p \|^2_2 - {\mathbf{v}}_p^\top {\mathbf{o}}_i) + {\boldsymbol{S}}_p^{\text{value}} + {\boldsymbol{S}}_p^{\text{key}} }. \label{app-eq:32score-joint} \end{align}\tag{15}\] This is consistent with the joint score expression in Equation ?? of Section 4.2.

7.5 Scores for Grouped-Query Attention↩︎

In models with Grouped-Query Attention (GQA) [30], multiple query heads share a single key and value head, reducing KV cache storage at the architectural level. To support cache eviction under this setting, we explicitly incorporate the head-group structure into our analysis and derive the corresponding OBCache scores.

We use superscripts \(g\) to denote the KV-head index, \(h\) to denote query-head index, and \(H(g)\) to denote the set of query heads associated with KV head \(g\). Our previous derivation omits head indices for simplicity. Here, we specifically include head indexing for multi-head attention (MHA). If the pruning-induced eviction error is defined on the per-head attention output, \[{\mathcal{L}}:= \| \widehat{\mathbf{O}}^h - {\mathbf{O}}^h \|_F^2, \;\text{where } \;{\mathbf{O}}^h,\widehat{\mathbf{O}}^h \in \mathbb{R}^{s \times d},\] the OBCache scores for the \(h\)-th head naturally extend to: \[\begin{align} {\boldsymbol{S}}_{p, h}^{\textrm{value}} &= \sum_{i} |{\mathbf{A}}_{i,p}^h|^2 \cdot \|{\mathbf{v}}_p^h\|^2 \\ {\boldsymbol{S}}_{p, h}^{\textrm{key}} &= \sum_{i} |{\mathbf{A}}_{i, p}^h|^2 \cdot |{\mathbf{Z}}_{i, p}^h|^2 \cdot \|{\mathbf{v}}_p^h - {\mathbf{o}}_i^h\|^2_2 \\ {\boldsymbol{S}}_{p, h}^{\textrm{joint}} &= 2 \sum_{i} |{\mathbf{A}}_{i,p}^h|^2 \cdot {\mathbf{Z}}_{i,p}^h \cdot (\| {\mathbf{v}}_p^h \|^2_2 - {{\mathbf{v}}_p^h}^\top {\mathbf{o}}_i^h) + {\boldsymbol{S}}_{p,h}^{\textrm{value}} + {\boldsymbol{S}}_{p, h}^{\textrm{key}} \end{align}\] In standard MHA, each output head \({\mathbf{O}}^h\) depends on its own \({\mathbf{K}}^h\) and \({\mathbf{V}}^h\). In GQA, however, a group of output heads \(\{{\mathbf{O}}^{h}\}_{h \in H(g)}\) shares the same key and value \({\mathbf{K}}^{g}\) and \({\mathbf{V}}^{g}\), i.e., \[{\mathbf{V}}^g \equiv {\mathbf{V}}^h, \forall ~h\in H(g), \quad {\mathbf{K}}^g \equiv {\mathbf{K}}^h, ~\forall ~h \in H(g)\] If we now modify the objective to be the sum of output perturbations across heads within a group: \[{\mathcal{L}}:= \sum_{h \in H(g)} \| \widehat{\mathbf{O}}^h - {\mathbf{O}}^h \|_F^2,\] the OBCache scores for the \(g\)-th KV head introduce an additional summation over its associated query heads: \[\begin{align} {\boldsymbol{S}}_{p, g}^{\textrm{value}} &= \sum_{h \in H(g)} {\boldsymbol{S}}_{p, h}^{\textrm{value}} = \sum_{h \in H(g)} \sum_{i} |{\mathbf{A}}_{i,p}^h|^2 \cdot \|{\mathbf{v}}_p^g\|^2 \\ {\boldsymbol{S}}_{p, g}^{\textrm{key}} &= \sum_{h \in H(g)} {\boldsymbol{S}}_{p, h}^{\textrm{key}} = \sum_{h \in H(g)} \sum_{i} |{\mathbf{A}}_{i, p}^h|^2 \cdot |{\mathbf{Z}}_{i, p}^h|^2 \cdot \|{\mathbf{v}}_p^g - {\mathbf{o}}_i^h\|^2_2 \\ {\boldsymbol{S}}_{p, g}^{\textrm{joint}} &= \sum_{h \in H(g)} {\boldsymbol{S}}_{p, h}^{\textrm{joint}} = 2 \sum_{h \in H(g)} \Bigg( \sum_{i} |{\mathbf{A}}_{i,p}^h|^2 \cdot {\mathbf{Z}}_{i,p}^h \cdot (\| {\mathbf{v}}_p^g \|^2_2 - {{\mathbf{v}}_{p}^{g}}^\top {\mathbf{o}}_i^h) \Bigg) + {\boldsymbol{S}}_{p,g}^{\textrm{value}} + {\boldsymbol{S}}_{p, g}^{\textrm{key}} \end{align}\] Note that objective formulations other than summation can also yield valid per-KV-head scores. In all experiments, we use the above GQA-aware formulation to gather per-query-head scores without retaining KV cache for all query heads, which creates redundant KV memories as implemented by SnapKV [14].

8 Limitations and Future Work↩︎

OBCache presents several opportunities for further improvement:

8.0.0.1 Structural Bias to Earlier Tokens.

One limitation is the structural bias toward earlier tokens, a common drawback in many existing cache eviction methods that also emerges in OBCache, particularly as the perturbation window increases. Although heuristics such as retaining a fixed-size recent window help mitigate this bias by leveraging the locality of attention patterns, such strategies remain static and empirically motivated. Future work should investigate the design and adaptability of perturbation windows more systematically to further improve saliency estimation.

8.0.0.2 Choice of Perturbation Window in Decoding.

Additionally, OBCache’s current dynamic cache eviction strategy in the decoding phase directly follows H\(_2\)O by accumulating saliency scores over time, effectively minimizing perturbations over the full sequence history. However, under our pruning-based formulation, alternative objectives could also be explored. For example, saliency scores could instead be accumulated only within a recent attention window. Such variants may yield improved adaptability in long-context generation.

8.0.0.3 Extended Real-World Evaluation.

The evaluation settings considered in this work can also be extended to broader real-world scenarios. For example, OBCache could be further studied under dynamic cache eviction beyond language modeling tasks. In decoding-centric tasks such as mathematical reasoning, dynamic cache eviction may play a more important role than prefill-stage eviction. In addition, OBCache has not yet been evaluated on multi-round agentic tasks, which may more directly reflect real-world deployment settings.

8.0.0.4 Extension to other KV Tasks.

On the theoretical side, OBCache introduces a flexible perturbation-minimization framework that can potentially be extended to a broader class of KV cache compression strategies. For instance, due to the flexibility of structured pruning, the OBCache formulation could be adapted to key-channel pruning by redefining the pruning units from tokens to channels. By leveraging suitable approximation techniques, one could derive corresponding closed-form saliency scores from an output-aware perspective, enabling informed channel-level pruning decisions. Moreover, our approach relies on the diagonal Hessian assumption used in Optimal Brain Damage, which naturally supports token removal. In contrast, the classical Optimal Brain Surgeon framework relaxes this assumption by allowing compensation across pruning units, adjusting the remaining parameters to reduce loss. This perspective aligns naturally with recent cache merging methods and may inspire more theoretically grounded cache management techniques for long-context LLMs.

Overall, OBCache provides a flexible foundation for future research on KV cache compression. Its structured and theoretically motivated formulation opens promising directions toward more principled KV cache management.

9 Experimental Setup↩︎

All experiments are conducted on NVIDIA A100-80GB GPUs. We adapt the KVPress library² [36] with PyTorch [39] to implement the cache eviction algorithms. We use two representative instruction-tuned large language models as backbones: LLaMA-3.1-8B-Instruct [10] and Qwen-2.5-7B-Instruct [3]. Both models employ Grouped Query Attention (GQA) and natively support context windows of up to 128K tokens. All model inference is performed in half precision (bfloat16) without quantization.

For tasks with long prefill prompts, we use FlashAttention-2 [37] for prefill attention computation to reduce memory overhead. Since FlashAttention does not materialize intermediate attention scores, we follow the standard workaround [14], [24], [25] and recompute attention weights for selected query positions when needed to compute saliency scores for token eviction.

9.1 Datasets↩︎

We evaluate OBCache and existing cache eviction methods on three benchmarks: RULER [32], LongBench [33], and the PG19 dataset [34], targeting prefill-stage static cache eviction and decoding-stage dynamic cache eviction, respectively.

9.1.0.1 RULER.

RULER is a synthetic long-context benchmark that expands upon the vanilla Needle-In-A-Haystack [31] test and comprises more comprehensive and challenging tasks. It includes 13 synthetic tasks across four categories: retrieval, multi-hop tracing, aggregation, and question answering. We follow the setup from the official repository³ using the default task configurations. For both LLaMA-3.1 and Qwen-2.5, context lengths of 4K and 32K tokens are evaluated. Each context-length setting contains 1,625 randomly generated samples, with 125 samples for each of the 13 synthetic tasks.

9.1.0.2 LongBench.

LongBench includes 16 datasets across six task categories: single-document QA, multi-document QA, summarization, few-shot learning, synthetic reasoning, and code completion. The average input length is 6,711 words, which is approximately 16K tokens. We do not perform input truncation during our experiments. Following SnapKV [14] and AdaKV [25], cache eviction occurs only during the prefill phase. We report task-specific evaluation metrics (e.g., Exact Match/F1 for QA tasks, ROUGE for summarization tasks) as the scores. Following LongBench’s official evaluation script⁴, we do not apply a chat template for TREC, TriviaQA, SAMSum, Lcc, and RepoBench-P tasks.

9.1.0.3 Language Modeling.

For decoding-phase cache eviction, we adopt the PG19 test set following the setup in StreamingLLM [12]. PG19 contains 100 full-length books, each averaging around 70K tokens. We evaluate on the standard test split⁵ and compute perplexity across varying context lengths (from 1 to 32K), using a fixed KV cache budget of 1024 tokens for all dynamic cache eviction methods.

9.2 Baseline Setups↩︎

9.2.1 Prefill Eviction↩︎

To evaluate baselines for prefill-stage cache eviction, we follow the experimental setup of SnapKV [14] and perform cache eviction only during the prefill phase. In these benchmarks, the prompt length dominates memory usage and serves as the primary bottleneck; therefore, applying decoding-stage eviction has limited impact on task performance. In the query-aware compression setting, the model has access to the task-specific downstream query when evicting tokens. In the query-agnostic setting, the model performs eviction based only on the context KV cache without access to the downstream query (e.g., a question about a document). This setting mimics a more challenging compression scenario under real-world cache reuse settings.

9.2.1.1 H\(_2\)O.

H\(_2\)O originally accumulates attention weights across all historical query positions to make eviction decisions. This is incompatible with FlashAttention-based prefill implementations, as it requires recomputing full attention matrices. To enable efficient implementation, we follow SnapKV’s implementation of H\(_2\)O and accumulate query positions only within a recent perturbation window. For all prefill-stage eviction tasks, the perturbation window size is set to 64. All tokens within the perturbation window are treated as fixed recent tokens, while the remaining cache budget is allocated to heavy hitters.

9.2.1.2 TOVA.

TOVA does not accumulate attention weights and instead makes eviction decisions solely based on the most recent attention distribution. In all prefill-stage eviction evaluations, the entire cache budget is allocated to heavy hitters. In its original implementation, TOVA averages scores across all attention heads and evicts the same token positions for all heads. For fair comparison with other baselines, our implementation instead performs per-head eviction, consistent with the other baselines.

9.2.1.3 SnapKV.

In SnapKV, the same 64-token perturbation window as in H\(_2\)O is used. The key difference is that SnapKV additionally applies a one-dimensional pooling filter to smooth the accumulated attention scores before eviction. For all SnapKV and OBCache-enhanced experiments, we follow the default implementation in KVPress [36], using an average-pooling function with a kernel size of 5.

9.2.1.4 AdaKV.

AdaKV determines adaptive head-wise budgets by globally ranking the top-\(k\) eviction scores across all attention heads. The eviction scores can originate from any prior scoring method. As the baseline configuration, we implement AdaKV using SnapKV scores (accumulated attention within a recent window together with pooling), which serves as our strongest baseline and the current state of the art. To integrate AdaKV with OBCache-enhanced counterparts, we replace the SnapKV scores in AdaKV with OBCache scores using the same 64-token perturbation window and average-pooling function. The safeguard hyperparameter in AdaKV is kept at its default value of 0.2.

9.2.1.5 VATP and CriticalKV.

VATP differs from OBCache-V only in the choice of norm. We replace the \(\ell_2\)-norm in OBCache-V with the \(\ell_1\)-norm to reproduce VATP. CriticalKV adopts a two-stage eviction pipeline. In the first stage, half of the KV budget is selected using a prior attention-only score (SnapKV). In the second stage, the remaining budget is allocated using the value-norm-scaled score. We reproduce CriticalKV using the implementation of KVPress.

9.2.2 Decoding Eviction↩︎

To evaluate baselines for decoding-stage cache eviction, we follow the setup of StreamingLLM [12], where eviction decisions are made at every decoding step. All eviction methods are evaluated under a 1024-token cache budget.

9.2.2.1 StreamingLLM.

For StreamingLLM, we retain the first 4 tokens as attention sinks and always preserve the most recent 1,020 tokens. Since attention sinks are essential for maintaining long-context generation quality, we also retain the first 4 sink tokens in all other methods.

9.2.2.2 H\(_2\)O.

At each decoding step, H\(_2\)O accumulates attention weights across all historical query positions to make eviction decisions. In all perplexity experiments involving H\(_2\)O and its OBCache-enhanced variants, a 256-token recent window is always reserved. The remaining 764 tokens are dynamically selected using their respective saliency scores.

10 Experimental Results↩︎

10.1 Efficiency Evaluation of OBCache Scores↩︎

To analyze the additional computation overhead introduced by OBCache scores, we compare their complexity against purely attention-based scores [11], [13], [14], as summarized in Table 2. For OBCache-V, the only additional operation beyond attention-based scores is computing the norm of the value state, which adds a negligible linear term in the per-head hidden dimension. OBCache-K and OBCache-V&K require computing the norm of the difference between the value vector and each attention output vector within the perturbation window, resulting in a complexity of \(O(W d_{\textrm{head}})\). This is more expensive than OBCache-V. However, in practical settings, the perturbation window \(W\) is often very small (e.g., \(W = 64\) in the prefill phase and \(W=1\) during decoding), so the additional overhead remains minor relative to the overall model computation.

We empirically benchmark peak memory usage and latency overhead during prefill and decoding with varying context lengths and batch sizes on a single A100-80GB GPU, as shown in Figure 6 and Figure 7. For each configuration, we measure: (1) the average time to first token (TTFT) in seconds, including the cache eviction operation, and (2) the per-token decoding latency where scoring and eviction are dynamically updated.

Across both single- and multi-batch settings, the results show that OBCache-V has nearly identical decoding latency (\(<\)​2ms) to attention-based methods across all context lengths. While OBCache-K and OBCache-V&K introduce additional latency (\(<\)​15ms), their cost remains substantially lower than full-cache decoding and does not scale linearly with context length, demonstrating practical efficiency. In the prefill phase, all eviction methods show negligible additional latency compared to the full-KV baseline. These empirical results align with the complexity analysis and demonstrate that OBCache scores can be efficiently computed. We also note that our implementation of score-based cache eviction is built on default PyTorch primitives (e.g., torch.gather and torch.cat) without customized kernels. As a result, at small context lengths, eviction methods exhibit higher latency than full-cache decoding. In the future, developing specialized kernels for score computation and cache eviction would yield additional speed improvements and further enhance practical efficiency.

10.2 RULER and LongBench Full Results↩︎

Due to space constraints, we present the full tables of results for RULER-4K and RULER-32K (both query-aware and query-agnostic) in Tables 3, 4, 5, 6. The full tables of results for LongBench (both query-aware and query-agnostic) are illustrated in Tables 7 and 8.

11 Implementation of OBCache↩︎

We provide a code implementation of OBCache in pseudo PyTorch style, as illustrated in Algorithms 8 and 9. These two algorithms demonstrate the computation of OBCache saliency scores and the cache eviction operation in the prefill phase.

References↩︎

---

1. For any two matrices \({\boldsymbol{C}}, {\boldsymbol{D}}\in \mathbb{R}^{s \times d}\), the output tensor via outer product is \({\boldsymbol{C}}\otimes {\boldsymbol{D}}\in \mathbb{R}^{s \times d \times s \times d}\).↩︎

2. https://github.com/NVIDIA/kvpress↩︎

3. https://github.com/NVIDIA/RULER↩︎

4. https://github.com/THUDM/LongBench/blob/main/LongBench/pred.py↩︎

5. https://huggingface.co/datasets/emozilla/pg19↩︎