# Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation
for Efficient LLM Inference

Yuan Feng1,3,$\dagger$, Junlin Lv1,3,$\dagger$, Yukun Cao1,3, Xike Xie2,3,*,
and S. Kevin Zhou2,3

###### Abstract

Large Language Models have excelled in various fields but encounter efficiency
limitations due to the substantial Key-Value (KV) cache required for long-
sequence inference. Recent efforts try to evict non-critical cache elements
during runtime, thereby reducing cache size within given memory budgets while
preserving generation quality. Our reexamination of foundational principles
reveals that prevailing methods aim to minimize an upper bound of eviction
loss, quantified as the L1 distance between the pre- and post-eviction outputs
of multi-head self-attention mechanisms. Moreover, our analysis indicates that
the common practices of uniformly assigning budgets across different attention
heads during cache eviction hinder their budget utilization, negatively
impacting generation quality. In light of these findings, we propose a simple
yet effective adaptive budget allocation algorithm. This algorithm not only
optimizes the loss upper bound in theory but also reduces the eviction loss in
practice by aligning with the intrinsic patterns of self-attention mechanisms.
Integrating this algorithm into two advanced methods, we develop Ada-SnapKV
and Ada-Pyramid. Extensive evaluations on 16 datasets and the Needle-in-a-
Haystack test confirm that they both significantly boost performance across
various tasks.

††footnotetext: $\dagger$Equal Contribution *Corresponding Author

##  1 Introduction

Autoregressive Large language models (LLMs) have achieved significant success
and are widely utilized across diverse natural language processing
applications, including dialogue systems (Yi et al. 2024), document
summarization (Laban et al. 2023), and code generation (Gu 2023). The
widespread deployments have propelled the development of their capacities to
process extended sequences. For instance, GPT-4 supports sequences up to 129K
(Achiam et al. 2023), Claude3 up to 200K (Anthropic 2024), and Gemini-Pro-1.5
(Reid et al. 2024) up to 1M tokens. However, transformer-based LLMs pose
significant efficiency challenges while processing long-sequence inputs.

Specifically, the inference process of LLMs consists of two stages: prefilling
and decoding. During prefilling, LLMs compute and store all Key-Value (KV)
cache pairs for the input tokens from the prompt. In the subsequent decoding
phase, LLMs autoregressively generate each output token using the most recent
token and all previous KV cache pairs. This stage continues until the maximum
designated length is reached or it is actively stopped. However, while dealing
with the long-sequence inputs, the size of the generated KV cache
significantly increases, easily exceeding the model’s parameter size (Sun et
al. 2024). This escalation leads to substantial storage challenges during both
the prefilling and decoding stages. Additionally, during the decoding phase,
each autoregressive step requires extensive KV cache I/O, which incurs
significant I/O latency. This latency now exceeds the time required for
computation and becomes the primary bottleneck in decoding speed (Tang et al.
2024; Dao et al. 2022). Consequently, the extensive KV cache associated with
long sequences poses significant efficiency challenges for the inference
process.

In response to the challenges posed by large KV cache size, a range of KV
cache eviction methods have been developed, as referenced in recent literature
(Zhang et al. 2024b; Xiao et al. 2023; Ge et al. 2023; Yang et al. 2024; Zhang
et al. 2024a; Li et al. 2024). As shown in Figure 1, these methods discard
most non-critical cache pairs after the prefilling stage of each layer in
LLMs. Therefore, by tailoring the cache sizes to specified storage budgets,
they not only reduce the storage burden but also boost the decoding speed
while maintaining the post-eviction generation quality. These eviction methods
are designed with plug-and-play capabilities, allowing for straightforward
integration into any LLM setup without the need for fine-tuning. They
typically utilize various strategies to selectively evict the majority of KV
cache pairs. The latest leading methods (Li et al. 2024; Yang et al. 2024;
Zhang et al. 2024a), employ a Top-K based selection scheme based on attention
weights. These approaches effectively distinguish between critical and non-
critical cache pairs, deciding which to retain and which to evict. Despite
these advancements, the challenge of minimizing quality loss while employing
these eviction methods remains an unresolved issue in the field.

Figure 1: Cache Eviction Framework. Cache Eviction removes a large number of
non-essential tokens, marked by dashed borders, immediately after the
prefilling stage of each layer. This process substantially reduces storage
requirements and accelerates decoding, while preserving the quality of post-
eviction generations.

Our study begins by reexamining the underlying principles of eviction methods
from a theoretical perspective. We reveal that Top-K based methods aim to
minimize an upper bound of eviction loss, quantified as the L1 distance
between the pre- and post-eviction outputs of self-attention mechanisms.
Moreover, We find the common practice of uniformly assigning budgets across
different attention heads leads to a misallocation of budgets, thus limiting
budget utilization. Based on these insights, we propose a simple yet effective
adaptive allocation algorithm. As an illustrative example shown in Figure 2,
it effectively improves budget utilization by dynamically allocating the
budget across different attention heads within the same layer based on their
concentration degrees, thereby enhancing the quality of post-eviction
generation. We substantiate the advantages of adaptive allocation both
theoretically and empirically. Theoretically, as detailed in Section 3.4, the
adaptive allocation reduces the upper bound of eviction loss compared to the
current uniform allocation. Empirically, as detailed in Section 3.5, the
alignment of the adaptive allocation with inherent varied degrees in attention
concentration among the heads within self-attention mechanisms effectively
reduces eviction loss in practice.

By integrating this adaptive budget allocation algorithm into two advanced
methods (Li et al. 2024; Yang et al. 2024; Zhang et al. 2024a), we develop two
adaptive eviction methods: Ada-SnapKV and Ada-Pyramid, respectively. Extensive
evaluations across 16 datasets, spanning various tasks in LongBench (Bai et
al. 2023), demonstrate that both Ada-SnapKV and Ada-Pyramid significantly
improve the generation quality. Additionally, the Needle-in-a-Haystack test, a
widely used evaluation method, further confirms that adaptive allocation
enhances in-context retrieval ability. The main contributions are summarized
as follows:

  * •

Reexamination of Cache Eviction Principle: We reveal the principles of cache
eviction methods that aim to minimize an upper bound of eviction loss,
quantified as the L1 distance between the pre- and post-eviction attention
outputs.

  * •

Identifying and Addressing Limitations in Budget Allocations: We pinpoint the
inefficiency of existing uniform budget allocation across heads. To address
this, we introduce an adaptive allocation algorithm that not only optimizes
the upper bound in theory but also reduces the eviction loss in practice.

  * •

Developments of Adaptive Eviction Methods: By integrating adaptive allocation,
we develop two novel methods: Ada-SnapKV and Ada-Pyramid. These methods both
significantly improve the post-eviction generation quality in the extensive
evaluations.

  * •

Innovative Direction to Enhance Cache Eviction: Our research pioneers a new
direction by adaptively allocating budgets based on unique characteristics
among heads to optimize cache eviction. This innovation, guided by our
theoretical analysis of the upper bounds of eviction loss, seeks significant
further development in this field.

(a) Uniform Allocation

(b) Adaptive Allocation

Figure 2: An Illustrative Example.This example features a five KV cache pairs
with associated attention weights across three heads. Adaptive allocation,
through reallocating budgets from the concentrated Head2 to the dispersed
Head1, increases the aggregated weights of retained KV cache pairs from 2.33
to 2.46, which correlates to a reduced eviction loss in Section 3.4 and 3.5.

##  2 Related Works

In the context of long-sequence inference, the vast scale of the KV cache
pairs leads to a memory-bound situation, causing significant storage and I/O
latency costs (Wang and Chen 2023). Some studies employ efficient memory
management strategies that reduce I/O time without altering the size of the KV
Cache, such as Page Attention (Kwon et al. 2023) and Flash Attention (Dao et
al. 2022). These studies are orthogonal to our work and we incorporate the
Flash Attention technique in our implementation to achieve efficient
computation. Some works, called KV cache quantization, reduce the size of
cache pairs by lowering the precision of individual entries. However, its
compression ratio is inherently limited—ideally, quantization from 16-bits to
1-bit achieves a theoretical limit of 6.25% compression but results in
catastrophic quality degradation. However, cache eviction methods could
compress the KV Cache to any given small budget, like less than 1% of the
original size while still maintaining certain generation quality. Moreover,
the techniques of quantization and eviction are orthogonal, suggesting
potential for further combined optimization in the future. Another work (Tang
et al. 2024) attempts to reduce I/O overhead by only recalling KV cache pairs
relevant to the current query for computation. However, it is constrained by
substantial memory overheads, making deployment on GPUs with limited storage
capacities challenging.

Recently, KV cache eviction methods have gained attention and rapid
development due to their flexible compression ratios given any storage budgets
and the advantage of being plug-and-play without the need for fine-tuning.
Earlier StreamingLLM (Xiao et al. 2023) simply maintains the cache of 4
initial and the recent tokes, discarding all others to adapt to long sequence
inference. FastGen (Ge et al. 2023) searches and combines multiple strategies,
like maintaining the cache of special tokens, punctuation tokens, and recent
tokens, based on the characteristics of attention heads. H2O (2024b) has
developed an eviction algorithm that utilizes query states of all tokens to
identify important KV pairs based on the Heavy Hitter method. The most recent
SnapKV (Li et al. 2024) identifies important KV cache pairs using query states
of several tokens within a recent window, evicting the less important ones.
This method effectively mitigates the quality degradation in cache eviction.
The Pyramid (Yang et al. 2024; Zhang et al. 2024a) further adjusts the budget
allocation across different layers in SnapKV, improving the generation quality
in small-budget scenarios. However, to our best knowledge, current eviction
methods have never tried to adaptively distribute the total budget across
different heads. Building on theoretical analysis and observations of inherent
attention patterns in LLM heads, we identify and demonstrate the necessity for
adaptive budget allocation in cache eviction.

##  3 Framework

###  3.1 Overview

In this section, we reexamine how the existing cache eviction strategies
retain essential information in the past KV cache from a theoretical
perspective. Inspired by theoretical findings, we propose a simple yet
effective algorithm for adaptive budget allocation, which is proven to be
better than the previous uniform allocation method in cache eviction
procedures, both theoretically and practically. Upon further integrating this
algorithm into two leading existing methods, we have developed two adaptive
cache eviction techniques: Ada-SnapKV and Ada-Pyramid. The key findings and
insights are as below:

  * •

How do the eviction strategies maintain generation quality despite discarding
substantial KV cache? Theorem 2 formalizes an upper bound of the eviction loss
by using the L1 distance between the pre- and post-eviction outputs and
Theorem 3 demonstrates that Top-K based eviction strategies aim to minimize
this upper bound.

  * •

Why is adaptive budget allocation better from the theoretical perspective?
Theorem 4 demonstrates that the upper bound of eviction loss under adaptive
allocation maintains at or below that under previous uniform allocation in the
cache eviction procedure. Thus it implies that the adaptive allocation is
better in theory.

  * •

Why is adaptive budget allocation better from the empirical perspective? As
shown in Figure 3 and 4, we identify that the inherent disparities in
attention concentration across different heads reliably translate theoretical
advantages into practical results, leading to consistent reductions of
eviction loss in practice.

###  3.2 Preliminary

We begin by providing a formal description of the computational processes
involving Multi-head Attention and KV cache in a single-layer of LLMs to
alleviate the burden of notation. LLMs are characterized by their
autoregressive generation mode, where each step involves using the last token
to predict the next token. Define $X\in\mathbb{R}^{n\times d}$ as the
embedding matrix of all tokens in sequence, and $x\in\mathbb{R}^{1\times d}$
as the last token used as input at the current timestep. To clarify the
subsequent theoretical exposition, we adopt the notation system from (Liu et
al. 2023) under the assumption of $h$ attention heads in one layer. The
transformation matrices for each head $i\in[h]$, $W_{i}^{Q}$, $W_{i}^{K}$,
$W_{i}^{V}\,\in\mathbb{R}^{d\times d_{h}}$, map token embeddings to their
respective Query, Key, and Value and the final output matrix
$W_{i}^{O}\,\in\mathbb{R}^{d_{h}\times d}$ transform the intermediate result
to the output hidden states. At each timestep, the previous KV cache pairs on
head $i$ have been initialized as:

$K_{i}=XW_{i}^{K},V_{i}=XW_{i}^{V}$ (1)

Then, the input token $x$ is mapped to its corresponding query, key, value for
each head, and the previous KV cache is updated accordingly:

$q_{i}=xW_{i}^{Q},k_{i}=xW_{i}^{K},v_{i}=xW_{i}^{V}$ (2)
$K_{i}=[K_{i}:k_{i}],V_{i}=[V_{i}:v_{i}]$ (3)

Finally the final output $o\in\mathbb{R}^{1\times d}$ is computed as follow:

$o=\sum^{i\in[h]}A_{i}V_{i}W_{i}^{O}$ (4)

where $A_{i}\in\mathbb{R}^{1\times n}$ is the attention weight calculated
by111The scaling factor $\sqrt{d_{h}}$ is omitted for simplification.:

$A_{i}=\text{softmax}(q_{i}K_{i}^{T})$ (5)

###  3.3 Reexamining the Principle of Cache Eviction

Cache eviction is dedicated to reducing the size of KV cache to fit within a
constrained budget by evicting non-critical cache pairs strategically.
Eviction masks $\left\\{\mathcal{M}_{i}\in\mathbb{R}^{1\times n}\right\\}$ can
be employed to simulate the post-eviction output $o^{\prime}$ of self-
attention mechanism during the generation process:

$o^{\prime}=\sum^{i\in[h]}A_{i}^{\prime}V_{i}W_{i}^{O}\>\text{where}\>A_{i}^{\prime}=\text{softmax}(q_{i}K_{i}^{T}+\mathcal{M}_{i})$
(6)

where each element $\mathcal{M}_{i}^{j}$ 222Given the first dimension of
$\mathcal{M}_{i}$ is 1, $M_{i}^{j}$ is used to simplify the notation for
$\mathcal{M}_{i}(1,j)$. Similarly, $A_{i}^{j}$ is in the same manner. in mask
$\mathcal{M}_{i}$ indicates that whether the $j\text{th}\in[n]$ KV cache pair
has been evicted in $K_{i},V_{i}\in\mathbb{R}^{n\times d_{h}}$ on each head
$i$:

$\displaystyle\mathcal{M}_{i}^{j}$ $\displaystyle=\begin{cases}0&\text{if the
$j$th cache pair on head $i$ is retained }\\\ -\infty&\text{otherwise the
$j$th cache pair on head $i$ is evicted}\\\ \end{cases}$
$\displaystyle\text{given budget
allocation}\left\\{B_{i}\right\\}\>s.t.\>{\sum^{i\in[h]}B_{i}=B}$

Thus, the budget $B_{i}$ for head $i$ corresponds to the number of zero
elements in $\mathcal{M}_{i}$. Theorem 1 further simplifies the output
$o^{\prime}$ by eliminating the softmax function. A detailed proof is provided
in Appendix A.1.

######  Theorem 1.

The post-eviction output $o^{\prime}$ can rewrite as:

$\displaystyle
o^{\prime}=\sum^{i\in[h]}\frac{A_{i}\odot\mathcal{N}_{i}}{||A_{i}\odot\mathcal{N}_{i}||_{1}}V_{i}W_{i}^{O}$
(7) $\displaystyle\text{where}\>\mathcal{N}_{i}^{j}=\begin{cases}1&\text{if
$K_{i}^{j}\>\text{and}\>V_{i}^{j}$ are retained}\\\ 0&\text{otherwise, evict
$K_{i}^{j}\text{and}V_{i}^{j}$}\\\ \end{cases}$
$\displaystyle\text{and}\odot\text{indicates element-wise multiplication}$
$\displaystyle\text{given budget
allocation}\left\\{B_{i}\right\\}\>s.t.\>{\sum^{i\in[h]}B_{i}=B}$

The reduction in generation quality during cache eviction stems from changes
in the attention output. Thus, we quantify the eviction loss as the L1
distance between the pre- and post-eviction outputs of self-attention
mechanisms:

$\text{Eviction Loss}=||o-o^{\prime}||_{1}$ (8)

Utilizing the row norm of matrix, we derive an upper bound $D$ for the
Eviction Loss in Theorem 2. For a detailed proof, refer to Appendix A.2.

######  Theorem 2.

The eviction loss caused by cache eviction can be bounded by $D$ as follows:

$\displaystyle\text{Eviction Loss}\leq
D=2hC-2C\sum^{i\in[h]}\sum^{\text{retained}j}A_{i}^{j}$ (9)
$\displaystyle\text{given budget
allocation}\left\\{B_{i}\right\\}\>s.t.\>{\sum^{i\in[h]}B_{i}=B}$

where $C=Max\left\\{\lVert V_{i}W_{i}^{O}\rVert_{\infty}\right\\}$ is the max
value in the row norms of Matrices $\left\\{V_{i}W_{i}^{O}\right\\}$ among all
heads.

Cache eviction strategies typically presuppose the stability of critical cache
pairs (Zhang et al. 2024b; Li et al. 2024; Yang et al. 2024; Zhang et al.
2024a), thereby facilitating cache eviction during the autoregressive
generation process. Consequently, the advanced strategies leverage query
states from several recent tokens to calculate attention weights with all past
cache pairs, which, in conjunction with Top-K selections, approximate the
identification of cache pairs critical in subsequent generations. For
simplicity, we assume that the eviction procedure relies solely on a single
query state $q_{i}$ associated with the last token $x$ for critical cache
detection. As shown in Algorithm 1, the core idea of current leading eviction
strategies based on Top-K selections is only to retain cache pairs
corresponding to $B_{i}$ highest weights
$A_{i}^{j}\in\text{Top-K}(A_{i},B_{i})$, while evicting those considered non-
essential. Obviously, given the budget allocation $\left\\{B_{i}\right\\}$ the
Top-K based eviction strategies maximize the aggregated weights as follows:

$\displaystyle\text{Top-K
Eviction}=\operatorname*{arg\,max}_{strategy}\sum^{i\in[h]}\sum^{\text{retained}j}A_{i}^{j}.$
(10) $\displaystyle\text{given budget
allocation}\>\left\\{B_{i}\right\\}\>\text{s.t.}\>\sum^{i\in[h]}B_{i}=B$

Therefore, we establish Theorem 3, which demonstrates that the principle of
Top-K based eviction aims to minimize an upper bound $D$.

Algorithm 1 Eviction Based on Top-K Selection

Input: Allocation $\left\\{B_{i}\right\\}$,KV Cache
$\left\\{K_{i},V_{i}\right\\}$, Weights $\left\\{A_{i}\right\\}$
Output: Compressed cache $\left\\{\hat{K}_{i},\hat{V}_{i}\right\\}$

1: for $i\leftarrow 1$ to $h$ do

2: initialize empty cache $\hat{K}_{i},\hat{V}_{i}$ for head $i$

3: for $j\leftarrow 1$ to $n$ do

4: if $A_{i}^{j}\in\text{Top-K}(A_{i},B_{i})$ then

5: retain and append cache pair $K_{i}^{j},V_{i}^{j}$ to
$\hat{K}_{i},\hat{V}_{i}$

6: else

7: evict cache pair $K_{i}^{j},V_{i}^{j}$

8: end if

9: end for

10: end for

11: return compressed cache $\left\\{\hat{K}_{i},\hat{V}_{i}\right\\}$

######  Theorem 3.

Given a budget allocation result $\left\\{B_{i}\right\\}$, the principle of
eviction strategies based on Top-K selections is to minimizes the upper bound
$D$ of eviction loss.

$\displaystyle\text{Top-K Eviction}=\operatorname*{arg\,min}_{strategy}D$ (11)
$\displaystyle\>\text{given budget
allocation}\>\left\\{B_{i}\right\\}\>\text{s.t.}\>\sum^{i\in[h]}B_{i}=B$

###  3.4 Adaptive vs. Uniform:Theoretical Perspective

In existing studies on cache eviction, the budget $B$ is uniformly assigned
across all heads within a layer: specifically, ${B_{i}=B/h}$. Consequently,
the upper bound of eviction loss $D$ is modified to $D^{\prime}$ under uniform
allocation:

$\displaystyle
D^{\prime}=2hC-2C\sum^{i\in[h]}\sum^{j\in[n]}_{A_{i}^{j}\in\text{
Top-K($A_{i}$,$B_{i}$)}}A_{i}^{j}$ (12) $\displaystyle\text{given uniform
budget allocation}\left\\{B_{i}=B/h\right\\}$

In contrast, we suggest adaptive allocation of the overall budget among heads
and introduce a simple yet effective budget allocation algorithm that
dynamically distributes the total budget $B$ based on the attention weights
$A_{i}$ across heads, with the allocated results denoted as
$\left\\{B^{*}_{1},B^{*}_{2},...,B^{*}_{h}\right\\}$ subject to
$\sum^{i\in[h]}B^{*}_{i}=B$. As shown in Algorithm 2, it firstly selects the
$B$ largest attention weights from all heads within one layer. Based on the
times each head is selected during the above procedure, different budgets
${B^{*}_{i}}$ are allocated to each head. Under this adaptive allocation, the
upper bound of eviction loss $D$ is modified to $D^{\prime\prime}$:

$\displaystyle
D^{\prime\prime}=2hC-2C\sum^{i\in[h]}\sum^{j\in[n]}_{A_{i}^{j}\in\text{
Top-K($A_{i}$,$B_{i}^{*}$)}}A_{i}^{j}$ (13) $\displaystyle\text{given adaptive
budget allocation}\left\\{B_{i}=B_{i}^{*}\right\\}$

######  Theorem 4.

The upper bound $D^{\prime\prime}$ of eviction loss with adaptive budget
allocation consistently remains at or below the upper bound $D^{\prime}$
associated with uniform allocation.

$D^{\prime\prime}\leq D^{\prime}$ (14)

Algorithm 2 Adaptive Budget Allocation

Input: Total Budget $B$, Attention Weights in $h$ heads
$\left\\{A_{i}\in\mathbb{R}^{1\times n}\right\\}$;
Output: Allocated Budgets of $h$ heads $\left\\{B^{*}_{i}\right\\}$;

1: Concatenate across heads $A$ = Cat($\left\\{A_{i}\right\\}$,dim=1)

2: Create head indicator $I$ = [$1...1:...:h...h$] with each index
$\left\\{i\right\\}$ repeat n times

3: Identify top indices $T$ = Top-K($A,B$).indices

4: Select the corresponding head indicator $I^{*}=I[T]$

5: Count frequencies of each $i$ in $I^{*}$ to determine
$\left\\{B^{*}_{i}\right\\}$

6: return Allocated Results of $h$ heads $\left\\{B^{*}_{i}\right\\}$

According to Theorem 4, our adaptive allocation algorithm achieves equal or
smaller eviction loss upper bound compared to the previous uniform allocation,
thereby offering theoretical advantages. The detailed proof can be found in
Appendix A.3.

Figure 3: Disparities in Attention Concentration Across Heads. (Mistral-
instruct-v0.2 on the first sample of the Qasper, a single-doc QA dataset in
LongBench) We aggregate different proportions of the top attention weights,
$\sum^{\text{retained}j}A_{i}^{j}$, to analyze attention concentration in
different head $i$. Most concentrated heads require a small cache proportion,
e.g., $5\%$, to aggregate weights close to 1, whereas other dispersed heads
need significantly larger proportions, such as $50\%$.

###  3.5 Adaptive vs. Uniform: Empirical Perspective

According to Theorem 4, the upper bound of eviction loss of adaptive
allocation is equal to or less than that of the prevailing uniform allocation.
We further demonstrate different attention heads within each layer of LLMs
exhibit significant disparities in the degrees of attention concentration,
which effectively translates this theoretical advantage of a reduced upper
bound into practical reductions in eviction loss.

For visualization in Figure 3 (a), most concentrated heads in Layer 8 like
head 1 require only 1% of the original cache budget to effectively retain the
aggregated weights $\sum^{\text{retained}j}A_{i}^{j}$ of 0.95. Conversely,
other heads like head 18 require nearly 50% proportion to near 0.95. This
characteristic is closely related to the upper bound $D$ of eviction loss, as
detailed in Theorem 3. Under such circumstances, the previous uniform
allocation encounters a budget utilization dilemma, as an illustrative example
in Figure 2: it must either endure substantial eviction losses in dispersed
heads or waste excessive and unwarranted budgets on heads with concentrated
attention. This significantly undermines the trade-off performance between the
total budget and generation quality in the cache eviction. In contrast, the
adaptive allocation algorithm assigns large budgets to dispersed heads, while
controlling the budget sizes for other heads, effectively maintaining the
overall size and mitigating the decline in generation quality. In Figure 4, we
show the averaged relative eviction loss in all layers, $\frac{\lVert
o^{\prime}-o\rVert_{1}}{\lVert o\rVert_{1}}$, with uniform and adaptive
allocations separately. Results show that the adaptive allocation consistently
reduces eviction losses across all 200 samples in the Qasper dataset. This
indicates that adaptively allocating budgets among different heads enhances
the budget utilization significantly, thereby reducing the output eviction
loss.

Algorithm 3 Ada-SnapKV/Ada-Pyramid in One Layer

Input: Overall budget $B$, Past cache $\left\\{K_{i},V_{i}\right\\}$, Tokens
in the recent window $X^{rec}\in\mathbb{R}^{win*d}$
Output: Compressed cache $\left\\{\hat{K}_{i},\hat{V}_{i}\right\\}$

1: for $i\leftarrow 1$ to $h$ do

2: $Q_{i}^{rec}=X^{rec}W_{i}^{Q}$

3: $\bar{A_{i}}=softmax(Q_{i}^{rec}K_{i}^{T})$

4: $\bar{A_{i}}=\bar{A_{i}}.maxpooling(dim=1).mean(dim=0)$

5: end for

6: get $\left\\{B_{i}^{*}\right\\}$ by invoking Algorithm
2($B,\left\\{\bar{A}_{i}\right\\}$)

7:
$\left\\{B_{i}^{*}\right\\}=\alpha\times\left\\{B_{i}^{*}\right\\}+(1-\alpha)\times\frac{B}{h}$

8: $\left\\{\hat{K}_{i},\hat{V}_{i}\right\\}$ = Algorithm
1($\left\\{B_{i}^{*}\right\\},\left\\{K_{i},V_{i}\right\\},\left\\{\bar{A}_{i}\right\\}$)

9: return compressed cache $\left\\{\hat{K}_{i},\hat{V}_{i}\right\\}$

(a) $B$ = 5% original size

(b) $B$ = 10 % original size

Figure 4: Comparison of Eviction Losses. (Using Mistral-7B-instruct-v0.2 on
200 samples from the Qasper Dataset) Cache eviction is implemented under
Uniform and Adaptive allocation strategies respectively, compressing the
original cache size to 5% and 10%. The adaptive version consistently yields a
lower average relative eviction loss $\frac{\lVert
o-o^{\prime}\rVert_{1}}{\lVert o\rVert_{1}}$ across all layers on different
samples compared to the prevailing uniform allocation.

###  3.6 Implementation

The current two leading eviction strategies, SnapKV and Pyramid, both utilize
the several tokens $X^{rec}\in\mathbb{R}^{win*d}$ from a recent window
(typically window size is 32) to identify and evict the less important cache
pairs. SnapKV excels at scenarios with large budgets, while Pyramid is more
effective in smaller budgets. The key distinction lies in how Pyramid and
SnapKV allocate budgets within different layers in LLMs. Pyramid suggests that
information aggregation among layers takes a pyramidal form, thereby
allocating a larger budget to shallower layers and progressively reducing it
in deeper layers through pre-set hyper-parameters. In contrast, SnapKV
distributes the budget uniformly across all layers. Nonetheless, same as other
eviction methods, they both allocate the budget uniformly among all heads
within a single layer. Thus, for a specific layer with a overall budget of
$B$, their eviction algorithms are the same.

We incorporate our adaptive allocation algorithm into both of them, resulting
in the creation of two novel adaptive strategies: Ada-SnapKV and Ada-Pyramid,
respectively. As shown in Algorithm 3333Algorithm 3 describes the process
serially for simplicity, but eviction operations can be parallelized easily in
practice., the adaptive budget allocation can be seamlessly integrated into
any eviction strategy by invoking the adaptive allocation Algorithm 2 before
eviction process in each layer. Both of Ada-SnapKV and Ada-Pyramid utilize a
max pooling layer to facilitate cache selection following the previous SnapKV
(Li et al. 2024) and Pyramid (Zhang et al. 2024a). Besides, a hyper-parameter
$\alpha$ in line 6, set by default to 0.5, prevents assigning tiny budgets to
highly concentrated heads, thereby enhancing fault tolerance of critical cache
detection. Additionally, we incorporate the Flash Attention (Kwon et al. 2023)
technique to support efficient computation of variable-length head caches,
with the computational efficiency evaluations provided in Appendix A.4.

| Single-Doc. QA | Multi-Doc. QA | Summarization | Few-shotLearning | Synthetic | Code |
---|---|---|---|---|---|---|---
|  NrtvQA |  Qasper |  MF-en |  HotpotQA |  2WikiMQA |  Musique |  GovReport |  QMSum |  MultiNews |  TREC |  TriviaQA |  SAMSum |  PCount |  PRe |  Lcc |  RB-P |  Ave. Score
Full Cache | 26.63 | 32.99 | 49.34 | 42.77 | 27.35 | 18.77 | 32.87 | 24.24 | 27.10 | 71.00 | 86.23 | 42.96 | 2.75 | 86.98 | 55.33 | 52.87 |  42.51
B=128h
H2O | 21.19 | 21.66 | 38.60 | 30.63 | 20.65 | 12.19 | 20.65 | 22.42 | 21.81 | 39.00 | 82.52 | 40.68 | 2.98 | 79.56 | 49.13 | 46.76 |  34.40
StreamingLLM | 16.61 | 14.74 | 31.40 | 28.05 | 21.36 | 12.08 | 18.44 | 18.91 | 19.26 | 43.50 | 74.22 | 29.00 | 2.75 | 31.65 | 41.27 | 38.84 |  27.63
SnapKV | 19.17 | 21.40 | 42.93 | 36.76 | 22.44 | 15.86 | 19.16 | 21.84 | 21.55 | 47.50 | 84.15 | 40.24 | 2.30 | 68.26 | 50.69 | 47.13 |  35.09
Pyramid | 20.16 | 21.77 | 43.55 | 36.78 | 23.12 | 14.39 | 19.53 | 22.03 | 21.47 | 51.00 | 84.62 | 40.24 | 2.79 | 70.77 | 50.57 | 46.53 |  35.58
Ada-SnapKV | 20.63 | 22.58 | 45.68 | 37.90 | 23.49 | 16.55 | 19.99 | 22.28 | 21.55 | 59.50 | 85.00 | 40.62 | 3.09 | 69.36 | 50.98 | 48.17 |  36.71
Ada-Pyramid | 20.50 | 21.71 | 45.61 | 36.81 | 23.57 | 15.84 | 19.75 | 22.13 | 22.00 | 60.50 | 84.04 | 40.51 | 3.21 | 73.60 | 51.24 | 48.02 |  36.81
B=256h
H2O | 21.54 | 22.92 | 42.56 | 31.07 | 22.53 | 13.76 | 22.52 | 22.40 | 23.09 | 40.50 | 84.20 | 40.77 | 3.41 | 86.10 | 50.98 | 48.17 |  36.03
StreamingLLM | 17.93 | 16.01 | 33.36 | 30.71 | 21.30 | 10.08 | 20.66 | 19.47 | 22.89 | 53.50 | 73.59 | 29.22 | 3.00 | 27.77 | 42.30 | 39.87 |  28.85
SnapKV | 22.37 | 23.74 | 48.13 | 38.56 | 22.43 | 15.66 | 21.91 | 23.13 | 23.15 | 61.50 | 85.45 | 41.42 | 3.09 | 84.54 | 53.22 | 50.24 |  38.66
Pyramid | 20.09 | 24.00 | 47.33 | 38.24 | 22.48 | 16.02 | 21.40 | 22.45 | 22.63 | 63.00 | 84.93 | 40.98 | 3.40 | 82.48 | 52.78 | 49.36 |  38.22
Ada-SnapKV | 22.55 | 25.78 | 48.33 | 40.30 | 24.24 | 16.64 | 21.63 | 23.03 | 23.19 | 67.00 | 85.78 | 41.53 | 3.47 | 87.07 | 53.86 | 51.13 |  39.72
Ada-Pyramid | 22.64 | 24.64 | 47.40 | 40.25 | 23.62 | 16.83 | 21.82 | 23.34 | 22.70 | 66.50 | 84.99 | 41.34 | 2.78 | 86.90 | 53.17 | 49.52 |  39.28
B=512h
H2O | 21.72 | 26.03 | 44.81 | 32.33 | 23.16 | 14.86 | 23.65 | 22.84 | 24.70 | 42.00 | 85.22 | 41.57 | 3.40 | 86.45 | 53.04 | 49.68 |  37.22
StreamingLLM | 18.76 | 17.17 | 37.09 | 30.21 | 21.64 | 9.93 | 24.44 | 20.00 | 25.57 | 62.00 | 72.36 | 29.95 | 2.48 | 18.17 | 43.70 | 40.13 |  29.60
SnapKV | 24.60 | 27.81 | 48.98 | 39.46 | 25.25 | 16.98 | 23.70 | 22.96 | 24.37 | 67.00 | 85.88 | 41.26 | 2.78 | 86.56 | 54.81 | 51.71 |  40.26
Pyramid | 23.23 | 27.94 | 48.87 | 40.50 | 24.36 | 16.74 | 23.22 | 23.16 | 24.37 | 67.00 | 85.73 | 41.74 | 3.16 | 85.67 | 54.16 | 50.34 |  40.01
Ada-SnapKV | 23.39 | 28.72 | 48.96 | 40.60 | 25.20 | 17.25 | 23.15 | 23.48 | 24.41 | 68.00 | 86.39 | 41.69 | 2.73 | 88.92 | 54.69 | 51.51 |  40.57
Ada-Pyramid | 24.03 | 28.98 | 48.39 | 39.25 | 24.50 | 18.38 | 23.13 | 23.90 | 24.30 | 68.00 | 85.89 | 41.89 | 2.98 | 87.71 | 54.46 | 51.39 |  40.45
B=1024h
H2O | 23.90 | 28.62 | 46.46 | 37.03 | 24.74 | 15.04 | 25.30 | 23.11 | 25.92 | 46.00 | 85.93 | 41.80 | 3.24 | 86.57 | 54.46 | 51.01 |  38.70
StreamingLLM | 19.42 | 21.69 | 41.75 | 32.40 | 22.18 | 11.18 | 27.13 | 21.09 | 26.59 | 67.00 | 71.79 | 30.11 | 2.88 | 16.57 | 44.82 | 39.76 |  31.02
SnapKV | 25.47 | 29.57 | 49.33 | 40.90 | 25.53 | 19.01 | 25.94 | 23.89 | 26.21 | 69.50 | 86.48 | 42.10 | 2.98 | 88.56 | 55.57 | 51.92 |  41.44
Pyramid | 24.21 | 29.86 | 48.93 | 40.75 | 25.05 | 18.77 | 25.73 | 24.06 | 25.65 | 68.50 | 86.31 | 42.25 | 2.97 | 87.17 | 54.75 | 52.10 |  41.07
Ada-SnapKV | 24.79 | 31.94 | 48.45 | 40.73 | 26.22 | 19.11 | 25.61 | 23.92 | 26.03 | 70.00 | 86.32 | 42.35 | 2.91 | 88.31 | 55.44 | 52.55 |  41.54
Ada-Pyramid | 25.09 | 30.94 | 48.18 | 40.00 | 26.52 | 19.10 | 24.93 | 23.71 | 25.86 | 70.00 | 86.34 | 42.64 | 2.56 | 86.92 | 54.93 | 51.90 |  41.23
Table 1: Comparison Based on Mistral-7B-Instruct-v0.2 Among 16 Datasets

(a) Mistral-7B-Instruct-v0.2

(b) LWM-Text-Chat-1M

Figure 5: Average Score Among 16 Datasets

(a) SnapKV

(b) Ada-SnapKV

(c) Pyramid

(d) Ada-Pyramid

Figure 6: Needle-in-a-Haystack Test. This test inserts a critical sentence
(the ’needle’) within a document (the ’haystack’) with extensive context, then
evaluates a model’s ability to retrieve the needle from the document. The
x-axis indicates the context length of the document, and the y-axis shows the
insertion depth of the needle. The Average Score is determined by averaging
the aggregated scores at various context lengths. Higher scores indicate an
improved capacity of the model for contextual retrieval.

##  4 Experiments

###  4.1 Settings

#### Datasets

Firstly, we carry out an comprehensive evaluation using 16 datasets, covering
domains of single-document QA (Kočiskỳ et al. 2018; Dasigi et al. 2021),
multi-document QA (Yang et al. 2018; Ho et al. 2020; Trivedi et al. 2022),
summarization (Huang et al. 2021; Zhong et al. 2021; Fabbri et al. 2019), few-
shot learning (Joshi et al. 2017; Gliwa et al. 2019; Li and Roth 2002),
Synthetic (Bai et al. 2023), and code generation (Guo et al. 2023; Liu, Xu,
and McAuley 2023), within LongBench (Bai et al. 2023), a benchmark for
evaluating multi-task performance with long-sequence inputs. These datasets
feature varying average input lengths from 1,235 to 18,409 tokens,
necessitating substantial KV cache size during generation, thereby rendering
them suitable for evaluating KV cache eviction strategies. Each dataset is
assessed using LongBench-recommended metrics, with quality scores up to 100.
Detailed dataset information is provided in Appendix A.6. We also utilize the
widely-used ’Needle-in-a-Haystack’ test, where key information is randomly
inserted into long texts to create prompts. This test evaluates whether LLMs
can extract this key information from extensive texts, specifically examining
the impact of proposed adaptive allocation on the models’ fundamental long
context retrieval abilities.

#### Baselines

We conduct a comparative analysis between our newly proposed methods, Ada-
SnapKV and Ada-Pyramid, and the previously established SnapKV and Pyramid,
acknowledged as two leading approaches under varying scenarios. Additionally,
we assess these methods against other earlier eviction strategies, including
StreamingLLM and H2O. StreamingLLM is designed to handle long-sequence inputs
by retaining caches of several initial tokens and others in a recent window.
And H2O compresses cache size by employing attention weight-based detection of
Heavy Hitters based on all queries and key states. In all experiments, the
hyper-parameter $\alpha$ in Algorithm 2 is set to 0.5. Both Ada-SnapKV and
Ada-Pyramid, as well as SnapKV and Pyramid, utilize the same configuration
settings as described in (Li et al. 2024), ensuring comparability with a
recent window size of 32 and a maximum pooling kernel size of 7. Parameters
for StreamingLLM and H2O conform to the default settings reported in the
literature (Zhang et al. 2024b; Xiao et al. 2023).

#### Base Models

In the experiments, we employ two open-source base models:
Mistral-7B-instruct-v0.2 (Jiang et al. 2023) and LWM-Text-Chat-1M (Liu et al.
2024). The Mistral 7B model features a context length of 32K and has been
adopted as the primary model in related studies (Li et al. 2024; Zhang et al.
2024a) due to its moderate parameter size and remarkable capability for long
text tasks. Meanwhile, LWM 7B model stands as the state of the art with its 1M
context length, facilitating performance evaluations for the Needle-in-a-
Haystack test under extreme context lengths.

###  4.2 Evaluations Among 16 Datasets

We assess all eviction strategies using cache budget $B\in\left\\{128\times
h,256\times h,512\times h,1024\times h\right\\}$ for each layer. Detailed
results for each dataset on the Mistral model are provided in Table 1, while
other results for the LWM model are placed in Appendix A.5 due to space
constraints. To demonstrate the efficacy of adaptive allocation, we take a
budget $B=128h$ as an example presented in Table 1. After integrating the
adaptive allocation algorithm, Ada-SnapKV enhances the quality scores in 15
out of 16 datasets compared to the original SnapKV, increasing the average
score from 35.09 to 36.71. Similarly, Ada-Pyramid surpasses the original
Pyramid in 14 of 16 datasets, boosting the average score from 35.58 to 36.81.

Figure 5 summarizes average scores of all strategies based on Mistral and LWM
across 16 datasets. Overall, SnapKV and Pyramid, as the current two leading
methods, exhibit closely matched performance, surpassing previous approaches
such as H2O and StreamingLLM. Furthermore, our Ada-SnapKV and Ada-Pyramid
strategies consistently improve the generated quality under varying budgets,
especially in small budgets. The two adaptive eviction strategies
alternatingly lead and surpass previous versions to become the new state-of-
the-art methods. This consistent improvement validates the necessity and
effectiveness of adaptive budget allocation, as previously demonstrated in
both theoretical derivations and empirical findings.

###  4.3 Evaluations on Needle-in-a-Haystack Test

As shown in Figure 6, we employ a Needle-in-a-Haystack test to demonstrate how
adaptive budget allocation can enhance long-context retrieval capabilities.
All configurations maintain a recent window size of 32 and a pooling kernel
size of 7 which is consistent with former experiments, where the maximum
inference length is limited to 37K in the full cache case on A800-80G. With a
cache budget of $B=128h$, the Ada-SnapKV and SnapKV extend the max length up
to 429K, while the Ada-Pyramid and Pyramid extend to 365K. Notably, Ada-SnapKV
and Ada-Pyramid both effectively improve long-text retrieval capabilities
compared to previous SnapKV and Pyramid. In particular, Ada-SnapKV and Ada-
Pyramid achieve near-lossless retrieval within the original 37K length, a feat
not replicated by the standard SnapKV and Pyramid. In terms of average score,
Ada-SnapKV improves from 94.84 to 95.99, while Ada-Pyramid increases from
96.02 to 96.69. Further details on decoding latency and maximum memory
consumption during inference within varied context lengths are available in
Appendix A.4.

##  5 Conclusion

In this study, we reexamine prevailing cache eviction strategies for efficient
inference in long sequences, revealing their goal to minimize an upper bound
of eviction loss, quantified as the L1 distance between pre- and post-eviction
outputs. By introducing an adaptive budget allocation among various attention
heads, we theoretically reduce this upper bound compared to previous methods.
Our empirical findings suggest that this adaptive approach significantly
benefits from the varied concentration degrees inherent in multiple heads
within the self-attention mechanism. The development of two novel adaptive
eviction methods, Ada-SnapKV and Ada-Pyramid, incorporating this adaptive
allocation, demonstrates significant improvements in comprehensive
evaluations. Furthermore, this research highlights the substantial potential
for advancing cache eviction strategies through our theoretical framework and
adaptive budget allocation, specifically designed to exploit the unique
characteristics of different attention heads in LLMs.

## References

  * Achiam et al. (2023) Achiam, J.; Adler, S.; Agarwal, S.; Ahmad, L.; Akkaya, I.; Aleman, F. L.; Almeida, D.; Altenschmidt, J.; Altman, S.; Anadkat, S.; et al. 2023.  Gpt-4 technical report.  _arXiv preprint arXiv:2303.08774_.
  * Anthropic (2024) Anthropic. 2024.  The claude 3 model family: Opus, sonnet, haiku.  Accessed: 2024-07-09.
  * Bai et al. (2023) Bai, Y.; Lv, X.; Zhang, J.; Lyu, H.; Tang, J.; Huang, Z.; Du, Z.; Liu, X.; Zeng, A.; Hou, L.; et al. 2023.  Longbench: A bilingual, multitask benchmark for long context understanding.  _arXiv preprint arXiv:2308.14508_.
  * Dao et al. (2022) Dao, T.; Fu, D.; Ermon, S.; Rudra, A.; and Ré, C. 2022.  Flashattention: Fast and memory-efficient exact attention with io-awareness.  _Advances in Neural Information Processing Systems_ , 35: 16344–16359.
  * Dasigi et al. (2021) Dasigi, P.; Lo, K.; Beltagy, I.; Cohan, A.; Smith, N. A.; and Gardner, M. 2021.  A dataset of information-seeking questions and answers anchored in research papers.  _arXiv preprint arXiv:2105.03011_.
  * Fabbri et al. (2019) Fabbri, A. R.; Li, I.; She, T.; Li, S.; and Radev, D. R. 2019.  Multi-news: A large-scale multi-document summarization dataset and abstractive hierarchical model.  _arXiv preprint arXiv:1906.01749_.
  * Ge et al. (2023) Ge, S.; Zhang, Y.; Liu, L.; Zhang, M.; Han, J.; and Gao, J. 2023.  Model tells you what to discard: Adaptive kv cache compression for llms.  _arXiv preprint arXiv:2310.01801_.
  * Gliwa et al. (2019) Gliwa, B.; Mochol, I.; Biesek, M.; and Wawer, A. 2019.  SAMSum corpus: A human-annotated dialogue dataset for abstractive summarization.  _arXiv preprint arXiv:1911.12237_.
  * Gu (2023) Gu, Q. 2023.  Llm-based code generation method for golang compiler testing.  In _Proceedings of the 31st ACM Joint European Software Engineering Conference and Symposium on the Foundations of Software Engineering_ , 2201–2203.
  * Guo et al. (2023) Guo, D.; Xu, C.; Duan, N.; Yin, J.; and McAuley, J. 2023.  LongCoder: A Long-Range Pre-trained Language Model for Code Completion.  arXiv:2306.14893.
  * Ho et al. (2020) Ho, X.; Duong Nguyen, A.-K.; Sugawara, S.; and Aizawa, A. 2020.  Constructing A Multi-hop QA Dataset for Comprehensive Evaluation of Reasoning Steps.  In Scott, D.; Bel, N.; and Zong, C., eds., _Proceedings of the 28th International Conference on Computational Linguistics_ , 6609–6625. Barcelona, Spain (Online): International Committee on Computational Linguistics.
  * Huang et al. (2021) Huang, L.; Cao, S.; Parulian, N.; Ji, H.; and Wang, L. 2021.  Efficient attentions for long document summarization.  _arXiv preprint arXiv:2104.02112_.
  * Jiang et al. (2023) Jiang, A. Q.; Sablayrolles, A.; Mensch, A.; Bamford, C.; Chaplot, D. S.; Casas, D. d. l.; Bressand, F.; Lengyel, G.; Lample, G.; Saulnier, L.; et al. 2023.  Mistral 7B.  _arXiv preprint arXiv:2310.06825_.
  * Joshi et al. (2017) Joshi, M.; Choi, E.; Weld, D. S.; and Zettlemoyer, L. 2017.  TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension.  arXiv:1705.03551.
  * Kočiskỳ et al. (2018) Kočiskỳ, T.; Schwarz, J.; Blunsom, P.; Dyer, C.; Hermann, K. M.; Melis, G.; and Grefenstette, E. 2018.  The narrativeqa reading comprehension challenge.  _Transactions of the Association for Computational Linguistics_ , 6: 317–328.
  * Kwon et al. (2023) Kwon, W.; Li, Z.; Zhuang, S.; Sheng, Y.; Zheng, L.; Yu, C. H.; Gonzalez, J.; Zhang, H.; and Stoica, I. 2023.  Efficient memory management for large language model serving with pagedattention.  In _Proceedings of the 29th Symposium on Operating Systems Principles_ , 611–626.
  * Laban et al. (2023) Laban, P.; Kryściński, W.; Agarwal, D.; Fabbri, A. R.; Xiong, C.; Joty, S.; and Wu, C.-S. 2023.  SUMMEDITS: measuring LLM ability at factual reasoning through the lens of summarization.  In _Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing_ , 9662–9676.
  * Li and Roth (2002) Li, X.; and Roth, D. 2002.  Learning question classifiers.  In _COLING 2002: The 19th International Conference on Computational Linguistics_.
  * Li et al. (2024) Li, Y.; Huang, Y.; Yang, B.; Venkitesh, B.; Locatelli, A.; Ye, H.; Cai, T.; Lewis, P.; and Chen, D. 2024.  Snapkv: Llm knows what you are looking for before generation.  _arXiv preprint arXiv:2404.14469_.
  * Liu et al. (2024) Liu, H.; Yan, W.; Zaharia, M.; and Abbeel, P. 2024.  World model on million-length video and language with ringattention.  _arXiv preprint arXiv:2402.08268_.
  * Liu, Xu, and McAuley (2023) Liu, T.; Xu, C.; and McAuley, J. 2023.  RepoBench: Benchmarking Repository-Level Code Auto-Completion Systems.  arXiv:2306.03091.
  * Liu et al. (2023) Liu, Z.; Wang, J.; Dao, T.; Zhou, T.; Yuan, B.; Song, Z.; Shrivastava, A.; Zhang, C.; Tian, Y.; Re, C.; and Chen, B. 2023.  Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time.  arXiv:2310.17157.
  * Reid et al. (2024) Reid, M.; Savinov, N.; Teplyashin, D.; Lepikhin, D.; Lillicrap, T.; Alayrac, J.-b.; Soricut, R.; Lazaridou, A.; Firat, O.; Schrittwieser, J.; et al. 2024.  Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context.  _arXiv preprint arXiv:2403.05530_.
  * Sun et al. (2024) Sun, H.; Chen, Z.; Yang, X.; Tian, Y.; and Chen, B. 2024.  Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding.  _arXiv preprint arXiv:2404.11912_.
  * Tang et al. (2024) Tang, J.; Zhao, Y.; Zhu, K.; Xiao, G.; Kasikci, B.; and Han, S. 2024.  Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference.  _arXiv preprint arXiv:2406.10774_.
  * Trivedi et al. (2022) Trivedi, H.; Balasubramanian, N.; Khot, T.; and Sabharwal, A. 2022.  MuSiQue: Multihop Questions via Single-hop Question Composition.  _Transactions of the Association for Computational Linguistics_ , 10: 539–554.
  * Wang and Chen (2023) Wang, K.; and Chen, F. 2023.  Catalyst: Optimizing Cache Management for Large In-memory Key-value Systems.  _Proceedings of the VLDB Endowment_ , 16(13): 4339–4352.
  * Xiao et al. (2023) Xiao, G.; Tian, Y.; Chen, B.; Han, S.; and Lewis, M. 2023.  Efficient streaming language models with attention sinks.  _arXiv preprint arXiv:2309.17453_.
  * Yang et al. (2024) Yang, D.; Han, X.; Gao, Y.; Hu, Y.; Zhang, S.; and Zhao, H. 2024.  PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference.  _arXiv preprint arXiv:2405.12532_.
  * Yang et al. (2018) Yang, Z.; Qi, P.; Zhang, S.; Bengio, Y.; Cohen, W. W.; Salakhutdinov, R.; and Manning, C. D. 2018.  HotpotQA: A dataset for diverse, explainable multi-hop question answering.  _arXiv preprint arXiv:1809.09600_.
  * Yi et al. (2024) Yi, Z.; Ouyang, J.; Liu, Y.; Liao, T.; Xu, Z.; and Shen, Y. 2024.  A Survey on Recent Advances in LLM-Based Multi-turn Dialogue Systems.  _arXiv preprint arXiv:2402.18013_.
  * Zhang et al. (2024a) Zhang, Y.; Gao, B.; Liu, T.; Lu, K.; Xiong, W.; Dong, Y.; Chang, B.; Hu, J.; Xiao, W.; et al. 2024a.  PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling.  _arXiv preprint arXiv:2406.02069_.
  * Zhang et al. (2024b) Zhang, Z.; Sheng, Y.; Zhou, T.; Chen, T.; Zheng, L.; Cai, R.; Song, Z.; Tian, Y.; Ré, C.; Barrett, C.; et al. 2024b.  H2o: Heavy-hitter oracle for efficient generative inference of large language models.  _Advances in Neural Information Processing Systems_ , 36.
  * Zhong et al. (2021) Zhong, M.; Yin, D.; Yu, T.; Zaidi, A.; Mutuma, M.; Jha, R.; Awadallah, A. H.; Celikyilmaz, A.; Liu, Y.; Qiu, X.; et al. 2021.  QMSum: A new benchmark for query-based multi-domain meeting summarization.  _arXiv preprint arXiv:2104.05938_.

##  Appendix A Appendix

###  A.1 Proof of Theorem 1

######  Theorem.

The post-eviction output $o^{\prime}$ can rewrite as:

$\displaystyle
o^{\prime}=\sum^{i\in[h]}\frac{A_{i}\odot\mathcal{N}_{i}}{||A_{i}\odot\mathcal{N}_{i}||_{1}}V_{i}W_{i}^{O}$
(15) $\displaystyle\text{where}\>\mathcal{N}_{i}^{j}=\begin{cases}1&\text{if
$K_{i}^{j}\>\text{and}\>V_{i}^{j}$ are retained}\\\ 0&\text{otherwise, evict
$K_{i}^{j}\text{and}V_{i}^{j}$}\\\ \end{cases}$
$\displaystyle\text{and}\odot\text{indicates element-wise multiplication}$
$\displaystyle\text{given budget
allocation}\left\\{B_{i}\right\\}\>s.t.\>{\sum^{i\in[h]}B_{i}=B}$

###### Proof.

Consider the softmax function as:

$\displaystyle softmax(x)=\frac{exp(x^{j})}{\sum^{j}exp(x^{j})}$ (16)

Thus, the attention weight after eviction procedure is:

$\displaystyle
A^{\prime}_{i}=softmax(s_{i}+\mathcal{M}_{i})\>where\>s_{i}=q_{i}K_{i}^{T}$
(17) $\displaystyle A^{\prime}_{i}$
$\displaystyle=\frac{exp(s_{i}^{j}+\mathcal{M}_{i}^{j})}{\sum^{j}exp(s_{i}^{j}+\mathcal{M}_{i}^{j})}=\frac{exp(s^{j}_{i})\odot
N_{i}^{j}}{\sum^{j}exp(s^{j}_{i})\odot\mathcal{N}_{i}^{j}}$ (18)
$\displaystyle=\frac{exp(s^{j}_{i})\odot
N_{i}^{j}}{\sum^{j}exp(s^{j}_{i})}\frac{\sum^{j}exp(s^{j}_{i})}{\sum^{j}exp(s^{j}_{i})\odot\mathcal{N}_{i}^{j}}$
(19)
$\displaystyle=\frac{A_{i}\odot\mathcal{N}_{i}}{||A_{i}\odot\mathcal{N}_{i}||_{1}}$
(20) $\displaystyle\text{where}\>\mathcal{N}_{i}^{j}=\begin{cases}1&\text{if
$K_{i}^{j}\>\text{and}\>V_{i}^{j}$ are retained}\\\ 0&\text{otherwise, evict
$K_{i}^{j}\text{and}V_{i}^{j}$}\\\ \end{cases}$

Thus:

$\displaystyle
o^{\prime}=\sum^{i\in[h]}A_{i}^{\prime}V_{i}W_{i}^{O}=\sum^{i\in[h]}\frac{A_{i}\odot\mathcal{N}_{i}}{||A_{i}\odot\mathcal{N}_{i}||_{1}}V_{i}W_{i}^{O}$
(22) $\displaystyle\text{where}\>\mathcal{N}_{i}^{j}=\begin{cases}1&\text{if
$K_{i}^{j}\>\text{and}\>V_{i}^{j}$ are retained}\\\ 0&\text{otherwise, evict
$K_{i}^{j}\text{and}V_{i}^{j}$}\\\ \end{cases}$ $\displaystyle\text{given
budget allocation}\left\\{B_{i}\right\\}\>s.t.\>{\sum^{i\in[h]}B_{i}=B}$

∎

###  A.2 Proof of Theorem 2

######  Theorem.

The eviction loss caused by cache eviction can be bounded by $D$ as follows:

$\displaystyle\text{Eviction Loss}\leq
D=2hC-2C\sum^{i\in[h]}\sum^{\text{retained}j}A_{i}^{j}$ (23)
$\displaystyle\text{given budget
allocation}\left\\{B_{i}\right\\}\>s.t.\>{\sum^{i\in[h]}B_{i}=B}$

where $C=Max\left\\{\lVert V_{i}W_{i}^{O}\rVert_{1}\right\\}$ is the max value
in the L1-norm of Matrices $\left\\{V_{i}W_{i}^{O}\right\\}$ among all head.

###### Proof.

By calculating the L1 distance between their outputs, we can obtain

$\displaystyle||o^{\prime}-o||_{1}=||\sum^{i\in[h]}(\textbf{1}-\frac{\mathcal{N}_{i}}{||A_{i}\odot\mathcal{N}_{i}||_{1}})\odot
A_{i}V_{i}W_{i}^{O}||_{1}$ (24)
$\displaystyle\leq\sum^{i\in[h]}||(\textbf{1}-\frac{\mathcal{N}_{i}}{||A_{i}\odot\mathcal{N}_{i}||_{1}})\odot
A_{i}V_{i}W_{i}^{O}||_{1}$ (25)
$\displaystyle\leq\sum^{i\in[h]}||(\textbf{1}-\frac{\mathcal{N}_{i}}{\lVert
A_{i}\odot\mathcal{N}_{i}||_{1}})\odot A_{i}\rVert_{1}\>\lVert
V_{i}W_{i}^{O}\rVert_{\infty}$ (26) $\displaystyle\leq
C\sum^{i\in[h]}||(\textbf{1}-\frac{\mathcal{N}_{i}}{\lVert
A_{i}\odot\mathcal{N}_{i}||_{1}})\odot A_{i}\rVert_{1}$ (27) $\displaystyle
where\>C=Max\left\\{\lVert V_{i}W_{i}^{O}\rVert_{\infty}\right\\}$

By expanding $A_{i}$, we can further simplify the expression.

$\displaystyle||o^{\prime}-o||_{1}\leq
C\sum^{i\in[h]}\sum^{\text{retained}j}(\frac{A_{i}^{j}}{\sum^{\text{retained}j}A_{i}^{j}}-A_{i}^{j})+\sum^{\text{evicted}j}A_{i}^{j}$
(28)
$\displaystyle=C\sum^{i\in[h]}\sum^{\text{retained}j}(\frac{A_{i}^{j}}{\sum^{\text{retained}j}A_{i}^{j}})-\sum^{\text{retained}j}A_{i}^{j}+\sum^{\text{evicted}j}A_{i}^{j}$
(29) $\displaystyle=C\sum^{i\in[h]}(2-2\sum^{\text{retained}j}A_{i}^{j})$ (30)
$\displaystyle=2hC-2C\sum^{i\in[h]}\sum^{\text{retained}j}A_{i}^{j}$ (31)

∎

###  A.3 Proof of Theorem 4

######  Theorem.

The upper bound $D^{\prime\prime}$ of eviction loss with adaptive budget
allocation consistently remains at or below the upper bound $D^{\prime}$
associated with uniform allocation.

$D^{\prime\prime}\leq D^{\prime}$ (32)

###### Proof.

$\displaystyle
D^{\prime}=2hC-2C\sum^{i\in[h]}\sum^{j\in[n]}_{A_{i}^{j}\in\text{
Top-K($A_{i}$,$B_{i}$)}}A_{i}^{j}$ (33) $\displaystyle\text{given uniform
allocation}\left\\{B_{i}=B/h\right\\}$ $\displaystyle
D^{\prime\prime}=2hC-2C\sum^{i\in[h]}\sum^{j\in[n]}_{A_{i}^{j}\in\text{
Top-K($A_{i}$,$B_{i}^{*}$)}}A_{i}^{j}$ (34) $\displaystyle\text{given adaptive
allocation}\left\\{B_{i}=B_{i}^{*}\right\\}$

Based on the operations of concatenation and Top-K in the first and second
lines of Algorithm 2, we can flatten the second $\sum$ in $D^{\prime\prime}$
as follow:

$\displaystyle
D^{\prime\prime}=2hC-2C\sum^{i\in[h],j\in[n]}_{A_{i}^{j}\in\text{Top-K}(A,B)}A_{i}^{j}$
(35) $\displaystyle\text{where}\>B=\sum^{i}B_{i}^{*}=\sum^{i}B_{i}$

Considering that $B=\sum{B_{i}}=\sum{B_{i}^{*}}$ , it is evident that
$\sum^{i\in[h],j\in[n]}_{A_{i}^{j}\in\text{Top-K}(A,B)}A_{i}^{j}\geq\sum^{i\in[h]}\sum^{j\in[n]}_{A_{i}^{j}\in\text{Top-K}(A_{i},B_{i})}A_{i}^{j}$.
This is because under the premise of identical total budget, the global Top-K
sum is greater than or equal to the sum of local Top-K sums of each head.
Thus:

$\displaystyle D^{\prime\prime}\leq D^{\prime}$ (36)

∎

| Single-Doc. QA | Multi-Doc. QA | Summarization | Few-shotLearning | Synthetic | Code |
---|---|---|---|---|---|---|---
|  NrtvQA |  Qasper |  MF-en |  HotpotQA |  2WikiMQA |  Musique |  GovReport |  QMSum |  MultiNews |  TREC |  TriviaQA |  SAMSum |  PCount |  PRe |  Lcc |  RB-P |  Ave. Score
Full Cache | 18.00 | 25.80 | 43.10 | 23.40 | 16.70 | 9.70 | 27.20 | 25.00 | 24.70 | 70.50 | 61.60 | 39.60 | 3.00 | 6.50 | 42.20 | 41.60 |  29.91
B=128h
H2O | 17.90 | 17.73 | 36.10 | 21.52 | 17.51 | 9.26 | 16.13 | 22.99 | 19.64 | 43.50 | 60.64 | 36.36 | 3.00 | 5.50 | 34.93 | 36.74 |  24.97
StreamingLLM | 12.81 | 11.32 | 29.04 | 17.24 | 13.67 | 6.91 | 16.34 | 20.25 | 17.35 | 41.00 | 52.74 | 25.77 | 0.50 | 3.00 | 28.38 | 30.98 |  20.46
SnapKV | 17.51 | 17.57 | 38.89 | 22.15 | 17.28 | 9.13 | 15.01 | 21.96 | 17.94 | 46.00 | 61.05 | 35.97 | 0.00 | 4.00 | 36.92 | 37.83 |  24.95
Pyramid | 18.17 | 17.58 | 39.08 | 22.05 | 16.78 | 8.13 | 14.74 | 22.24 | 17.88 | 47.50 | 60.11 | 37.02 | 0.50 | 3.50 | 36.96 | 38.73 |  25.06
Ada-SnapKV | 18.64 | 18.61 | 39.59 | 22.51 | 17.05 | 9.19 | 15.28 | 22.88 | 18.98 | 52.50 | 61.69 | 36.76 | 0.00 | 3.00 | 36.82 | 39.63 |  25.82
Ada-Pyramid | 18.35 | 18.93 | 39.49 | 22.57 | 16.83 | 8.61 | 15.05 | 23.22 | 18.85 | 55.50 | 60.93 | 37.39 | 0.50 | 3.50 | 36.55 | 39.79 |  26.00
B=256h
H2O | 18.99 | 19.18 | 38.78 | 21.88 | 17.33 | 9.16 | 16.88 | 23.29 | 20.51 | 48.00 | 60.36 | 38.07 | 3.00 | 5.50 | 37.18 | 37.94 |  26.00
StreamingLLM | 13.59 | 11.81 | 29.73 | 18.59 | 14.37 | 6.72 | 21.06 | 20.78 | 21.29 | 51.50 | 51.92 | 26.51 | 0.50 | 3.00 | 28.97 | 31.09 |  21.96
SnapKV | 19.27 | 20.61 | 40.78 | 22.81 | 16.83 | 9.89 | 16.23 | 23.17 | 20.07 | 53.50 | 61.75 | 38.41 | 0.00 | 4.00 | 38.25 | 40.57 |  26.63
Pyramid | 18.81 | 19.83 | 40.71 | 22.34 | 17.10 | 9.08 | 16.10 | 22.93 | 19.50 | 60.00 | 61.01 | 38.65 | 0.50 | 5.00 | 38.23 | 39.13 |  26.81
Ada-SnapKV | 18.99 | 21.08 | 41.18 | 22.89 | 17.64 | 9.52 | 16.71 | 23.05 | 20.48 | 67.00 | 61.27 | 38.74 | 0.00 | 3.50 | 39.60 | 40.96 |  27.66
Ada-Pyramid | 18.78 | 20.32 | 40.50 | 22.73 | 17.01 | 9.37 | 16.05 | 23.60 | 19.93 | 69.00 | 61.43 | 39.07 | 2.00 | 5.00 | 38.40 | 40.08 |  27.70
B=512h
H2O | 18.61 | 20.07 | 39.82 | 22.08 | 17.21 | 10.13 | 17.62 | 23.65 | 21.41 | 54.50 | 61.84 | 38.74 | 3.00 | 5.50 | 39.23 | 40.08 |  27.09
StreamingLLM | 13.94 | 13.13 | 33.06 | 18.26 | 14.44 | 7.41 | 25.24 | 21.00 | 23.78 | 60.50 | 52.04 | 26.31 | 1.00 | 3.00 | 30.10 | 31.76 |  23.44
SnapKV | 18.45 | 21.96 | 42.01 | 23.25 | 17.42 | 9.88 | 17.68 | 23.62 | 21.30 | 68.00 | 61.77 | 39.02 | 1.00 | 4.50 | 40.09 | 40.79 |  28.17
Pyramid | 18.46 | 22.85 | 42.24 | 23.27 | 16.75 | 9.45 | 17.41 | 24.62 | 21.20 | 70.00 | 60.61 | 39.32 | 3.00 | 6.50 | 39.63 | 40.78 |  28.51
Ada-SnapKV | 18.83 | 22.39 | 42.15 | 23.52 | 18.27 | 9.63 | 17.66 | 23.99 | 21.23 | 70.00 | 61.72 | 38.93 | 2.00 | 4.50 | 40.11 | 41.28 |  28.51
Ada-Pyramid | 18.64 | 22.86 | 41.81 | 23.61 | 16.67 | 9.45 | 17.35 | 23.75 | 20.79 | 70.00 | 60.66 | 39.61 | 3.00 | 5.50 | 39.87 | 40.99 |  28.41
B=1024h
H2O | 17.11 | 22.34 | 41.26 | 22.09 | 17.47 | 9.60 | 18.82 | 23.94 | 22.49 | 61.00 | 62.33 | 38.68 | 3.00 | 5.50 | 41.23 | 41.18 |  28.00
StreamingLLM | 14.78 | 16.77 | 37.64 | 18.77 | 14.63 | 7.39 | 26.43 | 21.47 | 24.21 | 67.00 | 53.00 | 25.99 | 0.50 | 3.00 | 31.51 | 32.31 |  24.71
SnapKV | 18.45 | 24.18 | 42.50 | 23.53 | 17.32 | 10.23 | 19.00 | 24.26 | 23.04 | 69.50 | 62.22 | 39.88 | 3.00 | 5.50 | 41.15 | 41.91 |  29.10
Pyramid | 18.48 | 24.87 | 42.11 | 23.45 | 16.97 | 9.84 | 18.93 | 24.50 | 22.77 | 69.50 | 61.65 | 39.73 | 2.50 | 5.00 | 41.07 | 41.27 |  28.91
Ada-SnapKV | 18.94 | 23.68 | 43.27 | 23.28 | 17.15 | 9.89 | 18.58 | 23.46 | 22.65 | 70.00 | 62.24 | 39.83 | 2.50 | 5.50 | 41.68 | 42.88 |  29.10
Ada-Pyramid | 19.00 | 23.83 | 43.36 | 23.48 | 17.03 | 9.32 | 18.70 | 24.11 | 22.61 | 69.50 | 61.83 | 39.75 | 2.50 | 6.00 | 40.85 | 41.80 |  28.98
Table 2: Comparison Based on LWM-Text-Chat-1M Among 16 Datasets

###  A.4 Computational Efficiency

We assess the computational efficiency of Ada-SnapKV and Ada-Pyramid under
adaptive allocation, as well as the original versions of SnapKV and Pyramid,
using a budget of 128h across various context lengths within needle-in-a-
haystack test. As shown in Figure 8, the peak memory footprint during
inference for Ada-SnapKV and Ada-Pyramid, as well as SnapKV and Pyramid,
remains almost consistent as sequence length increases, significantly lower
than that of vanilla Full Cache. Consequently, this allows the original
sequence length of 37K to be extended to most 429K, achieving a 10.59-fold
improvement. In terms of speed, as shown in 8, the decoding latency of the
four strategies remains almost consistent and is independent of the context
length, which is significantly lower than the decoding latency under Full
Cache. This is primarily due to cache eviction, which greatly reduces the size
of the KV Cache, thereby significantly alleviating the IO latency bottleneck
in the autoregressive decoding stage.

Figure 7: Peak Memory Footprint

Figure 8: Decoding Latency

###  A.5 Detailed results for LWM model Among 16 Datasets

The table 2 presents quality scores of different eviction strategies based on
the LWM model across 16 datasets. Overall, the results are consistent with
those of Mistral, and the adaptive allocation also leads to quality
improvements after cache eviction.

###  A.6 Detailed Information of Datasets

Table 3 provides a comprehensive description of information pertaining to 16
datasets.

###  A.7 Detailed Visualization of Head Concentration

Figure 9 supplements Figure 3 in the main paper by presenting the
visualization results across all layers. It can be observed that in all
layers, different heads exhibit significant variations in attention
concentration. This indicates that the adaptive allocation algorithm has great
potential to reduce the eviction loss in practice.

###  A.8 Code

We are in the process of making our source code available to the public.

Label | Task | Task Type | Eval metric | Avg len | Language | Sample Num
---|---|---|---|---|---|---
NrtvQA | NarrativeQA | Single-Doc. QA | F1 | 18,409 | EN | 200
Qasper | Qasper | Single-Doc. QA | F1 | 3,619 | EN | 200
MF-en | MultiFieldQA-en | Single-Doc. QA | F1 | 4,559 | EN | 150
HotpotQA | HotpotQA | Multi-Doc. QA | F1 | 9,151 | EN | 200
2WikiMQA | 2WikiMultihopQA | Multi-Doc. QA | F1 | 4,887 | EN | 200
Musique | MuSiQue | Multi-Doc. QA | F1 | 11,214 | EN | 200
GovReport | GovReport | Summarization | Rouge-L | 8,734 | EN | 200
QMSum | QMSum | Summarization | Rouge-L | 10,614 | EN | 200
MultiNews | MultiNews | Summarization | Rouge-L | 2,113 | EN | 200
TREC | TREC | Few-shotLearning | Accuracy | 5,177 | EN | 200
TriviaQA | TriviaQA | Few-shotLearning | F1 | 8,209 | EN | 200
SAMSum | SAMSum | Few-shotLearning | Rouge-L | 6,258 | EN | 200
PCount | PassageCount | Synthetic | Accuracy | 11,141 | EN | 200
PRe | PassageRetrieval-en | Synthetic | Accuracy | 9,289 | EN | 200
Lcc | LCC | Code | Edit Sim | 1,235 | Python/C#/Java | 500
RB-P | RepoBench-P | Code | Edit Sim | 4,206 | Python/Java | 500
Table 3: Details of 16 Datasets Figure 9:  Visualization of Heads’
Concentrations