Locret: Enhancing Eviction in Long-Context LLM Inference with Trained Retaining Heads

## Abstract

Abstract Large language models (LLMs) have shown remarkable advances in supporting long-context comprehension and processing tasks.
However, scaling the generation inference of LLMs to such long contexts incurs significant additional computation load, and demands a substantial GPU memory footprint to maintain the key-value (KV) cache of transformer-based LLMs.
Existing KV cache compression methods, such as quantization, face memory bottlenecks as context length increases, while static-sized caches, such as selective eviction, suffer from inefficient policies.
These limitations restrict deployment on consumer-grade devices like a single Nvidia 4090 GPU.
To overcome this, we propose Locret , an efficient framework for long-context LLM inference that introduces retaining heads to evaluate the causal importance of KV cache units, allowing for more accurate eviction within a fixed cache size. Locret is fine-tuned on top of the frozen backbone LLM using a minimal amount of data from standard long-context SFT datasets.
During inference, we evict low-importance cache units along with a chunked prefill pattern, significantly reducing peak GPU memory usage.
We conduct an extensive empirical study to evaluate Locret , where the experimental
results show that Locret outperforms the recent popular and competitive approaches, including InfLLM , Quantization, SirLLM , and MInference , in terms of memory efficiency and the quality of generated contents — Locret achieves over a 20 × 20\times 20 × and 8 × 8\times 8 × KV cache compression ratio compared to the full KV cache for Phi-3-mini-128K and Llama-3.1-8B-instruct .
Additionally, Locret can be combined with other efficient inference methods, such as quantization and token merging.
To the best of our knowledge, Locret is the first framework capable of deploying Llama-3.1-8B or similar models on a single Nvidia 4090 GPU, enabling 128K long-context inference without compromising generation quality, and requiring little additional system optimizations.
The experimental code and the tool kit of Locret is provided at https://github.com/huangyuxiang03/Locret

## 1 Introduction

Large language models (LLMs) have revolutionized AI development and deployment (Zhao et al., ; Minaee et al., ).
Recent advancements in LLMs’ ability to handle long-context tasks have further unlocked the potential of generative AI.
State-of-the-art LLMs now support significantly extended context lengths, with GPT-4o (OpenAI, ) and Llama-3.1 (Dubey et al., ) handling 128K tokens, Yi (Young et al., ) and Claude-3 (Anthropic, ) supporting 200K tokens, and Gemini-1.5 (Reid et al., ) reaching 10 million tokens.
These advances enable LLMs to tackle complex tasks like multi-hop reasoning (Li et al., ; Schnitzler et al., ), solving needle-in-a-haystack problems (Guerreiro et al., ; Wang et al., ), and powering advanced LLM agents (Qin et al., ; Wang et al., ) and AI-driven operating systems (Mei et al., ).
However, deploying generative inference under long-context settings on consumer-grade GPUs requires innovative algorithmic and system optimizations to handle this new paradigm efficiently.

Compared to traditional short-context LLM inference, long-context LLM inference shifts the computing paradigm in two key ways:
i) increased computational overhead for attention mechanisms: as context length grows, the computation required for obtaining attention scores increases quadratically, which results in a higher ratio of the computational budget in a transformer block;
ii) higher memory footprint for key-value (KV) caching: longer contexts require larger KV caches, which dramatically increases the peak memory usage.
These shifts demand innovative techniques to mitigate computational costs and manage memory usage effectively for long-context LLM inference.

Although various efforts have been made to overcome the
bottleneck
in LLM inference, these approaches fail to enable long-context inference on consumer-grade GPUs.
Models with compact architectures, e.g. MiniCPM-128K (Hu et al., ) and Phi-3-mini-128K (Abdin et al., ), reduce the computational load and memory usage, but cannot alleviate the KV cache burden in long-context scenarios.
Similarly, techniques like LLM model weight quantization (Frantar et al., ; Lin et al., ; Ma et al., ), activation quantization (Dettmers et al., ; Xiao et al., ; Zhang et al., ), or sparsification (Liu et al., ; Zhang et al., ) also fall short in effectively reducing the memory usage to a level supported by consumer-grade GPUs.

Recently, some specific optimizations have been proposed for long-context LLM inference.
For example, sparse attention mechanisms (Jiang et al., ; Ge et al., ; Lou et al., ) attempt to reduce runtime memory
through conduct selected few calculation,
and KV cache quantization (Liu et al., ; Hooper et al., ; Zandieh et al., ) reduces cache size by applying low-bit storage.
These methods can only offer limited compression rate, of which
the core issue is that the KV cache grows linearly with sequence length, and the above methods do not adequately address this problem.
On the contrary, combining chunked prefill with token-dropping techniques (Xiao et al., ; Yang et al., ) could offer a more effective solution, as it maintains a static-sized cache where the memory usage can be bounded.
However, current token-dropping and cache eviction methods (Zhang et al., ; Liu et al., ; Yao et al., ), whose token importance manually designed according to the inner statistics,
suffer from accuracy loss and performance degradation due to inaccuracies in estimating token importance — The weakening correlation between local and global importance as sequences grow exacerbates this issue.
To address these limitations, we propose a lightweight training-based paradigm that provides more accurate token importance scoring to tackle the long-context LLM inference problem. We highlight our contributions below:

Contribution 1:
We propose a lightweight training-based paradigm for selective KV cache eviction for long-context LLM inference, with an offline training cost $<$1 GPU hours.
We tackle the problem of KV cache eviction by a learning based approach.
We introduce the retaining heads, with a small number of additional parameters, fine-tuned on top of the frozen backbone LLM using a minimal amount of data from standard long-context SFT datasets to estimate the causal importance of each cache unit.
Such a training paradigm is able to provide accurate token importance scoring prediction and can be integrated with other efficient inference algorithms, e.g., quantization and token merging.

Contribution 2:
We provide an efficient inference system implementation for Locret.
We integrate the retaining head mechanism into a chunked prefill inference framework, where we maintain a static-size cache set through evicting cache units with low predicted importance to limit the GPU memory usage.
Locret is able to preserve the most important cache units with the trained retaining heads, enabling precise attention approximation without compromising the inference latency.
Locret is also applicable to all transformer-based LLMs and various hardware, as it requires minimal modifications to the model’s inference process and only utilizes dense operators.

Contribution 3:
We conduct an extensive evaluation of Locret, which illustrates that Locret can not only obtain a comparable performance but also maintain inference efficiency.
Locret achieves over a $20\times$ and $8\times$ KV cache compression ratio
for Phi-3-mini-128K and Llama-3.1-8B-instruct, enabling full comprehension of long contexts on consumer-grade devices.
To the best of our knowledge, Locret is the first framework capable of deploying Llama-3.1-8B or similar models on a single Nvidia 4090 GPU, enabling 128K long-context inference without compromising generation quality, and requiring little extra system optimizations.

## 2 Related Work

Efforts in long-context LLM inference can be categorized by algorithm and system optimizations:

Algorithm optimizations. Optimizations aimed at reducing the size of the KV cache can generally be classified into three categories: quantization-based methods, token dropping methods, and sparsity-based methods. Quantization-based methods (Liu et al., ; Hooper et al., ; Zandieh et al., ; Zhang et al., ), which store the KV cache in low-bit representations, require hardware support for these formats and may slow down inference due to the overhead of dequantization. Token dropping methods typically follow two main strategies: eviction or the use of an attention pool. Eviction-based approaches, such as H2O (Zhang et al., ), ScissorHands (Liu et al., ), and SirLLM (Yao et al., ), rank tokens by certain statistical metrics to identify the most influential ones, discarding others to reduce memory usage. Attention pool-based methods (Nawrot et al., ; Rajput et al., ), such as StreamingLLM (Xiao et al., ) and LoCoCo (Cai et al., ), compress multiple adjacent KV cache units into a single unit using a specially designed transformation. Sparsity-based methods (Ge et al., ; Jiang et al., ; Yang et al., ; Lou et al., ; Lv et al., ) focus on leveraging the sparsity patterns of attention heads to reduce both computation and I/O. The combination of these approaches can be further enhanced by identifying specific attention patterns for each head and layer (Ge et al., ). For surveys of these methods, please refer to (Yuan et al., ; Kang et al., ; Shi et al., ).

System optimizations. The challenge of long-context inference can also be alleviated from a system-level perspective. Offloading-based methods (Sheng et al., ; Xiao et al., ; Wu et al., ; Lee et al., ) store the KV cache in CPU memory, retrieving only the most relevant chunks to the GPU before computing a new chunk. This approach reduces peak GPU memory usage, though at the cost of slower inference. Hardware-aware algorithms, such as flash attention (Dao et al., ; Dao, ; Shah et al., ) and page attention (Kwon et al., ), exploit GPU architecture (Ghorpade et al., ) to enable more efficient runtime memory management. In addition, reimplementing inference infra-structure in a more efficient programming language (; ; ), or adopting disaggregated inference (Jiang et al., ; Zhong et al., ; Qin et al., ; Hu et al., ), can greatly enhance inference efficiency. Algorithmic optimizations can be seamlessly integrated into such systems (Agrawal et al., ; Lee et al., ). For instance, KTransformers (KVCache.AI, ) adopts the chunked offloading technique from InfLLM (Xiao et al., ). However, system optimizations primarily focus on extending context length by leveraging hardware resources, rather than directly reducing the size of the KV cache.

## 3 Locret

### 3.1 Preliminaries

Transformer architecture.
We define the model inference of transformer-based LLMs as follows.
Given a token sequence $t_{1},t_{2},\cdots,t_{n}$, we denote the output hidden state of layer $i$ as $\textbf{H}^{(i)}$ and $\textbf{H}^{(0)}$ is the embeddings.
Each transformer layer consists of an attention layer and an MLP layer.
We assume the model follows a grouped-query attention (GQA) architecture (Ainslie et al., ), with $h$ query heads and a group size of $g$.
For multi-head attention (MHA), $g$ is set to 1.
The attention score of layer $i$’s head $j$ is calculated by
$\mathbf{A}_{j}^{(i)}=\text{softmax}\left(\mathbf{Q}_{j}^{(i)}{\mathbf{K}_{
\lceil{j/g}\rceil}^{(i)T}}/\sqrt{d_{m}}\right)\cdot\mathbf{V}_{\lceil{j/g}
\rceil}^{(i)}$,
where $d_{m}$ represents the hidden size for each head and
$\left[\mathbf{Q}_{j}^{(i)},\mathbf{K}_{j}^{(i)},\mathbf{V}_{j}^{(i)}\right]=
\mathbf{H}^{(i-1)}\cdot\left[{\mathbf{W}_{j}^{Q(i)}},{\mathbf{W}_{j}^{K(i)}},{
\mathbf{W}_{j}^{V(i)}}\right]$.
Next, we compute
$\mathbf{A}^{(i)}=\left[\mathbf{A}_{1}^{(i)},\cdots,\mathbf{A}_{h}^{(i)}\right]
\cdot\mathbf{W}^{O(i)}$, finally followed by
$\mathbf{H}^{(i)}=\text{MLP}(\mathbf{A}^{(i)})$.

KV cache and chunked prefill.
During the prefill stage, all prompt tokens are processed in a single forward pass, where $\mathbf{Q}^{(i)}$, $\mathbf{K}^{(i)}$, and $\mathbf{V}^{(i)}$ each have a sequence length of $n$. In the decoding stage, only a single token is processed across layers, utilizing the KV cache units to reduce computation. Chunked prefill is a method for reducing peak memory consumption by processing tokens in chunks over multiple passes, with the assistance of the KV cache. Taking both KV cache and chunked prefill into account, the attention calculation can be modified as Equation 1, where $B$ represents the number of tokens processed in a single model pass. For decoding, $B=1$, while for chunked prefill, $B$ corresponds to the chunk size. We denote the attention output for tokens $n+1,n+2,\cdots,n+B$ as $\mathbf{A}[n+1\colon n+B]$, and the attention output for the $k$-th token as $\mathbf{A}[k]$.

$$ $\displaystyle\mathbf{A}{[n+1\colon n+B]}_{j}^{(i)}=\text{softmax}\left(\frac{ \mathbf{Q}{[n+1\colon n+B]}_{j}^{(i)}{\mathbf{K}{[1\colon n+B]}_{\lceil{j/g} \rceil}^{(i)T}}}{\sqrt{d_{m}}}\right)\cdot\mathbf{V}{[1\colon n+B]}_{\lceil{j/ g}\rceil}^{(i)}.$ (1) $$

Cache eviction.
Cache eviction in long-context inference is defined as follows.
Here, we slightly abuse the notation of heads and layers, and treat the key-value vector pair of a single token within one head as the smallest cache unit.
We denote the cache unit for the $k$-th token as $c_{k}=(\mathbf{K}[k],\mathbf{V}[k])$.
Assume a memory budget $b$, representing the maximum number of cache units that can be stored in GPU memory at any given time.
The abstract form of attention can then be written as $c_{k}=f(c_{1},c_{2},\cdots,c_{k-1})$.
With limited cache capacity, this calculation can only be approximated by $\tilde{c}_{k}=f(\tilde{c}_{p_{1}},\tilde{c}_{p_{2}},\cdots,\tilde{c}_{p_{b^{
\prime}}})$, where $b^{\prime}\leq b$, and $p_{1},p_{2}\cdots,p_{b^{\prime}}\in\{1,2,\cdots,k-1\}$.
Intuitively, the number of prior cache units involved cannot exceed the memory budget.
When $b^{\prime}=b$, indicating the cache is full, one cache unit must be evicted.
We select the unit to be evicted using a policy $p_{v}=\text{Policy}(\tilde{c}_{p_{1}},\cdots,\tilde{c}_{p_{b}};\tilde{c}_{k})$.
In such stated problem, the key challenge of cache eviction is to develop an effective policy function that minimizes the approximation error $\|\tilde{c}_{k}-c_{k}\|$.

Figure: Figure 1: The framework of Locret. “$\mathbf{R}$” represents the retaining head. $\text{P}_{i}$ and $\text{A}_{i}$ correspond to the $i$-th prompt token and answer token. “t” represents the time step in chunked prefill, “$b$” represents the budget size, and “$n_{s}$” represents the length of the stabilizers.
Refer to caption: extracted/5896273/pattern.png

### 3.2 Locret Framework

Locret is a training-based KV cache compression framework that works in conjunction with chunked prefill.
As illustrated in Figure 1, Locret operates in two stages: training and inference.
In the training stage, we modify the original LLM by appending a retaining head $\mathbf{R}$ to each attention module.
We then train the retaining heads $\mathbf{R}$ while keeping the LLM backbone frozen.
During chunked prefill inference, the retaining heads $\mathbf{R}$ are used to calculate the importance of each cache unit in the chunk.
We retain the cache units with higher scores, along with stabilizers (i.e., the last tokens), in the cache pool located in GPU memory.
Through this process, the retaining heads $\mathbf{R}$ learn and predict the patterns discovered by existing methods, as detailed in Appendix G.

From a mathematical perspective, cache eviction is performed by assigning each cache unit an importance score that reflects its influence on the calculation of subsequent tokens.
We refer to this estimation as the causal importance score (CIS) since it is computed in a causal manner.
The CIS of cache unit $k$ is calculated as $s_{k}=S(c_{1},c_{2},\cdots,c_{k})$.
By applying top-$b$ sparse attention based on the CIS, we can ensure that the trace (i.e. retaining and eviction) of each cache unit can be fit within a cache with a given memory budget.
Further details can be found in Appendix F.

However, since not all tokens can be stored in the cache simultaneously, calculating the actual CIS on-chip is impractical.
Instead, we use a heuristic approximation for CIS, defined as follows:
$\tilde{s}_{k}=S(\tilde{c}_{p_{1}},\tilde{c}_{p_{2}},\cdots,\tilde{c}_{p_{b^{
\prime}}})$,
where $b$ is the cache budget, all $\tilde{c}_{p_{i}}$ are approximated cached units, and $b^{\prime}\leq b$.
We hypothesize that if the scoring function for causal importance is sufficiently accurate, it will consistently select the most critical cache units, resulting in a negligible difference between the heuristic and the actual score.
Thus, we use the terms heuristic CIS and actual CIS interchangeably.

### 3.3 Training the Retaining Heads

In this section, we introduce Locret’s model architecture modifications and the corresponding training recipe.
We add additional parameters to compute the CIS $s_{k}$ (or $\tilde{s}_{k}$ for on-chip inference) with respect to all previous cache units.
Specifically, we inject a retaining head, consisting of a small MLP, into each layer.
The retaining head predicts the CIS for each head of the corresponding layer based on the concatenation of $[\mathbf{Q},\mathbf{K},\mathbf{V}]$.
Formally, with a slight abuse of notation, let the retaining head for layer $i$ be denoted as $\mathbf{R}$. The CIS at head $j$ of layer $i$ is then calculated as:
$\mathbf{\tilde{S}}=\mathbf{R}([\mathbf{Q},\mathbf{K},\mathbf{V}])=\sigma([
\mathbf{Q},\mathbf{K},\mathbf{V}]\mathbf{W}_{1})\mathbf{W}_{2}$.
In this equation, $\mathbf{W}_{1}\in\mathbb{R}^{(d_{m}+2d_{kv})\times d_{\mathbf{R}}}$ and $\mathbf{W}_{2}\in\mathbb{R}^{d_{\mathbf{R}}\times\frac{h}{g}}$ are the tunable parameters of $\mathbf{R}$, $\sigma$ is the activation function and $\mathbf{\tilde{S}}[k]_{j}$ is the predicted CIS of the $k$-th token at head $j$ of layer $i$.
This architecture implies that the importance estimation for a single head is not performed in isolation but rather considers all heads together.
Note that for GQA models, there are only $h/g$ output values corresponding to the number of heads in the KV cache.

We train the retaining head $\mathbf{R}$s on a small Question-Answer (QA) supervised fine-tuning (SFT) dataset, where each entry consists of a single prompt and one answer.
We define the CIS $s_{k}$ for the $k$-th token as the maximum attention score, before softmax, from all the answer tokens toward the $k$-th token.
Formally, for the $k$-th token at head $j$ of layer $i$, we approximate the predicted value $\mathbf{\tilde{S}}[k]_{j}^{(i)}$ to the ground truth $\mathbf{S}[k]_{j}^{(i)}:=\max_{p}\left(\mathbf{Q}_{j}^{(i)}\mathbf{K}_{j}^{(i)
T}\right)_{p,k}$, where $n_{q}(d)\leq p\leq n_{q}(d)+n_{a}(d)$, and $n_{q}(d)$ and $n_{a}(d)$ represent the lengths of the prompt and answer in data $d$, respectively. For an MHA model with $L$ layers and $h$ heads, the training objective is described in Equation 2.
For GQA models, we take the maximum attention score before softmax across different query heads within the same group as the ground truth for the corresponding KV head.

$$ $\displaystyle\operatorname*{arg\,min}_{\mathbf{W}_{1}^{(i)},\mathbf{W}_{2}^{(i )},i=1,2\cdots,L}\mathbb{E}_{d\in\mathcal{D}}\left[\sum_{i=1}^{L}\sum_{j=1}^{h }\sum_{k=1}^{n_{q}(d)}\mathcal{L}\left(\mathbf{\tilde{S}}[k]_{j}^{(i)},\mathbf {S}[k]_{j}^{(i)}\right)\right]$ (2) $$

The training loss consists of a regression loss and a smoothing loss.
We apply the Smooth-$\mathcal{L}_{1}$ norm between the predicted values and the ground truth.
Since important segments in natural language typically consist of adjacent tokens, we also apply the $\mathcal{L}_{2}$ norm between each pair of adjacent predicted values to enforce smoothness.
The complete training loss for Locret is given by Equation 3.

$$ $\displaystyle\mathcal{L}\left(\mathbf{\tilde{S}}[k]_{j}^{(i)},\mathbf{S}[k]_{j }^{(i)}\right)=\text{Smooth-}\mathcal{L}_{1}\left(\mathbf{\tilde{S}}[k]_{j}^{( i)},\mathbf{S}[k]_{j}^{(i)}\right)+\alpha\mathcal{L}_{2}\left(\mathbf{\tilde{S }}[k]_{j}^{(i)},\mathbf{\tilde{S}}[k+1]_{j}^{(i)}\right)$ (3) $$

Figure: Figure 2: L-Eval scores with different intermediate size of retaining head $d_{\mathbf{R}}$.
Refer to caption: x1.png

From our observations, the training of Locret exhibits strong robustness.
The performance variations shown in Figure 2 and Table 2 are minimal, despite changes in $d_{\mathbf{R}}$ and the dataset.
Details can be found in Appendix D.
Training statistics, including loss dynamics, are recorded in Appendix H.

### 3.4 Inference Implementation of Locret

Figure: Algorithm 1 Locret Inference

During the inference stage, we use the chunked prefill pattern and perform cache eviction based on the predicted CIS.
Since the predicted value $\mathbf{\tilde{S}}[k]_{j}^{(i)}$ depends solely on $\mathbf{Q}[k]_{j}^{(i)}$, $\mathbf{K}[k]_{j}^{(i)}$, and $\mathbf{V}[k]_{j}^{(i)}$, and because attention in decoder-only models is causal, $\mathbf{\tilde{S}}[k]_{j}^{(i)}$ remains consistent once calculated.
Thus, we store the KV cache units along with their corresponding causal importance values.
When the cache is full, we evict the units with lower causal importance values, as they are deemed less useful for future computations.
Cache eviction introduces context discontinuity, meaning some cache units at certain positions may be absent.
This can degrade generation quality and increase the error between the predicted and accurate CIS, as LLMs are typically not trained on such contexts.
To mitigate this, we retain the last tokens of the current chunk at each step of the chunked prefill process, ensuring a local and continuous context to minimize errors.
To demonstrate the effectiveness of this design, we perform an ablation study on the length of stabilizers $n_{s}$, shown in Figure 3.
Smaller $n_{s}$ results in severe performance degradation, and the model fails entirely when stabilizers are absent, as context discontinuity leads to instability in CIS prediction, causing errors in cache eviction and amplifying errors in hidden states.
More details are discussed in Appendix E.

Figure: (a)
Refer to caption: x2.png

We maintain a cache pool with a capacity of $b$ cache units, discarding units that exceed this limit.
For each chunk, the model processes the chunked input tokens alongside the current cache pool.
The newly generated KV pairs and their predicted scores are then concatenated with the existing cache.
Once the cache pool is full, only the $b$ cache units with the highest CIS values are retained.
At each chunked prefill step, except for the final step, we retain the stabilizers, i.e. the last $n_{s}$ cache units.
Additionally, we do not compress the last $n_{loc}$ tokens of the prompt, as they are critical for maintaining high generation quality due to their strong correlation with the query.
Finally, the answer is generated according to the compressed KV cache.
Algorithm 1 provides the pseudocode for Locret, where we formally describe the Locret inference process.

The GPU memory usage during Locret inference can be effectively bounded.
GPU memory for KV cache storage is limited to $\mathcal{O}(b+n_{loc})$, and the runtime memory usage of the attention mechanism is bounded by $\mathcal{O}(B\times(b+B+n_{loc}))$.
For comparison, while processing an input with $n$ tokens, full attention prefill requires $\mathcal{O}(n)$ for KV cache storage and $\mathcal{O}(n^{2})$ for runtime memory, whereas chunked prefill requires $\mathcal{O}(n)$ for KV cache storage and $\mathcal{O}(nB)$ for runtime memory consumption.

## 4 Experiments

In this section, we present the experiments conducted to evaluate the proposed framework, Locret, aiming to address the following questions:

(Q1) Can Locret obtain better end-to-end task performance compared to popular and competitive long-conetext inference approaches within similar or less peak memory?

(Q2) Can Locret improve inference speed compared to other approaches?

(Q3) Can Locret be combined with other efficient inference methods?

### 4.1 Experimental Setup

Models and training dataset.
We evaluate Locret on two long-context LLMs: Phi-3-mini-128K (Abdin et al., ) and Llama-3.1-8B-instruct (Dubey et al., ).
Both models can process up to 128K context tokens and follow MHA and GQA architectures, respectively.
Phi-3-mini-128K has 3.8B parameters, while Llama-3.1-8B-instruct contains 8B parameters, making both suitable for deployment on consumer-grade devices.
We inject retaining heads $\mathbf{R}$ into each layer, setting the intermediate size $d_{\mathbf{R}}$ to 1024 for both models.
The retaining heads are trained on the LongAlpaca dataset (Chen et al., ) for 3000 steps.
Training Locret is lightweight, with the tunable parameters comprising 8% and 2.5% of the total for the two models, respectively.
The complete training process takes 0.47 and 0.80 GPU hours on a single A800 GPU for each corresponding model.
The cache budget $b$ is set to 6000 for Phi-3-mini-128K and 16384 for Llama-3.1-8B-instruct.
More details on hyperparameters for both training and inference, as well as the system environment, can be found in Appendix A.

Benchmarks.
We evaluate Locret on selected subsets of $\infty$Bench (Zhang et al., ) and L-Eval (An et al., ).
For $\infty$Bench, we select R.PassKey, R.Number, E.Sum, E.QA, E.MC, Z.QA, E.Dia, C.Debug, and M.Find.
All selected subsets, except Z.QA, have an average length of approximately 100K tokens, while Z.QA has a longer average length of around 2000K tokens.
We exclude R.KV because it can be easily handled by calling a Python interpreter.
We also exclude C.Run and M.Calc due to their complexity for all methods, including full attention inference.
For L-Eval, we filter out all tasks with an average length shorter than 16384 tokens and evaluate on CodeU, NQ, CUAD, NarrativeQA, QMSum, and SPACE.
L-Eval has a relatively shorter context length compared to $\infty$Bench, allowing us to assess the effectiveness of Locret across different context lengths.
Metrics are reported according to the recommendations of the two frameworks, with further details provided in Appendix A.
We also report the peak memory usage, i.e. the average peak memory measured for the first entry of each task in the corresponding dataset, for reference.

Baselines.
As discussed in Section 2, existing algorithms for memory-efficient long-context inference can be categorized into offloading-based, quantization-based, token-dropping, and sparsity-based methods.
For each category, we select one representative method as the baseline.
We compare Locret against full attention inference (denoted as FullAttn), InfLLM (Xiao et al., ), KV cache quantization (Turganbay, ), SirLLM (Yao et al., ), and MInference (Jiang et al., ).
For quantization, we use Hugging Face Quanto () implementation, referring to the 2-bit quantization method as HF-2bits.
We omit HF-4bits and benchmark this combination in Section 4.4.
We do not include attention pool-based token-dropping methods in this benchmark, as they are orthogonal to our approach; further discussion is provided in Section 4.4.
Detailed introductions to the selected baselines can be found in Appendix A.
We also discuss the comparison between the trained Locret and the randomly initialized retaining heads $\mathbf{R}$ in Appendix B.

### 4.2 End-to-end Benchmark

We compare our method with the baselines on both $\infty$Bench and L-Eval to address Q1.
As shown in Table 2, Locret outperforms all baselines in terms of end-to-end performance.

In the $\infty$Bench tests, while all methods experience performance degradation compared to FullAttn inference, Locret, InfLLM, and MInference exhibits better performance than other methods, with only a modest drop in performance given the reduced memory usage.
Quantization, on the other hand, shows significant degradation and fails on all tasks.
SirLLM performs well on comprehensive tasks such as E.Sum and E.MC, but struggles with tasks that require precise memory, such as R.PassKey and R.Number.
Locret not only excels in context retrieval tasks but also achieves strong results in comprehensive tasks, earning the highest overall score among all competitors.

In the L-Eval tests, all methods show some degree of performance degradation.
Nevertheless, Locret achieves the best overall result, obtaining the highest score on most tasks.
L-Eval is a shorter but more complex dataset, where SirLLM performs particularly well.
Quantization fails on most tasks, resulting in the lowest overall score.
Both InfLLM and MInference suffer significant performance drops compared to FullAttn inference.
Locret consistently surpasses all competitors.

We also report memory consumption in Figure 4.
In the extreme long-context scenario ($\infty$Bench), Locret uses relatively less memory while achieving the best overall performance.
InfLLM performs well with limited GPU memory usage, but it requires a significant amount of CPU memory to store the full KV cache.
Quantization and SirLLM can achieve low memory consumption in some settings, but quantization introduces severe performance degradation.
MInference employs sparse attention patterns but does not compress the KV cache.
As a result, its minimum memory requirement equals the sum of the model weights and the full KV cache.
In the shorter context scenario (L-Eval), a similar phenomenon is observed.
For Phi-3-mini-128K, which has a larger KV cache, InfLLM and MInference exhibit higher memory consumption due to the need to store the full KV cache.
Other methods have similar memory footprints, with Locret achieving the best overall performance while using the least memory.
For Llama-3.1-8B-instruct, whose full KV cache is smaller, the memory bottleneck shifts to runtime computational memory for attention and other calculations.
All methods exhibit similar memory footprints, with Locret delivering the best overall performance.
Our experiments demonstrate that Locret is both effective and efficient, outperforming all baselines on multiple datasets and models while using less GPU memory.

Figure: Figure 4: Memory Statistics vs. Task Performance. The red lines correspond to the theoretical size of the model weights, while the blue lines represent the total theoretical size of the model weights and the full KV cache without any compression. The purple lines indicate the accuracies of FullAttn. “Total Memory” represents the total memory usage of both GPU and CPU memory.
Refer to caption: x5.png

**Table 2: The experimental results of Locret compared with all the baselines on $\infty$Bench and L-Eval, where higher score represents better performance. “Avg.” represents the average score across all tasks. The highest score in each column is marked in bold, and the second highest is underlined. Locret achieves the highest overall score among all competitors in every setting.**
| Method | R.PassKey | R.Number | E.Sum | E.QA | E.MC | Z.QA | E.Dia | C.Debug | M.Find | Avg.$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Phi-3-mini-128K on $\infty$Bench |  |  |  |  |  |  |  |  |  |  |
| FullAttn | 98.64 | 97.12 | 17.92 | 11.16 | 55.46 | 14.83 | 8.00 | 23.10 | 17.43 | 38.18 |
| InfLLM | 100.00 | 97.12 | 14.35 | 4.97 | 38.86 | 11.04 | 3.50 | 25.38 | 15.14 | 34.48 |
| HF-2bits | 0.00 | 0.00 | 13.80 | 1.44 | 1.75 | 0.20 | 0.50 | 0.00 | 0.57 | 2.03 |
| SirLLM | 3.39 | 3.39 | 21.06 | 6.32 | 44.98 | 11.99 | 5.00 | 22.34 | 21.71 | 15.58 |
| MInference | 99.32 | 95.93 | 14.44 | 8.11 | 40.61 | 10.60 | 9.00 | 15.48 | 15.43 | 32.25 |
| Locret | 100.00 | 97.46 | 16.82 | 7.61 | 46.29 | 11.31 | 10.00 | 27.92 | 29.71 | 34.73 |
| Llama-3.1-8B-instruct on $\infty$Bench |  |  |  |  |  |  |  |  |  |  |
| FullAttn | 100.00 | 99.32 | 26.79 | 15.06 | 68.12 | 13.40 | 17.00 | 20.56 | 34.00 | 43.81 |
| InfLLM | 100.00 | 100.00 | 24.24 | 14.21 | 51.97 | 10.76 | 11.00 | 26.25 | 35.71 | 41.57 |
| HF-2bits | 36.78 | 6.95 | 8.77 | 4.05 | 27.95 | 3.09 | 5.50 | 13.20 | 22.00 | 14.25 |
| SirLLM | 1.69 | 1.69 | 25.60 | 8.95 | 55.46 | 10.38 | 9.50 | 23.10 | 3.71 | 15.56 |
| MInference | 100.00 | 98.47 | 20.64 | 14.35 | 59.83 | 12.20 | 20.50 | 25.89 | 35.43 | 43.03 |
| Locret | 100.00 | 99.49 | 27.28 | 20.90 | 58.82 | 11.85 | 13.00 | 27.16 | 32.86 | 43.48 |

### 4.3 Speed Test on Real Consumer-Grade Devices

In this section, we examine the processing speed to demonstrate that Locret achieves its strong performance without compromising inference speed, addressing question Q2. We evaluate the inference speed on the R.PassKey task from $\infty$Bench and compare Locret against all the baselines introduced in Section 4.1, using a single Nvidia 4090 GPU with 24GB of memory, which is typical for consumer-grade AI devices.
We report the inference speed as the total number of tokens in the input and output sequences divided by the processing time, along with the accuracy of the measured task.
Since the original settings of some algorithms might lead to Out Of Memory (OOM) errors, we remove some tokens from the middle of the input sequence in those cases, marking these settings with ^∗, and report the valid context length in such scenario.
For settings without ^∗, we maximize the chunk size for higher speed when the method utilizes the chunked prefill pattern.

R.PassKey is a task where the model retrieves a 5-digit number from a large amount of irrelevant text, a task we believe to be relatively simple for humans.
Thus, we consider the task to have failed if the accuracy falls below 95%.
As shown in Table 3, aside from the settings that fail on this task, Locret achieves the highest inference speed among all methods that can correctly process R.PassKey.
Due to its MHA architecture, Phi-3-mini-128K has a larger KV cache, which leads to failures for both HF-2bits and MInference.
Storing the full KV cache on a single 4090 GPU is infeasible, as it requires 48GB of memory.
Although the quantized KV cache is reduced to 6GB, the converting processes between representations requires significant GPU memory for intermediate states, resulting in the failure of HF-2bits.
While InfLLM can run in memory-limited scenarios, its offloading process slows down inference, with I/O becoming the bottleneck in attention calculation.
SirLLM fails due to its inaccurate eviction policy, which cannot correctly identify the 5-digit number.
In the GQA model (Llama-3.1-8B-instruct), which has a smaller KV cache, the quantized cache can fit within the GPU memory.
However, the quantization and dequantization processes become the bottleneck, leading to significantly slower speeds.
The performance of InfLLM, SirLLM, and MInference is similar to that seen with Phi-3-mini-128K.
Although MInference benefits from faster encoding speeds, it fails on this task because it cannot process the entire input sequence at once.
Locret strikes a balance between inference speed and performance, making it a far more suitable solution for long-context scenarios on consumer-grade devices.

**Table 3: Executing R.PassKey on an Nvidia 4090. “tok/s” represents the inference speed, “C.Len” stands for the context length after truncation, and “Acc.” represents task accuracy. “Success” shows whether a setting has $\geq$95% accuracy. The highest score among 128K context is marked in bold.**
| Method | FullAttn | InfLLM | HF-2bits | SirLLM | MInference | Locret | HF-2bits^∗ | MInference^∗ |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Phi-3-<br>mini-128K | tok/s$\uparrow$ | - | 2276.38 | - | 2352.20 | - | 5080.85 | 1098.51 | 4099.92 |
| C.Len.$\uparrow$ | 128K | 128K | 128K | 128K | 128K | 128K | 30K | 14K |  |
| Acc.$\uparrow$ | OOM | 99.83 | OOM | 1.69 | OOM | 100.00 | 0.00 | 13.56 |  |
| Success | ✗ | ✓ | ✗ | ✗ | ✗ | ✓ | ✗ | ✗ |  |
| Llama-3.1-<br>8B-instruct | tok/s$\uparrow$ | - | 2287.66 | 1365.51 | 1589.75 | - | 3209.10 | 3680.06 | 5135.74 |
| C.Len.$\uparrow$ | 128K | 128K | 128K | 128K | 128K | 128K | 30K | 25K |  |
| Acc.$\uparrow$ | OOM | 100.00 | 35.59 | 1.69 | OOM | 100.00 | 26.78 | 20.34 |  |
|  | Success | ✗ | ✓ | ✗ | ✗ | ✗ | ✓ | ✗ | ✗ |

### 4.4 Orthogonality to Other Methods

We evaluate the combination of Locret with quantization and token merging to address Q3.

Figure: Table 4: Quantization with FullAttn and Locret. “M” represents Method and “$-\Delta$” represents the gap of average L-Eval score.

KV cache quantization. According to Zhang et al. (), eviction-based methods like H2O struggle with compatibility when combined with KV cache quantization.
Quantization introduces significant disturbance in the estimation of heavy-hitters, leading to severe performance degradation.
However, Locret is not affected by such issues and can be combined with quantization while maintaining most of its performance.
Here, we compare the performance degradation caused by quantization on Locret with that of the full attention method using the same metrics.
We use Quanto as the quantization backend and report the average L-Eval score with Llama-3.1-8B-instruct as the model backbone.
Table 4.4 shows that the performance drop caused by quantization on Locret is only slightly higher than that observed with the full attention method, indicating that Locret is a quantization-friendly approach. More details of the experiment are provided in Appendix C.1.

Token merging. As described in Section 2, token dropping can also be implemented through an attention pool.
Attention pool-based methods (Xiao et al., ; Cai et al., ; Mu et al., ; Munkhdalai et al., ) merge adjacent tokens or cache units into an attention pool, maintaining a static cache size.
These methods are orthogonal to Locret , as the evicted tokens can be merged into a small cache pool and retained in GPU memory.
We conduct the following experiment to demonstrate that Locret can serve as an effective plug-in scoring function within such frameworks, enhancing performance without increasing memory budget.
We select LoCoCo (Cai et al., ) as a representative of the latest attention pool-based methods.
LoCoCo maintains a cache set consisting of two parts: the heavy hitters and the convolved non-heavy hitters.
During each chunked prefill step, LoCoCo first identifies a set of heavy hitters according to H2O (Zhang et al., ), then applies 1-D convolution to the non-heavy hitters to compress them into a static size.
By replacing H2O’s heavy-hitter scoring function with Locret, we retain the cache units with high CIS and convolve the others.
We compare this combination with standalone LoCoCo and Locret on L-Eval using the Llama-3.1-8B-instruct backbone and report the average score across all selected tasks.
As shown in Table 4.4, Locret achieves a higher score than LoCoCo, and the combined algorithm outperforms both standalone methods.
This suggests that Locret provides a more accurate scoring function compared to H2O, and the two methods complement each other, demonstrating their orthogonality. Further details of the experiment are provided in Appendix C.2.

## 5 Conclusion

We propose Locret, a lightweight training-based method that enables memory-efficient inference of long contexts on consumer-grade devices.
Locret introduces a few additional parameters (retaining heads) to predict the CIS of each cache unit during chunked prefill and performs cache eviction based on the predicted CIS.
We conduct extensive experiments across different models and multiple datasets to compare Locret with major efficient inference techniques, including offloading, quantization, eviction, and sparse attention.
Results show that Locret outperforms all baselines, using less GPU memory and without requiring offloading to CPU memory.
Furthermore, Locret can be combined with other efficient inference methods, such as KV cache quantization and attention pool-based token dropping.
Locret maintains high inference throughput on consumer-grade devices and demonstrates robustness to variations in the intermediate size of retaining heads and training data.
The complete framework of Locret, including both training and inference, highlights its suitability for low-resource computing scenarios.
Locret can be applied to various application scenarios, such as end-side multi-modal model inference and context compression during disaggregated inference, which we plan to explore in future work.
We also aim to test Locret on additional model architectures, including multi-latent models and models of varying sizes.
Additionally, we are highly interested in integrating Locret into end-side deployment infra-structures to enhance inference efficiency, particularly for devices like smartphones and personal computers.

## Acknowledgment

Yuxiang Huang is supported by Tsinghua University Initiative Scientific Research Program (Student Academic Research Advancement Program). We acknowledge the discussion with Ruisi Cai for LoCoCo implementation, Xinrong Zhang for details of $\infty$Bench, Weilin Zhao for implementation issues, Chenyang Song for model sparsity, Shuo Wang for long-context training recipe and Yuan Yao for future works on multimodal LLMs.

## Appendix A Hyperparameters, Environment and Baselines

### A.1 Training

During the training stage, we first insert retaining head $\mathbf{R}$s to each layar.
A retaining head is a small FFN consist of two linear transformations, and the non-linear function is aligned with other non-linears of the conresponding model, with an intermediate size of 1024.
We train the appended retaining head $\mathbf{R}$s on the LongAlpaca for 3000 steps with batch size set to 1 and maximum sequence length set to 10240.
We use the AdamW scheduler (Loshchilov, ) and the learning rate is set to 5e-4.
We conduct the training with a linear learning rate scheduler, whose warmup step number is set to 2000.
The balance factor between two training loss $\alpha$ is set to 0.0025.

### A.2 Inference

**Table 6: Hyperparameters in Locret’s inference stage.**
| Model | $b$ | $B$ | $n_{s}$ | $n_{loc}$ |
| --- | --- | --- | --- | --- |
| Phi-3-mini-128K | 6000 | 3072 | 2500 | 100 |
| Llama-3.1-8B-instruct | 16384 | 1024 | 2500 | 100 |

The inference hyperparameters of Locret is listed in Table 6.
Here, we follow the notations in Algorithm 1.
$b$ stands for the cache budget, $B$ is the chunk size of chunked prefill, $n_{s}$ is the length of stabilizers, and $n_{loc}$ represents the length of locally retained tokens at the end of the input sequence.

Hyperparameters of other baselines are as follows.
For InfLLM, we use the recommended settings for Llama-3 to evaluate Llama-3.1.
Since there is no recommendations of Phi-3-mini-128K, we use the settings for MiniCPM, whose architechture and size is similar to Phi-3-mini-128K, to conduct all the experiments.
For Quantization, we use the official implementation (Quanto backend) of Hugging Face.
For SirLLM, we set the start size to 4, recent size to 1000 for both models.
We set the token entropy size to 6000 and 16384 for Phi-3-mini-128K and Llama-3.1-8B-instruct respectively.
The chunk size of chunked prefill is also 3072 and 1024 for the corresponding model.
For MInference, we utilize the recommended settings for both models.

### A.3 System Environment

For all the experiments except the 4090 experiments in Section 4.3, we use a workstation with 8$\times$Nvidia A800/H800 GPUs and 104 Intel(R) Xeon(R) Platinum 8470 CPUs. The device has 1.0 TB CPU memory. The operating system is Red Hat 4.8.5. We conduct all experiments except the full attention full KV cache inference on a single GPU, and 2 GPUs for full attention settings.

For Section 4.3, we conduct the experiments on a single Nvidia 4090 GPU. The device has 512 AMD EPYC 9754 128-Core Processors and 1.0 TB CPU memory. GPUs and CPUs are connected through PCIe Gen 4, which has 16GT/s transmission speed. The operating system is Ubuntu 9.4.0.

### A.4 Baselines

We compare Locret with full attention inference, InfLLM, Quantization, SirLLM and MInference.
FullAttn inference is performed using vllm (Kwon et al., ), which includes automatic tensor parallelism.
InfLLM is a representative of the offloading-based methods, where the full KV cache is offloaded to CPU, and the most relavant blocks are retrieved to GPU during inference.
For quantization method, we use the Hugging Face implementation of 2-bits KV cache quantization, which is inspired by Liu et al. (), where quantization is conducted along channels instead of tokens.
We denote this method as HF-2bits.
SirLLM is an eviction-based token dropping algorithm, where tokens with low token-entropy is evicted once the cache is fullfilled.
MInference is a typical method of reducing peak GPU memory consumption through rule-based sparse attention, but it does not reduce the size of KV cache.
Note that InfLLM, HF-2bits and SirLLM does not have official implementation on Phi-3-mini-128K, thus we implement these three methods according to the original algorithm.
We only use the short factor of RoPE for InfLLM, and no further model modification is conducted for HF-2bits and SirLLM.

## Appendix B The Effect of Training

**Table 7: The results of Locret compared with randomly initialized retaining head $\mathbf{R}$s on $\infty$Bench and L-Eval.**
| Phi-3-mini-128K on $\infty$Bench |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Method | R.PassKey | R.Number | E.Sum | E.QA | E.MC | Z.QA | E.Dia | C.Debug | M.Find | Avg. |
| Random | 0.00 | 34.00 | 5.09 | 2.68 | 18.34 | 1.54 | 0.00 | 13.71 | 2.57 | 4.92 |
| Locret | 100.00 | 97.46 | 16.82 | 7.61 | 46.29 | 11.31 | 10.00 | 27.92 | 29.71 | 34.73 |

We compare the trained Locret to appending randomly initialized retaining head $\mathbf{R}$s on $\infty$Bench.
The results in Table 7 show that Locret training is effective.
Randomly initialized of retaining heads give random predictions and evict arbitary cache units at each step, resulting the failure on all tasks.

## Appendix C Orthogonality to Other Methods

### C.1 Combination with Quantization

**Table 8: L-Eval scores of FullAttn, FullAttn-4bits, Locret and Locret-4bits. (Detailed)**
| Llama-3.1-8B-instruct on L-Eval |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Method | CodeU | NQ | CUAD | NarrativeQA | QMSum | SPACE | Avg.$\uparrow$ |
| FullAttn | 10.0 | 66.84 | 38.91 | 23.11 | 18.76 | 16.86 | 29.08 |
| FullAttn-4bits | 7.78 | 66.64 | 38.25 | 22.76 | 18.85 | 16.84 | 28.52 |
| Locret | 8.89 | 63.03 | 37.21 | 23.59 | 18.17 | 16.87 | 27.96 |
| Locret-4bits | 4.44 | 63.22 | 36.95 | 22.80 | 18.43 | 16.81 | 27.11 |

We compare the combination of Locret and HF-4bits quantization with the full attention method and the standalong HF-4bits quantization.
We utilize the official implementation of Hugging Face, with Quanto as the backend of quantization.
Other hyperparameters are kept same as described in Section 4.1.
We conduct the experiment on L-Eval and report the average score, with Llama-3.1-8B-instruct backend.
The results in Table 8 shows that the degradation caused by quantization is not significantly high, showing that Locret exhibits good robustness on data representation and it is friendly to quantization.

### C.2 Combination with LoCoCo

**Table 9: L-Eval scores of LoCoCo, Locret and the combination LoCoCo+Locret. (Detailed)**
| Llama-3.1-8B-instruct on L-Eval |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Method | CodeU | NQ | CUAD | NarrativeQA | QMSum | SPACE | Avg.$\uparrow$ |
| FullAttn | 10.0 | 66.84 | 38.91 | 23.11 | 18.76 | 16.86 | 29.08 |
| LoCoCo | 4.44 | 61.10 | 35.84 | 19.83 | 18.15 | 16.71 | 26.01 |
| Locret | 8.89 | 63.03 | 37.21 | 23.59 | 18.17 | 16.87 | 27.96 |
| LoCoCo+Locret | 7.78 | 66.33 | 38.01 | 24.85 | 18.31 | 16.92 | 28.70 |

We compare the combination of LoCoCo and Locret with the standalone methods.
For LoCoCo, we train the convolution head with the size of convolved cache set to 2048.
We extend the context length through chunked prefill training to 64K, which is longer than all tasks’ average input length.
The convolution kernel is set to 21, and we train the newly-added convolution and layer norms for 200 steps, following the original setting.
Since the original Llama-3.1-8B-instruct supports 128K context length, we do not modify its positional embedding.
During Inference, we keep a cache budget size of 16384. In the standalone LoCoCo setting, there are 2048 cache units are convolved, while the others are the heavy-hitters selected by H2O.
In the combined algorithm, we replace H2O to Locret.
We select 14336 cache units with the highest CIS, and convolve the other evicted tokens into 2048 cache units.
In all methods, we set the local length to 0, following the original setting.

## Appendix D Training Robustness

Locret demonstrates high robustness to the training settings, suggesting that there is no need for careful tuning of training hyperparameters or meticulous selection of datasets.
Here, we ablate the intermediate size of the retaining heads $d_{\mathbf{R}}$ and train the retaining head $\mathbf{R}$s on various long-context tuning datasets to demonstrate the stability of results across different training settings.

### D.1 Intermediate Size of the retaining head

We align all the training settings as described in Section 4.1 and only change the intermediate size of retaining heads $d_{\mathbf{R}}\in\{256,512,1024,2048,4096\}$ with the backbone model Phi-3-mini-128K.
The trained model is evaluated on L-Eval and we report the average L-Eval score corresponding to each intermediate size.
Results are listed in Figure 2.
The performance variations among all the settings are minimal compared to the changes in the intermediate size, surpassing all baselines in Table 2.
This indicates that out method exhibits good performance stability regardless of the intermediate size of the retaining head $\mathbf{R}$s.

**Table 10: L-Eval scores with different intermediate size of the retaining head $d_{\mathbf{R}}$. (Detailed)**
| Phi-3-mini-128K on L-Eval |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| $d_{\mathbf{R}}$ | CodeU | NQ | CUAD | NarrativeQA | QMSum | SPACE | Avg.$\uparrow$ |
| 256 | 8.89 | 51.52 | 23.05 | 16.21 | 15.26 | 13.77 | 21.45 |
| 512 | 6.67 | 50.61 | 23.33 | 16.67 | 15.02 | 14.23 | 21.09 |
| 1024 | 8.89 | 51.49 | 22.23 | 16.42 | 14.86 | 14.06 | 21.33 |
| 2048 | 7.78 | 54.09 | 21.91 | 16.46 | 15.00 | 13.89 | 21.52 |
| 4096 | 10.00 | 52.33 | 23.52 | 16.15 | 14.81 | 14.02 | 21.81 |

We train different retaining head $\mathbf{R}$s with $d_{\mathbf{R}}\in\{256,512,1024,2048,4096\}$.
We keep all the other hyperparameters same, and train on the same dataset.
From Table 10, Locret shows stability to the intermediate size, in both overall performance and the performance of each single task.
While increasing the intermediate size, we observe very slight overall performance enhancement.
However, the performance variance is negligible compared to the increase of parameter size, thus we choose to maintain the intermediate size in a small scope to take balance of performance and efficiency.

### D.2 Training Data Insensitivity

We also consider the sensitivity of the training data, which leads us to ablate the training dataset by training on LongAlign (Bai et al., ) and Anti-Haystack (Pan, ), comparing these results with those from LongAlpaca (Chen et al., ) in the original training setting.
We also align other settings to the original setting and choose the backbone model to be Phi-3-mini-128K.
We report the average L-Eval score for each training dataset.
The results in Table 2 shows that Locret has high insensitivity towards different training data.
The performance impact of different data recipes is minimal, indicating that our method can be trained on any long-context tuning dataset.

**Table 11: L-Eval scores of Locret trained on various dataset. (Detailed)**
| Phi-3-mini-128K on L-Eval |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Dataset | CodeU | NQ | CUAD | NarrativeQA | QMSum | SPACE | Avg.$\uparrow$ |
| LongAlpaca | 8.89 | 51.49 | 22.23 | 16.42 | 14.86 | 14.06 | 21.33 |
| LongAlign | 10.00 | 55.13 | 21.34 | 16.40 | 15.01 | 14.09 | 22.00 |
| Anti-Haystack | 8.89 | 52.91 | 20.87 | 13.73 | 13.84 | 14.10 | 20.72 |

We conduct training on various datasets and benchmark the trained weights on L-Eval with Phi-3-mini-128K backbone, to show the stability towards training datasets.
For each datasets, we set the training hyperparameters same and truncate the context to 10240 tokens.
We train the first 3000 steps of LongAlpaca and LongAlign.
Since Anti-Haystack is a relatively smaller dataset, we utilize the whole dataset, which consist of 2424 entries.
The results in Table 11 shows that different training dataset recipe exhibits minor effect towards the overall performance.
Locret can obtain competitive performance without delicately selecting the training data.

## Appendix E Discontinuous Context and Stablizers

Evicting cache units results in context discontinuity, which causes unstable CIS prediction and inaccurate calculation of later tokens.
Thus, we always retain the stabilizers, which are consist of the last $n_{s}$ cache units in each chunked prefill step.
We ablate $n_{s}$ on R.Number of $\infty$-Bench in the proposed algorithm to demonstrate the necessity of incorporating stabilizers in the design.
The results in Figure 3(a) show that lower stabilizer length $n_{s}$ causes severe performance degredation and the model fails completely when the stabilizers are absent.
We report the maximum absolute error of the last hidden state of the input prompt across different layers in Figure 3(b).
Large errors can be observed when the stabilizers are short or absent.
We also report the mean absolute error of the predicted causal importance values with different stabilizer lengths, compared to the case without evicting any cache units, in Figure 3(c).
We also observe high errors when the stabilizer length is limited.
This explains the reason for failure when the stabilizers are short or absent: context discontinuity leads to instability in the prediction of CIS, resulting in errors during cache eviction and amplifying errors in the hidden states.

## Appendix F Causal Importance Score Simulates a Cache Problem

In this section, we show that assigning each cache unit a CIS and calculate each cache units with top-$b$ cache units simulates a cache problem, i.e. the calculation process can be done in a cache. Thus, Locret mathmatically equals to top-$b$ sparse attention.

###### Definition F.1 .

(Causal Calculation)
Given a sequence of objects $c_{1},c_{2},\cdots,c_{n}$, if

$\forall 1\leq i\leq n$, $c_{i}=f(c_{1},c_{2},\cdots,c_{i-1})$

then $f$ is a causal calculation. $c_{1},c_{2},\cdots,c_{n}$ is the generated sequence respective to $f$.

For all causal calculations, we can easily split the function into two parts: a selection function and a another function. Formally,

$\forall\text{ causal calculation }$f$,~{}\exists\text{ function }g,Sel,$

$g:2^{\{c_{1},c_{2},\cdots,c_{n}\}}\rightarrow\{c_{1},c_{2},\cdots,c_{n}\},$

$Sel:2^{\{c_{1},c_{2},\cdots,c_{n}\}}\rightarrow 2^{\{c_{1},c_{2},\cdots,c_{n}
\}};~{}X\mapsto Y\subseteq X,$

$\text{s.t. }f=g\circ Sel.$

###### Definition F.2 .

(Causal Importance Score)
Given a causal calculation $f$ and $c_{1},c_{2},\cdots,c_{n}$ is the generated sequence of $f$. $s_{1},s_{2},\cdots,s_{n}\in\mathbb{R}$ is a sequence of numbers. If

$s_{i}=h(c_{i}),$

then $\{s_{i}\}$ is a CIS of sequence $\{c_{i}\}$. $h$ is a causal importance scoring function.

###### Definition F.3 .

(Cache Problem)
Given a causal calculation $f=g\circ Sel$, its generated sequence $\{c_{i}\}$ and a positive number $b\in\mathbb{Z}_{+}$, if $f$ satisfies the following two condion, then $(f,b,\{c_{i}\})$ is a cache problem with budget $b$.

- •
$\forall 1\leq i\leq n,~{}|Sel(c_{1},\cdots,c_{n})|\leq b,$
- •
$\forall 1\leq m_{1}<m_{2}\leq n,~{}Sel(c_{1},\cdots,c_{m_{2}})\backslash Sel(c
_{1},\cdots,c_{m_{1}})\subseteq\{c_{m_{1}+1},\cdots,c_{m_{2}}\}.$

###### Theorem F.1 .

(Calculating cache units with Top-$b$ CIS is a cache problem.)
Given a causal calculation $f=g\circ Sel$, and its generated sequence $\{c_{i}\}$, a CIS $s_{i}=h(c_{i})$ and a positive number $b\in\mathbb{Z}_{+}$, if the selection function $Sel$ satisfies the following condition,

$Sel(c_{1},c_{2},\cdots,c_{i})=\{c_{p_{1}},c_{p_{2}}\cdots,c_{p_{b^{\prime}}}\}
,~{}s_{p_{1}},s_{p_{2}}\cdots,s_{p_{b^{\prime}}}\in\text{Top-}b(s_{1},s_{2},
\cdots,s_{i})$

then $(f,b,\{c_{i}\})$ is a cache problem with budget $b$.

Proof.
(1) For all $i$ of $1\leq i\leq n$, $|Sel(c_{1},\cdots,c_{i})|=|\{c_{p_{1}}\cdots,c_{p_{b^{\prime}}}\}|=|\{s_{p_{1}
},\cdots,s_{p_{b^{\prime}}}\}|$. Since $s_{p_{1}},s_{p_{2}}\cdots,s_{p_{b^{\prime}}}\in\text{Top-}b(s_{1},s_{2},\cdots
,s_{i})$, $|\{s_{p_{1}},\cdots,s_{p_{b^{\prime}}}\}|\leq b$. Thus $|Sel(c_{1},\cdots,c_{i})|\leq b$.

(2) For all $1\leq m_{1},<m_{2},\leq n$,

$$ $\displaystyle Sel(c_{1},\cdots,c_{m_{2}})\backslash Sel(c_{1},\cdots,c_{m_{1}})$ $\displaystyle\subseteq\{c_{m_{1}+1},\cdots,c_{m_{2}}\}$ $\displaystyle\iff\{s_{p_{1}},\cdots,s_{p_{m_{2}}}\}\backslash\{s_{q_{1}}, \cdots,s_{q_{m_{1}}}\}$ $\displaystyle\subseteq\{s_{m_{1}+1},\cdots,s_{m_{2}}\}.$ $$

Assume $\exists s\in\{s_{p_{1}},\cdots,s_{p_{m_{2}}}\}\backslash\{s_{q_{1}},\cdots,s_{
q_{m_{1}}}\}$ but $s\notin\{s_{m_{1}+1},\cdots,s_{m_{2}}\}$. Since $s_{p_{1}},\cdots,s_{p_{m_{2}}}=\text{Top-}b(s_{1},\cdots,s_{m_{2}}),~{}s\in\{s
_{1},\cdots,s_{m_{2}}\}$. Thus $s\in\{s_{1},\cdots,s_{m_{1}}\}$. $s$ is not in the Top-$b$ values of first $m_{1}$ scores, thus there exists $b$ values larger than $s$, denote as $s_{l_{1}},\cdots,s_{l_{b}}$. Then, $s_{p_{1}},\cdots,s_{p_{m_{2}}}=\text{Top-}b(s_{l_{1}},\cdots,s_{l_{b}},s_{m_{1
}+1},\cdots,s_{m_{2}}).$ From this, we can obtain that $\min\{s_{p_{m_{2}}}\}\geq\min\{s_{l_{1}},\cdots,s_{l_{b}}\}>s$, $s\notin\{s_{p_{1}},\cdots,s_{p_{m_{2}}}\}$. Contradiction. Finally, there must be $s\in\{s_{m_{1}+1},\cdots,s_{m_{2}}\}$.
From (1)(2), $f$ satisfies the two conditions of cache problem. Thus, calculating cache units with Top-$b$ CIS is a cache problem.

###### Definition F.1 .

###### Definition F.2 .

###### Definition F.3 .

###### Theorem F.1 .

## Appendix G Retained Patterns of Locret

We investigate the retained patterns of Locret.
We trace the cache units at each attention head through the chunked prefill on R.Number, M.find and E.MC of $\infty$Bench with backbone Phi-3-mini-128K, and investigate the pattern variation among different layers on R.Number.
We display the results in Figure 6 and Figure 6.
The yellow parts are the retained cache, where the y-axis represents cache position and x-axis is the time axis.

Figure 6 shows that the pattern is mostly decided by the tasks, where both heads shows similar pattern in the same task.
In R.Number, we are able to observe a strong signal between token 10000 and 15000, which is the position of the inserted number string, indicating that Locret can identify the potentially answer-related parts by giving high predicted values of CIS.
In M.Find, we can observe the StreamingLLM (Xiao et al., ) pattern, where the tokens at the beginning of the sequence are always important.
This is also mentioned as the $\Lambda$-pattern in MInference.
We can also discover the vertical lines in the middle of the sequence.
This pattern is also approached by MInference (Jiang et al., ) by the pattern “vertical-and-slash”.
In E.MC, H2O (Zhang et al., ) and ScissorHands (Liu et al., ) pattern can be observed, following the assumption that if a token is activated at some point, it will continue to be activated in the consequencing process.
Noticing that the vertical lines always come in groups, which is the fundament of InfLLM (Xiao et al., ) retrieving blocks to calculate. The comparison between two heads also shows that different heads exhibits different features.
Head 22 of layer 11 shows stronger vertical lines at some point, where retained pattern of head 14 layer 11 is more even.
Head 14 of layer 11 also gives stronger signal to the initial tokens, where this effect is less strong in head 22 layer 11.
We also conduct experiments to investigate the patterns across layers.
In Figure 6, we show that the pattern variance of the same head in different layers can be large.
In shallow layers, e.g. layer 1 and 5, the retained cache units appears to be periodical and semantic independent.
However, in middle layers, e.g. layer 13 and 17, the position of the inserted number string is strongly highlighted, indicating that semantic takes over to be the dominant factor.
In the deepest layers, e.g. 21, 25 and 29, the highlighted vertical line at the position of the inserted string becomes more accurate.

The retained pattern at different layers shows various features, which might be a good handle to investigate how LLMs understand and process natural language queries.

Figure: Figure 5: Head patterns across multiple tasks.
Refer to caption: extracted/5896273/patterns_heads.png

## Appendix H The Training Process of Locret

Here, we present changing trend of loss and accuracy during training in Figure 7.

Figure: Figure 7: Training loss and accuracy during the training process.
Refer to caption: extracted/5896273/training_dynamics.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. Could the authors clarify how LOCRET’s retaining heads would handle extremely high-context lengths (e.g., 10 million tokens)? Would additional constraints or modifications be required to manage the scoring of cache units in such contexts?
2. While SIRLLM performs poorly on memory-demanding tasks, it performs well on comprehension tasks. Could the authors comment on potential reasons LOCRET outperforms SIRLLM in these scenarios, particularly when both approaches manage memory through cache eviction?
3. Could the authors provide further insights or examples where the heuristic scoring might diverge significantly from the true causal importance? This would clarify the potential trade-offs in LOCRET's eviction policy.
soundness: 2
strengths: 1. The paper presents a framework combining trained retaining heads with chunked prefill, contributing a distinctive approach to KV cache management in long-context inference. Unlike previous methods, LOCRET’s retaining heads learn a heuristic for cache importance, adapting to specific model architectures and sequence types, which provides greater flexibility across transformer-based LLMs.
2. The empirical evaluation is rigorous, with comparisons across a diverse set of baselines, including INFLLM, Quantization, SIRLLM, and MINFERENCE. The experiments cover both long and shorter context scenarios, supporting the paper’s claims of LOCRET’s superiority in maintaining performance while reducing memory usage.
3. LOCRET offers a good solution for deploying long-context LLM inference on consumer-grade hardware by significantly reducing the KV cache size without compromising quality. This contribution is valuable given the rising importance of long-context LLM applications in various fields.
4. The paper is well-organized, providing a clear explanation of LOCRET's architecture, training process, and the underlying intuition behind retaining heads. Diagrams effectively illustrate the framework and its mechanisms, enhancing reader understanding of the complex process of cache unit scoring and selective eviction.
summary: The paper proposes LOCRET, an framework designed to enhance memory efficiency in long-context large language model (LLM) inference by using retaining heads to score and selectively retain key-value (KV) cache units. The primary challenge addressed is the high computational and memory demands posed by long-context LLM inference, which often limits deployment on consumer-grade devices. LOCRET introduces a trained retaining head mechanism that evaluates and prioritizes cache units based on their causal importance, offering a scalable and efficient approach that maintains inference quality on devices such as Nvidia 4090 GPUs. The paper conducts a comprehensive evaluation, comparing LOCRET with various memory-efficient inference baselines, demonstrating notable improvements in memory compression and inference quality without sacrificing speed.
weaknesses: 1. While the use of retaining heads to score and retain cache units is a valuable idea, the approach may benefit from further differentiation from existing token-dropping and quantization-based methods. Some parts of the scoring approach appear to overlap with traditional token importance estimation techniques (e.g., heavy-hitter approaches). A more comprehensive analysis highlighting LOCRET’s distinctions from similar heuristics in cache management would strengthen the contribution.
2. The results indicate promising efficiency gains but lack granular performance data on how LOCRET’s accuracy scales with different cache budgets across various architectures. Additionally, while the framework shows reduced memory requirements, further evidence on latency and computation trade-offs associated with retaining heads would be beneficial for practitioners evaluating deployment feasibility.
3. Although LOCRET is tested across two LLM architectures, the applicability of this approach to a broader set of LLMs with diverse attention mechanisms (e.g., sparse attention) is not explored in depth. Discussing potential limitations or adjustments required for alternative models would enhance the generalizability of the method.

---

confidence: 5
contribution: 3
initial_rating: 3
presentation: 3
questions: - Regarding the use of a static-sized KV cache in selection-based KV cache eviction methods, can you explain why "the weakening correlation between local and global importance as sequences grow exacerbates this issue"?
- During training, the first loss term merely learns the maximum value of each column in the attention score. How effective would it be to directly use the maximum value of each column as a metric during inference?
- The paper mentions that methods like H2O cannot be effectively combined with KV quantization approaches. What are the actual performances of these methods?
- There are many papers similar to H2O that use attention score statistics for pruning, such as SnapKV and PyramidKV [2]. How does the method proposed in this paper compare with these approaches?
- Is the Stabilizer used only for selecting recent tokens?
- Is the performance improvement in this paper due to the SFT? What would be the effect if SFT were directly applied to the model?
- Should the number of heads in a retaining head be the same as in Query, or should it match Key/Value? If it matches Query, in structures like Grouped-Query Attention where each head's Key/Value corresponds to multiple heads' Query, how did you train this setup?

[1] LLM Knows What You are Looking for Before Generation
[2] Dynamic KV Cache Compression based on Pyramidal Information Funneling
soundness: 2
strengths: - The method proposes a lightweight training-based selective key-value cache eviction paradigm for long-context language model inference, with an offline training cost of less than 1 GPU hour.
- Extensive validation on various datasets confirms the superiority of our proposed method over the baselines discussed in the paper.
- An efficient inference system implementation is provided, integrating a retaining head mechanism into a segmented pre-filling inference framework. It maintains a fixed-size cache set by evicting cache units with low predicted importance, thereby controlling GPU memory usage.
- The paper discusses the inadequacies of existing methods such as KV quantization, which fail to address the overhead caused by linear growth in KV size. Our selection-based KV cache eviction method utilizes a static-sized KV cache and outperforms previous strategies in preserving important KV cache entries.
summary: To address the substantial overhead of KV cache in long-context reasoning with large language models, this paper introduces a novel method named LOCERT for KV cache pruning. LOCERT utilizes a more precise pruning metric called the causal importance score (CIS) to preserve the most significant KV cache entries.
weaknesses: - The proposed method requires additional training, and although the authors claim it only needs one hour, it also utilizes an eight-card A800 server, which is still resource-intensive.
- The novelty of the proposed method is modest. It is unclear why the training of heads to perform KV cache eviction, predicting each KV's importance, and using the causal importance score (CIS) for pruning, is superior to existing methods like H2O.
- The paper lacks a detailed analysis of the causal importance score (CIS) and needs a deeper discussion to explain why this metric effectively reflects the importance of KV cache.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: Please refer to weaknesses.
soundness: 3
strengths: 1.	This paper proposes a training-based KV cache compression framework LOCRET for selective KV cache eviction for long-context LLM inference. The proposed framework on two LLMs outperforms related methods on two LLMs and two benchmarks.
2.	The paper is easy to follow.
summary: This paper proposes a training-based KV cache compression framework LOCRET for long-context LLM inference. The framework introduces retaining heads to evaluate the causal importance of KV cache units, allowing for more accurate eviction within a fixed cache size. The proposed framework is evaluated with two LLMs on Nvidia 4090 GPU.
weaknesses: 1.	The paper claimed “LOCRET is also applicable to all transformer-based LLMs and various hardware”. However, the proposed method is only evaluated with two LLMs (Phi-3-mini-128K and Llam-3.1-8B-instruct) and one hardware platform (Nvidia 4090 GPU).
2.	The proposed framework is validated with ∞Bench and L-Eval. How is the performance on other long-context benchmarks, such as longBench, et al. ?

---

confidence: 5
contribution: 3
initial_rating: 5
presentation: 2
questions: - Could you clarify why there is a significant difference in performance between SirLLM and LOCRET in Table 3? If both methods operate under the same KV budget, the latency bottleneck should primarily stem from the attention operation. What factors contribute to LOCRET’s superior performance despite this similarity?
- Why is it necessary to keep the last $n_s$ caches? Could the retaining head detect and manage these recent tokens effectively? Does this indicate that the retaining head’s predictions are not sufficiently accurate for recent tokens, and if so, what improvements could address this limitation?
- How does LOCRET handle noisy datasets, such as conversational data with inconsistent or off-topic turns? Are there cases where retaining incorrect KV pairs causes irreparable errors during generation, and if so, how does the method mitigate such risks?
soundness: 3
strengths: The motivation of the paper is well-articulated, and the experiments are thoughtfully designed. Specifically:

- The claims are strongly supported by comprehensive experimental results. The framework addresses the core issue of KV cache growth through the use of retaining heads, with detailed benchmarks comparing LOCRET against several existing methods.
- The selective eviction strategy, guided by the use of CIS, is convincingly motivated. The experiments are well-structured, thoroughly exploring various datasets, models, and baselines, providing strong evidence of LOCRET’s effectiveness.
- The empirical evaluations comprehensively assess memory usage, inference speed, and performance across a diverse set of tasks. The results are consistently underpinned by sound theoretical analysis. Additionally, LOCRET facilitates long-context inference on GPUs like the Nvidia 4090, significantly enhancing the accessibility of advanced LLMs on consumer-grade hardware.
summary: This paper proposes LOCRET, a novel framework for long-context LLM inference aimed at reducing GPU memory usage through trained retaining heads. Unlike existing static cache eviction methods, LOCRET uses lightweight training to estimate the causal importance of KV units, achieving more accurate cache eviction. The experimental results demonstrate memory efficiency and competitive generation quality with models like Llama-3.1-8B to perform 128K token inference on a single RTX 4090 GPU.
weaknesses: The core idea of this paper is to develop an effective eviction policy through training retaining heads. However, several weaknesses need to be addressed:

- SirLLM is not an appropriate baseline for evaluating token eviction strategies. SirLLM is designed primarily for multi-turn conversations and is not tested on benchmarks like InfiniteBench or L-Eval. A more suitable baseline for eviction-based methods would be SnapKV [1]. Although chunk prefilling may not align perfectly with SnapKV, the authors could still avoid OOM errors and reduce GPU peak memory usage by employing layer-by-layer token dropping during prefilling.
- The benchmark suite lacks depth, particularly for information retrieval tasks. The retrieval task within InfiniteBench is overly simplistic, comprising repeated sentences that can be trivially discarded. I recommend that the authors incorporate experiments on RULER [2], following the MInference settings, to provide a more meaningful evaluation of retrieval performance.
- Token eviction based methods may struggle in multi-turn conversation scenarios. For example, in key-value retrieval tasks, if the user queries a different key-value pair during a subsequent turn, the model’s accuracy could degrade significantly due to missing context or prematurely evicted tokens.

If the authors address these concerns, I would consider raising my score.

[1] SnapKV: LLM Knows What You are Looking for Before Generation

[2] RULER: What's the Real Context Size of Your Long-Context Language Models?

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 2
questions: 1. Stabilizer Length: Could the authors provide more insight into how the stabilizer length ns affects performance across different models and datasets? Is there an optimal range for ns?
2. Theoretical Justification: Can the authors elaborate on the causal importance score's theoretical properties and explain how it ensures minimal approximation error during cache eviction?
3. Hyperparameter Sensitivity: Have the authors conducted ablation studies on the cache budget b and chunk size B? How do these parameters impact performance and memory usage?
4. Generalization: It might be out-of-scope, but how well does LOCRET generalize to other transformer architectures, such as encoder-decoder models or those with different attention mechanisms?
5. Limitations: Are there specific tasks or contexts where attention pool-based methods might outperform LOCRET? How does LOCRET handle scenarios with severe context discontinuity?
6. Quantization Methods: You mention KV cache quantization techniques, mentioning the computation overhead as their limitation. Could you compare these techniques, e.g., KVQuant, with sparse attention methods such as FastGen?
7. Combination: You mention the possibility of combining your approach with other efficient inference methods. Could you expand on this with results?
soundness: 3
strengths: Novelty: The introduction of retaining heads for estimating causal importance is a novel approach to KV cache management.
Practical Impact: Enables deployment of large LLMs on consumer-grade GPUs without significant performance loss.
Comprehensive Evaluation: Extensive experiments across multiple datasets and models validate the effectiveness of LOCRET.
Compatibility: LOCRET can be integrated with other efficient inference methods like quantization and token merging.
Lightweight Training: Requires minimal additional training time and resources
summary: This paper presents LOCRET, a framework that enhances long-context large language model (LLM) inference on consumer GPUs. LOCRET introduces "retaining heads," lightweight components added to a frozen LLM backbone to estimate the importance of each key-value (KV) cache unit. LOCRET optimizes cache eviction and reduces GPU memory usage during inference by predicting which cache units are crucial. Combined with chunked prefill, it outperforms methods like InfLLM and SirLLM in memory efficiency and generation quality, enabling models like Llama-3.1-8B to run 128K context inference on a single Nvidia 4090 GPU without performance loss.
weaknesses: Clarity of Presentation: The paper contains grammatical errors and unclear notations, hindering understanding.
Theoretical Depth: The theoretical underpinnings, particularly regarding the causal importance score and its properties, could be more thoroughly developed.
Hyperparameter Analysis: Limited discussion on the impact of key hyperparameters (e.g., cache budget, chunk size) on performance.
Limited Discussion of Limitations: The paper does not sufficiently explore potential drawbacks or scenarios where LOCRET may underperform.
Reproducibility: Some essential details for reproducing results are located in the appendix rather than the main text.

Meta-review: The paper introduces LOCRET, a framework for memory-efficient long-context inference in LLMs. LOCRET employs lightweight retaining heads trained to predict the causal importance of KV cache units, enabling selective cache eviction. Using this mechanism with `chunked prefill`, LOCRET achieves efficient inference on consumer-grade GPUS while maintaining strong performance on long-context tasks.

## Strengths

- LOCRET introduces a query-independent mechanism for cache eviction, which addresses the bottleneck of KV cache growth in long-context inference (Reviewer WKXD). The paper enables long-context inference on consumer-grade hardware. 
- The paper includes detailed comparisons across multiple benchmarks and baselines (e.g., SnapKV, SirLLM, and H2O) and explores various hyperparameters like cache budget, stabilizer length, and chunk size (Reviewer WKXD, Reviewer V3dC).
- LOCRET consistently demonstrates strong memory efficiency and performance on a variety of tasks (Reviewer V3dC).
- The framework requires minimal training resources (less than 1 GPU hour) and makes only minor modifications to existing LLM architectures, ensuring diverse applicability (Reviewer xs3N, Reviewer WKXD).

## Weaknesses

- LOCRET has been evaluated primarily on decoder-only architectures and does not generalize to encoder-decoder models or alternative attention mechanisms, such as sparse attention (Reviewer WKXD, Reviewer j7cb).
- Poor performance on complex query-driven tasks like RULER indicates that LOCRET struggles without query-awareness, limiting its applicability in certain retrieval scenarios (Reviewer KmJZ, Reviewer xs3N).
- While the experiments are extensive, some reviewers noted a lack of focus on challenging benchmarks or real-world datasets (Reviewer KmJZ). For example, InfiniteBench’s simplicity and the lack of off-topic turns in multi-turn conversation datasets reduce the impact of the results.
- SnapKV and similar baselines were tested under chunked prefill scenarios, which are not compatible with their design, potentially inflating LOCRET’s relative performance (Reviewer KmJZ).
- LOCRET exhibits performance degradation on complex retrieval tasks like multi-key RULER subtasks, raising concerns about its robustness and generalizability (Reviewer KmJZ).

Based on the reviews and the rebuttal, I recommend **Reject** for this paper. The paper tackles an important problem, however, the paper in its current form has few issues that seem not to be properly addressed in the rebuttal:

- LOCRET fails to generalize beyond decoder-only architectures, struggling with encoder-decoder models and alternative attention mechanisms (Reviewer WKXD, Reviewer j7cb).
- Its poor performance on query-driven benchmarks like RULER, especially on complex multi-key tasks, raises concerns about its applicability in real-world scenarios (Reviewer KmJZ). The results suggest that LOCRET lacks the generalizability needed to handle tasks where query-awareness is critical.
- While the authors run extensive evaluations, some of the chosen benchmarks (e.g., InfiniteBench) are overly simplistic and fail to challenge the proposed method sufficiently (Reviewer KmJZ).
- The comparisons with baselines like SnapKV under chunked prefill scenarios are unfair, as these methods are not designed for such setups (Reviewer KmJZ). 
- The proposed method, while practical, does not demonstrate significant contribution over existing approaches like SnapKV and H2O. The core idea of training retaining heads to predict cache importance lacks sufficient theoretical development and differentiation from prior work (Reviewer xs3N, Reviewer WKXD).
- The paper does not convincingly address challenges in query-driven or multi-turn conversation tasks. LOCRET struggles with off-topic or inconsistent input contexts, and the experiments lack robust evaluation in such scenarios (Reviewer KmJZ).

Decision: Reject