Locret: Enhancing Eviction in Long-Context LLM Inference with Trained Retaining Heads on Consumer-Grade Devices
Authors: Yuxiang Huang
Binhang Yuan
Xu Han
Chaojun Xiao
Zhiyuan Liu
Venue: Accepted by TMLR ()

Abstract: Scaling the input context length of a large language model (LLM) incurs a significant increase in computation cost and memory footprint to maintain the attention key-value (KV) cache.
Existing KV cache compression methods suffer from inefficient compression strategies and limited memory reduction effects, making it difficult for LLMs to conduct long-context inference on consumer-grade devices, especially when inferring long-context stream input.
Such obstacles prevent consumer-grade devices from supporting more complex applications, creating challenges for the democratization of LLMs.
To overcome this, we propose Locret, a framework to create an eviction policy compatible with chunked prefill. By evaluating the causal importance of KV cache units using \textit{retaining heads}, Locret enables precise eviction of cache units, facilitating efficient long-context inference.
In our empirical studies, Locret outperforms the recent popular and competitive approaches in terms of memory efficiency and generation quality
--- Locret achieves up to $20\times$ of KV cache compression ratio within less than $10\%$ performance loss.
Furthermore, Locret achieves 128K+ long-context inference on a single NVIDIA 4090 GPU without compromising generation quality and only costs $<1$ GPU hour of additional training.

Published in Transactions on Machine Learning Research (12/2025)
Locret: Enhancing Eviction in Long-Context LLM Inference with Trained Retaining Heads on Consumer-Grade Devices
Yuxiang Huang'
huang- Baanggya2/0 S wew.nnghuueeuu.cr.. Dept. of Computer Science and Technology, BNRIST, IAI, Tsinghua University, Beijing, China
Binhang Yuan'
biguan
Department of Computer Science and Engineering, HKUST, Hong Kong, China.
Xu Han', Chaojun Xiao, Zhiyuan Liu Dept. of Computer Science and Technology, BNRIST, TAI, Tsinghua University, Beijing, China.
fhan-mu, zej, luay) @tsinghua edu.cn
Reviewed on OpenReview: https: http////pperreeeeee/mmme net/ formm? id= YPVECTBQHE
Abstract
Scaling the input context length of a large language model (LLM) incurs significant increase in computation cost and memory footprint to mainlain the attention key-value (KV) cache. Existing KV cache compression methods suffer from inefficient compression strategies and limited memory reduction effects, making it difficult for LLMs to conduct long-context inference on consumer-gres -eruume--grrde devices, especially when inferring long-context stream input. Such obstacles prevent consumer-grade devices from supporting more complex applications, creating challenges for the democratization of LLMs. To overcome this, we propose LOCRET, framework to create an eviction policy compatible with chunked prefill. By evaluating the causal importance of KV cache units using relaining heads, LOCRET enables precise eviction of cache units, facilitating efficient long-context inference. In our empirical studies, LOCRET outperforms the recent popular and competitive approaches in terms of memory efficiency and generation quality - LOCRET achieves up to 20x of KV cache compression ratio within less than 10% performance loss. Furthermore, LOCRET achieves 128K+ long- context inference on single NVIDIA 4090 GPU without compromising generation quality and only costs < GPU hour of additional training.
1 Introduction
In recent years, large language models (LLMs) have revolutionized generative Al (Zhao et al., 2023; Minaee et al., 2024), and the advancements of LLMs in handling long-context tasks have further unlocked the potential of generative AI. As result, the context lengths supported by state-of-the-arl LLMs have been significantly extended, such as GPT-4o (OpenAI, 2024) handling 128K tokens, Claude-3 (Anthropic, 2024) supporting 200K tokens, and Gemini-1.5 (Reid et al., 2024) even reaching 10M tokens. These improvements enable LLMs to tackle complex applications with extremely long or streaming inputs, such as multi-hop reasoning (Ti et al., 2024a; Schnitzler et. al., 2024), L.I.V-rivven agents (Qin et al., 2024b; Wang et al., 2024), and Al-powered operating systems (Mei et al., 2024). Some recent efforts (Hu et al., 2024b; Abdin al., 2024) have successfully deployed LLMs consumer-grade end-side devices instead of cloud servers and conducted inference with limited context. We envision that unleashing the potential of long-context inference on consumer-grade devices will revolutionize the development. of personalized AI applications and the democratization of LLMs. However, conducting long-context LLM inference on consumer-grude devices remains a challenging problem that requires algorithmic innovations and systematic optimizations.
"Work doue during internship HKUST.
I Findicates corresponding authors.

---

Published in Transactions on Machine Learning Research (12/2025)
As context length scales, the challenge of long-context LLM inference arises from two major aspects: the in- creased computational cost of the attention mechanism and the higher memory footprint due to the key-v (KV) cache. This leads to the failure of traditional optimizations targeting model backbones to provide a suf- ficient solution. Specifically, backbone-targeted optimizations, such as compacting model architectures (Hu et al., 2024b; Abdin et al., 2024) and quantixing model weights (Frantar et al., 2023; Detimers et al., 2022; Xiao et al., 2023; Lin al., 2024), fail to improve the efficiency of attention patterns or the KV cache, as attention's 's quadratic complexity concerning sequence length remains unaddressed. To this end, recent efforts target optimizing attention patterns and the KV cache achieve efficient long-conters inference.
Recent eatentiionttagetingg optimizations, including sparse attention (Jiang et al., 2024a; Ge et al., 2024; Lou et al., 2024) and KV cache quantization (Liu et al., 2024b; Hooper et al., 2024; Zandieh et al., 2024), show promising results to accelerate attention computation and reduce memory footprint. However, they fail fundamentally address the core challenge: the KV cache grows linearly with contert length. Layer-wise chunked prefill combined with attention sparsity (c.g., SNAPKV (Li et al., 2024b)) can alleviate this problem to certain exlent. This technique lypically performs cache eviction after the precise attention computation for each layer and requires access to the entire sequence. It can theoretically support longer sequences by limiting the maximum memory usage to a single layer's KV cache, but it cannot handle streaming input whose length grows continually. The combination of cache eviction methods (Xiao et al., 2024c; Yang et al., 2024) and chunked prefill offers a more promising approach by maintaining a static cache size and supporting streaming input. Yet, existing eviction techniques like H2O (Zhang et al., 2024d) and SNAPKV show significant discrepancies between local importance estimation and global importance estimation, i.c., it is hard to estimate the importance of each token only based its previous tokens. Instead, these methods require large number of subsoquent tokens to make accurate cstimation. Other methods like SIRLLM (Yao et al., 2024) show great local-global estimation consistency but suffer from performance degradation. We present a detailed analysis of the background in Section 3.2.
Here, we list our contributions below:
We propose LOCRET, a light training-based paradigm for selective KV cache eviction in long-context LLM inference. It introduces learnable retaining heads to estimate the causal importance score (CIS) for token selection, with offline training cost of <1 GPU hour. Additionally, we present LOCRET-Q, query-aware variant of LOCRET, slightly modified to handle query-driven lasks (e.g., retrieving value of a given key in a long JSON string).
Contribution 1:
Contribution 2: We provide an efficient inference system implementation for LOCRET, integrating retaining heads into a chunked prefill framework. This integration limits GPU memory usage by evicting low-CTS KV cache units during the prefill process, thereby accelerating the prefill time. LOCRET compatible with all decoder-only LLMs and imposes minimal additional hardware requirements.
Contribution 3: We extensively evaluate LOCRET, demoustrating its ability to achieve comparable perfor- mance with full KV cache while maintaining inference efficiency. LOCRET achieves over 20x and 8x KV cache compression ratios for Phi-3-mini-128k and Llama-3. 1-----instruct,, respectively. Additionally, LOCRET-Q Q accelerates prefill by over 2x on query-driven tasks without significant performance degradation. This framework enables full comprehension of long contexts on consumer-gras grade devices without compromis- ing the generation quality, and introduces minimal additional system optimizations.
2 Related Work
This paper focuses on optimizing long-context. LLM inference. Existing efforts can be categorized into algorithm and system optimization. For more details about LLMs, please refer to the surveys (Zhao et al., 2023; Lu et al., 2024).
A Algorithm Optimizations aim to reduce the size of the KV cache and can generally be classified into three categories: quantization-bast methods, sparsity- based methods, and token dropping methods. Quantization-basal methods (Liu et al., 2024b; Hooper et al., 2024; Zandich et al., 2024; Zhang et al., 2024a) use low-bit values to represent the KV cache, reducing cache memory overhead and improving cache comput- ing efficiency. These quantization-based methods suffer from hardware-oriented operator customization and
2

---

Published in Transactions on Machine Learning Research (12/2025)
additional inverse quantization overhead. Sparsity-based methods (Ge et al., 2024; Jiang et al., 2024a; Yang et al., 2024; Lou et al., 2024; Lv et al., 2024; Xu et al., 2025) leverage the sparsity patterns of attention heads to reduce both computational and 1/O costs. Combining different patterns can yield further optimization by identifying specific patterns for each head (Ge et al., 2024; Jiang et al., 2024a; Xiao et al., 2024b). Trainable sparse altention methods, such as NSA (Yuan et al., 2025) and MoBA (Lu et al., 2025), adapt sparsity pal- terns during training to improve task performance. For more details on sparsity-based mothods, please refer to the surveys (Yuan et al., 2024; Kang et al., 2024; Shi et al., 2024). Although quantization-beln and sparsity- based methods have achieved promising results, they cannot address the issue that the KV cache memory overhead increases linearly with the context length. Eviction-based methods, such as H2O (Zhang et al., 2024d), SCISSORHANDS (Liu et al., 2024a), and SIRLLM (Yao et al., 2024), rank KV cache units by certain statistical metrics to identify the most influential ial units, discarding others to reduce memory usage and speed up attention computation. INFINIPOT (Kim et al., 2024) introduces context distillation combined with chunked prefill, enabling effective long-context processing even under memory-constranng conditions. Pooling-based methods (Nawrol al., 2024; Rajput et al., 2024), especially STREAMINGLLM (Xiao et al., 2024c) and LoCoCo (Cai et al., 2024a), compress multiple adjacent KV cache units into single unit using pre-designed transformations. More important units will merge into compressed units with higher weights. Eviction- Dvietion-based and mooling-based methods drop or merge tokens to maintain static cache size, but struggle with accurate victim selection and optimal pooling function design.
methods
System Optimizations alleviate the challenge of long-context inference from a system-level perspective, by fully considering hardware features. Offloading- based methods (Sheng et al., 2023; Xiao et al., 2024a; Wu et al., 2024; Sun et al., 2024) use CPU memory to store the KV cache and retrieve only the most relevant chunks to GPU memory before inferring a new chunk. Those methods reduce maximum GPU memory usage at the cost of introducing CPU-GPU communication overhead. Hardware-aware methods, such ALAS-AATTENTION (Dao et al., 2022; Dao, 2024; Shah et al., 2024) and PAGE- AOGE-ATTENTION (Kwon et al., 2023), enable more efficient runtime memory management by considering GPU architectures (Ghorpade et al., 2012). In addition, building inference infrastructures with a more efficient programming language (llama.cpp; llama2.c; rust rostforme))), or adopting disaggregated inference (Jiang et al., 2024b; Zhong et al., 2024; Qin et al., 2024a; Hu al., 2024a), can also greatly improve long-context inference efficiency. Since system optimiza- tions primarily enhance efficiency by leveraging hardware resources rather than directly optimizing attention patterns or the KV cache, relying solely on them cannot adequately address the challenges of long-context LLM inference. Several efforts have integrated algorithm optimizations into system optimizations (Agrawal et al., 2023; Lee et al., 2024), such as KTRANSFORMERS (KVCache.AI, 2024) leveraging offloading based on INFLLM (Xiao et al., 2024a), and STARATTN (Acharya et al., 2025) and APB (Huang et al., 2025) integrate approximate attention with sequence parallelism to achieving faster inference.
3 Background
3.1
Preliminaries
Transformer Architecture. Given a loken sequence (4,,.. (4n as the input prompt of the transformer- based LLM, we denote the output hidden states of the i-th layer as ( and denote the input embeddings of the first layer as ()) For each transformer layer, it consists of an attention block and a feedforward neural network (FFN) block. Attention blocks often follow the grouped-query attention (GQA) architecture ( (Ainslie et al., 2023), with & query heads and h/g KV heads, where g is the group size, i.e., heads share the same K V heads. The multi-head attention (MHA) architecture adopted in the original transformer can be regarded as a special GQA (g- = 1). In the i-th layer, the attention score of the j-th query head is formalized as
Q
()
W OD w(OK
(iv [i/s]]
1//l 13/1]
D/ol
A (4) = softmax (M O o
(TT [2/9]/
dm
(()
[i/ol
where dm represents the hidden size for each head. After obtaining the attention score, the output of the i-th atlention block is A = A (4)
, Ald w((0, and the output hidden states is H FFN (0).
3

---

Published in Transactions on Machine Learning Research (12/2025)
Top 101-00%% Consistency (9)
100
Table 1: ooBench scores of H2O, SNAPKV and LOCRET. H₂O and SNAPKV are exccuted in a chunked prefill manner.
H2O
SnapKV
60
SIrLLM
Phi-3-mini-t on ooBench
Locret
Method
R.Number
E.Sum
E.MC
C.Debug
Avg-T
40
FULLArTN
97.12
17.92
55.46
23.10
48.40
0.5K1K1.5K2K2 StW.://.........5544..5......... KKKKKKKKKKK
Prefix Length
3.39
15.35
45.41
20.57
21.18
SNAPKV
2.54
15.44
41.92
21.43
20.33
Locret
97.46
16.82
46.29
29.71
47.57
Figure 1: The Po.1n consistency,
KV Cache and Chunked Prefill. Given the input prompt sequence (t,,... ,tn), during the prefill stage, all prompt tokens are processed in a single forward pass. After the prefill, K()) | ,.. K
h/0 and
Q are stored as the KV cache, whose sequence length is n. During the decoding stage, cach time token is decoded, forward pass is conducted only for this token and decode the next token. In this process, the KV cache is used to avoid redundant attention computation. Chunked prefill is a method for reducing peak memory usage by segmenting input sequence and prefilling tokens chunk by chunk. Considering both the KV cache and chunked prefill, the attention block can be modified as:
Q ......... (i)
Q 1: BJ K[1 +
( [3/9]
A[n + 1: n B1S sof
M
V(L: + 13/00'
(2)
dm
where A/n + 1: + B] denotes the attention output for the tokens {n + 1,..,........................1 B), and B is the number of tokens processed in single forward pass. For decoding, B 3= 1, while for chunked prefill, B corresponds to the chunk size. For the k-th token in the context, its attention output [A)), its key and value vectors are KKk and Vik].
Cache Eviction. In the cache eviction process, we treat the KV vector pair of a single token within one attention head as the smallest cache unit. We denote the cache unit of the k-th token c. (K(A,, V(k)). Assuming memory budget b, representing the maximum number of cache units that can be stored, the rast form of the attention block can be written as ((1, C22 Ck-1). With limited cache capacity, this process can only be approximated by F = f(2p) Ep2. ), where < b, and P1,P2 ... ,.. E {1,2, ..........--- 13. When the cache is full, one unit must be evicted. We select the victim using some policy Po = Policy((p..""" '*" E 2 and the key challenge is develop a policy that minimizes the error l| - cxll-
3.2 The Global and Local Discrepancy of Scoring Functions in Existing Methods
Conducting KV cache eviction in conjunction with chunked prefill offers potential solution for reducing peak GPU memory usage during long-context prefill; such combination utilizes scoring function of KV pair importance to identify essential KVs. Cache importance scoring functions can generally be categorized into two types: causal and non-causal. Existing causal functions, e.g., SIRLLM (Yao et al., 2024), predict cache importance without relying on future information, which, as shown Section 5, leads to subopti- mal performance. Non-causal functions require information from subscquent cache units to determine the importance score of given cache unit, making them dependent on prefilling the entire sequence. Since their requirement for complete sequence information cannot be satisfied in the chunked prefill process, this introduces a significant discrepancy in importance estimation when based on global versus local context. In this section, we analyze two major KV cache eviction methods, H2O (Zhang et al., 2024d) and SNAPKV (Li et al., 2024b), to illustrate this discrepancy and its impact on task performance.
sequence
To quantify the global and local discrepancy, we define the consistency as follows. Given a sequence of n KV pairs C1, Cg. .Cn and scoring function S(c)lc).." , (.m) (representing the importance of KV pair e, with respect to "* , ,,m)), calculate the global importance scores 81, "., Sn, and the local importance
4

---

Published in Transactions on Machine Learning Research (12/2025)
Training
Inference
Old Cache
New Cache
A
0.5
4.1
7.9
Updats
58
QKT
Pg
Pa
softmax(
V
A
49
cd
A
2.6
13.6
3.0
Traning Loss
Pie
28
76
6.5
7.3
4.8
(Qp, Kp. Vp
Q
Vi
Smooth-L,
+arl
25
2.5
7.
3.8
Attention
Wi
W
W
MLP
Attention 23
[Hp.Ha)
5.
Exitiom
Attention with Retaining Head
P. P
A,
Stabilours
The framework of LOCRET. "R" represents the retaining head. P. and A, correspond to the i-th prompt token and answer token. "b" represents the budget size, and "n," represents the length of the stabilizers. For simplicity, our notation here does not reflect the concept of layers.
Figure 2:
scores 81, '2, ... by
$ s S(
....., ca), 5i Slale,,e........
(3)
Then, the consistency of two scores are defined as
k (s, .. (s' ,,,,), [0.
(4)
We show the the consistency of the top 10% essential KV pairs, i.e. Po. and the lask performance evaluated on four representing tasks of ocBench (Zhang et al., 2024b) in Figure and Table 1.
The results highlight that scoring functions requiring future information (H₂O and SNAPKV) suffer from significant discrepancies when subsequent cache units are not considered. The top 10% essential KV identifi- cation of the first 0.5K tokens can only achieve less than 80% and 40% accuracy compared with considerating the complete sequence for H2O and SNAPKV V correspondingly. Such discrepancy leads the failure of H₂O and SNAPKV in accurately retrieving information from the context, particularly in R.Number. Specifically, the model is unable to identify the importance of certain cache units at the time they are first encountered. Our proposed LOCRET, however, avoids such inconsistencics and achioves superior performance.
4 Locret KV Cache Eviction with Causal Importance Score
4.1 Framework of Locret
LOCRET is training-based KV cache eviction framework that works in conjunction with chunked prefill. As illustrated in Figure 2, LOCRET operates in two stages: training and inference. In the training stage, we modify the original LLM by appending retaining head R to each attention module. We then train the retaining heads R while keeping the LLM backbone frozen. During the chunked prefill inference stage, the retaining heads R can obtain the importance each cache unit. 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 R learn and predict the patterns discovered by existing methods, c.g. attention sink in Xiao et al. (2024c) and vertical lines in MINFERENCE (Jiang et al., 2024a), as detailed in Appendix J.
The eviction policy assigns each cache unit importance score reflecting its influence comprehending subsequent context. This estimation is causal, termed the causal importance score (CIS). The CIS of the k-th unit depends only on the preceding units and the k-th unit itself. Due to memory const raints, calculating the exact CIS on-chip impractical. Intuitively, discarding certain KV cache during the attention calcu- lation introduces errors, manifesting as biases in the CIS. However, as long as the importance estimation is sufficiently accurate, the core information can be preserved, and the loss of information in the attention process be minimized. Therefore, the error introduced in the CIS be neglected.
5

---

Published in Transactions on Machine Learning Research (12/2025)
100
Max Abs, Error Lase Hidden State
3.5
*- Layer 8
- Layer 16
9.75
Mean Abs. Error of Predicted as
0.4
Head 8
Head 16
0.50
3.0-
0.45
80
9.50
0.3
Accuracy
2.5-
0.40
60
2.0
9.25
0.2
0.35
0,30
40
- Layer 24
-X- Layer 32
---
Layer
Layer 32
0.5
0.28
20-
0.4
0.26
15.
500 1000 1500 2000 2500 3000 0
0
2000
0
2000
2000
0
2000
n
(a)
(b)
(c)
Figure 3: R. Number with different stabilizer lengths nts (a) Task accuracy under different Ng. (b) Maximum absolute error of the last hidden state.
(c) Mean absolute error of the predicted CIS. We conduct this
experiment on entries 101-120 of R.Number using the Phi-3-mini-128%
4.2 Training 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] for the k-th cache unit. Specifically, we inject a retaining head, consisting of small MLP, into each layer. This design stems from the observation that a lightweight. MLP can effectively learn essential KVs and minimize training loss, eliminating the need for more complex architectures. From observation, such small MLPs do not slow down model inference, with details claborated in Appendix I. The retaining head predicts the CIS for cach head of the corresponding layer based on the concatenation of [Q, K, VJ. Formally, with a slight abuse of notation, let the retaining head for layer i be denoted as R.. The CTS at head j is then calculated as: $ = R(Q,,, VT) = ([Q, K, 1)W1)W2 the tunable parameters of R. and is the activation function, S[A] = [S..],,.... S[k)a/g E R where S[k], is the predicted CIS of the k-th token. This architecture implics that the importance estimation for a single head is not performed in isolation but considers all heads.
In this equation, W1 (2dm .dR and W2 IR &
We train the retaining head Rs small Question-Answer (QA) supervised fine-tuning (SFT) dataset, where each entry consists of a single prompl and one answer. We define the CIS S(k], for the k-th token at head j as the maximun attention score, before softmax, from all the answer tokens toward the k-th token, Formally, given a training instance d, for the k-th token at head j of layer i, we approximate the predicted value S[A] (2) to the ground truth S[K] == max, (Q. ((T") p.k" where ng(d) < S ng + nd(d),, and
(d) and na(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 Equation 3. 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. This design leverages that semantically relevant tokens yield higher attention scores, while answer tokens (e.g., "the") are not. Thus, we use max pooling to filter them out.
L h (d)
ΣΣ A (SIN)) sIN!!)
argmin w(( w 1,2.....
EdED
The training loss consists of a regression loss and smoothing loss. We apply the Smooth-C, norm between the predicted values and the ground truth. Since important segments in natural language typically consist of adjacent tokens, we also apply the L2 norm between each pair of adjacent predicted values to enforce smoothness. The complete training loss for LOCRET is given by Equation 6.
c (|a]())= Smooth-Ca
[ ($1 ($IH)),((H)))) + aL2 ( 11(5)
(6)
From our observations, the training of LOCRET exhibits strong robustness. Despite changes in dg and the dataset, the performance variations shown in Figure 6 and Table 8 are minimal.
6

---

Published in Transactions on Machine Learning Research (12/2025)
Algorithm 1: LOCRET Inference
Input: Model M, Prompl tokens x, Local length Mote Stablizer length nl. Budget b, Chunk size B
Output: Generated tokens Zgen
// Leave the last niz out make sure they chunk posstppsssittstt + spliss_pluullsaea _htitk_h_ x.length() -Roee B) score_cccche cache f- [] 0.
evicted.
K_ K_ _cache, V_ __cache, score for chunk E chunk _ooiik__ostioosseecc do begin_ end pos <- chunk.begin_ pos, chunk.end pos pos
K _chunk, V __chunk, score chunk <- M(æ]bbeinn boctbeeginpppseddd__e pos], K cache, V_ V_cache)
K cache + Sancat((__aahhaa K <ccunn)
__cchee + Concat(V_cachh V "_chunk)
score_ _cache + Concat(score cache, score chunkk
if chunk is not the last chunk then
// Keep the last n. caches to maintain higher context continuity.
score waree__ccessree_cccheeeeee wacche..nnnghhhhnnnscorr scr..scccc_cccheelleeghh)))) +00
indices +- top-b(score top-h(@core_ccach))..in cache).indices
K. cache, cache, score score_eacchee cache K cache indices], V wacchae score sccre__ccche______e____eeeeeeeee indices]
K_ cache, V_ c_ache, score score_cachhe t- (Meele.leggth(()mmmiii......
K cache, V y_cache)
Eyen + M.generate cache, V cccche
return Tyen
4.3 Inference Implementation of Locret
During the inference stage, use the chunked prefill pattern and perform cache eviction based on the predicted CTS. Since the predicted CTS does not rely on subsequent tokens, it remains consistent once calculated. Thus, we store the KV cache units along with their corresponding causal importance values. When the cache full, we evict the units with lower causal importance values, as they are deemed less useful for future computations. Such eviction performed during chunked prefill. When processing new chunk, we first compute its KV cache, concatenate it with the previously retained cache, and evict redundant units to adhere the budget size. Note that we cache the Dre- KV cache and reassign continuous position embeddings from the beginning to enhance context continuity.
Cache eviction introduces context discontinuity, mcaning some cache units at certain positions may be absent, which can degrade generation quality. To mitigate this, we retain the last n, tokens of the current chunk, named as the stabilizers, at each step of chunked prefill, ensuring a local and continuous context to minimize errors. As shown Figure 3, smaller ns 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 evietion and amplifying errors in hidden stales. More details discussed in Appendix H. We provide pseudocode of LOCRET inference in Algorithm 1.
4.4 Processing Query-Driven Tasks
Query-driven lasks are characterized by highly sparse yet query-correlated critical regions within the context. A representative example is the MultikeyNIAHH task from RULER (Hsich al., 2024). Such contexts are inherently challenging compress elfectively without query information. To effectively handle such tasks, introduce LOCRET-Q. Q, a query-aware variant of LOCRET. When training the retaining heads, we prepend the last lg query tokens to the sequence and gather CIS labels, to avoid dismatch between training and inference. At inference, the query is inserted at the sequence start, ensuring its visibility across all eviction actions. This adaptation enables LOCRET-Q to perform query-aware eviction.
5 Experiments
We conduct experiments to evaluate whether LOCRET address the following questions:

---

Published in Transactions on Machine Learning Research (12/2025)
Table 2: Benchmark results on xeBench. " "Avg." represents the average score across all tasks. The highest score in each column is marked in bold, and the second highest is underlined.
Method
R. KassKey
R.Number
E.Sum
E.QA
E.MC
Z.QA
E.
C.Debug
M.Flnd
Avg.t
Phi-3-mini-128 on xx Bench
FULLATTX
98.6420.00
97. 97.1220.00
17.9220.00
11.1610.00
53.4620.00
14.8320.00
8.00g0.00
23. 10.1020.00 oo
17-43 on
38.18 18
INFLLM
100.00
97.12
14.35 uo
4.97.0.00
38.8610.00
11.04.0.00
25.38 3.50
00
15.14 ou
34.48 20.00
HF-2BITS
2.00
0.00
13.8 23.00
1.44g0.00
1.75_0.00
0.20..o.0
0.50g0.00
0.00g0.00
0.57
2.03_00.0
SIILLM
2.09
3.39
21.06
6.32
44.98
11.99
5.00
22.3440.00
15.58 20.71
MINFERENCE
39.32
95.93
14.4420.00
8.11
40.61.0.00
10.6020.00
9.00
16. 15.4820.00
15.43-20.00
34.32220.00
Locret
99.8820.10
91.30
16.7520.85
8.0320.35
43.38
11.11 20.21 20.21
9.00 73
26.65
30.19 19
37.3721..1
Llama-3.
1.mma-33--8BB-nttuuct coBench
FULLATTN
100.00 ou
99.3240.m uc
26.791..u0 on 15.0620.00
68.1220.00
13.40
17.00
20.561 oo
34.00 on
43.81
INFLLM
100.00
100.00+6.00
24.2440.00
14.21+0.00
51.97
10.76
40.00
26.25+0.00
35.71 70.71=o.0
41.57=0.00
HF-2BITS
36.7820.00
6.95
8.7720.00
4.05g0.00
27.95.0.00
3.09-0000
5.50g0.00
13:20
22.00 20:00
14.25.0.00
SIRLLM
1.00
1.600
25 20.60+0.00
8.95
55.46+0.00
10.38+0.00
23.10+0.00 9.50
15.56+0.00 3..0
MINFERENCE
100.00 100000.0 ou
98. 47 10.00
20.64 oo
14.3510.u0
59.8320.0m oo 12.20 uo 20.50
20.00
25.89 oo
35.43 20 ou
43.03
Locret
100.00
99.49
27.28
20.90
58.82
11.85
13.00
27.16
32.86
43.48
Table 3: Benchmark results on L-Eval. "Avg." represents the average score across all tasks. score in each column is marked in bold, and the second highest is underlined.
The highest
Method
CodeU
NQ
CUAD
NarrativeQA
QMSum
SPACE
Avg.1
Phi-3-aini-122 on L-Eval
FULLATIN
8.89
59.14
30.34+0.00
17.59+0.00
16.05=0.00
14.51
24.42=0.00
INeLLM
5.56+0.00
34.32=0.00
14.53
14.80+0.00
13.31
14.81+0.00
16. 16.22:00.0
HF-2BITS
0. 0.00 no
1.69 uo
6 6.40
2.0420.00
2.73
3.34.0.00
2.70
SIRLLM
20.00
37.92=20.0
20.89
40.00
13.70.0.00
14.46g0.00
18. 18.40=0.00
MINFFRENCE
7.78
25.21 20.00
26.64
15.14
15.78
14.87
17.57
Locret
7.78
52.37 20.98
22.34.0.33
15.01
15.53 ra
13.88
21.15
Llama-3. 1--8--nnttrrc on L-Eval
FULLATTN
10.00 ou
66.84_u.0z
38.91
23. 10.00
18.76gu.00
16.86
29.08_u.or
INFLLM
6.67 20.00
54.77
33.7620.00
20.35
17.62
16.73.0.00
24.98_0.00
HF-2BITS
1.10
29. 20.00
18.98
9.46=0.00
14.02
13. 00.00000000
14.52=00.0
SIRLLM
5. 5.5620.00
58. 20.00
25.00
21. 20:00
17.3220.00
16.44.0.00
25.66_0..0
MINFERENCE
4.78+0.00
31.80=0.or
36.93+0.0
19.4440.00
18.14+0.00
16.76 16.70
21.81=0.00
Locret
7.78
63.90 g0.04
27.10
23.5440.0
18.54_o0 g0.32
16.94.0.07
27.99 20.33
QQ1) Can LOCRET obtain better end-to-end task performance compared to popular and compotitive long- context. inference methods using similar or less peak memory?
(22 is LOCRET able to achieve a faster inference speed on consumer-grade grade devices?
Q3) How can LOCRET process query-driven tasks?
Q4) Can LOCRET achieve stable performance across various hyperparameter settings and training recipes?
5.1
Experimental Setup
Models and Datasets. We conduct experiments on two long-context LLMs: Phi-3-mini-1288 (Abdin et al., 2024) and Llama-3. 1-8B-instruct (Dubey et al., 2024). Both models can process up to 128K context tokens and follow MHA and GQA architectures, respectively. The parameter sizes of these two models are also suitable for deployment these two models, and the intermediate size dR is 1,024. The retaining heads are trained on the LongAlpaca dataset (Chen et al., 2024) for 3,000 sleps , with a 5e-4 learning rale, 10,240 sequence length, and a set lo 2.5e-3. 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 an A800 GPU for each corresponding model. More important hyperparameters are listed in Table More details on hyperparameters and system environments can be found in Appendix A.
consumer-grade devices. We inject retaining heads R into each layer of
8

---

Published in Transactions on Machine Learning Research (12/2025)
<Becch
L-Eval
Phi-3-mini-1286
Liama-3 11..1.88---ntttee
Phi-3-minl-11ak
1amm-3.1-88--nnst
Accuracy (%)
40
Accuracy (%)
Accuracy (%6
Accuracy (%)
30
40
20-
10
20
40
60
80
20
30
40
10
20
30
15
20
25
30
Total Memaryy (GB)
Total Memmary (GB)
Total Memory (GB)
Total Memory (GB)
Accuracy (%)
Accuracy (%)
Accuracy (%)
Accuracy (%)
20-
20
10
20
20
10
30
15
20
25
GPU Memary (GB)
GPU Mernary (GB)
GPU Memory (GB)
GPU Memory (GB)
HF-2biits
SirLLM
Minderence
28
SMLLM
Minference
Model Weightss
Model Weiggts Full KV Cashe
FullARtn Accuracy
Model Weights
Model Weights Full KV Cache
FullAttn Accuracy
The red lines correspond to the theoretical size of the model weights, while the blue lines represent the total 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.
Figure 4: Memory Statistics vs. Task Performance.
Benchmarks. We evaluate LOCRET on selected subsets of ooBench (Zhang et al., 2024b) and L-Eval (An et al., 2024). For query-driven tasks, we evaluate LOCRET on RULER (Hsieh et al., 2024). For oeBench, we select R.PassKey, R.Number, E.Sum, E.QA, E.MC, Z.QA, E.Dia, C.Debug, and M.ndnd. 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 casily handled by calling 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 16K tokens and evaluate models on CodeU, NQ, CUAD, NarrativeQA, QMSum, and SPACE. Metrics are reported according to the recommendations of the two datasets, with further details provided in Appendix A. Although temperature sampling can be useful in real-world applications, we adopt greedy generation aligned with baseline methods. Since LOCRET does not modify sampling, we believe this setup is representative. We also report the peak memory usage, i.e. the average peak memory measured for the first entry of each task in oBench and L-Eval, for reference.
Baselines. As discussed in Section 2, existing algorithms for memory-efficient long-context inference can be basee,, and token-dropping methods. For each
categorized into offloading- based, sparsity-based, category, we select one representative method as the baseline. We compare LOCRET against full attention inference (denoted as FULLATTN), INFLLM (Xiao et al., 2024a), MINFERENCE (Jiang et al., 2024a), KV cache quantization (Turganbay, 2024), and SIRLLM (Yao et al., 2024). For quantization, we use HuggingFace Quanto (Hugging-Face implementation, referring to the 2-bit. quantization method as HF-28rTs. We omit HF-4BITS and benchmark the combination with LOCRET in Appendix E. We do not include attention pooling-based methods in this benchmark, as they are orthogonal to our approach, and further discussion about this is provided in Appendix E. When conducting query-driven tasks, we compare LOCRET and LOCRET-Q with H₂O (Zhang et al., 2024d), SNAPKV (Li et al., 2024b), SIRLLM, and INFINIPoT (Kim et al., 2024). Environmental setup and details of the selected baselines can be found in Appendix A.
5.2 End-to-end Benchmark
We compare all the methods on Beench and L-Eval to address Q1. In Table and Table 3, LOCRET outper- forms baselines in terms of end-to-end performance, showing: (1) On oc ocBenhh, while all methods experience performance degradation compared FULLATTN, LOCRET, INFLLM, and MINFERENCE exhibit better per- formance than other methods, with only modest drop in performance given the reduced memory usage. Quantization 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. PassKay and R.Number. LOCRET not only excels in context retrieval tasks but also achieves strong results in com- prehensive tasks, earning the highest overall score among all competitors. (2) On L-Eval, all methods show performance degradation. Nevertheless, LOCRET achieves the best overall performance, obtaining the high- est scores on most tasks. L-Eval is a shorter but more complex dataset, where SIRLLM performs particularly well. Quantization fails on most tasks. Both INFLMM and MINPERENCE, suffer significant, performance drops
9

---

Published in Transactions on Machine Learning Research (12/2025)
Table 4: Executing Rayasskey on an NVIDIA 4090. "tok/s" represents the inference speed, "C.Len" stands for the context length after truncation, and "Ace." represents task accuracy. The highest score among 128K context is marked in bold.
Method
FULLATTN
INFLLM
HF-2BITS
SMLLM MINFERENCE
Locret
HF-2BITS*
MINFERENCE"
tok/st C.Len.t
2276.38
2352.20
5080.85
1098.51
4099.92
Phi-3-
128K
128K
128K
128K
128K
128K
30K
14K
mini-128K
Acc.t
OOM
99.83
0OM
1.69
QOM
100.00
0.00
13.56
tak/st
2287.66
1365.51
1689.75
3209.10
3680.06
5135.74
Llama-3.1- 1- 8B-instruct
C.Len.t
128K
128K
128K
128K
128K
128K
30K
25K
Ace.t
OOM
100.00
35.59
1.69
OOM
100.00
26.78
20.34
compared to FULLATTN inference. LOCRET consistently surpasses all competitors. Notably, LOCRET can even surpass FULLATTN on certain tasks, such as R.Number and C.Debug in oeBench. This is because the critical regions in the input context that contain the correct answers can be distinguished from the surround- ing noise. Evicting irrelevant tokens reduces imput noise, thereby enhancing the model's ability to retrieve the correct information. We believe this phenomenon highlights a beneficial property for long-context processing.
We report memory consumption in Figure 4, showing: In the extreme longecontext scenario (ocBench), LOCRET uses relatively less memory while achieving the best overall performance. INFLLM performs well with limited GPU memory usage, but it requires significant amount of CPU memory to store the full KV cache. HF-2BrTs 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-conlext scenario (I-Eval), memory usage is mainly determined by the KV-cache size. For 1ii13-minn-128K, which employs MHA and thus has a larger cache, MINFERENCE still consumes substantial GPU memory due to the lack of KV reduction, while INFLLM lowers GPU memory by offloading at the cost of higher total memory. SIRLLM achieves memory usage comparable to LOCRET, and HF-2BrTS consumes similar memory but suffers severe performance degradation. For Llama-3. 1-8B-instruct,, whose KV cache is smaller, the bottleneck shifts to runtime memory for attention and other computations, so MINFERENCE and HF-2BIT exhibit notably lower memory usage but significant performance loss. Among baselines with similar memory usage, LOCRET performs the best.
In summary, our experiments demonstrate that LOCRET both effective and efficient, outperforming all basclines on multiple datasets and models while using less GPU memory.
5.3 Processing Speed on Real Consumer-Grade Devices
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 RaPassKey task from cBench and compare LOCRET against all the baselines, using single NVIDIA 4090 GPU with 24GB of memory, which is typical for consumer-grade Al devices. We report the inference speed (the total number of tokens within the input and output sequences divided by the processing time) and the task accuracy. Since the default settings of some baselines can cause out-of-memory (OOM) errors, we simulate a common fallback used in GPU-comstrained scenarios: truncating the long input from the middle. This approach incurs substantial performance degradation due to significant information loss, which also reflects real-world behavior when applicd in practice. We mark these settings with and report the corresponding valid context length. For settings without. *, we maximize the chunk size for higher inference speed.
R. PasaskKe is task where the model retrieves 5-digit number from large amount of irrelevant text, 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 4, aside from the settings that fail on this task, LOCRET achieves the highest inference speed among all methods that can correctly process R. R.PassKey: (1) Due to its MHA architecture, Phi-3-mini- -12KK has larger KV cache, which leads to failures for both HF-2BITS and MINFERENCE. Storing the full KV cache on a single 4090 GPU infeasible, as requires 48GB of memory. Although the quantized KV cache is reduced to 6GB, converting representations requires significant GPU memory for
10

---

Published in Transactions on Machine Learning Research (12/2025)
Table 5: Performance, prefill speed, and decode speed on BULER-128K, The best and second-highest scores among wriction-based methods in cach column are highlighted in bold and underlined, respectively. We randomly generate 20 entries for each task for all settings. Prefill and decode speed are reported in tok/s. FULLATTN is implemented using FTASHATTN.
Method
SG1
SG2
SG3
MK1
MK2
MK3
MV
MQ
VT
CWE
FWE
QA1
QA2
Avg.
Prefill
Decode
FULLATTN
100.00
100.00
100.00
95.00
95.00
75.00
91.25
100.00
60.00
72.00
66.67
75.00
40.00
82.30
4319.95
12.32
MINFERENCE
100.00
95.00
100.00
95.00
15.00
10.00
81.25
93.75
77.00
14.50
70.00
40.00
30.00
63.19
7205.06
2.61
SNAPKV
100.00
95.00
0.00
65.00
40.00
10.00
62.50
90.00
58.00
6.00
46.67
50.00
30.00
50.24
4203.34
36.49
H₂O
20.00
0.00
0.00
0.00
0.00
0.00
3.75
2.50
1.00
0.00
53.33
100.00
15.00
15.04
464.73
44.70
SIRLLM
0.00
5.00
5.00
0.00
5.00
5.00
6.25
8.75
12.00
0.00
80.00
30.00
15.00
13.23
9717.41
40.86
INFINIPOT
100.00
40.00
0.00
15.00
0.00
0.00
15.00
31.25
93.00
26.50
51.67
20.00
15.00
31.34
9099.02
34.36
LOCRET
100.00
45.00
35.00
10.00
5.00
0.00
20.00
17.50
69.00
46.50
73.33
20.00
34.33
9587.09
37.38
LOCRET- Q
100.00
100.00
100.00
70.00
100.00
100.00
70.00
85.00
60.00
34.50
80.00
50.00
35.00
75.78
9587.84
40.11
its intermediate states, resulting in the failure of HF-2BITS. While INFLLM can run in memory-limited scenarios, its offlodding process slows down inference. SIRLLM fails due to its inaccurate eviction, which cannot correctly identify the 5-digit number. (2) In the GQA model (Llama-3. 1.lama33.1-BB -----innttuutt which has a smaller KV cache, the quantized cache can fit. within the GPU memory. However, this method introduces no sparsity in attention computation or memory access, resulting in relatively lower inference speed compared with other methods. The performance of INFLLM, SIRLLM, and MINFERENCE is similar to that seen with Mmi-3-mini-128h Although MINFERENCE benefits from faster encoding speed, it fails this task because it cannot process the entire input sequence at once. LOCRET strikes balance between inference speed and performance, making it far more suitable solution for long-context scenarios on consumer-grade grrsdde devices.
5.4 Locret-Q: Supporting Query-Drive Tasks
We mainly address question Q3 in this section. As introduced by Sun et al. (2024), existing eviction-based techniques exhibit significant. performance degradation when applied to query-driven tasks. Due to this, we evaluate LOCRET against classical eviction- based baselines using the RULER benchmark (Hsieh et al., 2024).
Table 5 shows LOCRET-Q against SNAPKV (Li et al., 2024b), H₂O (Zhang et al., 2024d), SIRLLM (Yao et al., 2024), and INfINIPOT (Kim et al., 2024) on RULER with a 128K context length. For reference, we also include results for FULLATTN and MINFERENCE. Metrics include prefill speed, decode speed, and task performances are reported. We randomly generate 20 entries for each task for all methods. As shown in Table 5, all elficient inference methods exhibit performance degradation to some extent. However, LOCRET- Q outperforms other eviction-based methods and even surpasses MINFERENCE, demonstrating its effective- ness on query-driven tasks. SNAPKV shows performance degradation, while H₂O and SIRLLM fail com- pletely on RULER. To further compare methods that do not fail completely, we conduct a higher-preciste RULER benchmark (Table 7), generating 500 random entries per task. SNAPKV. For speedup, methods combining eviction and chunked prefill (LOCRET-Q, LOCRET, SIRLLM, and INFINIPOT) significantly reduce prefill time, achieving > 2x speedup over FULLATTN. SNAPKV cannot accelerate prefill due to no computation reduction. H₂O suffers from extremely slow prefill as it relies on full- sequence attention scores, incompatible with efficient implementations like FLASH-ATTN. Decoding speeds are similar across eviction-based methods, and all of them are faster than FULLATTN. Notably, LOCRET fails on RULER, showing a gap compared to LOCRET-Q, highlighting the necessity of query-awareness for query- centric tasks. A simple modification to LOCRET unlocks its potential for such tasks. To ablate the effect of placing the query at the beginning of the input sequence, we conduct an ablation study presented in Table 6. The results show that not all methods benefit. from this modification; in contrast, LOCRET benefits the most, significantly outperforming the others. This indicates that the superior performance of LOCRET-Q is not due to early access to the query, but rather stems from its query-aware KV cache eviction mechanism.
outperforms
5.5 Hyperparameter and Training Robustness Analysis
To adderss Q4, we first examine the inference hyperparameters: budget, stabilizer length, and chunk size.
11

---

Published in Transactions on Machine Learning Research (12/2025)
Table 6: Ablation on putting the query the front of the input. We ran- Table 7: RULER-1280 of 500 domly generate entries for each task for all settings. The average accuracy entries for each task. of RULER-12SK is reported.
Method
RULER-I2SH
Query-Frout FULLATTN MINFERENCE SNAPKV HzO SIRLLM INFINIPOT LOCRET
FULLATIX
82.20
82.30
63.19
50.24
15.04
13.23
31.34
34.33
MINFERENCE
72.97
75.66
54.19
48.09
8.89
13.09
32.47
75.73
SNAPKV
48.76
AVV
-6.64
-9.00
-2.15
-6.15
-0.14
1.13
41.40
LOCRET-Q
75.54
Budget. To evaluate the robustness of LOCRET under different budget constraints, we compare the proposed method with SNAPKV (Ti al., 2024b) using chunked prefill on LongBench (Bai et al., 2024b). As shown in Figure 5a, when the budget size increases, LOCRET demonstrates a faster performance improvement compared to SNAPKV, and consistently outperforms SNAPKV.
Stabilizers Length. As discussed in Figure 3, stabilizers play a crucial role context retrieval tasks. However, NLU tasks, the stability of n, remains relatively high. We evaluate QMSum with different stabilizer lengths na, with the budget set at 6000. As illustrated in Figure 5b, performance remains consistent when n, is small. The observed performance degradation at larger ng values is due to the reduced space available for other cache units.
Chunk Size. Executing long-context inference on hardware with varying GPU memory limitations choices of chunk size. When the chunk size changes, LOCRET shows stable performance. We test on the NQ dataset from L-Eval using multiple chunk sizes ranging from 256 4096. The results, shown in Figure 5c, highlight the stability under various chunk sizes.
longBeench Score
60
40
SnapiV
Locret
30
1000 1500 2000
3000 Budget
2500
3500
1000 2000 3000
5000
6000
1000
2000
3000
4000
(b)
(c)
Figure 5: Scores of LOCRET under (a) various budgets; (b) various ni5 (c) various chunk size.
We further investigate LOCRET'S training stability by experimenting with various intermediate sizes dR of the retaining head and conducting training under different datasets. LOCRET demonstrates high robustness these variations in both intermediate size and training recipes, indicating that it does not require careful hyperparameter luning or meticulous dataset selection.
25.0
25,22
Table 8: L-Eval scores of LOCRET trained on various datasets.
22.5
21.09
21333
21.81
Lobt
200
175
Dataset LongAlpaca LongAlign Anti-Haystack
255
512
1924
2048
4095
L-Fval
21.33
22.00
20.72
Figure L-Eval scores with different dR.
5.6 Additional Experiments
Additional experiments are included in the appendices. We evaluate LOCRET on LongBench in Appendix B, explore the combination of LOCRET with other methods in Appendix E, challenge LOCRET with extremely long inputs in Appendix F, and benchmark LOCRET in multi-turn conversation scenario Appendix G.
6 Conclusion
We propose LOCRET, a lightweight training- based method that enables memory-efficient long-context LLM inference on consumer-grade devices. LOCRET introduces retaining heads predict the CIS of each cache unit during chunked prefill and performs accurate cache eviction. We conduct. extensive experiments across
12

---

Published in Transactions on Machine Learning Research (12/2025)
different models and multiple datasets to compare LOCRET with major efficient inference techniques, and results show that LOCRET outperforms all baselines, using less GPU memory and without requiring offloading to CPU memory. LOCRET-Q, a query-aware variant of LOCRET, can further process query-centric tasks without significant performance degradation. Future work will involve testing LOCRET on other model architectures, e.g. encoder-decoderss and evaluating LOCRET on more personal devices, e.g. NVIDIA Jetson.
Acknowledgment
This work is supported by the HKUST startup grant R9895 from CSE,RGC-ECS project 26218024, RGC- NSFC project CRS HKUST601/24 and the high-quality development project of MIIT. Yuxiang Huang is supported by Beijing National Science Foundation (No. QY24253) and 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 ooBench, Weilin Zhao for implementation issues, Chenyang Song for model sparsity, Shuo Wang for long-context training recipe and Yuan Yao for future works multimodal LLMs. We also thank Minsoo Kim for discussions on InfiniPot.
References
Marah Abdin, Sam Ade Jacobs, Ammar Ahmad Awan, Jyoti Aneja, Ahmod Awadallah, Hany Awadalla, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Harkirat Behl, et al. Phi-3 technical report: highly capable language model locally on your phone. arXiv:2404.14 14219, 2024.
Shantanu Acharya, Fei Jia, and Boris Ginsburg. Star attention: Efficient Ilm inference over long sequences. Proceedings of ICML, 2025.
Amey Agrawal, Ashish Panwar, Jayashree Mohan, Nipun Kwatra, Bhargav S Gulavani, and Ramachandran Ramjee. Sarathi: Efficient Ilm inference by piggybacking decodes with chunked prefills. orXiu:2308. 16369, 2023.
Joshua Ainslic, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. Proceedings of EMNLP, 2023.
Chenxin An, Shansan Gong, Ming Zhong, Xingjian Zhao, Mukai Li, Jun Zhang, Lingpeng Kong, and Xipeng Qiu. L-eval: Instituting standardized evaluation for long context language modcls. Proceedings of ACL, 2024.
Anthropic.
The claude 3 model family: Opus, sonnet, haiku, 2024. URL https://www-cod antoroointoiiic... htt.aaaaabbb.aaaaaaaaaaaaaaaaae. Mod.:e//00.006667721ddee. /atp:s2277/odee_Carerrr _aa____aardd________ 2laude_3.pde
Yushi Bai, Xin Lv, Jiajic Zhang, Yuze Hc, Ji Qi, Lei Hou, Jie Tang, Yuxiao Dong, and Juanzi Li. Longalign: A recipe for long context aligument of large language models. Proceedings of EMNLP, 2024a.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. Proceedingss of ACL, 2024b.
Ruisi Cai, Yuandong Tian, Zhangyang Wang, and Beidi Chen. Lococo: Dropping in convolutions for long context compression. Proceedings of ICML, 2024a.
Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, et al. Pyramidkv: Dynamic kv cache compression based on pyramidal information funncling. arXiv:2406.0000 02066,, 2024b.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient. fine-tuning of long-context large language models. Proceedings of ICI.R, 2024.
13

---

Published in Transactions on Machine Learning Research (12/2025)
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. Proceedings of ICLR, 2024.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-ofhcient exact attention with io-awareness. Proceedings of NeurIPS, 2022.
Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm.int8(): 8-bit matrix multiplication for transformers at scale. Proceedings of NewrIPS, 2022.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv:2407.21788, 2024.
Elias Frantar, Saleh Ashkboos, Torsten Hoefter, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained translormers. Proceedings of TCLR, 2023.
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what discard: Adaptive ky cache compression for llms. Proceedings of ICLR, 2024.
Jayshree Ghorpade, Jitendra Parande, Madhura Kulkarni, and Amit Bawaskar. Gpgpu processing in cuda architecture. Proceedings of ACIJ, 2012.
Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length Ilm inference with kv cache quantization. arXim:2401.0 18079, 2024.
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv:2404.066 066544 2024.
Cunchen Hu, Heyang Huang, Liangliang Xu, Xusheng Chen, Jiang Xu, Shuang Chen, Hao Feng, Chenxi Wang, Sa Wang, Yungang Bao, et al. Inference without interference: Disaggregate Ilm inference for mixed downstream workloads. 2rX1::441.11188, 2024a.
Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang. Yuxiang Huang, Weilin Zhao, el. al. Minicpm: Unveiling the potential of small language models with scalable training strategies. Proceedings of COLM, 2024b.
Luyang Huang, Shuyang Cao, Nikolaus Parulian, Heng Ji, and Lu Wang. Efficient attentions for long document summarization. arXiv preprint arXiv:2104.0211 2021.
Yuxiang Huang, Mingye Li, Xu Han, Chaojun Xiao, Weilin Zhao, Sun Ao, Hao Zhou, Jie Zhou, Zhiyuan Liu, and Maosong Sun. Apb: Accelerating distributed long-context inference by passing compressed context blocks across gpus. Proceedings of ACL, 2025.
Hugging-Face. URL httpa://githhbe http://tww....coggggggg////t/aaa
Huiqiang Jiang, Yuchong Li, Chongruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin-Yew Lin, et al. Minference 1.0: Accelerating pre-filling for long-context 1lms via dynamic sparse attention. Proceedings of ICML, 2024a.
Youhe Jiang, Ran Yan, Xiaozhe Yao, Yang Zhou, Beidi Chen, and Binhang Yuan. Hexgen: Generative inference of large language model over heterogencous environment. Proceedings of ICML, 2024b.
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly super- vised challenge dataset for reading comprehension. arXiv preprint arXie:170s 101117055555.11 03551, 2017.
Hao Kang, Qingru Zhang, Souvik Kundu, Geonhwa Jeong, Zaoxing Liu, Tushar Krishna, and Tuo Zhao. Gear: An efficient ky cache compression recipefor near-lossless generative inference of Ilm. 073::22403.05527, 2024.
14

---

Published in Transactions on Machine Learning Research (12/2025)
Minsoo Kim, Kyuhong Shim, Jungwook Choi, and Simyung Chang. Infinipot: Infinite context processing memory-consttrinee llms. Proceedlings to EMNLP. 2024.
KVCache.AI. Ktransformers: A flexible framework for experiencing cutting-edee Ilm inference optimizations, http://gghabbbco////aa/aaaaa/a/a
2024. URL
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonza- lez, Hao Zhang, and Ton Stoica. Efficient memory management for large language model serving with pagedattention. Proceedings of SOSP, 2023.
Wonbcom Loe, Jungi Lee, Junghwan Sco, and Jacwoong Sim. Infinigon: Efficient generative inference of large language models with dynamic kv cache management. Proceedings of OSDI, 2024.
Yanyang Li, Shuo Liang, Michael R Lyu, and Liwei Wang. Making long-context language models better multi-hop reasoners. Proceedings of ACL, 2024a.
Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Llm knows what you are looking for before generation. 2rXiv:2404.14400 2024b.
Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Wei-Ming Chen, Wei-Chen Wang, Guangxuan Xiao, Xingyu Dang, Chuang Gan, and Song Han. Awq: Activation-aware weight quantization for on-device Ilm compression and acceleration. Proceedings of MLSys, 2024.
Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for Ilm kv cache compression at test time. Proceedings of NewrIPS, 2024a.
Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu, Kivi: A tuning-fre asymmetric 2bit quantization for ky cache. Proceedings of 1CML, 2024b.
1lama2.c. URL
llama.cpp. URL
cpp.
I Loshchilov. Decoupled weight decay regularization. Proceedings of ICLR, 2019.
Chao Lou, Zixia Jia, Zilong Zheng, and Kewei Tu. Sparser is faster and less is more: Efficient sparse attention for long-range transformers. arXiv:2406.100 16747, 2024.
Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du, Tao Jiang, Chao Hong, shaowei Liu, Weiran He, En- ming Yuan, Yuzhi Wang, et al. Moba: Mixture of block attention for long-context context 1lms. arXiv preprint W:Xiv:2502.3318 2025.
Zhenyan Lu, Xiang Li, Dongqi Cai, Rongjie Yi, Fangming Liu, Xiwen Zhang, Nicholas D Lane, and Mengwei Xu. Small language models: Survey, measurements, and insights. 1xii:2409.11990 2024.
Junlin Lv, Yuan Feng, Xike Xie, Xin Jia, Qirong Peng, and Guiming Xie. Critiprefill: A sogment-wise criticality-based approach for prefilling acceleration in llms. arXiv:2409.12/ 2024.
Kai Mci, Zelong Li, Shuyuan Xu, Ruosong Ye, Yingqiang Ge, and Yongfeng Zhang. Aios: Llm agent operating system. arXiv:2403.1000 16971, 2024.
Shervin Minacc, Tomas Mikolov, Narjes Nikzad, Mcysam Chenaghlu, Richard Socher, Xavier Amatriain, and Jianfeng Gao. Large language models: A survey. arXiv:2402.06196, 2024.
Jesse Mu, Xiang Li, and Noah Goodman. Learning to compress prompts with gist tokens. Proceedings of NewrIPS, 2024.
Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave no context behind: Efficient infinite context transformers with infini-attention. 4rX0v:2404.07114, 2024.
15

---

Published in Transactions on Machine Learning Research (12/2025)
Piotr Nawrot, Adrian Lańcuuci, Marcin Chochowski, David Tarjan, and Edoardo M Ponti. Dynamic memory compression: Retrofitting llms for accelerated inference. arXim2203.096336 2024.
OpenAL. Openai gpt-4o, 2024. URL
Wenbo Pan. Anti-haystack, 2024. URL
Ruoyu Qin, Zheming Li, Weiran He, Mingxing Zhang, Yongwei Wu, Weimin Zheng, and Xinran Xu. Moon- cake: Kimi's keeacce-centric architecture for Ilm serving. (rxfo2240.000075 2024a.
Yujia Qin, Shengding Hu, Yankai Lin, Weize Chen, Ning Ding, Ganqu Cui, Zheni Zeng, Yufei Huang, Chaojun Xiao, Chi Han, Yi Ren Fung, Yusheng Su, Huadong Wang, Cheng Qian, Runchu Tian, Kunlun Zhu, Shihao Liang, Xingyu Shen, Bokai Xu, Zhen Zhang, Yining Ye, Bowen Li, Ziwei Tang, Jing Yi, Yuzhang Zhu, Zhenning Dai, Lan Yan, Xin Cong, Yaxi Lu, Weilin Zhao, Yuxiang Huang, Junxi Yan, Xu Han, Xian Sun, Dahai Li, Jason Phang, Cheng Yang, Tongshuang Wu, Heng Ji, Zhiyuan Liu, and Maosong Sun. Tool learning with foundation models. ACM Computing Surveys, 2024b.
Shashank Rajput, Ying Sheng, Sean Owen, and Vitaliy Chiley, Inference-friendly models with mixattention. 00f:::2409.501111 2024.
Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean-baptiste Alayrac, Radu Soricut, Augeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. Gemini 1.5: Unlocking multi- modal understanding across millions of tokens of context. XXiu:2403.05530, 2024.
rustformers. URL attps://ge thu...uu.b.o http:///soooooomrs///i..a
Julian Schnitzler, Xanh Ho, Jiahao Huang, Florian Boudin, Saku Sugawara, and Akiko Aizawa. Morehopqa: More than multi-hop reasoning, arXiv:2406. 13397, 2024.
Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. arXi::247.00800 08608, 2024.
Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christo pher Ré, Ton Stoica, and Ce Zhang. Flexgen: High-throughput generative inference of large language models with a single gpu. Proceedings of ICML, 2023.
Luohe Shi, Hongyi Zhang, Yao Yao, Zuchao Li, and Hai Zhao. Keep the cost down: A review on methods to optimize I1m's kv-cache consumption. Proceedings of COLM, 2024.
Hanshi Sun, Li-Wen Chang, Wenlei Bao, Size Zheng, Ningxin Zheng, Xin Liu, Harry Dong, Yuejie Chi, and Beidi Chen. Shadowky: Ky cache in shadows for high-throughput long-context Im inference. arXiv:2410.211 21465,,...5.... 2024.
Raushan Turganbay. Unlocking longer generation with key-value cache quantization, 2024. URL https: htw.://g......oooo////rr--r/
//huggingface co/blog/kv- httbblog//--ccceeeeee cache- quattizataniioiin....m
Lei Wang, Chen Ma, Xueyang Feng, Zeyu Zhang, Hao Yang, Jingsen Zhang, Zhiyuan Chen, Jiakai Tang, Xu Chen, Yankai Lin, et al. survey on large language model based autonomons agents. Frontiers of Computer Science, 1866)1,,3345, 2024.
Jianbo Wu, Jie Ren, Shuangyan Yang, Konstantinos Parasyris, Giorgis Georgakoudis, Ignacio Laguna, and Dong Li. Lm-offload: Performance model-guided generative inference of large language models with parallelism control. Blog of PASA Lab, 2024.
Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, Song Han, and Maosong Sun. Infllm: Unveiling the intrinsic capacity of llms for understanding extremely long sequences with trainingffree memory. Proceedings of NeurIPS, 2024a.
Guangxuan Xiao, Ji Lin, Mickael Seznee, Hao Wu, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient. post-training quantization for large language models. Proceedings of ICML, 2023.
16

---

Published in Transactions on Machine Learning Research (12/2025)
Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Junxian Guo, Shang Yang, Haotian Tang, Yao Fu, and Song Han. Duoattention: Efficient long-context I1m inference with retrieval and streaming heads. arXiv:24 10.10199, 2024b.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficientt streaming language models with attention sinks. Proceedings of ICLR, 2024c.
Ruyi Xu, Guangxuan Xiao, Haofeng Huang, Junxian Guo, and Song Han. Xattention: Block sparse attention with antidiagonal scoring. arxiv preprint or Xiw:22033.16228 2025.
Dongjie Yang, XiaoDong Han, Yan Gao, Yao Hu, Shilin Zhang, and Hai Zhao. Pyramidinfer: Pyramid kv cache compression for high-throughput Ilm inference. Proceedings of ACL, 2024.
Yao Yao, Zuchao Li, and Hai Zhao. Sirllm: Streaming infinite retentive 1lm. Proceedings of ACL, 2024.
Jiayi Yuan, Hongyi Liu, Yu-Neng Chuang, Songchen Li, Guanchu Wang, Duy Le, Hongye Jin, Vipin Chaud- hary, Zhaozhuo Xu, Zirui Liu, et al. Kv cache compression, but what must we give in return? compre- hensive benchmark of long context capable approaches. Proceedings of EMNLP, 2024.
Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, YX Wei, Lean Wang, Zhiping Xiao, et al. Native sparse attention: Hardware-aligned and natively trainable sparse attention. Proceedings to ACL, 2025.
Amir Zandich, Majid Daliri, and Insu Han. Qjl: 1-bit quantized jl transform for kv cache quantization with zero overhead. arXiv:2406.cm 004.44882, 2024.
Tianyi Zhang, Jonah Yi, Zhaozhuo Xu, and Anshumali Shrivastava. Kv cache is 1 bit per channel: Efficient large language model inference with coupled quantization. (tf:::24050391111 2024a.
Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang, Zhiyuan Liu, et al. do bench: Extending long context evaluation beyond 100k lokens. Proceedings of ACL, 2024b.
Zhenyu Zhang, Shiwei Liu, Runjin Chen, Bhavya Kailkhura, Beidi Chen, and Atlas Wang. Q-hitter: A better token oracle for efficient 1lm inference via sparso-quantized kv cache. Proceedings MLSSs, 2024c.
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Proceedings of NewrIPS, 2024d.
Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolci Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, al. A survey of large language models. arXiv:2303. 18223, 2023.
Yinmin Zhong, Shengyu Liu, Junda Chen, Jianbo Hu, Yibo Zhu, Xuanzhe Liu, Xin Jin, and Hao Zhang. Distserve: Disaggregating prefill and decoding for ceedings of OSDI, 2024.
large language model serving. Pro-
A Hyperparameters, Environment and Baselines
A.1 Training
During the training stage, first insert retaining head Rs to each layar. A retaining head is a small FFN consist of two lincar transformations, and the non-lincar function is aligned with other non-linears of the conresponding model, with an intermediate size of 1024. We train the appended relaining head Rs 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, 2019) and the lcarning rate is set to 5c-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 is set to 0.0025.
17

---

Published in Transactions on Machine Learning Research (12/2025)
Table 9: Hyperparameters in LOCRET's inference stage. "b" is cache budget, "B" refers to chunk size of chunked prefill, "n," " refers to stabilizers length and "noo." is local length.
Model
b
B
Rlo
Phi Thi---minii-221
6000
3072
2500
100
Llama-3 3.
16384
1024
2500
100
A.2 Inference
The inference hyperparameters LOCRET listed in Table Here, we follow the notations in Algorithm 1. stands for the cache budget, is the chunk size of chunked prefill, n, is the length of stabilizers, and nose represents the length of locally retained tokens at the end of the input sequence.
Hyperparameters of other basclines Llama-3 to evaluate Llama-3.1. Since there is no recommendations of Phi-3-mini--22K, for MiniCPM, whose architechture and size is similar to Phi- Phi-3-mini--a -12XX 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-1281 and Llama-3. 1--8--insttrctt respectively. The chunk size of chunked prefill also 3072 and 1024 for the corresponding model. For MINFERENCE, we utilize the recommended settings for both models.
follows. For INFLLM, we the recommended settings for
use the settings
A.3 System Environment
For all the experiments except the 4090 experiments in Section 5.3, we use a workstation with 8xNVIDIA A800/HSO GPUs and 104 Intel(R) Xcon(R) Platinum 8470 CPUs. We only use GPU from the cluster for training, the GPU requirements are less than 80GB for all training procedures. 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 single GPU, and 2 GPUs for full attention settings.
For Section 5.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., 2023), which includes automatic tensor paral- lelism. INFLLM is a representative the offloading- based methods, where the full KV cache 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 al. (2024b), where quantization is conducted along channels instead of tokens. We denote this method as HF-2BITS. SIRLLM is an eriction-based token dropping algorithm, where tokens with low token-entropy is evicted once the cache is fullfilled. We use the official implementation of SirLLM, which includes some CPU operations including importance sorting. MINFERENCE is typical method of reducing peak GPU memory consump- tion through rule-based sparse attention, but it does not reduce the size of KV cache. Note that INFLLM, thus we implement
HF-2BITS and SIRLLM does not have official implementation on Phi-3-mini-1221 these three methods according the original algorithm. We only use the short factor of RoPE for INFLLM, and no further model modification is conducted for HF-2ArTS and S.R.L.M
18

---

Published in Transactions on Machine Learning Research (12/2025)
Table 10: LongBench scores of LOCRET compared with baselines.
Method
gov
triviaqa
narrative qa
qusumm
musique
2wikimqa
multifield
repobench
report
qu_
-p
FULLATTN
33.35
86.38
18.21
19.51
19.82
33.37
49.82
58.02
MINFERENCE
32.94
86.87
19.46
19.57
18.85
33.30
49.14
SIRLLM
32.92
85.61
21.08
21.59
24.32
34.97
48.52
59.15
INFLLM
25.96
81.87
20.83
19.61
13.63
27.43
41.29
55.73
Locret
33.46
82.39
24.56
23.35
25.12
35.93
52.77
57.16
Method
qasper
hotpotqa
multi
trec
passage retrieval
passage count
samsum
lcc
Avg.1
FULLATTN
41.07
43.06
26.57
67.00
93.50
2.97
23.15
51.86
41.73
MINFERENCE
43.56
26.35
68.00
89.00
2.10
25.58
53.68
41.73
SIRLLM
40.17
47.00
26.44
63.00
3.00
23.11
40.51
INFLLM
30.51
38.05
25.36
64.50
10.00
7.50
0.28
61.59
32.95
Locret
40.17
48.70
26.41
62.00
83.00
3.00
26.37
52.61
42.31
A.5 RULER Benchmark
To evaluate LOCRET-Q's performance query-centric lasks, compare it with selected eviction-based baselines: SNAPKV, H2O, SIRLLM, INFINIPOT, and vanilla LOCRET. We also include FULLATTN (im- plemented with FLASH-ATTN) and MINFERENCE for reference. The RULER benchmark consists 500 syuthetic queries per task, each with a context length of 128K tokens. All methods are tested on
For LOCRET-Q and LOCRET, set the budget size b 6000, chunk size B to 4096, stabilizers length n, to 2500, and local length Roc to 100. For SNAPKV, the voting window size is set to 100, with the last 100 tokens retained. For H2O, due to its reliance on full-sequence attention scores, we use a layer-wise chunked prefill pattern with a chunk size of 1024. Larger chunk size would result in an out-of-memory error. For SIRLLM, we configure the start size to 4, recent size to 1000, and budget size to 6000. For INFINIPOT, we set the budget size to 6000 with 50% NuC ratio. All evaluations are conducted on single NVIDIA A800-80GB GPU.
For prefill and decode speed testing, all methods except H₂O are implemented with FLASH-ATTN; H₂O uses Toror'''s 's vanilla attention due to its incompatibility with efficient attention implementations. The speeds are averaged on the first 5 entries of NIAH-Simple-1.
B Evaluation on LongBench
We conduct additional experiments to evaluate Locret LongBench (Bai et al., 2024b), comparing it with baselines such as Full Attention, MInference, InfLLM, and SirLLM. For this evaluation, used Ahi-3-mini-128t with a retained head trained on LongAlign. To ensure fair comparison, we excluded all Chinese subtasks from LongBench and focused solely on the English subtasks, as thi---min was not specifically trained on Chinese corpora. The results are presented below. For LOCRET we follow the hyperparameters presented in Table
We also report the maximum memory usage, including the GPU memory, the CPU memory, and the total maximum memory, alongside the average score LongBench. For FULLATTN, exclude the maximum memory usage, aligning with Figure 4.
From the experiments above, LOCRET demonstrates the best overall performance and excels in the majority of subtasks. It outperforms all the baselines without any noticeable performance degradation while consuming less memory. Although MInference also avoids performance drops, it requires more GPU memory compared LOCRET. SirLLM achieves comparable memory usage but shows some performance decline compared to FULLArTN and LOCRET. InfL.M exhibits the most significant, performance drop, and its offloading
19

---

Published in Transactions on Machine Learning Research (12/2025)
Table 11: Comparison methods on LongBench and memory usage.
Method
LongBench
Max GPU Memory
Max CPU Memory
Total Max Memory
FULLATTN
41.73
MINFERENCE
41.73
27.63
0.17
27.80
SIRLLM
40.51
18.29
0.05
INFLLM
32.95
20.03
8.95
28.98
Locret
42.31
17.71
0.15
17.86
mechanism results in the highest CPU memory usage, making the method with the largest total memory consumption. These results highlight LOCRET outstanding approach for evaluation LongBench.
C
Training Robustness
C.1 Intermediate Size of the retaining head
We align all the training settings as described Section 5.1 and only change the intermediate size retaining heads dR E {256, 512, 1024, 2048, 4096) with the backbone model Phi-3-mini-193 The trained model is evaluated on L-Eval and report the average L-Eval score corresponding each intermediate size. Results are listed in Figure 6. The performance variations among all the settings minimal compared to the changes in the intermediate size, surpassing all baselines in Table 2 and Table 3. This indicates that out method exhibits good performance stability regardless of the intermediate size of the retaining head Rs..
Table 12: L-Eval scores with different intermediate size of the retaining head dR- (Detailed)
Phi-3-nini-2 on Eval
du
CodeUU
NQ
CUAD
NarrativeQA
QMSum
SPACE
Avg.t
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 Rs with dR E {256, 512, 1024, 2048, 4096). We keep all the other hyperpa- ramotors same, and train on the same datasct. From Table 12, LOCRET shows stability to the intermodiato size, in both overall performance and the performance each single task. While increasing the intermediate size, observe very slight overall performance enhancement. However, the performance variance is negligi- ble compared to the increase of parameter size, thus we choose to maintain the intermediate size in small scope to take balance of performance and efficiency.
C.2 Training Data Insensitivity
We also consider the sensitivity of the training data, which leads to ablate the training dataset by training LongAlign (Bai et al., 2024a) and Anti-Haystack (Pan, 2024), comparing these results with those from LongAlpaca (Chen et al., 2024) in the original training setting. We also align other settings to the original setting and choose the backbone model to be Phi-3-mini-1288 We report the average L-Eval score for each training datasct. The results Table 8 shows that LOCRET has high insensitivity towards different training data. The performance impact of different data recipes minimal, indicating that our method can be trained any long-context. tuning dataset.
20

---

Published in Transactions on Machine Learning Research (12/2025)
Table 13: L-Eval scores of LOCRET trained on various dataset.
(Detailed)
Phi-3-mini- L-Eval
Dataset
CodeU
NQ
CUAD
NarrativeQA
QMSum
SPACE
Avg."
Lonaalpaca
8.89
51.49
22.23
16.42
14.86
14.06
21.33
LongAlign
10.00
55.13
21.34
16.40
14.09
22.00
Anti- Hay-tHaystcck
8.89
52.91
20.87
13.73
13.84
14.10
20.72
Table 14: L-Eval scores of LOCRET trained on different tasks.
Phi-3-nini-122 L-Eval
Code CodeU
QA
Summarization QMSum
Dataset
Avg.1
NQ
CUAD
NarrativeQA
SPACE
LCC (Code)
7.78
56.87
14.01
20.93
15.10
13.87
21.43
Trivia-QA (QA)
56.77
14.15
20.49
14.75
14.03
21.51
Gov- Rov-Report (Summ.)
4.44
57.18
14.16
21.66
15.53
13.40
21.06
We conduct training on various datasets and benchmark the trained weights on I-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 relatively smaller dataset, we utilize the whole dataset, which consist of 2424 entries. According to (Chen et al., 2024; Bai et al., 2024a; Pan, 2024), all three datasets are general long-context SFT datasets, consisting of mixed collection of diverse long- context tasks such as long-document question answering, summarization, and code generation. Therefore, the primary differences among these datasets lie the relative proportions data quality of individual documents. The results in Table 13 shows that different training dataset recipe exhibits minor effect towards the overall performance. LOCRET can obtain competitive performance with- out delicately selecting the training data. We further perform cross-task evaluation to demonstrate the generalizability of LOCRET, as shown in Table 14. Specifically, we consider three representative long-context tasks: coding, question answering (QA), and summarization. The model is trained on LCC2 for coding, TriviaQA (Joshi et al., 2017) for QA, and GovReport (Huang et al., 2021) for summarization. As presented in Table 14, a model trained on one task can be effectively transferred to others. For QA and summariza- tion, the transfer is nearly scamless, while the coding task shows modest performance drop when trained without code data. However, models trained on coding data still generalize well to other lasks. Overall, LOCRET exhibits strong cross-task transferability and adaptability across diverse long-context domains.
various task types and the
D Ablation the Smoothing Loss
We incorporate a smoothing loss, denoted as L2, into the training objective of LOCRET (Equation 6), which is applied to the predicted CIS between adjacent positions. This design encourages locality in predictions, a property commonly observed in natural language. To examine the effect of the smoothing loss, we evaluate models trained with and without L2 on the Numube task xeBnnh,, where the target number appears in a locally concentrated region within noisy document. As shown in Table 15, removing Le leads to a significant performance degradation, confirming the effectiveness of the smoothing constraint. One illustrative failure case is when the model predicts "810007779" instead of the correct "8100077779", missing one "7", This example highlights that without enforcing locality, the model tends to lose information from crucial regions.
21

---

Published in Transactions on Machine Learning Research (12/2025)
Table 15: Ablation on the smoothing loss L2.
Phi-3-mini-120 R.Number
Dataset
R.Number
LOCRET
97.46
LOCRET w/o La
79.15
E
Orthogonality to Other Methods
Table 16: Quantization with FULLATTN and LOCRET. "M" represents Method and "-A" rep- resents the gap of average L-Eval score.
Table 17:
The average L-Eval scores of Lo- CoCo, LOCRET, and the combination of Lo- CoCo and LOCRET.
Setting
M
M-4bits
-A
Method
LoCoCo
LOCRET
Combination
M-FULLATTN
29.08
28.52
0.56
L-Fval
26.01
27.96
28.70
M--LCRET
27.96
27.11
0.85
KV cache quantization. According to Zhang et al. (2024c), eviction-based methods like H2O struggle with compatibility when combined with KV cache quantization. Quantization introducee 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--8--B 8a---Btituuut caused by quantization on LOCRET' is only slightly higher than that observed with the full atiention method, indicating that LOCRET is a quantization-frienlly approach. More details of the experiment are provided in Appendix E.1.
as the model backbone. Table 16 shows that the performance drop
Token merging. As described Section 2, token dropping can also be implemented through an attention pool. Attention pool-based methods (Xiao et al., 2024c; Cai et al., 2024a; Mu et al., 2024; Munkhdalai et al., 2024) merge adjacent tokens 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 al., 2024a) as a representative of the latest allention 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 set of heavy hitters according to H₂O (Zhang et al., 2024d), then applies 1-D convolution to the non-heavy hitters to compress them into static size. By replacing HzO's '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 1.1.....0 ---ttinstruct backbone and report the average score across all selected tasks. As shown Table 17, 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 E.2.
Hcad-wise Budget Allocation. Since LOCRET evict cache units across the attention heads independently, it is compatible with head-wise budget allocation. Here, we combine LOCRET with PYRAMIDKY (Cai et al., 2024b). PYRAMIKVV assumes that identifing the important cache in deeper layers are simpler than shallow layers, thus allocates more budget to the shallow layers. We evaluate LOCRET+Pye PYRAMIDKY on the following subtasks of coBecch using Phi-3-mini 128K. Results presented in Figure 18 shows the compatibility of the two methods.
22

---

Published in Transactions on Machine Learning Research (12/2025)
Table 18: BeBhnch scores of the combination LOCRET and PYRAMIDKY V.
Phi-3-mini--e ccBench
Method
R.Number
E.Sum
E.MC
C.Debug
Avg.t
LOCRET
97.46
16.82
46.29
29.71
47.57
LOCRET PYRAMIDVV
99.66
15.82
48.03
30.00
48.38
E.1
Combination with Quantization
Table 19: L-Eval scores of FULLATTN, FULLATTN-4bits, LOCRET and LOCRET-4bits, (Detailed)
11ama-3.1-88-ine 1-88-insttuct on L-Eval
Method
CodeU
NQ
CUAD
NarrativeQA
QMSum
SPACE
Avg.
FULLATIN
10.0
66.84
38.91
23.11
18.76
16.86
29.08
FULLArIN-lbits
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
LocnET-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 the backend of quantization. Other hyperparamelers are kept same described in Section 5.1. We I.-Fval and report the average score, with Llama-3. 1--B-instttuct backend.
conduct the experiment The results in Table 19 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.
E.2 Combination with LoCoCo
(a) L-Eval scores of LoCoCo, LOCRET and the combination LOCOCO+LOCRET. (Detailed)
--st-ttuettii on T. Eval
Method
CodeU
NQ
CUAD
NarrativeAA
QMSum
SPACE
Avg-1
FULLATIN
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+LoCRRT
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 2048. We extend the context length through
train the convolution head with the size of convolved cache chunked prefill training to 64K, which is longer than all tasks' average input length. The convolution kernel is set to 21, and 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 noi. modify its positional embedding. During Inference, we keep cache budget size of 16384. In the standalone LoCoCo setting, there are 2048 cache units are convolved, while the others are the heavy-hitters selocted by H2O. In the combined algorithm, we replace H₂O 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.
F Extremely Long Context Evaluation
We create a dataset similar ocBench's N.Number, with an average length of 10 million tokens. Each data point contains 10-digit number string inserted into irrelevant. context, and the task is to retrieve
23

---

Published in Transactions on Machine Learning Research (12/2025)
the inserted number. The dataset consists of 50 examples, with the number strings uniformly distributed throughout the context. We used the hyperparametens from Table 9, with the exception of setting the chunk size to 10240 to speed up inference. The results, presented below in Table 21, show that Locret can efficiently process extremely long contexts. In this experiment, the cache budget is set to 6000, and the compression ratio 1747.6x.
Table 21: Inference speed with Retaining Heads.
Phi-3-mini-128K on 10M context
Dataset Passkey 10M
LOCRET 100.00
G Compressing Multi-turn Conversations
Compared to query-aware eviction methods, such as SNAPKV (Li et al., 2024b). LOCRET is more suit- able solution for multi-turn conversation scenarios. This is because the evaluation of cache importance in LOCRET based on the cache itself, rather than being dependent on the subsequent query. To demonstrate this, we evaluate LOCRET on the Kock-Paper-Scissee
benchmark introduced in SIRLLM (Yao et al., 2024). Since SIRLLM specifically designed for such scenarios, we use it our basclinc in this benchmark. Results in Table 22 show that Locret is also effective in multi-turn conversation contexts.
The hyperparameters
aligned with those used in SIRLLM, with the cache budget set to 1024, and no stabilizers are retained, as SIRLLM does not retain local tokens in this benchmark. We perform 2000 turns as same as the original SIRLLM settings. The results are presented below.
Table 22: Rock- Lock-Paper--e Paper-Scissors Scissors scores of LOCRET and SIRLLM.
Phi-3-mini-1284
Rock-Papor--Scisory
Rock tie
Paper tie
Scissors tie
Avg.
Preference
win
lose
win
lose
win
lose
win"
lose.
SIRLLM
40.00
31.75
28.25
27.5
36.55
35.96
29.35
25.15
45.50
32.28
36.57
Locret
50.00
31.05
30.35
19.45
50.20
27.25
20.70
33.78
33.98
H
Discontinuous Context and Stablizers
Evicting cache units results in context discontinuity, which causes unstable CIS prediction and inaccurate caleulation of later tokens. Thus, always retain the stabilizers, which consist of the last n, cache units in each chunked prefill step. We ablate n, R.Nubber of oc-Bench in the proposed algorithm to demonstrate the necessity of incorporating stabilizers in the design. The results in Figure 3a show that lower stabilizer length n. 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 Figure 3b. 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 3c. We also observe high errors when the stabilizer length is limited. This explains the reason for failure when the stabilizers short absent: context discontinuity leads to instability in the prediction of CTS, resulting in errors during cache eviction and amplifying errors in the hidden states.
Retaining Heads Do not Slow Down Inference
We evaluate the model's forward throughput under varying context lengths, both with and without retain- ing heads. The results are summarized below in Table 23. "R" represents the retaining heads, and the throughput is reported in tokens per second (tok/s) in the format "Avg. Std."
24

---

Published in Transactions on Machine Learning Research (12/2025)
Table 23: Inference speed with Retaining Heads.
Context Length
1024
2048
3072
4096
w/o R Speed
18674 / 443
19743
/ 464
19982
/ 402
20304
187
w/ R Speed
17118
1117
18503
546
19054
283
19153
174
From the results, significant latency increase is observed when using retaining heads. The numerical differences are attributed to systematic variations rather than additional overhead introduced by retaining heads during inference.
J 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 Benhh with backbone investigate the pattern variation among different layers on N.Number. We display the results in Figure 7 and Figure 8. The yellow parts are the retained cache, where the y-axis represents cache position and x-axis is the time axis.
and
Figure 7 shows that the pattern is mostly decided by the tasks, where both heads shows similar pattern in the same task. In R.Nubeer, we are able to observe 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., 2024c) pattern, where the tokens at the beginning of the sequence are always important. This also mentioned as the A-pattern sequence. This pattern is also approached by MINFERENCE (Jiang et al., 2024a) by the pattern "vertical-and- slash". In F.MC, H2O (Zhang et al., 2024d) and ScissorHands (Liu et al., 2024a) pattern can be observed, following the assumption that if 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., 2024a) 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 investigate the patterns across layers. Tn Figure 8, 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.
MINFERENCE. We can also discover the vertical lines the middle of the
The retained pattern different layers shows various features, which might be a good handle to investigate how LLMs understand and process natural language queries.
25

---

Published in Transactions on Machine Learning Research (12/2025)
R.Number
M.Find
E.MC
Layer 11
1904
3550-
300
Head 14
2008
3009
2080
2004
3540
2540
3000
3000
10900 20000
1000 2009
1908
(2:00
2008
1500 2000 CVRV
2000
Layer 11
Head 22
2503
2102
2000
.000
15006 11300 20100
7556
10000
1000
10800
29300
1966
600 00000 19d10 2000 2600
3000
Figure 7: Head patterns across multiple tasks.
Layer 1
Layer
Layer 9
Layer 13
000
moop
PANM PA6E
1908
100
19000 4000
MA0 1900
Layer 17
Layer 21
Layer 25
Layer 29
Figure 8: Layer patterns of Numbbr
26


Reviews:
[{"review_id": "8TxYJIZlZF", "reviewer": "Reviewer_4ga9", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a novel eviction-based strategy for key-value (KV) cache optimisation with the goal of making long-context LLM inference more efficient. This strategy is called Locret and is based on training a separate set of retaining heads that predict an importance score of elements in the KV cache. The retaining heads are trained on a smaller, shorter-context datasets by supervising the desired importance score with the maximum attention score over the answer produced from a KV cache prefilled on the prompt. The heads are then used at generation time to prefill the KV cache up to its memory limit by removing those pairs with a low importance score. Consequently, combined GPU and CPU memory never exceeds the desired memory budget once the budget is set. Locret is empirically verified on a number of long-context benchmarks against other memory reduction techniques. \n#### Strengths\n1. **The explanation of the method is clear.** The entire pipeline that Locret is building is clearly explained. In particular, the exposition is well-written throughout. The problem that is being solved is clearly demarcated as is its solution. I also appreciate a sufficient amount of preliminaries for the reader to catch up on the specifics of the topic as well as related work.\n   \n2. **The method is intuitive and easy to understand, and the authors highlight they had to solve some issues.** The idea of training a separate module that learns to predict a meaningful score for retaining key-value pairs is certainly interesting and makes intuitive sense. However, it is certainly not intuitive that this strategy of building a KV cache *without* LLM retraining still works. From their exposition, the authors clearly had to optimise this idea for practical success and had to deal with a number of engineering challenges, e.g. adding stabilisers.\n#### Weaknesses\n1. **The writing could be augmented with more explanation of why choices were made.** Locret makes some choices in its pipeline, e.g. an MLP as retaining head or training on QA datasets with maximum attention score as target, without explaining why. This is a minor weakness and aimed to further improve the exposition.\n2. **The experimental section does not provide convincing evidence for Locret.** See my specific concerns on the matter below.\n\nI want to highlight here that Locret does *seem* to show promise with reasonable and interesting results, but that this can not be made certain yet. The method is interesting and only requires a more consistent and complete experimental section to show that it is impactful.", "claims_and_evidence_yn": "No", "claims_explanation": "While the method is explained well, it still hinges on a very convincing experimental section in order to be of any impact. However, I have a number of concerns with the experimental section.\n1. **No variability metrics are reported.** All reported metrics are a single, average value without any notion of how variable those averages are. Please use variability metrics in order for the reader to gain an understanding of how statistically significant these results are. *Without variability metrics, reviewers could be looking at statistical noise instead of significance.*\n   \n2. **No consistent choice of baselines across experiments and missing baselines/ablations.** I am failing to see why there is not consistent choice of baselines across all experiments. For example, the evaluation on $\\infty$Bench and L-Eval compares to InfLLM, 2-bit quantisation, SirLLM and MInference, while the query-driven experiment does not consider InfLLM, SirLLM or quantisation and adds SnapKV, H$_2$O and InfiniPot. Why are not all these methods used across the board? Additionally, it seems two important baselines are missing; higher-bit quantisation (at least 4-bit) and an ablation baseline that only utilises stabilisers. First, a comparison to 4-bit quantisation is important because it is much more commonly used than 2-bit quantisation. Even more, to avoid the overhead from inverse quantisation, why not quantise the LLM as well as, again, is more commonly done in practice? Second, a comparison to a baseline with a KV cache that *only* uses stabilisers (so only the last tokens fitting in memory) is important because of the importance stabilisers show to have for Locret too. In particular, Figure 5b shows that when the stabilisers fill the entire budget (6000) and there is no selection made by the retaining heads, the performance suddenly increase substantially. *Without these two baselines, the impact of Locret is not strongly supported.*\n   \n3. **Overly strong or wrong experimental conclusion in some places.** First example at the end of Section 5.2, paragraph 2: \"In the shorter context scenario (L-Eval), a similar phenomenon is observed.\" A similar phenomenon is most certainly not observed from Figure 4 as all methods, except perhaps HF-2bits on Llama-3.1 actually *increase* both the GPU and total memory cost. This observation is very counter-intuitive. *Could you please elaborate?*\n   Second example in the second paragraph of Section 5.3: \"the quantization and dequantization processes become the bottleneck, leading to significantly slower speed performance\". It does not lead to slower speeds. Table 4 shos that quantisation with HF-2bits is even quicker for Llama-3.1, it is the performance that tanks for quantisation. *Can you elaborate why this is the case for Llama-3.1 but not for Phi3?*\n   \n4. **Reporting of strange or incomparable metrics as well as weird testing practices.** Example 1 is Table 4; there is no need for an arbitrary \"success\" checkmark. The accuracy speaks for itself. Example 2 is the \"solution\" to OOM of some methods in Section 5.3: \"we remove some tokens from the middle of the input sequence in those cases, marking these settings with ∗,\". Since the task here is to recover a passkey from a large amount of irrelevant text, the removing of tokens might remove the passkey to find. Additionally, no precise explanation is given *how* it is decided which tokens get removed. Example 3, perhaps the most incorrect practice, is the *limiting of testing to the first 20 entries of the test set if bad performance is observed*. This way of testing only produces numbers that are *incomparable* as they are computed from different test sets. Poor performance is no reason to not report comparable metrics throughout (Tables 5 and 6).", "audience_yn": "Yes", "audience_explanation": "I think this work can be interesting for those members of the community that aim to optimise the efficiency of LLM inference as well as their practical ramifications.", "requested_changes": "In accordance with the above raised issues, I request the authors to\n1. Add variability metrics. Preferably those with less distributional assumptions, e.g. quartiles/quantiles.\n2. Make set of baselines consistent across all experimental sections. Add 4-bit quantisation (with quantised LLM) as well as baseline with only stabilisers.\n3. Check the experimental section for errors and rewrite to be more faithful.\n4. Solve or explain the use of strange testing practices.\n5. Locret-Q should be part of the method (Section 4) and should not suddenly appear in the experimental section.", "broader_impact_concerns": "I do not have any broader ethical concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "4ic1dmLWV8", "reviewer": "Reviewer_Fwwc", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces Locret, a lightweight, training-based framework for key-value (KV) cache eviction in long-context LLM inference. It tackles the fundamental challenge of scaling LLMs on consumer-grade devices, where memory constraints limit long-context processing.\nThe core innovation lies in the use of trainable “retaining heads” that learn a causal importance score (CIS) for each cache unit, enabling precise and causal cache eviction compatible with chunked prefill. Locret requires minimal additional training (<1 GPU hour) and can be seamlessly integrated into existing decoder-only architectures.\nA variant, Locret-Q, extends the approach to query-driven tasks by introducing query-aware eviction. Extensive experiments on ∞Bench, L-Eval, LongBench, and RULER show that Locret achieves:\n\n- Up to 20× compression of the KV cache with <10% quality loss.\n\n- 128K+ context inference on a single NVIDIA 4090 GPU.\n\n- State-of-the-art tradeoffs between accuracy, memory, and inference speed, outperforming strong baselines like H2O, SnapKV, SirLLM, and MInference.\n\nOverall, Locret offers a practical, principled, and efficient approach to democratizing long-context LLM inference.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The paper provides rigorous empirical validation and clear theoretical motivation:\n\n- It systematically analyzes the local-global discrepancy in existing scoring functions (e.g., H2O, SnapKV) and demonstrates Locret’s superiority in maintaining consistency without access to future tokens.\n\n- The architectural design (retaining heads, causal CIS training) is well-motivated, with ablations isolating each component (budget size, stabilizer length, chunk size, training dataset, intermediate size).\n\n- Benchmark results across multiple long-context datasets consistently show Locret’s performance advantages in both efficiency (memory/speed) and accuracy, especially under realistic consumer-grade settings.\n\n- Experimental transparency is high: training time, hardware settings, and hyperparameters are all provided, and comparisons include a wide range of state-of-the-art baselines. The empirical evidence thus fully supports the paper’s claims, and the methodology is both replicable and technically sound.", "audience_yn": "Yes", "audience_explanation": "- The method is algorithmically novel (introducing trainable retaining heads and causal importance scoring).\n\n- In addition, it is systemically practical (lightweight integration, low compute overhead).\n\nOverall, the proposed framework bridges the gap between model design, memory optimization, and system deployment, making it of strong interest to readers working on efficient inference, compression, and scalable reasoning in LLMs.", "requested_changes": "None", "broader_impact_concerns": "There are no major ethical red flags in this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Gg2FGidw9v", "reviewer": "Reviewer_otMS", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "### Summary\n\nThe paper addresses the problem of efficiently evicting KV caches for long-context inference in large language models in causal settings, such as streaming or chunked prefill. The authors argue that most existing eviction strategies (e.g., H2O, SnapKV) implicitly rely on global importance estimates that are not available in causal settings. To quantify this mismatch, they present a consistency metric that compares global (full-context) and local (causal) token importance scores, demonstrating that existing strategies are unstable across the two regimes. To address this, they propose Locret, a causally trained importance scorer (CIS) that predicts future token relevance using only past attention, implemented via Retaining Heads. The method aims to minimize the discrepancy between local cache scores and ideal full-context scores. \n\nIn the experiment, the paper trained Locret on the LongAlpaca dataset. It evaluated the proposed method across multiple benchmarks. The results show that Locret not only matches or exceeds full-attention baselines in accuracy but also provides significant memory savings and speed up.\n\n\n### Strengths\n\n1. Clear and well-motivated diagnosis of why prior eviction methods fail in streaming/chunked inference.\n2. The method requires minimal training and is relatively lightweight.\n3. Strong and extensive empirical results. The paper experiments on two models and two benchmarks, comparing with many baselines.\n\n\n### Weaknesses\n\n1. Generalization of the trained Retaining Heads across tasks is implied but not extensively validated. \n    1. The causal training objective relies on future attention during supervision. It would be clearer to discuss whether this introduces bias towards the datasets used to train the Retaining Heads.\n    2. Table 12 shows that there are some inconsistencies between training datasets. Extending this experiment and using statistics to quantify the \"insensitiveness\" would strengthen the paper.\n2. Although it is a good sign, some results outperforming full attention lack discussion. These results introduce a contradiction between the consistency (Figure 1) and the benchmark performances (Tables 2 and 3).\n3. The CIS training objective includes a smoothing / L2 regularization term, but no ablation is provided to isolate its contribution. Is the performance gain primarily due to the Retaining Head architecture, or does the smoothness constraint dominate? A comparison of CIS trained with and without the regularization term would clarify this.\n4. In general, the paper is well organized and provides a brief introduction to the extra materials in the appendix. However, Section 3 would be clearer if the paper introduced preliminary (4.1) and included the limitations of the causal inference cache.", "claims_and_evidence_yn": "No", "claims_explanation": "There are several claims made by the paper. The majority of them have experimental results to support their claims in the chunked prefilled setting.\n\n1. Locret uses less memory while maintaining performance (Section 5.2)\n2. Locret achieves a higher speed-up on a small GPU (Section 5.3)\n3. Locret can handle query-driven task (Section 5.4)\n4. Locret requires stabilizer (Figure 3)\n\nHowever, some implicit claims would need more evidence:\n\n1. Locret robustness (Section 5.5)\n    > From our observations, the training of Locret exhibits strong robustness. Despite changes in $d_R$ and the dataset, the performance variations shown in Figure 6 and Table 7 are minimal.\n    The robustness of different training data is not clearly shown in Table 7 (as there are only three data points). Upon closer examination of Table 12, several inconsistencies are evident from task to task across different training datasets.\n2. Locret requires $L_2$. While the stabilizer is shown to be important, it is unclear whether the $L_2$ is needed.", "audience_yn": "Yes", "audience_explanation": "This paper addresses a crucial problem in making LLMs more efficient.", "requested_changes": "1. Clarify the training and generalization behavior of Retaining Heads. The authors can include one of the following changes:\n    - Provide a cross-task generalization experiment (e.g., train on dataset A, evaluate on dataset B without retraining) to support the claim that Locret is robust to training distribution shifts.\n    - Expand Table 12 into a broader study with mean/variance statistics to better quantify robustness rather than relying on a few manually selected data points.\n2. Some benchmark results exceed full attention, which appears contradictory given that consistency metrics suggest Locret approximates rather than improves upon full attention. Please provide a short discussion explaining this effect.\n3. Include an ablation on the $L_2$ smoothing.\n4. Improve clarity in Section 3. The definition of the consistency metric is clear. Still, the narrative would be more straightforward to follow if Section 4.1 (preliminary formulation of local vs. global scoring) were introduced earlier or briefly referenced. Additionally, explicitly stating the limitations of causal scoring before presenting consistency results would help motivate why inconsistency matters.", "broader_impact_concerns": "No broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)