OmniKV: Dynamic Context Selection for Efficient Long-Context LLMs
Authors: Jitai Hao
Yuke Zhu
Tian Wang
Jun Yu
Xin Xin
Bo Zheng
Zhaochun Ren
Sheng Guo
Venue: ICLR 2025 Poster (2025)
Keywords: Efficient LLMs
KV cache
Long Context LLMs
Primary area: foundation or frontier models, including LLMs
TL;DR: We proposed a drop-free and training-free KV cache saving inference method for long context LLMs.

Abstract: During the inference phase of Large Language Models (LLMs) with long context, a substantial portion of GPU memory is allocated to the KV cache, with memory usage increasing as the sequence length grows. To mitigate the GPU memory footprint associate with KV cache, some previous studies have discarded less important tokens based on the sparsity identified in attention scores in long context scenarios. However, we argue that attention scores cannot indicate the future importance of tokens in subsequent generation iterations, because attention scores are calculated based on current hidden states. Therefore, we propose OmniKV, a token-dropping-free and training-free inference method, which achieves a 1.68x speedup without any loss in performance. It is well-suited for offloading, significantly reducing KV cache memory usage by up to 75% with it. The core innovative insight of OmniKV is: Within a single generation iteration, there is a high degree of similarity in the important tokens identified across consecutive layers. Extensive experiments demonstrate that OmniKV achieves state-of-the-art performance across multiple benchmarks, with particularly advantages in chain-of-thoughts scenarios. OmniKV extends the maximum context length supported by a single A100 for Llama-3-8B from 128K to 450K. Our code is available at https://github.com/antgroup/OmniKV.git.

Published as a conference paper at ICLR 2025
OMNIKV: DYNAMIC CONTEXT SELECTION FOR EFFI- CIENT LONG-CONTEXI LLMs
Jitai ? Yuke Zhu'; Tian Wang', Jun Yu², Xin Xin?, Bo Zheng', Zhaochun Ren", Sheng Guo
"MYbank, Ant Group Shandong University "Harbin Institute of Technology
LLeiden University ffelix. yk, guosheng gwwshenggemmyann..... on ji sitaihac@mail sdu, edu. en,
2 on, rene 1 liacs. le leidenuniv.nl.
ABSTRACT
During the inference phase of Large Language Models (LLMs) with long con- text, substantial portion GPU memory is allocated to the KV cache, with memory usage increasing the sequence length grows. To mitigate the GPU memory footprint associate with KV cache, some previous studies have discarded less important tokens based the sparsity identified in attention scores in long context scenarios. However, we argue that attention scores cannot indicate the future importance of tokens subsequent generation iterations, because attention scores are calculated based on current hidden states. Therefore, we propose Om- niKV, a saken-dropping-free and training-free free inference method, which achieves 1.68x speedup without any loss in performance. It is well-suited for offloading, significantly reducing KV cache memory usage by up to 75% with it. The core innovative insight of OmniKV is: Within a single generation iteration, there is a high degree of similarity in the important tokens identified across consecutive layers. Extensive experiments demonstrate that OmniKV achieves state-of-the performance across multiple benchmarks, with particularly advantages in chain- of-thoughts scenarios. OmniKV extends the maximum context length supported by a single A100 for Llama-3-8B from 128K to 450K. Our code is available at
INTRODUCTION
Large Language Models (LLMs) have demonstrated their profound impact across multitude of
applications, such as chatbots (Achiam et al., 2023 Metal 2024), agents (Wu et al.) 2023 Chan
et al. 2023 Liu et al. 2023b), and embodied robotics (Mai
2023 Zhang al.) 2024b).. These
applications confirm the demand for LLMs with strong long context processing capabilities to tackle
complex tasks based on prior interaction histories or given information materials. Currently, some
studies have successfully extended the maximum context length of LLMs (Peng et al. 2023 Yang
al. 2024 Young et al. 2024 Meta 2024
Chen et al., 2023
In long-context scenarios, LLMs inference incurs more massive GPU memory usage. Particularly, a significant portion of this attributed to the KV cache, which is proposed to accelerate LLMS generation and reduce redundant computation. Moreover, the memory occupied by the KV cache increases linearly with the sequence length. For instance, in the case of Llama-3-8B,, with a batch size of 8 and a context length of 128K tokens, the KV cache alone occupies over 134GB of GPU memory, presenting significant challenge to inference systems. To alleviate the high memory usage of the KV cache in long-context scenarios, previous studies (Zhang et al. 2024c Li et al.; 2024 Liu et al. 2024a) have attempted to identify and discard less important tokens based on sparsity in attention in long-context scenarios (Wang et al. 2021 Ribar et al., 2024). Tokens with lower cumulative attention scores are discarded, thereby reducing GPU memory occupation.
However, we argue that in multi-step reasoning scenarios, the important tokens vary depending on the reasoning step. This variation arises because attention scores are calculated based on the
Equal contribution. 'Corresponding author. Work done during internship at MYbank.

---

Published as a conference paper at ICLR 2025
current hidden states, meaning that the attention scores of tokens only reflect their relevance to current reasoning step. Consequently, some tokens with low scores may be recalled as important and relevant tokens in subsequent reasoning steps as shown in Figure lb Discarding tokens may lead to the loss of crucial information for completing subsequent reasoning steps.
Motivated by this, propose OmniKV, novel inference method that retains the KV cache for all tokens. This approach not only achieves performance comparable to that of the original model but also accelerates decoding efficiency when context is longer than 32K. OmniKV is mainly built upon an innovative insight: For a specific context, the tokens with high attention scores are very similar across various layers. We refer to this phenomenon as Inter-Layer Attention Similarity. In other words, similar set of tokens is identified as important across multiple layers.
Specifically, OmniKV offloads most layers' KV cache to CPU memory when processing the input prompt (i.e., the prefil stage) but retains few "filter" layers' entire KV cache. During the next token prediction (i.e., the decode stage), OmniKV first exploits the sparsity of attention map, selecting a few top-scoring tokens using top-k in these few "filter" layers. Next, other layers directly use the subset of tokens chosen by the preceding filter layers as context. This way, OmniK V only needs to load small subset of tokens from CPU memory into GPU memory for attention computation. For decoding iteration, since many layers share the same index of tokens, we only perform & 3 transfers between GPU and CPU. By using asynchronous transfer to overlap computation with transfer time, the efficiency of decoding is accelerated due to the shorter context computed in most layers.
We conducted extensive experiments on various LLMs, including 1aamma---8--6223- Yi-9B-200K
(Young et al.) (024) and Llama-3. 1.mma--3.1-7 3-1B (Meta 2024) across benchmarks LongBench (Bai et al.] 2023] and InfiniteBench (Zhang al. 2024a). We tested both single-step reasoning and multi- step reasoning (Chain of Thoughts) setups. The results showed that OmniK achieved the best performance in both settings, especially in the multi-step reasoning setup. This demonstrates the necessity of drop-free and the effectiveness of dynamic context selection. Moreover, OmniKV can accelerate inference when context lengths exceed 32K. By using only a single A100 GPU, compared to the original model's 's efficiency, OmniKV achieves a 1.7x speedup with 128K context. When using three A 10s for origin model to pipeline inference with 450K context, OmniKV still achieves a 1.87x speedup while using only single A100.
RELATED WORK
Token Dropping and Offloading. Most similar studies to ours focus on discarding unimportant tokens based on accumulated attention scores after the prefill stage. These tokens are then com- pletely dropped in the decode stage (Zhang et al. 2024c| Liu et al. 2024a Li et al. [0024] Ge et al.] 2024a). However, these methods might discard tokens that could become important in future rea- soning steps. To ensure the lossless of information, we dynamically select sparse subset of the KV cache for each generation iteration to guarantee performance. In a related approach, Quest (Tang et al. 2024) recognizes the importance of dynamic selection. Nonetheless, it fails to reduce memory usage and may compromise recall, Due to representing a block with a single vector, Quest's capacity to retrieve relevant tokens may be compromised.
Another type of similar works involves KV cache offloading. Many studies offload layers' KV cache to CPU memory when VRAM is not sufficient (Sheng et al. 2023 Kwon et al. 2023b). However, these methods do not leverage sparsity in attention. The data transfer volume over PCIe 10 times more than ours in long-context scenarios. InfLLM (Xiao et al. (024a) similarly divides sequences into blocks, selecting few representative vectors as retrieval keys for block. Then offloads other data to CPU memory. However, the chosen few representative vectors may fail to fully capture the block's ' information, resulting in relatively low recall rate.
Other Efficient Methods. Considerable efforts have been made to minimize KV cache while incurring minimal performance loss the model. One class of work compresses KV cache, such as ICAE (Ge et al., 2024b) and Gist (Mu al. 2024), which utilize L.MMs as auto-encoders to compress the context shorter sequence. Additionally, ly, there works that directly compress prompts the language level, thereby indirectly compressing KV cache, such as LLMLingua (Jiang et al. 2023
"hanss//tugginnfaccccoogrartttna
2

---

Published as a conference paper at ICLR 2025
Similar to model weight quantization, there are also attempts quantize KV cache, such KIVI (Liu et al., 2024b) and SmoothQuant (Xiao et al. 2023). These compression or quantization works orthogonal to method and can be used in conjunction.
Sparsity in LLMs. The sparsity of attention long-context scenarios has been observed by previ- studies (Liu et al. 2023a; Ribar et al., 20224, Wang et al., 202I). For example, Minference (Jiang ctt al. 2024b) has demonstrated that with context of 128k, only 4k tokens are required to accumu- late 96.4% of the total attention score. However, the tokens with high attention scores vary across different generation iterations, indicating that the sparse pattern is dynamic (Tang et al.) 2024). This implies that it is necessary to compute the full attention in every generation iteration and every layer
to determine the sparse pattern. Quest (Tang et al. 2024) and SparQ (Ribar et al., 2024) employed approximate attention methods to circumvent the computationally expensive full attention. Infini- Gen (Lee et al. 2024) introduced cross-layer similarity between consecutive two layers, leveraging this characteristic pre-select the critical KV cache. However, the loading time may still exceed the computation time, resulting in GPU idleness. In contrast, we observe that the sparse patterns between different layers exhibit high similarity, and thus we only compute full attention for a few layers to obtain the sparse patterns for subsequent layers. To the best of our knowledge, we are the first to highlight this.
INSIGHTS
Layey Actn Sinllarty
Context about Apollo 11
<BOS=Apolo was manned mission and Buzz Aldinn ... They spent around 135 minutes on the lunar surface. collecting 55 kllngrams rock samples. Mearwhilo ..... spleshing down in the Pacific Ocean on July 24, concluding an dight-day missian.
by Ammsronn
interLayer
0.875
Duestions:
1. When did the first manned mission to the moon take place? Who the
0.850
Llama3-88 n-8
begins with
0.825
Llama3-86
Whet does play leam,
accomplished?
LLM Output (w. Full Attention)
0
5
10
15
20
Transformer Layer
place July 20,
3. He the second person to walk on the moon, and he collected- 2:-eecteddd221.655 kilograms rock samples.
Inter Attn Similarity
0.95
"cax as Current Token !!
"<B>" as Current Taken
0.90
(b)
Sammy Moodin Soores
Gam Meednn Srows
0.85
0.4
24
Llama3.1-708
0.80
Ulama311-708
n=16
9.2
62
0.0
20
40
60
20
200
Transformer ID
(c)
(a)
Figure 1: Analysis of attention. (a) Inter-Layer Attention Similarity. This shows high similarity of important tokens/sub-contex between layers even 16 layers apart (b) An example a multihop question demonstrates the variation in important tokens across different generation iterations. (c) Analysis of variety in important tokens. We retained the set of important tokens, then calculated the cumulative attention scores of missing tokens.
Our work is grounded in three pivotal insights, which verify through experiments utilizing pop- ular models -lama-3--B-IInsttee Yi-9B-200k, and Llama-3. 1-0....1------mmm---
Intra-Layer Attention Sparsity. Studies have consistently revealed that attention matrices within LLMs layers exhibit sparsity (Wang et al., 2021 Deng et al., 2024). This characteristic implies reduced subset of tokens. Some studies have enhanced inference speed or reduced GPU memory requirements based on sparsity
that LLMs can generate nearly equivalent outputs by focusing
(Zhag et al.
2024c Li et al. 2024 Tang al. 2024). Due to the presence of sparsity, OmniKV
3

---

Published as a conference paper at ICLR 2025
Context Bank
Visill KV Cashe
Context Filter
CPU Mem
GPU Mem
SV Cachee on GPU
Diu Window
Context Selector
Offloaded
LLM
KV Cache
X
Shanelirdiee
Bull Vett Select
GPU
select (fiter Lever)
Neasert
Figure 2: Overall Framework of OmniKV decode stage. There are three types of layers in Om- niKV. In prefill stage, all layers perform full attention and generate KV cache of context. Then OmniKV offloads the KV cache generated by green layers to CPU memory. In the decode stage, orange layers perform full attention because of lower sparsity inference efficiency. Purple (filter) layers perform full attention and a) select important tokens using context selector based on attention scores calculated
observation window. Green layers b) load the subset of KV cache selected by preceding filter layers to GPU and perform sparse attention. c) Only the KV Cache on the GPU is visible to the LLM for generating next token.
utilizes only small subset of tokens in most layers. In this way, we not only reduced computation but also decreased the communication volume between CPU and GPU.
Inter-Layer Attention Similarity. We introduce the concept of inter-layer attention similarity, which defined as, a fixed subset of tokens that receive significant attention in specific layer, maintains their prominence throughout successive layers. The value of similarity of layer can also be viewed as "filter" ability of this layer, and is calculated as the mean of the summation of fixed tokens subset's attention scores in subsequent layers. Figure Ia demonstrates that, beyond certain shallow layer, the similarity for one layer to successive layers becomes exceptionally high. Although the overall similarity between layers already high, some layers exhibit higher "filter" ability than others. We refer to these layers as "filter" layers. Subsequently. these layers function as context selectors within OmniKV, identifying crucial tokens for each generation iteration and thus facilitating sparse attention in subsequent layers.
Inter-Token Attention Variability. Intuitively, the important tokens should vary throughout the generative process of LLMs, particularly under multi-task or multi-reasoning scenarios such as the Chain of Thoughts (CoT) (Wei et al.] (003). As shown in Figure [bb for a multi-hop question, highlight 12 tokens with the highest attention scores for two decoding steps respectively in the CoT scenario. We can observe that, apart from the special BOS token, the other important tokens entirely different. Meanwhile, Multi-Hop QA task (Ho et al. 2020) also confirm this variability. For each generation iteration, we compute the cumulative attention scores for the missing tokens token set, which stores 25% most pivotal tokens during the prefill stage (i.e. Heavy Hitters in H20). The spikes observed in the figure indicate that some missing tokens have significant attention scores. This phenomenon substantiates our intuition that the subsets of critical tokens identified exhibit significant fluctuations across different generative steps. Motivated by this insight, OmniKV retains all KV cache to ensure that performance remains unaffected.
demonstrated in Figure Ic empirical studies the
METHOD
Based on the aforementioned insights, we propose OmniKV, a token lelen--rooping-ffiie free and training -ree inference method. This design empowers OmniKV to sustain the performance of LLMss in multi-
4

---

Published as a conference paper at ICLR 2025
reasoning settings. As depicted in Figure 2 OmniKV comprises two pivotal modules: the Context Bank and the Context Selector.
The inference of rego-regressive LLMs be divided into two stages: 1) Prefiil, which encodes the intermediate computational state of the input prompt as KV cache K, V to circumvent redundant KV vector computations, and outputs the next token the first input for decoding; 2) Decode, which takes the token predicted from the previous decode iteration as the current token, and predicts the next token.
During the prefill stage, OmniKV initialize the Context Bank to store most "non-filter" layers' KV cache in CPU memory based on inter-layer attention similarity. In the decode phase, OmniKV adopts plug-and-play Context Selector that dynamically identifies important subsets of KV cache K, V on few "filter" layers. Then the Context Bank propagates the selections the "non-filter" layers and load the corresponding subset of KV cache in pack to GPU memory, since these layers share a same index of tokens. this way, OmniKV reduces computation and data transfer costs.
4.1 CONTEXT BANK
The proposed Context Bank utilizes inter-layer attention similarity to prefetch important tokens. In scenarios of insufficient GPU memory, the Context Bank can also asynchronously preload the corresponding KV cache from CPU memory, thereby alleviating memory constraints.
For the sake simplicity in analysis, we will ignore the batch size here. In the prefill phase, an L- layer LLM creates KV caches, (K, V.) by applying attention projection matrices W and W to hidden states h², yielding tensors in 12 R Nxxx for both keys and values. Here, N represents the length of the tokenized prompt within the prefill context, H signifies number of attention heads, and d denotes the hidden size of per attention head.
Firstly, We need to determine which layers are more effective for identifying important tokens. These selected layers are referred as "filter" layers. As illustrated in Figure Tal although the similarity between adjacent 16 layers is high, a significant gap is observed between the mean simi- larities of 8 adjacent layers compared to 16 layers. To enhance performance by reducing the distance between "filter" layers, OmniKV utilizes set of hyperparameters IL, where the size (m, m S. 3). Compared to using a single "filter" layer, the sub-context used by "non-filter" layers theoretically exhibits a higher degree of similarity when using multiple "filter" layers.
Consequently, OmniKV performs full attention on the layers within L identify small subset of important tokens T, (i E L). Due to the shallow layers' reduced sparsity, OmniKV also performs full attention without selection on preceding layers up to layer Lo, (I Lo)). Then, OmniKV utilizes only these important tokens T; as sub-context for the sparse attention layers 1, where TL, < < T4++1- Here, h." represents the hidden states the observation window at layer i. The context selector identifies important tokens that have significant attention scores over the observation window.
ContextSelector((t K;) if E L
T T
for i 2 Lo
otherwise
To avoid unnecessary GPU waiting when loading KV cache, OmniKV also performs full attention on the L-adjacent layers, denoted as (1 + 1heL. This interleaves data transfer with computation. Here, h represents the hidden states of the last token. Finally, the entire attention mechanism can be formulated as follows:
out; = Attention, (h', K,, Vi) (ttention,,h', K, (T.). (T.)))
if E or - or < Lo otherwise
(2)
OmniKV significantly reduces the sequence length to less than 10% in sparse attention layers, which leads to a decrease in time complexity. Upon identifying the critical tokens T, in L,, (where Tl = i). OmniKV retrieves the corresponding subset of KV caches KT[T), [T. (where T, < T.41) for sparse layers as a sub-context from the CPU memory.
Packed Load. Since layers between "filter" layers share the same sub-contex, tokens' index T, the KV cache for series of consecutive sparse attention layers can be packed and loaded from the CPU to the GPU at the nearest preceding "filter" layer. Consequently, OmniKV conducts only
5

---

Published as a conference paper at ICLR 2025
m, I loads, significantly reducing the slow PCle transfer overheads compared to loading each layer.
4.2 CONTEXT SELECTOR
As described in Section 4. OmniKV selects important tokens T; in "filter" layers I.. Inspired by previous works (Li et al., 2024, Xiao et al. (2024)), we propose a unified framework for token selection. OmniKV selects important tokens based score vector S, E R The score S, is calculated using a observation window h Commencing with the local window as the query states, and the full context hf the key states, compute the attention scores A;:
Q.Kg Vd
Q wyh K, - b, A = Softmax
A, EERE E
Next, to get the score S, first apply reduce-max obtain the maximum score across attention heads. Subsequently, a weighted vector a is utilized to perform a weighted summation on the attention scores. Finally, we leverage topk to identify the important tokens T, on score S,.
1- S
Σ ==
Jj max A. OShcH
E
T. arg top k(S,), ogt<[bş]
T, T R
(4)
Different values assign varying weights to local window tokens and yield distinct selection patterns. To further investigate which tokens of the observation window possess stronger "fil- ter" capability to identify important tokens T, this study explores three methods: 1) Uniform:
a = 01) bution when weighted and summed attention scores. 2) Exponential: = (2'-1h²] 3
[h? This approach implies that each token in the window produces equivalent contri- }
hh"l This
approach implies that tokens closer to the end of the window produce higher contribution. 3) Last
Token: concat(a {0){- hh' {1)). This approach implies that only consider the last token's
attention score in the window.
EXPERIMENTS
To demonstrate the effectiveness of OmniKV, we conducted extensive experiments on Llama-3- 8B-262K, Yi-9B-200K and Llama-3.1
1---0B----00----- Instructt mainly using the datasets InfinitBBench (Zhang) cct al. 2024a) and LongBench (Bai et al. (023). OmniKV demonstrated its effectiveness on 8B and 70B through experiments. Furthermore, we have explored its efficacy Llama-3.1-405B, with the results presented in See D.8 Ablation studies and minor or detailed experiments can be found in Section
Implementation. For most tasks, we adopted greedy decoding. To prevent repetitive outputs, we employed top-p decoding with = 0.95, temperature = 0.8 for summarization tasks in In- finiteBench. All performance and latency experiments were conducted on Nvidia A100 GPUs. Llama-3. 3.BBBOB utilized 4-bit weight quantization via bitsandbytes (Dettmers et al.] (021). while other models employed float16 formatting. We makes minor modifications based on Huggingface's transformers (Wolf et al. 2020). For exponential and uniform context selectors, we set the local window size to 16. The "filter" layers L are set respectively {2, 8, 18), (6, 11, 30), (4, 19, 41} for Mama-3-8B-262Kk Yi-9B-200K and
To ensure fairness in comparison, we strictly set OmniKV to retain Mem of KV caches on GPU. This means we dynamically adjust the token budget for sparse attention layers based on the length of the context prompt length. Mem% consists of two parts: First, OmniKV retains all KV caches for full attention layers, which occupies budget for sparse layers set to OmniKV can be expressed as 2|Li+La + 1.
2|L|+Lo of the total KV caches. Second, OmniKV's token T.
[ 0% of KV caches. Therefore, the Mem% of KV cache for
[p]
[P] (1 )1.+ For example, when Mem is set to 30% and the memory usage of layers I. is 25%, then the token budget is set 6.7%.
[T]

---

Published as a conference paper at ICLR 2025
Table 1: Performance single-step reasoning on LongBench (Bai et al. 2023). Italics indicate that the model uses full attention baseline. Bold indicates the best performance under the same model. refers to InfLLM''s KV cache memory budget being roughly set to a specific value due to its highly integrated implementation. indicates dropping tokens within each chunk H2O due to incompatible with flash attention. Detailed results for every sub-task can be found in Table and
Methods
Mein Single Doc QA Multi Doc QA Summarize
Few Shot
Synthetic
Code
Llama 3-80 362K
100%
29.2
22.9
24.9
65.9
43.5
48.9
39.2
30%
62.9
InfLLM
~30%
28.1
15.3
19.1
62.7
36.5
48.5
35,0
StreamingLLM
30%
19.3
17.3
18.6
49.7
11.7
52.3
28.
OmmiKY w/ uni
30%
29.6
23.3
23.7
64.0
41.5
48.4
38.4
OmmiKV w/ exp
30%
29.5
22.9
23.9
65.1
41.5
35.1
38.3
OmniKV last
30%
29.5
22.9
24.4
64.9
41.2
48.4
38.5
11-9B-200K
100%
28.6
33.3
20.2
71.2
31.0
67.3
41.9
H20
30%
26.0
33.5
17.5
68.7
66.6
40.6
InfLLM
~30%
25.9
33.6
20.0
70.8
26.3
65.9
38.8
StreamingLLM
30%
12.4
11.3
13.5
56.3
2.9
60.7
26.2
OmniKV w/ uni
30%
28.2
34.3
20.0
71.0
29.1
65.0
41.3
OmniKV w/ exp
30%
28.1
34.0
20.0
71.2
30.6
65.5
41.6
OmniKV last
30%
33.6
70.8
65.3
Llama 3.1-708
100%
42.2
44.8
25.5
68.9
58.0
55.7
49.2
H20+
20%
36.6
38.9
24.1
61.2
24.5
$4.0
39.9
InfLLM
-0%%
36.1
62.5
41.1
39.8
39.5
Streamingl.IM I.M
20%
16.7
12.7
18.7
45.6
7.5
61.1
27.0
OmniKV w/ uni
20%
40.8
43.8
23.6
67.7
53.8
OmniKV w/ exp
20%
42.0
44.3
24.6
68.4
57.6
55.2
48.7
OmmiKV last
Baselines. We employed three state-of-the-ar methods for memory reduction baselines, well as Full Attention. For methods other than InfLLM and Full Attention, we strictly limited the KV cache size %. Specific settings and implementations
be found in the Section C 1) H2O (Zhang et al. 2024c), which discards tokens based on attention scores, has been proven to be superior in performance by previous work (Yuan al. 2024). + indicates dropping tokens within each chunk in H2O due to incompatible with flash attention. 2) InfLLM (Xiao et al. 2024a),, similar our method, does not discard any tokens. It divides the sequence into blocks, then chooses a few important tokens in each block as retrieval keys, and retrieves several relevant blocks as context when decoding. 3) StreamingLLMM (Xiao et al.) 2024b) observes that the starting part of the input prompt tokens occupies a large portion of the attention scores, and thus, in addition to the sliding window approach, retains an initial window as part of the context. 4) Full Attention (original model), which does not discard any tokens and uses all tokens context, serves the theoretical performance upper bound for comparison.
Datasets. To test OmniKV's performance in single-step reasoning, we primarily used two widely applied benchmarks: 1) InfiniteBench (Zhang et al. 2024a) with an average length 145.1K, covering multiple tasks. We uniformly adopted 128K context for testing and truncated inputs exceeding 128K at the middle. 2) We tested LongBench's ' 18 tasks across multiple categories, with most tasks' average length ranging from 5K 15K (Bai et al. 2023). During testing, all models supported a context length longer than the longest sample, eliminating the need for truncation.
To assess OmniKV's performance in multi-step reasoning. utilized 2WikiMQA (Ho al.) 2020) and HotpotQA (Yang et al.) (018)) from LongBench. However, QA tasks may be biased, as questions often focus on special information, such person's birthplace, birth date, or awards received. LLMs might prioritize this type of information during the prefill stage, potentially affecting the accurate evaluation of their long-tes text capabilities. Motivated by this, we propose benchmark called Reta This task consists of dictionary and equation adding two numbers. LLMs need to use the answer from the equation to search the dictionary, find the corresponding key, and output the value associated with that key, Example and details can be found in Section
5.1 PERFORMANCE
Performance with Single-Step Reasoning, We demonstrate the effectiveness of OmniKV in the Single-Step Reasoning scenarios, which is the standard evaluation method for both benchmarks. In this format, the model receives an input and directly answers the question without providing inter-

---

Published as a conference paper at ICLR 2025
Table 2: Performance of single-step reasoning on InfiniteBench (Zhang et al. 2024a).
Methods
|SMeem] En.Sum En.QA MnLMC Ea. ZA.QA Code Debug Math. Find Rass.as Ret..umber Retr.V Avg.
LLaMA-3.B
13.3
65.9
20.8
26.5
100.0
14.4
H20+
30%
21.8
13.8
65.5
7.0
12.5
23.0
M0.5
98.4
74.0
5.2
35.2
InfLLM
44.5
29.1
24.5
OmniKY w/ Inst
30%
22.5
12.7
65.5
5.0
12.6
20.0
26.5
100.0
100.0
9.6
37.4
11-98-200K
100%
19.3
11.1
66.3
2.0
15.6
24.1
24.8
100.0
100.0
20.2
38.
H20+
21.6
66.3
InfLLM
<030%
6.7
12.1
37.1
3,5
14.8
21.5
34,0
100.0
100.0
0
32,9
OmniKY w/
66.8
23.8
19.8
38.2
mediate reasoning steps. Results from LongBench and InfiniteBench presented in Table and 2] respectively. The results indicate that OmniKV achieves the best performance and consistently performs very close to the original full attention baseline across all tasks, showcasing OmniKV V's stability. Particularly in Llama-3. 1-70B, OmniKV significantly outperforms the baselines. In some task categories, OmniKV even surpasses the baseline results. On InfiniteBench, although OmniK performance on the Math. task is slightly lower than InfLLM, it remains very close to the orig- inal model. However, InfLLM performs poorly on the KV Retrieval task. Thus, the stability of OmniK V's performance, compared to the original model, facilitates its direct application in real- world scenarios without the need for additional testing.
7M0kMMGA Score
Holpot Scoree
RStageRattr Score
50
40
30
20
25
35
40
20
25
30
40
20
25
30
35
40
Token Budget (Percentage)
Taken Budget (Percentage)
Tokan Budget (Percentage)
120
Omn KV las:
Omniil w/
Dmniii wi eop
St StremingliM
Tull Eal
--- Till OaT
Figure 3: Performance of multi-step reasoning with different token budget ratio of KV cache on three multi-hop tasks.
Performance with Multi-Step Reasoning. OmniKV dynamically selects the necessary context for attention computation. To demonstrate the necessity of drop-free and effectiveness of dynamic context selection, we tested models on three multi-step reasoning tasks and adopted the CoT output format to solve the problems. We use Exact Match the metric, meaning that as long as the standard answer appears in the model's output, is considered correct. The reason for employing a simple metrie is due to more linear metric being better able to quantify performance gaps (Levy eet al., 2024; Schaeffer et al., 2023). The results of Llama-3.1-70FF are shown in Figure 3 Results show that OmniKy V achieved the best performance under all budgets, demonstrating the effectiveness of dynamic context selection. In the 2StageRetr task, we observe that H2O's accuracy cannot exceed its budget, indicating that without prior knowledge, H2O can only randomly retain key-value pairs in task dictionary. We also conducted experiments on the Yi-9B-200K, as detailed in Section D.4
Choice of Context Selectors. In Section 4.2 three context selectors proposed. From Table and Figure 3 it is evident that in single-step reasoning settings, the exponential and last selectors perform better, while the uniform selector lags slightly but not significantly. In multi-step reasoning settings, the last selector exhibits the best performance, while the uniform selector performs some- what poorly. The exponential selector can be viewed as an intermediate state between the last and uniform selectors. The fact that the last and exponential selectors outperform the uniform selector suggests that the last may be the most optimal context selector. Moreover, the last method aligns with the pretraining paradigm of LLMs, the computation of context scores remains entirely consistent with the original model.
From engineering practice and inference efficiency perspective, last is also the simplest and has the lowest latency, as this selector does not require maintaining cumulative scores for the window. Therefore, in subsequent analysis experiments, primarily focus the last context selector.
8

---

Published as a conference paper at ICLR 2025
Decode Lntenoylltten (ms)
250
Brutal Offioed
O0M
Brutal Offiodd
90M
Full
Full
200
gous
Prefill Latency
150
IMFLIM
InfL
Omiik
Omn KV
100
Ommikt officed
100
Omn offieed
50
450
256
450
Context Longth Takens)
Contoxt Longth Takens)
Figure 4: End-to-end latency results. The left and right figures show the latency of the decode and prefill stages respectively. OOM indicates the maximum context length supported by some methods on single A100 GPU.
5.2 LATENCY AND TRADE-OFF
We evaluated the end-to-end inference latency of OmniKV using a single NVIDIA A A100 80GB GPU and 12 cores of an Intel Xeon Platinum 8369B CPU at 2.90GHz. A 32-layer LLaMA-3 1--A---8--262-
model was utilized, with "filter" layers L = {2, 8, 18). applying flash attention (Dao et al. 2022). and a batch size of 1. For different context length settings, the token budget for sparse attention was set to 2048. During decoding, 50 tokens were generated, and the mean latency per token was calculated over all decoding steps. In the prefill stage, measured the Time To First Token (TTFT).
The results are presented in Figure 4 The "Brutal Offload" approach refers to offloading all KV cache of the final 20 layers CPU, and pre-loading all KV cache for each layer 4 layers in advance. InfLLMM does not perform full attention during prefill, instead utilizing sparse attention in chunks.
Consequently, it achieves better latency at 450K context. However, employing sparse attention in the prefill stage may impact the method's performance (Yuan et al. (024).. OmniKV exhibits identical latency to full attention during the prefill stage, which is attributed the offload process being covered by the full attention computation. We also tested the original model's latency of 450K context with 3 A100s. Overall, our method demonstrates the best latency performance.
Notably, when GPU memory is sufficient, i.e., within a 128K context for single A A100, OmniKV can store the entire KV cache in VRAM. At this point, OmniKV V can still perform sparse attention with- out offloading (OmniKV w/o offload,, thus achieving 1.68x decoding efficiency of 21.0 tokens/s at 128K context. Morcover, we observe that OmniKV achieves lower latency than Full Attention at 128K. This suggests that we can simply adopt segmentation strategy for OmniKV to accelerate decoding speed at any context length longer than 32K.
Under settings that save 70% of KV cache, Llama-3-8B with a 450K context on a single A100 80G
37.50
Perf Infgened
37.25
120
Decade Labency (ms)
GPU at speed of 7.5 tokens/s. We also validated the per- formance under 512K ultra-long context on the Passkey task using Maama---BB-0045) achieving perfect results as shown in Figure 7 Under 80% memory reduction, we can run Jlama-3.1-70p with a context of 150K tokens at a speed of 4.5 tokens/s.
37.00
100
36.50
36.25
36.00
128
256
512
1024
2048
4096
4000
Trade-off between Performance and Efficiency. The smaller the token budget used for sparse attention lay- ers, the lower the latency caused by loading KV cache from CPU. We evaluated the trade-off between perfor- mance and efficiency on Llama-3-8B-2622 for various token budgets on InfiniteBench with 128K context. The results are presented in Figure Here, 6400 represents the token budget under Mem% = 30%. This finding indicates that retrieving only 128 relevant tokens per sparse layer as context can yield satisfactory average score 35.9 higher than H20+. Furthermore, selecting 1024 tokens achieves a favorable balance between performance and latency. Detailed results could be found in Table
Token Budget Sparse Layers
Figure 5: Trade-off on token budgets for sparse layers. Left axis shows the aver- age score InfiniteBench, and the right axis shows the latency of decoding.

---

Published as a conference paper at ICLR 2025
39
Inter-Laer Alt
LongBence Bench
Inter-Layer
Guqaer
Gerper
Long
38
6.
(a)
Layer ID
Avg.
Hit Rate
0.04
36
0.050
0.02
0.025
Random
5
10
Layer ID
5
10
15 Layer
20
25
30
20 Layer ID
(c)
(b)
Figure 6: Analysis of "filter" ability in layers. For (a) and (b), left: Llama-3-8B-262K, right: Yi- 9B-200K. (a) The inter-layer attention similarity, also referred the "filter capability of layers, displays similar trends across various tasks. (b) Varying capabilities to capture important tokens. (c)
Varying performance on LongBench using 1lama-3-8B-262KK
5.3 ANALYSIS
As shown in Figure Ta sparse patterns identified by some "filter" layers have higher similarity with subsequent layers than other layers. Naturally, we have the following two research questions: 1) Are these "filter" layers' ability task-dependent layers have a greater ability to identify genuinely important tokens? 3) Is the performance varied with "filter" capabilities?
this more model characteristies? 2) Which
Yask-Independe Filter Layers. To answer the first research question, we conducted experiments on different tasks collect those layers' "filter" characteristics. The results, shown in Figure ba display the trends of similarity values for Llama-3-8B and Yi-9B across various tasks. We can ob- serve that these curves for different tasks generally follow same trends, suggesting that the strength of "filter" is not task-dependent. Instead, the "filter" ability is more likely an intrinsic characteris- tic of the layers themselves. This implies that once we select appropriate hyper-parameters L, our method can be adapted to any task.
Accuracy of Context Selection. We use the CLongEval dataset (Qiu al. 2024) to test whether the layers accurately assign higher attention scores the context containing the answer. This dataset provides the reference chunk where the standard answer is located. Although this chunk covers broad range,
still compute the hit ratio as the proportion of important tokens within the reference chunk. The results are shown in Figure 6b] We observed that the layers with stronger "filter" ability demonstrate higher hit ratio compared to neighboring layers. For example, the 8th layer in Llama3-8B exhibits a peak in both Figure 66a and 6b) Likewise, the 14th layer in Yi-9B displays similar characteristics. This suggests that certain layers develop stronger capacity for important token retrieval after training.
Performance Vary with Filter Ability. We conduct experiments using Llama3-8B-262K on Long- Bench. As shown in Figure 6c| results indicate that layers 8, 10, 11, and 13 exhibit relatively superior performance, which corresponds well with the higher "filter" ability observed in these lay- ers as shown in Figure [6a] Similarly, the sudden performance decline in layer 12 and the abrupt improvement performance from layer to layer are also aligned in Figure ba Detail results can be found in Section D.6
CONCLUSION
This paper proposes OmniKV, a token-dropping- and training-free inference method, delivering 1.7x improvement in inference efficiency without compromising performance in long-text sce- narios. Moreover, O is highly compatible with offloading techniques, significantly reducing KV cache memory consumption. The method is simple implement and has promising practical application prospects. In our future work, we plan to explore the integration of OmniKV with KV cache quantization techniques to further minimize the usage of KV cache and enhance efficiency.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
The hyper-parameters. hardware environment, decoding methods, and other pertinent details are presented in Section [ The core algorithm for OmniKV is provided Section B Detailed settings
and implementation of baselines can be found in Section C
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303. 087744000777... 2023.
Joshua Ainslie, James Lee- Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head check- points, 2023. BRLLtttp://
http///axvvv.... org/ abs / 2305. 13245.
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. arXiv preprint arXiv:2308. 14508, 2023.
Chi-Min Chan, Weize Chen, Yusheng Su, Jianxuan Yu, Wei Xue, Shanghang Zhang, Jie Fu, and Zhiyuan Liu. Chateval: Towards better Ilm-based evaluators through multi-agent dehate, arXiv preprint N.Xiv:2308.070000 2023.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-conteds large language models. arXiv preprint arXiv:2309. 2023.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022. URL htt.psss///eiii.org//////////
2205. 14135
Yichuan Deng, Zhao Song, and Chiwun Yang. Attention is naturally sparse with gaussian distributed input, 2024. URL httt:ttt::////t //tpsii///axv..////a wwa.xiiv.oog/aaae hbtp///ss/2244... 02690
Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wiss quantization. arXiv preprint arXiv:21 2021.
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for 1lms. In The Twelfth International Conference on Learning Representations, 2024a.
Tao Ge, Hu Jing, Lei Wang, Xun Wang, Si-Qing Chen, and Furu Wei. In-context autoencoder for context compression in a large language model. In The Twelfth International Conference on Learning Representations, 2024b.
Xanh Ho, Khoa Duong Nguyen, Saku Sugawara, and Akiko Aizawa. Constructing a multi- hop qa dataset for comprehensive evaluation of reasoning steps, 2020. URLhttss:///axt orq/ orc/abs// 2011. 01060
Huiqiang Jiang, Qianhui Wu, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. LImlingua: Compressing prompts for accelerated inference of large language models. In Proceedings of the 2023 Confer- ence on Empirical Methods in Natural Language Processing, pp. 13358-13376. 2023.
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin-Ye Yew Lin, et al. Minference 1.0: Accelerating pre-filling for long-context Ilms via dynamic sparse attention. arXiv preprint vrXiv:2407.02499 2024a.
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo. Surin Ahn. Zhenhua Han, Amir Abdi, Dongsheng Li, Chin- Yew Lin, et al. Minference 1.0: Accelerating pre-filling for long-context Ilms via dynamic sparse attention. arXiv preprint arXiv:2407.02490. 2024b.
11

---

Published as a conference paper at ICLR 2025
Woosuk Kwon, Zhuohan Li. Siyuan Zhuang. Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Prin- ciples, pp. 61 1-626, 2023a.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Prin- ciples, pp. 111-626, 2023b.
Wonbeom Lee, Jungi Lee, Junghwan Seo, and Jaewoong Sim. [InfiniGen]: Efficient generative inference of large language models with dynamic {KV] cache management. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), pp. 155-172, 2024.
Mosh Levy, Alon Jacoby, and Yoav Goldberg. Same task, more tokens: the impact of imput length on the reasoning performance of large language models, 2024. URL httpps://axxii..... abs/2400 14848
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, 2024. URL
www.ooggggbbbaaeeeee.eee 2404. 14469
Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Re, et al. Deja vu: Contextual sparsity for efficient Ilms at inference time. In International Conference on Machine Learning, pp. 22137-22176. PMLR, 2023a.
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 IIm kv cache compression at test time. Advances Neural Information Processing Systems, 36, 2024a.
Zijun Liu, Yanzhe Zhang, Peng Li, Yang Liu, and Diyi Yang. Dynamic Imm-agent network: An llm- agent collaboration framework with agent team optimization. arXiv preprint 10X0022310002170, 2023b.
Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu, Kivi: A tuning-free asymmetric 2bit quantization for kv cache. arXiv preprint arXiv:2402.022 02750, 2024b.
Jinjie Mai, Jun Chen, Guocheng Qian, Mohamed Elhoseiny, Bernard Ghanem, et al. Lim a robotic brain: Unifying egocentric memory and control. 2023.
Meta. The llama 3 herd of models, 2024. URL http://ss//aaaa...oo//////000000
ModelTC. Lightllm: An efficient lightweight IIm serving framework. ModelTC/Lightlnn 2023.
com/
Jesse Mu, Xiang Li, and Noah Goodman. Learning to compress prompts with gist tokens. Advances in Neural Information Processing Systems, 36, 2024.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint wwXiv:230900071, 2023.
Zexuan Qiu, Jingjing Li, Shijue Huang, Wanjun Zhong, and Irwin King. Clongeval: chinese benchmark for evaluating long-context large language models. arXiv preprint arXiv:2403.03514, 2024.
Luka Ribar, Ivan Chelombiev, Luke Hudlass-Galln Charlie Blake, Carlo Luschi, and Douglas Orr. Sparq attention: Bandwidth-efficient Ilm inference. In Forty-firss International Conference Machine Learning, 2024.
Rylan Schaeffer, Brando Miranda, and Sammi Koyejo. Are emergent abilities of large language models mirage?, 2023.
15004
12

---

Published as a conference paper at ICLR 2025
Ying Sheng, Lianmin Zheng. Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Ré, lon Stoica, and Ce Zhang. Flexgen: tiightthoughhut generative inference of large language models with single gpu. In International Conference Machine Learning, pp- 31094-31116. PMLR, 2023.
Jiaming Tang, Yilong Zhao. Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. Quest: Query-aware sparsity for efficient long-context I1m inference. In Forty-first International Confer- ence on Machine Learning, 2024.
Hanrui Wang, Zhekai Zhang, and Song Han. Spatten: Efficient sparse attention architecture with cascade token and head pruning. In 2021 IEEE International Symposium on Performafee Computer Architecture (HPCA), pp. 97-110. IEEE, 2021.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed Chi, Quoc Le, and Denny Zhou. Chain-of-thought prompting elicits reasoning in large language models, 2023. URL
org wwtpz///abb//22/0 abtp://..ss/22...1111
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowiez, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Huggingface 's transformers: State-of- the-art natural language processing, 2020. URL https:tp///cm /ttpp::///rxiia wttp://...r//ab//11..007777
Qingyun Wu, Gagan Bansal, Jieyu Zhang, Yiran Wu, Shaokun Zhang, Erkang Zhu, Beibin Li, Li Jiang, Xiaoyun Zhang, and Chi Wang. Autogen: Enabling next-gen I1m applications via multi- agent conversation framework. arXiv preprint arXiv:2308.088155 2023.
Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, and Maosong Sun. Infllm: Training-free wow.eeeooooocooootttteeteotttpem extrapolatio for Ilms with an efficient context memory, 2024a. URL
220.aaab./2220.02 04617
Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pp. 38087-38099, PMLR, 2023.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Rep- resentations, 2024b.
An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li. Dayiheng Liu, Fei Huang, et al. Qwen2 technical report. arXiv preprint arXiv:2407.1067. 2024.
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhutdinov, and Christopher D Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. arXiv preprint arXiv: 1809:180.00000 096000 2018.
Alex Young, Bei Chen, Chao Li, Chengen Huang, Ge Zhang, Guanwei Zhang, Heng Li, Jiangcheng Zhu, Jianqun Chen, Jing Chang, et al. Yi: Open foundation models by 01. ai. arXiv preprint 2403022003. 2024.
Jiyy Yuan, Hongyi Liu, Shaochen, Zhong, Yu-Neng Chuang. Songchen Li. Guanchu Wang, Duy Le, Hongye Jin, Vipin Chaudhary, Zhaozhuo Xu, Zirui Liu, and Xia Hu. Kv cache compression, but what must we give in return? comprehensive benchmark of long context capable approaches, 2024. URL https: http://////aa..o..///////0// 01527
Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang, Zhiyuan Liu, and Maosong Sun. sobench: Extending long context evaluation beyond 100k tokens, 2024a. URL PRRItttps:: /////xaxvvvv
wtapixxivv.orr.. org. abs/ 2402. 13718
Yang Zhang, Shixin Yang. Chenjia Bai, Fei Wu, Xiu Li, Xuelong Li, and Zhen Wang. Towards efficient Imm grounding for embodied multi-agent collaboration. arXiv preprint arXiv:2405. 2024b.
13

---

Published as a conference paper at ICLR 2025
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 gen- erative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024c.
14

---

Published as a conference paper at ICLR 2025
A 2STAGERETR
To mitigate the influence of these priors on the results, we propose the benchmark 2StageRetr. The main idea behind this benchmark is to construct a two-step reasoning problem, using the result of the first step to complete the second step. The task description is placed at the end of the prompt, preventing the model from "mentally calculating" the result of the first step. Thus avoids LLMs assigning high attention scores answer tokens based on mentally computed results. Specifically, the 2StageRetr consists of dictionary composed of multiple pairs of numbers and colors, followed by an addition equation. We ensure that the result of the addition is always key in the dictionary. Then the model needs to find the corresponding color in the dictionary based on the result of the addition, and output it. An example of 2StageRetr is shown below.
2StageRetr Example
Lets play game. You have a dict and a mathematical addition equation. The keys of dictionary can be any number. You need to find the corresponding key value the dictionary after performing the addition and output value correspouding to that key. The Dict {0: lime, yellow. 2: red, 3: black, 4: brown, 17: brown, 18: maroon, 19: teal, 20: red, ..... 28: brown. 29: violet) The equation is 8 + 10=? ? Answer the corresponding color based on the addition result.
Example output: Since
18, the corresponding color "maroon".
This design essentially ensures that LLMs, without prior knowledge of the problem, use pre-trained prior knowledge to assign higher attention scores to certain key-value pairs in the dictionary. Con- sequently, the H2O method is fundamentally akin to randomly discarding key-value pairs from the dictionary. Hence, its performance cannot surpass the proportion of its token budget. In contrast, our method selects the most relevant context based on the information within the observation window each time, thereby achieving superior performance.
We set the dictionary size to not exceed 200, with numbers arranged sequentially from to the maximum value. This arithmetic sequence should be relatively simple for current LLMs, thereby primarily assessing the model's 's retrieval capability.
The average length of this dataset is only 739 tokens, with the maximum length being 1382 tokens. However, current cache drop-based methods do not perform very well on this task.
B DETAILED IMPLEMENTATION
The actual code implementation required only a minor modification of the code sourced from Hug- gingface's 's Transformers library (Wolf al. (020).. Here we present the core pseudo code of Om- niKV in A This algorithm demonstrate the attention forward procedure of one layer.
C BASELINES SETTINGS
Here we provide detailed description of the settings for different baselines.
1) H20. As H2O requires the output of attention scores, which incompatible with flash attention, the intermediate activation values directly lead to out of memory errors when processing long se- quences. Therefore, when the context length exceeds 60K or when using models larger than 30B, must calculate attention in chunks and remove tokens based on the attention score within each chunk. We denote this approach with + in our results. To ensure fairness in comparison, we strictly set H2O to retain Mem% of KV caches. We also modified code for supporting GQA (Ainslie al. 2023 in Llama-3 and Yi.
2) InfLLM. Since InfLLM uses LRU-based block cache and has many hyper-parameten limiting the KV cache by percentage leads to efficiency degradation and performance decline. Therefore, InfLLM is configured
average of Mem% of KV caches.
15

---

Published as a conference paper at ICLR 2025
Algorithm 1: Attention Forward of OmniKV
Input: observation window's hidden states h,", filter layers I. context cache K, V, visible KV
cache K", V current layer i, attention query weight W7, token budget k, window
weights a
if / E IL then
Q
h
get query states
Q!'? Va
A
t Softmax
get attention scores
S, + Σド= 1h.1- 1-1 maxosA<H [,
> get context scores
T2 t arg top k(S,)
> select important tokens
ost<[bb]
1 < get the index of L
for j = L + L++1 do
Kj. V + LL V,TT))
load subsets of KV cache to GPU
end
end
E or i << Lo or i 1- E L then KJ V + V.
original cache for layers performing full attention
end
Finally, perform normal Attention with visible KV cache KY, VT.
3) StreamingL.LMM. We strictly set StreamingI.LMM LLM retain Mem% of KV caches with 1 %% for initial sink tokens and Mem 1% for local window tokens.
4) Full Attention (original model). To avoid excessive GPU memory consumption by intermediate activations in the MLP, we split the hidden states along the sequence length dimension into chunks of size 4000. This reduces peak memory usage without affecting efficiency.
D ABLATION STUDIES AND DETAILED EXPERIMENTS
D.1 NEEDLE-IN-A SEEDLE-INAA-HAYSTACE HAYSTACK
As OmniKY employs an 8B model capable of supporting 450K context on a single A100 GPU, conducted needle-n-aahhyssackk tests with maximum context length of 450K, achieving fully accurate performance. We continued to use LL = {2, 8, 18). with token budget of 1024 for sparse layers. We also conducted tests with an input length of 512K single NVIDIA H20, also achiev- ing entirely accurate results. The results are shown in Figure
D.2 DETAILED TRADE-OFF
Here, present the detailed trade-off between efficiency and performance. Our test results InfiniteBench are shown 3 The latency here refers end-to-end time of per token in decoding stage.
Table 3: Detailed results trade-off between performance and efficiency.
Budget
En.Sum
En.QA
En.MC
En.Daa
Zh.QA
Code.. Code..ehuug
Math. Find
RT.psskky
RT.Num
RTKVV
Avg.
Latency (ms)
256
19.4
13.8
62.9
60
11.8
20.6
26.6
100.0
100.0
2.8
36.4
$1.9
512
20.8
12.9
63.3
4.5
11.8
20.6
26.6
100.0
100.0
5.8
36.6
55.7
1024
2048
21.4
12.9
65.5
30
12.2
20.3
26.6
100.0
100.0
7.6
37.0
78.8
4096
20.7
12.6
65.5
4.5
12.4
20.3
26.6
100.0
100.0
8.6
37.1
95.5
16

---

Published as a conference paper at ICLR 2025
Needle in A Haystack Llama-3-88 3-8B-1M w/ OmniKV 512K Context
1.0
Depth Percent (%)
0.8
33
0.6
56
0.4
67
0.2
89
100
0.0
10k
16k
32K
64k
128K
256k
512K
Context Length
Figure 7: Needle-in-a-hen in-a-haysttcck test. Results indicate that OmniKV applied Llama-3-8B-In 3-8B-1M
achieves perfect retrieval results.
D.3 DETAILED LONGBENCH RESULTS
We present the detailed results on LongBench, showing the scores for each subtask, as shown in
Table 4 and
Table 4: Detailed results on LongBench (Part 1).
Single-Decament QA NarrativeUA
Mulli-Ducument QA HotpotQA
Summariration Gov Report
Model
MuttiFielQQA EN
Qusper
2WikiMQA
Durender
Musique
qmium
vesum
43.6
16.0
28.2
21.7
28.3
27.5
14.3
$4.5
25.1
15.1
HZO
39.8
16.8
26.5
19.3
23.9
24.5
13.2
$1.2
25.3
15.2
StreaminLLLM
22.1
13.4
22.5
20.1
12.9
26.]
12.3
26.5
21.4
8.1
IatLLM
41.2
19.8
8.4
24.2
0.2
OmniKV w/ exp
43.7
28.3
28.5
26.7
14.5
33.6
24.5
13.8
E--8-200K
37.1
13.6
35.3
35.6
19.4
50.
27.5
50.7
20.4
9.8
H20
34.4
14.9
29.1
35.5
20.1
50.
28.7
22.5
20,3
9.4
17.1
2.6
17.7
19.2
10.8
10.9
4.5
19.0
15.9
5.9
IntLLM
25.6
15.2
37.1
31.5
23.3
$3.3
26.5
20.7
9.2
20.4
OrwiKV w/exp
27.8
9.0
Llama 3.1.00
56.9
27.9
44.1
54.5
31.2
$8.8
34,8
24.0
17.5
H20L
45.1
29.5
35.4
42.7
27.3
52.7
32.9
$2.9
23.0
16.6
SLL
25.3
19.1
14.9
19.5
10.0
15.5
4.4
26.7
19.8
9.8
IrfLLM
502
24.7
43.2
47.3
32.7
43.2
21.4
19.9
20.1
15.3
OmniKV w/
54.8
28.0
33.9
24.2
16.5
w/exp
28.3
D.4 MULTI-STEP REASONING RESULTS OF Y1-9B-200K
To further validate the effectiveness of OmniKV in multi-step reasoning, we continued experiments using Yi-9B-200K. The average length of 2StageRetris 739, so the token budget 739 x0.067 49. To avoid an excessively low budget, while ensuring faimess in comparison, we set the "filter" layers
L (3, 11, 30} to allocate more token budgets for sparse layers.
As shown in Table 6 OmniKV achieved the best results across all three datasets, particularly under a constrained token budget. This further validates the effectiveness of our approach. Yi-9B-200K does not follow instruction to directly answer in 2StageRetr.
D.5 ABLATION STUDIES
OmniKV comprises Context Selector and a Context Bank, which are highly coupled modules. When the Context Selector is removed, implying the computation of entire context, the method
17

---

Published as a conference paper at ICLR 2025
Table 5: Detailed results on LongBench (Part 2).
Fow-Shot
Syntheic Tasks
Codc
Model
Isht
trec
TriviaQA
Passage Count
Passage Retrieval EN
Icc
RepoBench-p
Llama-3-88- 3-8B-262K
44.5
69.5
83.7
0.0
87.0
52.5
45,4
H2O
35.0
68.5
85,5
0.0
85.5
42.8
43.8
StreamingLLMM
20.5
60.0
68.7
1.7
21.8
55.8
48.8
InfLLM
37.0
68.5
82.7
6,5
66.5
52.3
44.8
OmniKV w/ last
44.5
69.5
80.9
0.0
82.5
51.4
45.4
OmniKY uni
41.5
68.5
82.0
OmniKV exp
43.5
69.5
82.5
0.0
37.5
32.9
Yi-9B-200K
48.5
78.5
86.8
3.0
59.0
71.6
63.1
H2O
42.0
77.5
86.8
2,8
59.5
70.7
62.6
StreamingLLM
24.3
69.5
75.2
3.1
2.8
64.4
$7.1
InfLLMM
48.0
77.0
87.6
3,2
49.5
69.4
OmniKY
48.5
78.5
69.5
OmniKV uni
48.0
78.5
2.8
55.5
68.9
OmniKV exp
48.5
78.5
86.6
3.3
58.0
69.7
61.3
Maama-3. 1.708
46.0
75.0
85.9
18.5
97.5
48.2
63.3
H2O+
30.0
68.5
50.4
StreamingLLM
4.0
54.5
78.4
1.5
67.3
InfLLM
36.0
65.5
86.0
5.5
76.8
32.3
47.4
OmniKV
44.5
74.3
85.9
18.0
97.3
49,2
60.5
OmniKV w/ uni
42.0
74.5
86.8
18.0
97.5
47.8
59.8
OmniKV w/ exp
44.0
75.0
86.4
18.0
97.3
48.7
61.8
Table 6: Multi-step reasoning results of Yi-9B-200K.
Model
CoT
%Mem
2WikiMQA
HotpotQA
2StageRetr
Yi-9B-200K
100.0%
27.0
40.0
Vi-9B-200K
100.0%
64.0
48.5
36.0
H2O
30.0%
55.5
45.5
13.0
H2O
40.0%
60.5
47.0
19.0
H2O
50.0%
64.0
50.0
26.0
OmniKV V
30.0%
58.0
51.0
14.0
OmniKV
40.0%
62.0
51.0
32.0
OmniKV
50.0%
64.0
50.0
31.0
becomes equivalent to full attention if GPU memory is sufficient. Otherwise, we have to offload some layers' KV cache, then the approach becomes identical to Brutal Offload. Although the per- formance of Brutal Offlood is entirely equivalent to that of the original model, the frequent loading and offloading during the decode stage incurs significant overhead.
Without the Context Bank, indicating that we do not apply inter-layer attention similarity, we can only utilize the Context Selector at each layer to choose the token set However, selecting impor- tant tokens/context itself requires full attention at the current layer, which means efficiency would be significantly reduced. The only potential benefit might using shorter context to avoid inter- ference from irrelevant information, possibly leading to better performance.
Table 7: Detailed results on trade-off between performance and efficiency.
Variant
En.Sum
Ea.MC
En.Dia
Zh.QA
Code Debug
Math.Find
passkey
KT.Nam
KT.KV
OmniXVV
OmniKV wio CB
21.8
13.5
61.5
13.6
21.5
26,2
100.0
100.0
9.8
37.2
184.4
OmniK
22.5
12.7
65.5
5.0
12.6
20.0
26.5
100.0
100.0
9.6
37.4
$4.3
D.6 PERFORMANCE OF FILTER LAYERS
Theoretically, if we select "filter" layers with stronger "filter capabilities, we can expect improved performance. However, ensuring completely fair performance comparison when testing different
18

---

Published as a conference paper at ICLR 2025
selected layers is challenging. This is due to the sequential nature of the layers; if ensure that the total number layers performing full attention remains consistent, there can still be issues with uneven spacing between "filter" layers. Attempting to evenly distribute the "filter" layer set introduces additional variables. Nonetheless, experimental results indicate that layers with enhanced "filter" capabilities tend to exhibit superior performance to some extent. We conduct experiments on both LongBench and InfiniteBench. The results are shown in Table
Table Performance on LongBench of different filter layers settings.
Filler Layers
Single- Doc QA
Multi-Doc QA
Summarizatiin
Few-Shot
Synthetic
Code
Avg.
2,4,18
27.0
19.9
21.4
58.1
46.0
35.8
2,5,18
29.3
22.8
23.4
63.4
42.5
47.1
38.1
2,7,18
29.3
22.1
23.6
65.2
41.8
47.2
38.2
2,8.18
29.6
22.9
24.4
65.0
41.3
48.4
38.6
2,9.18
28.6
22.5
24.1
65.1
41.8
47.4
38.2
2,10,18
30.2
23.0
24.0
65.1
42.8
47.5
38.8
2,11,18
31.0
22.7
23.9
65.4
43.0
48.2
39.0
2,12,18
29.8
21.2
23.0
63.2
42.1
47.8
37.8
2,13,18
30.1
22.5
23.7
65.4
42.5
48.2
38.7
D.7 AN EXAMPLE OF INTER-LAYER ATTENTION SIMILARITY MAP
To directly observe the "filter" capability of layers, we also demonstrated the similarity between any two layers. We visualized the cumulative attention scores for the top-2048 token set from Llama-3- 8B-262K on HotpotQA as a measure of similarity, consistently as before. shown in Figure X for layer& 8, despite being 12 layers apart from layer 20, the token set Ts selected by layer still achieves a cumulative attention score 0.87 at layer 20. This substantially demonstrates the effective ability of layer 8 to select important tokens.
D.8 EFFECTIVENESS OF OMNIKV ON LARGER MODELS
We first visualized the Inter-Layer Attention Similarity Llama-3.1 405B, as illustrated in Figure The 405B model continues to exhibit remarkably high inter-layer similarity, indicating that OmniK can be effectively applied to it. Additionally, we conducted further evaluations on two tasks from LongBench, Qasper and Qmsum. And the results presented in Table 9
Table 9: OmniKV in 405B.
Setting
qasper
qmsum
Llama 3.1 405B
50.0
25.5
OmniKV 48.5
25.9
D.9 COMPATIBILITY OMNIKV
In this paper, most of the experiments are conducted using the Huggingface Transformers. However, this framework is typically used inference, such vLLM (Kwon et al. 2023a) and LightLLM (ModeITC) (023), are much faster than those based on Hugging Face Transformers. Therefore, we adapted OmniKV one of them, LightLLLM,, demonstrate that OmniKV has excellent usability in real-world scenarios. The results are shown Table IO
inference engine. Current engines used for large model
Continuous Batching, Continuous Batching is a technique used in the training and inference of Large Language Models (LLMs) to optimize computational efficiency, particularly in scenarios where the model processes sequences of varying lengths. The core idea behind continuous batching to dynamically group input sequences into batches that can processed in parallel, even the sequences have different lengths. OmniKV can be natually integrated with Continuous Bathcing. Based on lightLLM, we find that the continuous batching technique decrease the average latency of multiple 6en-thh requests from 41.7s 7s to 37.8s.
19

---

Published as a conference paper at ICLR 2025
10
15
20
25
30
1:0
cat
3:88
CRT
031
CA3
C40
C.0
535
048
C30
AJR
037
1M0
671
672
058
5.4
5.4
0.70 8:2
5.95 955
0.5
0.7
6.0
6.6
5:4
06B
0.99
649
SMM
0.11
0.0
0.43
071
579
631
5:00
0.4
9.41
0.8
5.0
0.22
5.5
GMI
034
5:0
082
0.29
5:0
0.6
OST
035
0.4
0355
6:00
C10
1:04
045
045
0.00
4:75
0:04
6.0
0.0
Layer
3.07
032
2.84
SNI
5:09
09.
082
5.07
5:N
0.4
20
032
0:11
0:50
25
2.00
6.35
0.5
30
Layer ID
Figure 8: An Example of Inter-Layer Attention Similarity Map.
Table 10: Average latency of per token for each request (ms/token, tp=4).
Setting
LightllmOmnikV
LightIlm+origin
vLLM+origin
128k, bs=16
46.2
73.5
72.3
256k, bs=8
44.9
75.4
73.1
512k, bs=4
44.9
78.1
75.6
Tensor parallelism. Tensor parallelism(TP) is a technique used in large language models (LLMs) and other deep learning models to distribute the computation of tensors (multi-dimensional arrays) across multiple devices, such as GPUs. This approach particularly useful for training and inference with very large models that cannot fit into the memory of single device. Using LightLLM, we have implemented OmniKV with tensor parallelism. Experimental shown in TTble[TTresulle indicate that with TP enabled, OmniKV V can still achieve considerable acceleration compared the original model. Another observation is that Lightllm+OnnikY experiences reduced decoding latency as sequence length increases. This is because the "non-filter layers" have fixed sequence length of 2048. As the batch size decreases, the computation also decreases for mnnn-filee layers".
More technical details of TP. The attention heads is distributed across different GPUs in tensor parallelism, and the outputs are aggregated using all-reduce, making inter-GPU communication OmniKV maintains indices of important tokens for
key factor in efficiency. For tp each card. Each index of important tokens is computed solely by the card's own attention heads, thereby eliminating the need for inter-GPU communication. Since different attention heads may
20

---

Published as a conference paper at ICLR 2025
Inter-Layer Attn Similarity
0.95
0.90
0.85
Llama3.1-405B n
0.80
Jlama3.1-405B nn
25
50
75
100
Transformer Layer ID
Figure 9: Inter-Layer Attention Similarity of Llama-3.1 405B.
Table 11: Throughput results of decoding (tp=4), using
Setting
Lightllm+OnnikY
Lightllm+origin
vLLM+origin
128k, bs=16
345.9
217.7
221.3
256k, bs=8
177.5
106.1
109.4
512k, bs=4
89.1
51.2
52.9
have distinct indices, the system intuitively exhibits greater flexibility. We have also substantiated this through experiments(Shown in Table 12
Table Performance comparison across datasets.
Setting
2WikiMQA
Qasper
HotpotQA
lcc
gmsum
tp=8
21.7
28.0
27.3
52.3
25.9
tp=1
21.8
28.6
27.0
51.4
24.9
Pipeline parallelism. Pipeline parallelism(PP) is technique used parallelize the execution of LLMs across multiple devices or processors. This approach particularly useful for deep neural networks, where the model can be divided into smaller, sequential stages, and each stage can be processed on different device. OmniKV is also natually compatible with pipeline parallelism. The experiments of Llama-3.1-70B shown in Table I is conducted using pipeline parallelism.
Context parallelism. Context parallelism(CPP) is technique used in large language models (LLMs) improve training and inference efficiency by distributing the processing of input sequences across multiple devices or processors. This approach leverages the fact that many LLMs, such as transform- ers, process input sequences in a way that allows for parallel computation. Indeed, when OmniK V is used, there is no need to enable context parallelism. This is due to OmniKV only needs around 2048 tokens to achieve excellent results. Nevertheless, OmniKV can still integrated with context parallelism. In this case, the communication latency in multi-GPU parallelism often exceeds the computational latency a single GPU. Hence, we can employ context parallelism only in dense (full) attention layers, while utilizing single GPU for sparse attention layers. This means sim- ply need to configure which layers will implement context parallelism. can allocate samples to different GPUs when the batch size is larger than to prevent resource wastage.
21

---

Published as a conference paper at ICLR 2025
D.10 PREFILL ACCELERATION
OmniKV is designed to enhance decoding speed. However, given the importance of prefill acceler- ation, have modified OmniKV optimize the prefill stage well. Experimental results indicate that OmniK V-prefill achieves 1.90x reduction in latency without performance loss compared to the original model with an input length of 256K. The code here is implemented using the Transformers library, enhanced with Flash Attention.
Specifically, during the prefill stage, we divide the sequence dimension into multiple chunks for computation. For each chunk, we also use the "filter layer" to select important tokens. Since the sequence length of the query chunk is no longer 1, we employ the uniform computation method introduced in Section 4.2 to select important tokens in the "filter layer".
Table 13 and Table 14 show latency and quality preserving results. Here, "r" denotes the size of the selected tokens, with the ratio of token size to chunk size being "r". Thus, the larger the value of r, the less information lost. The results indicate that r=2 has achieved fine balance between efficiency and performance.
Table 13: Prefill latency (s) single NVIDIA-H20.
Setting
256k
128k
64k
32k
16k
OmniKV-prefill, r=1
105.9
36.7
14.3
6.2
2.8
OmniKV-prefill, r=2
107.2
38.8
15.9
6.5
2.9
OmniK r=4
112.9
42.8
17.1
6.8
3.0
OmniK V-prefill, r=
133.3
49.2
19.0
6.9
2.9
MInference Llama-3-8B-261 (origin)
61.6
28.1
13.0
6.5
3.7
203.8
58.6
18.3
6.5
2.6
MInference (Jiang et al. 2024a) specifically designed for Prefill acceleration and optimizes mul- tiple CUDA kernels. But OmniKV solely relies on the Huggingfaceetrrassffoness library, conse- quently less efficient.
Table 14: Quality preserving results on LongBench.
Setting
Mem
Single-Doe QA
Multi-Doc QA
Summarization
Few-Shot
Synthetic Tasks
Code
Llama-3-8D 3-m---3---22K (origin)
100%
29.2
22.9
24.9
65.9
43.5
48.9
Minference
30%
22.4
25.3
65.8
49.8
OmniK V-prei---efillrrr
30%
29.8
22.0
24.0
63.5
50.0
OmnikV-prefil r=2
30%
30.0
23.1
24.3
65.1
40.0
49.6
Omnik K V.
30%
29.8
22.9
24.3
42.0
49.1
Omni V-prefill,
30%
29.8
22.9
24.3
65.9
49.0
OmniK V-prefl-reffil,,, all
30%
29.6
23,3
23.7
64.0
41.5
48.4
22


Reviews:
[{"review_id": "trj2cZuJ9J", "reviewer": "Reviewer_diVc", "summary": "The inference of large language models (LLM) is memory intensive. This paper proposes a new method for reducing the GPU memory usage when doing LLM inference on GPUs. The main insight of this paper is the similarity of attention score matrix between different consecutive layers. By leveraging this insight, the author proposes a new method to reduce the GPU memory usage of the KV cache  by over 75% via offloading to CPU memory, while not significantly affect the accuracy.", "questions": "Thanks for submitting the marvelous paper to ICLR. In general, I believe the paper is trying to solve a timely and important problem. While I appreciate the nice results the proposed methods can achieve, my major concern of this paper is that it introduces some hyper-parameters that could be non-trivial to tune. So, I am wondering how would you tune the filter layers in your designed algorithm?\n\nBesides, since the author does not provide any intuition or mathematical guarantee in the paper, I would be concerned about the generalizability to bigger LLMs such as Llama 3.1 405B. So I am wondering if you have some intuition to explain why attention score matrix could be similar across different layers? While I understand it could be expensive to run experiments with larger LLMs, do you have any intuition that the proposed method will work even with larger LLMs?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The author aims at a timely problem of reducing memory usage of KV cache for LLM inference.\n* The proposed method is effect yet simple to implement.\n* The author demonstrate real performance gain and capability to run LLM inference with longer sequence length with real workloads.", "weaknesses": "* There is no mathematical guarentee that the propose method is universally applicable to all LLMs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "vHPptjg3TT", "reviewer": "Reviewer_ztbP", "summary": "This paper introduces a KV cache selection algorithm based on cross-layer similarity and shows reasonable downstream performance and wall clock speedups.", "questions": "As stated in weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper explains the cross-layer KV cache selection similarity well and designs a real-world system based on this. The evaluations include the 70B-level model's results, and the decoding speed is also very impressive.", "weaknesses": "1 Inference-time pipeline parallelism, especially in the same node, is uncommon.   \"which is 1.87 times faster than the original model running on three A100 GPUs in pipeline\" does not make sense to me.  A tensor parallelism baseline is expected in this case, even if I can understand it is very hard to beat.      \n2 This paper uses cross-layer KV cache selection similarity. However, this idea was introduced by InfiniGen in a previous paper. No proper discussion is seen in this paper. If this is solved, I will raise my score.     \n3 The baselines shown in this paper are not proper. Only InfLLM is a dynamic KV selection method, but it is mainly used for extending context windows. I will ask for comparisons with Quest or Loki, as they are also dynamic KV selection methods. Static KV selection methods like H2O can not only save inference time but also save total memory without CPU offloading. If this is solved, I will raise my score.\n\n[1] Lee, W., Lee, J., Seo, J., & Sim, J. (2024). {InfiniGen}: Efficient generative inference of large language models with dynamic {KV} cache management. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24) (pp. 155-172).    \n[2] Tang, J., Zhao, Y., Zhu, K., Xiao, G., Kasikci, B., & Han, S. (2024). Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference. arXiv preprint arXiv:2406.10774.          \n[3] Singhania, P., Singh, S., He, S., Feizi, S., & Bhatele, A. (2024). Loki: Low-Rank Keys for Efficient Sparse Attention. arXiv preprint arXiv:2406.02542.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "0TQ2N4Zxdh", "reviewer": "Reviewer_eZxa", "summary": "The paper proposes an inference optimization framework which leverage inter-layer attention similarity to offload most of the KV into CPU and select relevant context dynamically to alleviate the memory bound on GPU while maintaining the integrity of contexts.", "questions": "1. Related to weakness 1, how does this method compare to vLLM/SGLang with FlashInfern/FlashAttention backend, in terms of throughput and latency? If slower, is OmniKV compatible with practical serving framework like vLLM? \n2. Related to weakness 2, can you benchmark TP=4,8 for <10b scale mdoels? Also, tp=8, pp=4, for any 70B models? I'd like to see (1) latency and throughput for prompt length of 128k, 256k, 512k(probably need ctx parallel) and output length of 512, (2) quality test on needle-in-a-haystack on 128k, 256k, 512k ctx.\n3. Could you benchmark against [1], for both speed and quality preserving?\n\n[1] MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The method can keep full context, which was reflected in the needle and LongBench tasks. \n2. The method can achieve preserve better model quality while having more speed-up over other methods in this class.", "weaknesses": "1. The memory manipulation doesn't seem to be compatible with fundamental serving mechanisms like continuous batching and prefix caching. It's questionable how such algorithm can work in real world.\n2. It will be quite challenging to make the proposed method work with tensor parallel,pipeline parallel, and context parallel, which are also essential in production.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "jB4zZc91rx", "reviewer": "Reviewer_UJXu", "summary": "The paper presents OmniKV to reduce GPU memory usage when serving Large Language Models (LLMs) with long contexts. OmniKV is based on the observation that the importance of history tokens cannot be determined by the current stage attention score. Thus, permanently discarding some of past KV tokens may incur accuracy loss. OmniKV is a token-dropping-free and training-free inference method that reduces KV cache GPU memory usage by over 75% without performance degradation, which also accelerates  inference efficiency in long-text scenarios. The key insight of OmniKV is that important tokens are very similar across consecutive LLM layers. So that the KV Cache, which is offloaded to CPU, can be loaded to GPU in the meantime for different layers. Results demonstrate that OmniKV can handle a 450K context with a decoding latency of 7.52 tokens/s, which is 1.87 times faster than the original model running on three A100 GPUs in pipeline mode.", "questions": "1) What is the dense baseline implementation used for efficiency evaluation in Figure 4?\n2) In the last sentence of abstract, the authors mention that “By using a single A100 and Llama-3-8B, OmniKV can handle a 450K context with a decoding latency of 7.52 tokens/s, which is 1.87 times faster than the original model running on three A100 GPUs in pipeline.” I would appreciate if the authors could elaborate i) what is the baseline setting for “original model running on three A100 GPUs in pipeline”; ii) what will be the latency for “original model” if the model is served with tensor parallelism rather than pipeline parallelism.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1) This paper focuses on a timely and important problem: Reducing the KV Cache memory usage of long-context LLMs. \n\n2) OminKV greatly reduces the GPU memory consumption when serving long-context LLM through offloading. It reduces the LLM inference latency by only loading important history tokens to GPU for attention computation. OmniKV well preserves model’s performance on long-context benchmarks such as LongBench and InfiniteBench.\n\n3) OmniKV notices that the important tokens for consecutive layers are similar, merging the CPU-GPU memory loading procedure for multiple layers to reduce memory transaction and important token identification overheads.", "weaknesses": "1) The optimization techniques proposed in this paper cannot accelerate the prefilling speed of LLM. \n\n2) In Figure 4, it seems that OmniKV (w/ offload) is slower than the dense baseline until the sequence length reaches 128K. For OmniKV w/o offload, the speed also lags behind the dense baseline when sequence length < 32K.\n\n3) The dense baseline (Full) seems not to be strong enough. For Llama-3-8B model on A100, TensorRT-LLM and vLLM take less than 20ms to decode a token (sequence length = 128K, batch size = 1) . However, the dense baseline takes ~80ms,  according to Figure 4.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)