Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs
Authors: Suyu Ge
Yunan Zhang
Liyuan Liu
Minjia Zhang
Jiawei Han
Jianfeng Gao
Venue: ICLR 2024 oral (2024)
Keywords: Large Language Model
Efficient Inference
Generative Inference
Key-Value Cache
Primary area: representation learning for computer vision, audio, language, and other modalities
TL;DR: We introduce adaptive KV cache compression, a plug-and-play method that reduces the memory footprint of generative inference for Large Language Models (LLMs) and accelerates its generation throughput.

Abstract: In this study, we introduce adaptive KV cache compression, a plug-and-play method that reduces the memory footprint of generative inference for Large Language Models (LLMs). Different from the conventional KV cache that retains key and value vectors for all context tokens, we conduct targeted profiling to discern the intrinsic structure of attention modules. Based on the recognized structure, we then construct the KV cache in an adaptive manner: evicting long-range contexts on attention heads emphasizing local contexts, discarding non-special tokens on attention heads centered on special tokens, and only employing the standard KV cache for attention heads that broadly attend to all tokens. Moreover, with the lightweight attention profiling used to guide the construction of the adaptive KV cache, FastGen can be deployed without resource-intensive fine-tuning or re-training. In our experiments across various asks, FastGen demonstrates substantial reduction on GPU memory consumption with negligible generation quality loss. We will release our code and the compatible CUDA kernel for reproducibility.

Published as a conference paper at ICLR 2024
MODEL TELLS YOU WHAT TO DISCARD: ADAPTIVE KV CACHE COMPRESSION FOR LLMs
Suyu Ge Yunan Zhang" Liyuan ,Minjia , Jiawei Jianfeng Gao² University of Illinois Urbana- Sahhan-Chammaiin, Miicrosoft (suyuge2, yunanz2, janijeilli edu [lucliu, minjiaz, jfgao wwtpe///cicoooootcooommo
ABSTRACT
In this study, we introduce adaptive KV cache compression, plug-and-play method that reduces the memory footprint of generative inference for Large Lan- guage Models (LLMs).. Different from the conventional KV cache that retains key and value vectors for all context tokens, we conduct targeted profiling to discern the intrinsic structure of attention modules. Based on the recognized structure, we propose FastGen, which constructs the KV cache in an adaptive manner: evict- ing long-range contexts on attention heads emphasizing local contexts, discarding non-special tokens on attention heads centered on special tokens, and only em- ploying the standard KV cache for attention heads that broadly attend to all tokens. Moreover, with the lightweight attention profiling used to guide the construction of the adaptive KV cache, FastGen can be deployed without resource-intensive fine-tuning or re-training. In our experiments across various asks, FastGen demon- strates substantial reduction on GPU memory consumption with negligible gener- ation quality loss.
INTRODUCTION
Based on the Transformer architecture, autoregressive language models have attracted extensive at- tention (OpenAI, 2023; Touvron et al., 2023b). Along with the increase of model size, these models present significant challenges in terms of computational complexity and GPU memory consump- tion (Shazeer et al., 2017). Since these models achieve remarkable success across diverse applica- tions, there is pressing need for serving these models in an economically feasible manner.
The generative inference of LLMs usually involves using the KV Cache mechanism to improve the generation speed. KV cache stores previously computed Key/Value vectors in attention calculation and reuses those values for the current token generation. As such, it avoids recalculations of previous tokens at each token generation step at the cost of extra memory consumption. Despite being prominent technique, the memory consumption of KV cache increases rapidly as the model size and generation length increase, drastically increasing the pressure of on-device memory.
When memory usage exceeds GPU capacity, the generative inference of LLMs typically resort to offloading (Aminabadi al., 2022; Sheng et al., 2023). While these methods help mitigate the pressure on the scarce GPU memory from using KV cache, offloading KV cache to CPU/NVMe can still add non-trivial overhead to generative inference performance due the limited PCIe bandwidth between the GPU and CPU on many devices. Therefore, it becomes a crucial task to reduce the memory footprint of KV cache without costly retraining or fine-tuning.
Our study starts from the observation (Figure 1) that there are abundant structures observed in at- tention modules (Michel et al., 2019: Voita et al., 2019; Clark et al., 2019; Wang et al., 2020; Child et al., 2019), and not all attention modules need to attend to all tokens (Liu et al., 2023b; Zhang al., 2023; Liu et al., 2023a). Intuitively, harvesting such structures and compressing cached vec- tors could substantially reduce memory consumption and accelerate text generation.
Based on this intuition, we propose FastGen to accelerate the generative inference by adaptively compressing the KV cache on the fly. First, we employ efficient profiling algorithm to recognize
Authors contributed equally to this research. Code is available http://gthha.ccmmmaeillll/aaaaaa

---

Published as a conference paper at ICLR 2024
the structural patterns for attention modules. Under the guidance of this profiling, then struct the KV cache for various modules adaptively. With this diagnose- compresssssess approach, FastGen effectively reduces the memory footprint of KV cache while preserving the model quality.
Algorithm 1: Promnt Encoding.
Algorithm 2: FastGen-Token Generation.
Input: Feasible Policy Set (C), Prompt
Input: Adaptive KV cache ((C², K V*)
Output: Adaptive KV Cache
Output: Generated Text
for Attention Head H; in LLM do
last prompt token
K Q' V + (Prompt)
2 for E {1,... ,Max Generate Length} do for Attention Head H; in L.M do
A' < softmax(Q'K**)
C + apply Equation to A'
K K' V H1 K Vi) K K Vi t V*,
/* C²: optimal policy
*/
K V + =(KK'V',C
K +- Kc., Va.
K K' t Ko.a
7
f sample from LLM prediction
return {C', K.,V
return {zy}
Accmulative Attention at Layer Head
Accmulative Attention at Layer Head
Accmulative Attention at Layer Head
3mm
1210
Special Takens
Punctustion
Loca ity
Othars
Figure 1: Different attention heads usually have different structures. Left: Four common attention structures (more details elaborated in Section and Section 4). Right: Attention map composi- tions of three attention heads that are in the same layer.
Prequency-
Frequency
Fine-tuned LiaMa 7h
Fine-tuned LiMMa 136
Fine-tured LiaMa 30h
tune-tuned LaMa 65h
50%
Win-late Ooe Carhe
Wi--ate Over Eull Cathe
50%
Tul Certre
50%
30%
30%
30%
30%
Wi-datte
20%
20%
20%
20%
10%
1.0%
10%
10%
40%
80%
0%0%%
80%
0%
70%
0%.
6.5%
Figure 2: Performance Adaptive KV Cache (FastGen) and Fixed KV Cache (Frequency, Local, and Frequency+I Local; Zhang al., 2023 and Liu et al., 2023a) on AlpacaEval.
In our study, FastGinn recognizes five fundamental attention structures and applies them correspond- ingly. Specifically, some attention modules mostly attend to local contexts, for which construct a KV cache that evicts long-range contexts; some primarily attend to specific tokens/punctuations, for which create KV cache that retains only special tokens/punctuations: some have attention maps that are column-we -iiie sparse, for which we discard the least frequently attended tokens; and some broadly attend all tokens, for which employ the standard KV cache and store all tokens.
In this way, FastGen is able to compress the KV cache while retaining the original functionality of attention modules. Remarkably, FastGen does not require any fine-tuning and can be applied in
2

---

Published as a conference paper at ICLR 2024
plug-and-pin manner. This is big advantage of FastGen, because the training cost on extra-large models (Brown et al., 2020), can hardly be afforded by many research labs or practitioners.
We evaluate FastGen Llama 1 (Touvron et al., 2023b) with suite of major benchmarks cover- ing generative tasks in math, code, knowledge, and common sense reasoning. FastGen effectively performs KV cache compression with negligible generation quality loss (i.e., recover over 95% of attention scores with 35% cache compressed). Notably, as to the 30b model in Figure 2, FastGen (50% cache compressed) surpasses all fixed KV compression methods (15% cache compressed).
RELATED WORK
Token Dropping and KV Cache Compression. Many efforts have been made to improve the model efficiency for LLMs. For recurrent neural networks, method is to skip multiple tokens at a given time step (Campos et al., 2017; Seo et al., 2017; Hansen et al., 2019). Since Transformer models quickly attracted lots of attention, Goyal et al. (2020) proposes eliminate redundant words in BERT (Devlin et al., 2019) based on their attention scores, while Dai et al. (2020) compresses the input sequence by adding pooling layers to the encoding modules of the transformer architecture. Recently, Huang et al. (2022) adds a token selection task to the original BERT model that learns to select performance-cre -rruiall tokens, and Kim et al. (2022) designs a learnable threshold detect unimportant tokens to prune. Mcanwhile, many efforts have been made to explore the possibility of compressing the hidden state of tokens rather than explicitly reducing the sequence length (Guan et al., 2022; Sun et al., 2022; Zhou et al., 2020).
Nevertheless, these methods can only be applied to aenaautoregressive models and typically require an additional re-training phrase, making them less suitable for auto-regressive LLMs like ChatGPT and Llama. Recognizing this gap, researchers started examining the potential of pruning tokens within the KV cache of auto-regressive LLMs. Mu et al. (2023) learns to compress the prompts into a few special tokens to reduce memory pressure during caching. However, the token prediction requires model re-training and could be expensive overhead during inference. Meanwhile, sev- eral concurrent methods propose to leverage accumulated attention score as the criteria to identify important tokens in the KV cache (e.g., Sheng al., 2023; Zhang al., 2023; Liu al., 2023a). In- stead of investigating specific eviction policy, this study aims to synergistically coordinate diverse eviction policies better align with model-specific attributes.
Underlying Structure of Attention. Inspired by the success of Transformer, extensive studies have been conducted to explore the underlying mechanism of different self-attention heads. Voita et al. (2019) analyzed the self-attention heads in BERT using LRF (Bach et al.. 2015) and characterized them into interpretable roles, one of which attending adjacent tokens all the time. Michel et al. (2019) demonstrated that heads in the same layer could have different impact on the performance while the importance of each head changes across tasks. Clark et al. (2019) and Kovaleva et al. (2019) identified such patterns some heads primarily attend to separator tokens, adjacent tokens and a combination of these. While most previous studies mainly considered encoder models, Fast- Gen is motivated by consistent patterns we have observed in decoder-only models. Like previous studies, FastGen also explores the structure of the attention mechanism to improve inference effi- ciency, but focusing characterizing the KV cache of different attention heads.
ADAPTIVE KV CACHE COMPRESSION
In this section we first introduce the problem formulation, and then present attention profiling and adaptive KV cache compression.
3.1 GENERATIVE INFERENCE OF AUTOREGRESSIVE LLMS
A typical generative model inference involves two steps: prompt encoding and token generation.
Prompt Encoding. When an autoregressive transformer-basee LLM generates the i-th token, the attention module needs to refer to all the preceding (KV vectors) of these tokens. To cireumvent redundant KV vector computations when generating succeeding tokens, all KV vectors are stored in the KV cache once they are generated.
1 tokens, i.e., the key and value vectors
3

---

Published as a conference paper at ICLR 2024
Token Generation. Once prompt encoding is finished, the LLM generates the output token by token. At each generation step, the LLM needs to encode the new token(s) generated in the previous step. After token is generated, its associated KV vectors appended to the current KV cache. Thus, the size of KV cache increases linearly with the number of tokens being generated.
3.2 FASTGEN
As described in Section 2, many previous studies of compressing KV cache for improving inference efficiency do not leverage the intricate attention structure in LLMSs. As to be detailed in Section 4, at- tention heads L.MSs often function distinctively, indicating the need for tailoring the compression strategy to each individual attention head.
With these insights, we introduce FastGen: dual-phase algorithm for crafting an adaptive KV cache. During the prompt encoding phase, model profiling is conducted to discern the behavior of various attention heads, so that we can choose the most appropriate compression strategy for each head. Then, in the token generation phase, instead of indiscriminately appending new KV vectors for each newly generated token, we manage the KV cache based on the selected compression strategies.
3.3 MODEL PROFILING
Model profiling conducted based on the result of prompt policy C, we mark the corresponding KV cache compression as Kc, Vc f(K, V C), where K and Ve are the compressed KV cache. Then, for attention map A softmax(Qk' ) we pick the optimal policy that can recover A with recover ratio T with the minimum memory cost:
p encoding. Specifically, for a compression
C arg min CacheMemoryCost(C) s.t. A - softmaax Cec
K -1
(1)
where C is the set of all feasible compression policies, CacheMemoryCost(C is the target KV cache budget of the compression policy C, and T is predefined hyper-parameter representing how much want the policy to recover A. As to be discussed in Section 5, FastGen is able recover +95% of the attention map with +40% compression ratio for a 65B model. The final prompt encoding algorithm that includes model profiling is presented in Algorithm
Intrinsically, our method assumes that the structure of the attention map for a head is stable through the generation process. So, it is sufficient to use only the encoded prompt select proper com- pression policy. It is worth noting that existing literature has provided theoretical justifications for using solely encoded prompts capture attention structures for the full contexts (Zhang al., 2023; Liu et al., 2023a). In our study, we also empirically verified this, to be elaborated in Section 4.
3.4 KV CACHE COMPRESSION POLICIES
In our experiments we observe that a large number of attention heads closely follow certain patterns, to be detailed in Section 4. Thus, in addition to the conventional full KV cache policy, also consider four fundamental KV cache compression policies. While we mainly use these four fundamental KV cache compression policies for evaluation in this study, it is easy for FastCien to use numerous other strategies. The four KV cache compression policies are:
Special Tokens. We keep in KV cache only special tokens, such as the begin-of-the-sentence token 
, the instruction token 
, and on. This policy is referred to as Special-
Punctuation. We keep in the KV cache only punctuation tokens like "" "." "?"" This policy is referred to as Cpunct.-
Locality This policy evicts long-range contexts. Once the relative distance between the context token and the current token exceeds a threshold, the KV cache of the context token will be evicted. The threshold is determined by pre-defined ratio rı of the length budget of local context over the input sequence length. This policy referred as Clocal-
Frequency (Heavy Hitter) This policy has been used in multiple previous studies (e.g., Sheng et al., 2023; Zhang et al., 2023; Liu et al., 2023a). We monitor for each token its cumulative sum of attention score, then treat these scores token frequency and only keep the most frequent
4

---

Published as a conference paper at ICLR 2024
tokens in the KV cache. The length budget of frequent tokens the current sequence length is controlled by a ratio ry. This policy is referred frequent-
Hybrid Policies. In practice, it is often necessary to use hybrid policies that combines the afore- mentioned compression policies. Since the total number of hybrid policies is hugh, in our study use a greedy method to construct small set of hybrid-policies as follows
C = (Cspcciaal- Cspecial+punct.
ppocial-punct.-ffreuueee Cspecial-punct--frqqueee
Ciull:: (2)
where the sum of two compression strategies to compute the union of their compressed KV cache, and Cfull refers to full KV cache without compression.
We use Special as a component in all hybrid policies for two reasons: 1) We observe that high attention scores are usually allocated towards Special- as to be detailed in Section 4, indicating that Special are crucial for attention map recovery; 2) the compressed cache of C Special is memory- efficient since there are usually less than special tokens in a sentence. In other words, brings little-tonoo extra memory cost by always including Special Similarly, Cpunct. often used component form hybrid policies due its memory-efficiency, i.e., the number of punctuations in sentence is small. The final algorithm for token generation is presented in Algorithm 2.
4 DIVERSITY AND STABILITY OF ATTENTION STRUCTURES
In this section we present an empirical study to show the effectiveness of adaptive KV cache com- pression. First, demonstrate that different attention heads typically possess distinct structures. Then, we show that these attention head structures remain relatively consistent across different at- tention heads different positions. We do by analyzing the attention scores of Llama 1 65B using random samples from GSM8k (Cobbe et al., 2021).
4.1 HEAD DISTINCTIVE ATTENTION STRUCTURE
Layer
Layer 10
Layer 20
layer 30
Layer 40
Layer 50
Layer 60
Layer 70
Layer 80
Cra
Cypers pute, +deguntt=io
Cpaccus pund. Pieguant
Capcca
Figure 3: Attention profiling result distribution across different layers.
Setting. We perform model profiling with recover threshold of 0.95 and compute the distribution of profiling results for (1, 10, 20, 30, 40, 50, 60, 70, 80) layers. The result is shown in Figure
Observation. Figure 3 shows that attention heads in different layers have vastly different struc- tures. Specifically, for initial and final layers, they have more attention heads assigned to the full KV cache, indicating attention heads in these layers are likely to attend to all tokens. Meanwhile, for middle layers, the attention map focuses on special tokens, indicating that most attention heads of these layers primarily attend to special tokens (i.e., the accumulated attention score special to- kens is higher than 0.95 for these attention heads). Figure shows the structure different attention heads in the same layer. We see that attention structures differ across different layers and heads.
5

---

Published as a conference paper at ICLR 2024
These results indicate that it is suboptimal to apply the same KV cache to all layers without adapta- tion, and that it is beneficial to detect the structure of each attention head so as to select the optimal compression policy construct the KV cache.
4.2 PROFILE TENDS TO BE CONSISTENT IN ONE SEQUENCE
The previous section demonstrates the great potential for constructing adaptive KV cache in accor- dance with the structure of different attention heads. Here, we show for each instance, it is sufficient conduct one-shot model profiling, outlined in Section 3.3. Specifically, for given prompt, show that the attention structure of each head remains consistent through the decoding process.
Spedal
Punctuction
localty
Layer Head 9
Laver Head
Layer Head
Layer 39 hend
Acom Burddd ARDI (011
Decoding Step
Decoding
Decoding Step
Decoding Stop
Layer
Sccumn
Decoting
Becoding
Decoding
Figure Accumulated attention score at 1st (prompt encoding), 10th, 20th, 30th decoding steps.
Setting. Following Figure 1, we compute the accumulated attention score for attention heads in different layers of Llama 1 65B multiple decoding steps (i.e., 1st, 10th, 20th, 30th). We visualized the resulting accumulated score in Figure 4.
Observation. Despite some fluctuations of accumulated attention scores across time steps, the pattern of the attention maps remains relatively stable. For example, Layer 33 Head 0) and Layer 23 Head almost only attend to the special token, while the locality and punctuation plays an important role in Layer 23 Head 0. As to Layer 23 Head 3, more than 10% of the attention score is allocated to the others portion, making suitable for a uncompressed KV cache Cfull-
In addition, we observe that large portion of attention scores are on special tokens in all cases. This justifies the greed method we used to construct hybrid policies, as described in Section 3.4.
5 EXPERIMENT
We conduct comprehensive experiments to demonstrate the effectiveness of FastGen on memory footprint reduction and generation quality preserving. First, we report the trade-off between mem- ory reduction and end-to-end generation quality in Section 5.1, and discuss the compression ratio of FastGen in Section 5.2. To demonstrate the superiority of FastGen on real-world systems, we demonstrate the end-to end-to-end latency change in Section 5.3 and the profiling overhead in Section 5,4. Finally, we present ablation studies and discussions in Section A.I.
5.1 TRADE-OFF BETWEEN PERFORMANCE AND MEMORY REDUCTION
Backbones. We conduct experiments with both Llama (Touvron et al., 2023a) and its fine-tuned variants, with model sizes ranging from 7B to 65B. For fined-tuned variants, do not choose

---

Published as a conference paper at ICLR 2024
FestCinn
Frequency-Local
Cache
GSMBK LLMe 7b
GSMBL Bk LaMa 13b
GSMBK LLaMas 30b
GSMBK LLaMa 65b
16
32
14
1.30
35
10-
26
30-
407
60%
80%
100%
40%
60%
80%
100%
40%
60%
80%
100%
40%
60%
80%
100%
Budget
Cache Budget
Budget
Busiget
Human Eval LLaMa
Humman Evel ULaMa 1.3
Human Evel LLaMt 30b
Hummn Evel LLaMa 65b
10
16
20
Pars81
1.
10
100%
60%
80%
100%
46%
IV Carhe Burdget
KV Cache Budget
KV Carce Burdeet
KV Cache Budget
NQ M
NQ LLaMa 136
LLaMa
NQ LLaMa 65b
29
34
42
46
28
41
33
26
42
38
40
40%
60%
80%
100%
40%
60%
80%
100%
40%
60%
80%
100%
20%
45%
60%
80%
100%
KV Cache Budget
KV Cache Budget
KV Cache Budget
KV Cathe Budget
TQA LLaMa 7b
TQA LLaMa 136
TQA LLa 30b
TQA LLaMa 65b
60
67
76
66
56
65
64
72-
407
60%
80%
100%
40%%
80% Cache Budget
60%
100%
40%
60%
80%
100%
40%
65%
80%
100%
Budget
Budget
Cache Budget
Figure 5: Performance of Adaptive KV Cache (FastGen) and Fixed KV Cache (Frequency+Local; Zhang et al., 2023 and Liu et al., 2023a) of Llama 1 on GSM8k, HumanEval, NQ, and TQA.
the open-sourced Llama 2-chat (Touvron al., 2023b) model due its grouped-query attention techniques. Instead, we use the original multi-head attention architecture in this study and leave the integration of grouped-query attention future work. To prepare comparable instruction- following model for analysis, we fine-tuned the Llama I model with open-sourced instruction-tuning datasets. Specifically, the fine-tuned variants are trained LIMA' data (Zhou al., 2023) and Open Assistant (Köpf et al., 2023) data.
Tasks. We use standard generation tasks to evaluate Llama and our fine-tuned Llama models. For Llama 1, we choose different tasks, including HumanEval (Chen al., 2021), GSM8K (Cobbe et al., 2021), ), NQ (K (Kwiatkowski et al., 2019) and TQA (Kembhavi et al., 2017) to evaluate models' abilities on different domains (code, math, question answering and reading comprehension). Note that in the four tasks, each testing sample is in generative format, where answers extracted after model generation finishes. This is crucial for a fair comparison on model's generation quality. We evaluate the instruction finetuned LLaMa model the instruction tuning benchmark AlpacaEval (Li et al., 2023), which consists of 805 question prompts from diverse domains.
Experiment Setup. The evaluation of the Llama model follows the default setting and evalu- ation metrics on each benchmark. We calculate Fl scores for GSM8K, NQ and TQA, and use the code execution Pass@1 I rate for HvmanEval. While evaluating an instruction-tuninng model remains challenging, follow previous work (Zhou al., 2023; Touvron al., 2023b) to use GPT4 as an evaluator for wair-wise comparison between two different model generations. For each prompt,
'httss//hugginngfcc.cc http://iiiinnccttlataatcsa//aaaa
7

---

Published as a conference paper at ICLR 2024
input the FastGen generation and the generation from the model with Full KV Cache as a pair, and ask GPT4 to judge which one better. We then calculate the win rate of FastGen over Full Cache. Hypothetically, the win rate of lossless method should be around 50%. Aside from full-cache models, we also include non-adaptive KV cache methods for comparison. Specifically, all attention head without any adaptation, as
apply Clocal- frequeett and Clocal+frequent baselines. is worth mentioning that Coocal+ftequent is very strong baseline as it identical the H₂O method (Zhang et al., 2023) and the Scissorhands method (Liu et al., 2023a). We set r 0.3, ry = 0.3 in FastGen, and only change the recovery ratio T to control the pruned KV cache ratio. For generation, we use nucleus sampling (Holtzman al., 2019) with temperature 0.6, p 0.9. Experiments are conducted on 8 NVIDIA A100 80GB GPUs.
Main Results, In Figure and Figure 5, we present the model quality as a function of KV cache budget increasing from 30% to 100%. For 30B models, FastGen (50% cache compressed) surpasses all non-adaptive KV compression methods (15% cache compressed) Also, we can see FastGen achieves more KV cache reduction ratio the model size increases, while preserving the same model quality. For example, achieving 45% win rate, FastGen can get as much as 44.9% pruned ratio on Llama 1-65B, compared 16.9% pruned ratio on Llama 1-7B. In all settings, FastGen shows consistent and significant improvement over non-adaptive compression methods. The results validate the effectiveness of adaptive KV cache compression using FastGen, despite its simplicity.
KV Cache
Model Size
Win rate
Full
FastGen
Pruned ratio
1.9Gb
56.6%
91%
30.8%
7B
4.3Gb
2.6Gb
39.8%
95%
37.7%
3.6Gb
16.9%
98%
47.4%
3.1Gb
53.4%
91%
32.0%
13B
6.7Gb
4.1Gb
39.0%
95%
39.9%
5.5Gb
18.3%
98%
48.7%
5.7Gb
56.7%
93%
37.0%
30B
13.1Gb
6.7Gb
48.8%
95%
42.5%
9.5Gb
27.4%
98%
47.5%
9.4Gb
56.3%
93%
40.9%
65B
21.5Gb
11.8Gb
44.9%
95%
44,2%
13.8Gb
36.0%
98%
49.8%
Table 1:
Memory footprint reduction by FastGen. We compared the memory consumption between models with full KV cache, and models compressed by FastGen fine-tuned Llama 1.
5.2 MEMORY FOOTPRINT REDUCTION ANALYSIS
We report the KV cache memory footprint reduction in Table 1. For all the evaluated 7B-65B models, we evaluate the memory consumption with fixed batch size of 16, sequence length of 512, and model weights fp16 format. We observe that FastGen substantially reduces the KV cache memory footprint across all model sizes, with more significant reductions for larger models. Taking a win rate over 45% as little-to-no quality regression, FastGen can achieve 440% memory reduction in Llama 1-65B, ~00% in Llama 1-30B, ~00 in Llama 1-13B and Ilama 1-7B.
5.3 END-TO-END LATENCY IMPROVEMENT
Table 2: End-to-end latency comparison on Llama 1-7B.
Batch size
16
[prompt gen lca)
(32,512)
[32.2018]
[32.8192]
[32,16384]
[512.32]
[512,512]
[4096.4096]
[512.512]
[4096,1096]
[512.512]
13.35
57.37
299
799.14
1.12
19.16
167.64
23.44
OOM
OOM
DS
11.58
47.12
201.23
435.74
0.79
10.45
91.04
12.93
127.94
OOM
FastCien
44.6
Speed-upe Suddd-pp (8) over HE
16.03%
22.30%
40.00%
55.009
34.80%
49.309
54.10%
54.90%
OOM
Speed-upe (3) over DS
3.20%
5.35%
10.83%
17.42%
7.59%
7.08%
15.50%
18.25%
35.78%
OOM
8

---

Published as a conference paper at ICLR 2024
To analyze the end-to-end speedup of FastGen, present the end-to-end latency improvement over full-cache setting and a strong model acceleration baseline in Table 2. In the experiment, we record the total duration in seconds, measured from the start of prompt encoding, until the end of generation as the end-to-end latency. For the full-cache baseline, we adopt the widely used Hugging Face Accelerate (HF) (Gugger al., 2022), denoted as HF in Table 2. For FastGen, customized kernel to handle the KV cache pruning operation. Specifically, we adapt the kernel from Deepspeed (DS) (Aminabadi et al., 2022) by adding the KV cache sparsity operation. We include the Deepspeed performance for fair comparison, denoted as DS in Table 2. All methods are tested on the same Nvidia V100 GPUs.
implemented
As shown in Table 2, we can observe that FastGen achieves significant end-to-end speed-up across all the generation settings. For the least significant case, FastGen can have decent 16.04% latency improvement over the HF baseline achieve up to 55.0% latency reduction over HF with FastGen a generation length of 16k. We can also observe that the relative speedup is greater with longer generation length. For example, given batch size 1, FastGen's s relative speed-up rises from 16.04% to 55.0%, as the generation length grows from 512 16k. When comparing FastGen to DeepSpeed, we still observe significant speed-up that gets bigger with batch size and generation length. Considering DeepSpeed a full- stack optimized inference system, where only attention computation is optimized, there still much room to further improve FastGen by polishing the sparsity kernel. We leave this unique research and engineering challenge to future works.
short generation length of 512. In the best cases, we can
5.4 PROFILING COST
Table 3: Profiling time of Llama 1-65B. The Overall Generation Duration is measured from the start of decoding to the end of the generation length. The Profiling Duration measured from the start of the decoding until Fastgen finishes the policy search.
Generation Length
Overall Generation
Profiling Duration (s)
Decoding Time Per Token
Profiling/Overaal (%)
Duration
128
30.98
0.11
0.35%
256
50.1
0.11
0.10
0.21%
512
94.98
0.11
0.10
0.12%
1024
157.43
0.11
0.10
0.07%
To better understand the overhead of profiling step, we compare the profiling time with the total generation time across different generation lengths. We present the result in Table 3.
We can observe that the profiling time only accounts for a very small percentage the total genera- tion duration, up to 0.35% in our tested cases. Also, the overhead decreases as the generation length increases, dropping to 0.07% when the generation length comes to 1024.
In terms of extra memory usage, it's mainly introduced by one of the compression strate- gies, Cfrequent- which needs to store an extra cumulative sum of attention scores for each attention head. To provide a detailed analysis, for each layer, the dimension of on),
the KV cache is (batch-sizn sittcch-iie,,. num.of f fuead..hhaad, sequence _len, while the dimension of extra memory for the cumulative attention scores is (bat ch.s ze, num head, sequence _een). Considering hidden hidden..immnssi 128 for all model sizes, the memory overhead is 1/128=0. 78% compared to storing KV cache only, which is a negligible cost.
www.dsiiiinnn
CONCLUSION
We have presented FastGen, novel method that significantly improves the inference efficiency of LLMs, with no visible quality loss, using lightweight model profiling and adaptive key-value caching. Areas for future explorations include combining FastGen with other model compression techniques, such as quantization and distillation, and other efficient attention architectures, such as coouped-query attention.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
Research was supported in part by US DARPA KAIROS Program No. FA8750-19-2-1004 and IN- CAS Program No. HR001 121 12IC0165, National Science Foundation 19-5555515,, and the Molecule Maker Lab Institute: An AI Research Institutes program supported by NSF under Award No. 2019897, and the Institute for Geospatial Understanding through an Integrative Discovery Envi- ronment (I- (I-GUIDE) by NSF under Award No. 2118329, Any opinions, findings, and conclusions or recommendations expressed herein are those of the authors and do not necessarily represent the views, either expressed or implied, of DARPA or the U.S. Government.
REFERENCES
Reza Yazdani Aminabadi, Samyam Rajbhandari, Minjia Zhang, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng. Jeff Rasley, Shaden Smith, Olatunji Ruwase, and Yuxiong He. Deepspeed- inference: Enabling efficient inference of transformer models at unprecedented scale. SC22: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-1 2022.
Sebastian Bach, Alexander Binder, Grégoire Montavon, Frederick Klauschen, Klaus-Rober Müller, and Wojeiech Samek. On pixel-wise explanations for non-linear classifier decisions by layer-wise relevance propagation. PLoS ONE, 10, 2015. URL https: https////pp...saai...oo.omm semanticscholar:
/trg/Corpus 9327892.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss. Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Pro- cessing Systems 2020, NeurIPS 2020, December 6- 2020, virtual, 2020.
Víctor Campos, Brendan Jou, Xavier Giró Nieto, Jordi Torres, and Shih-Fu Chang. Skip rn:: Learning to skip state updates in recurrent neural networks. ArXiv, abs/1708.00634, 2017. URL https: / /api. wemanticschiie. ar. www.ooaa.rrrrrroroeeeeeoooooee
wttg//orpssTiii TD: 1859294.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint 10Xiv:2100700374: 2021.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. CoRR, abs/1904. 10509, 2019. URL http:: wwwpsiiivorrr.oooooooooooo abs /tsss199 1904. 10509.
Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. What does BERT look at? an analysis of BERT'S attention. In Proceedings of the 2019 ACL Workshop Black- boxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 276-286, Florence, Italy, August 2019. Association for Computational Linguistics. doi: 10. 10.18653/1//1194440 URL
org/ //wp///// W19-4828.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Rejichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 14168, 2021.
Zihang Dai, Guokun Lai, Yiming Yang, and Quoc Le. Funnel-transformer: Filter- ing out sequential redundancy for efficient language processing. In Hugo Larochelle, Mare Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neu- ral Information Processing Systems 2020, NeurIPS 2020. December 6-12. 2020, vir- tual, 2020. URL https:// nww.eeddd.......i.....ee.eee.... wwtppapee/222//aaaaaaaaeeaeeeaea 2cd2 2e029115665540000445554aaaaaaa Abst html.
10

---

Published as a conference paper at ICLR 2024
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre- vre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume (Long and Short Papers), pp. 4171- 4186. Association for Computational Linguistics, 2019. doi: 10.18653/v 10653/vl/n19-11 1423. URL
https: //tp::///do..eee/ org/orggg/oooo /10 w1w11//1///nnt
1423.
Saurabh Goyal, Anamitra R. Choudhury, Saurabh Raje, Venkatesan T. Chakaravarthy, Yogish Sab- harwal, and Ashish Verma. Power-bert: Accelerating bert inference via progressive word- vector elimination. In International Conference Machine Learning, 2020. URL https:
//api. semant cwcccccclaseo www.laar.rr/ooopuuooooee 2197299799 92793.
Yue Guan, Zhengyi Li, Jingwen Leng, Zhouhan Lin, and Minyi Guo. Transkimmer: Transformer learns to layer-wise skim. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022, pp. 7275-7286. Association for Computational Linguistics, 2022. doi: 10. 1/2022/al/0222cllll org/10. 10.118653/v wtwp///0// 2022.. acl 5022
URL https: /toppp:///ii... //ooi.
Sylvain Gugger, Lysandre Debut, Thomas Wolf, Philipp Schmid, Zachary Mueller, Sourab Man- grulkar, Mare Sun, and Benjamin Bossan. Accelerate: Training and inference scale made sim- ple, efficient and adaptable. https::/gg thub www.ccm/////oooo www/cogg/oiiioccccccco higginggacce /acce aaaaaleeeeaeeeeeeeeeeeee lerate, 2022.
Christian Hansen, Casper Hansen, Stephen Alstrup, Jakob Grue Simonsen, and Christina Lioma.
Neural speed reading with structuralljjump-ssm.. ArXiv, 190400061, 2019. URL https:
//api. www.ssescatttccchooeoeeo war.chhhhhhho.ocooooccooccc..a.a wttg/Corpueeeeeeee ID: 90258012.
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text 2019.
degeneration. arXiv preprint
Xin Huang, Ashish Khetan, Rene Bidart, and Zohar Karnin. Pyramid-bert: Reducing complexity via successive core-set based token selection. In Smaranda Muresan, Preslav Nakov, and Aline Villav- icencio (eds.). Proceedings of the 60th Annual Meeting of the Association for Computational Lin- guistics (Volume 1: Long Pupers), ACI. 2022, Dublin, Ireland, May 22-27, 2022, pp. 8798-8817. Association for Computational Linguistics, 2022. doi: 10. 16335//1/2222.2llloooo btttps://diis http://i..ooogg///// /ttp//6/////222.2/l/c2 awwp///2222....1ooggg. 602.
URI.
Aniruddha Kembhavi, Minjoon Seo, Dustin Schwenk, Jonghyun Choi, Ali Farhadi, and Han- naneh Hajishirzi. Are you smarter than a sixth grader? textbook question answering for multi- modal machine comprehension. 2017 IEEE Conference on Computer Vision and Pattern Recog- nition (CVPR), pp. 5376-5338 5384, 2017. URL https: /ttps://app.... semant iommantiissn wesccccllaaac oorg/
CorpusID::3l1
Sehoon Kim, Sheng Shen, David Thorsley, Amir Gholami, Woosuk Kwon, Joseph Hassoun, and Kurt Keutzer. Learned token pruning for transformers. In Aidong Zhang and Huzefa Rangwala (eds.), KDD '22: The 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, Washington, DC, USA, August 14- 18, 2022, pp. 784-794. ACM, 2022. doi: 10.1 145/333333668 3539260.
3539260. URL https: / doi org/ 10. wet..1115533333
Andreas Köpf, Yannic Kilcher, Dimitri von Rütte, Sotiris Anagnostidis, Zhi-Rui Tam, Keith Stevens, Abdullah Barhoum, Nguyen Minh Duc, Oliver Stanley, Richárd Nagyfi, Shahul ES, Sameer Suri, David Glushkov, Arna Dantuluri, Andrew Maguire, Christoph Schuhmann, Huu Nguyen, and Alexander Mattick. Openassistant conversations democratizing large language model alignment. CoRR, abs/2304.07321, 2023. doi: 10448550WaXXxx22304..722 URL https: //dos org/ 10. 004888500aaiiii 2304. 07327.
Olga Kovaleva, Alexey Romanov, Anna Rogers, and Anna Rumshisky. Revealing the dark secrets of BERT. In Kentaro Inui, Jing Jiang, Vincent Ng, and Xiaojun Wan (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing. EMNLP-IJCNLP 2019, Hong Kong, China, November 3-7, 2019, pp. 4364 4373. Association for Computational Linguistics, 2019. doi: 10. 10.19965// URL https::////.. : /doi. wrgg/001..0.. 10111.....13.1111 /p 1001191144 1445.
11

---

Published as a conference paper at ICLR 2024
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Al- berti, Danielle Epstein, Illia Polosukhin, Matthew Kelcey, Jacob Devlin, Kenton Lee, Kristina N. Toutanova, Llion Jones, Ming-Was Wei Chang, Andrew Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural questions: a benchmark for question answering research. Transactions of the Association of Computational Linguistics, 2019.
Xuechen Li, Tianyi Zhang, Yann Dubois, Rohan Taori, Ishaan Gulrajani, Carlos Guestrin, Percy Liang, and Tatsunori Hashimoto. Alpacaeval: An automatic evaluator of instruction-folloviing models. https:/ github...... www.oo/ttttoooeteeooooooooe lab/ aaaaa//laaaaaaaaalaaaa 2023.
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 LLM KV cache compression at test time. CoRR, abs/2305.17118, 2023a. doi: 00.48550/arXiv.2355.. 17118. URL https::
2305 111...7111111 17118.
Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivas- tava, Ce Zhang, Yuandong Tian, Christopher Ré, and Beidi Chen. Deja vu: Contextual spar- sity for efficient ilms at inference time. Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.). International Confer- ence on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 22137-221766. PMLR, 2023b. URL https:: pttpppoocediiii ings..gggg.. ml www.mmmrrsrrsseeeemmmmmmmmmmmmme press /22022122111 200221ll2ammm html.
Paul Michel, Omer Levy; and Graham Neubig. Are sixteen heads really better than one? In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https: httts:://ppoceeedngg...eeee
cc wwtp//ppeeeieeeeppeeeeeeeeeeeeee /tt.ess/pppeee/01/////////////// pdf.
Jesse Mu, Xiang Lisa Li, and Noah D. Goodman. Learning to compress prompts with gist tokens. CoRR, abs/2304.08467, 2023. doi: 104485500arxxi.2344.46666 URL https: //t//////////////////////// wwd.doorrggooooooo /rg/ 10. (0.48550/ark/ 2304. 08467.
OpenAl. Gpt-4 technical report, 2023.
Minjoon Seo, Sewon Min, Ali Farhadi, and Hannaneh Hajishirzi. Neural speed reading via skim-rnn. ArXiv, abs/1711.02085, 2017. URL https:////// /api. semanticscholar orgg
Noam M. Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe V. Le, Geoffrey E. Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of- experts layer. ArXiv, abs/1701.06533. 2017. URL https:/ /ttp:///ppiiai www.app.ssaaaiicccccceseocccoeee
org/CorpusTD: 12T::224422
Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Daniel Y. Fu, Zhiqiang Xie, Beidi Chen, Clark W. Barrett, Joseph Gonzalez, Percy Liang, Christopher Ré, Ioan Cristian Stoica, and Ce Zhang. theouhhhou generative inference of large language models with single gpu. In International Conference on Machine Learning, 2023. URL https://api. semant www.tticcchla.... org/ wwagg//rppussc ID: 257495837.
Tianxiang Sun, Xiangyang Liu, Wei Zhu, Zhichao Geng, Lingling Wu. Yilong He, Yuan Ni, Guo- tong Xie, Xuanjing Huang, and Xipeng Qiu. A simple hash-based early exiting approach for Smaranda Muresan, Preslav Nakov, and Aline Villav-
language understanding and generation. icencio (eds.), Findings of the Association for Computational Linguistics: ACL 2022, Dublin, Ireland, May 22-27, 2022, pp. 2409-2421. Association for Computational Linguisties, 2022. doi: 10. 10t2:/1.1122222.nnn..g1 acl.189. URL https:/ https://ddo wtt..///.......//////... sindings-acl.. 189.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023a.
12

---

Published as a conference paper at ICLR 2024
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023b.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned, July 2019, URL org/p1gg. P19- 1580.
Ziheng Wang. Jeremy Wohlwend, and Tao Lei. Structured pruning of large language models. In Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu (eds.). Proceedings of the 2020 Conference on Empirical Methods Natural Language Processing, EMNLP 2020, Online, November 16-20, 2020, pp. 6151-6162. Association for Computational Linguistics, 2020. doi: 10.18653/v1/2020. emnlp-main main.aai.......... URL https: /tt. org/ 10 18653/v1/2220 sannp-maing 496.
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark W. Barrett, Zhangyang Wang, and Beidi Chen. H20: Heavy-hitter oracle for efficient generative inference of large language models. CoRR, abs/2306. 14048, 2023. doi: 10.485500arXX.2236644444 URL https: //t::////oo..ee org/10. 4855 48550/arXiv.. 2306. 14048.
Chunting Zhou, Pengfei Liu, Puxin Xu, Srini Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, Lili Yu, Susan Zhang, Gargi Ghosh, Mike Lewis, Luke Zettlemoyer, and Omer Levy. LIMA: less is more for alignment. CoRR, abs/2305. 1206. 2023, doi: 10.485500arXiv22300. 11206. URL https:://///
bttt:::://ooo //o.i///o///////// wew.oo.rro.......e...eoo. woog/00...... 48550/ /a..///rxivvv. .2055.230...0 11206.
Wangchunshu Zhou, Canwen Xu, Tao Ge. Julian J. McAuley, Ke Xu, and Furu Wei. BERT loses patience: Fast and robust inference with early exit. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020. URL https:/ http:////occeeeenn...oo
ca/ htc/paper/2220/heee Abstract html.
A APPENDIX
A.1 ABLATIONS
For all the ablations, we use a fixed targeted recovery ratio T = 0.98.
How one policy affect all the other policies? We study the complementary effects of each policy on the combination of all other policies in our framework. We examine changes in pruned KV cache and win rate while fixing the targeted recovery ratio T. We take the full policy set as our control set C. For each ablation, we remove one of the policies from all policy combination in C. We summarized the results in Table 4, which suggests the Cieqquent- and the Cspecial are the most important policies. Removing them will ineur 3.67% and 2.11% win rate drop respectively. We can also observe from the pruned cache ratio that Creqquent and Clocal reduce more KV caches than the others. However, their standalone ada-aivet deployment yields suboptimal performance, as depicted in Figure 2, further verifying the importance of adapting different compression policies.
13

---

Published as a conference paper at ICLR 2024
Feasible Policy Set
Pruned KV Ratio
Win Rate
C
36.04%
49.75%
(Cpunct. C punct.+frrquuent +unnet.ffreuuuenttolcalll Ciull
31.16%
47.64%
sSppciall C speciall C feeccilll Culll
34.23%
49.56%
(ppeciall C special-punne. C "speciall-umnet--ffeqeeeee Cful))
30.18%
49.06%
speecial C Special++punct. special-puncct-oocll Cuul))
21.26%
46.08%
Table 4: Complementary effects of each policy. We display the win rate each method over full
cache setting. We evaluate the fine-tuned Llama 1-65B on AlpacaEval with the same parameters.
Cache Order
Pruned KV Ratio
Win Rate
Ospecial Opunct. Crequent
Clocal
36.04%
49.75%
Ospecial Cirequent
Ooocl
Cpunct.
36.40%
47.64%.
Table 5: Policy order ablation fine-tuned Llama 1-65B with AlpacaEval.
Pruned KV Cathe
50%
50%
50%
50%
Win Bate
40%
30%
30%
30%
20%
0.1
0.2
0.3
0.4
20%
0.1
0.2
0.3
0.4
20%
0.1
0.2
0.3
0.4
20%
0.1
0.2
0.3
0.4
Locality Ratio
Locality Rabio
Frequency Ratio fy
Frequency Ration
Figure 6: Hyper-parameter ablation fine-tuned Llama 1-65B with AlpacaEval.
Which policy should we add first (and last)? As in Section 3.4, we use greed method to construct adaptive KV cache. Here, examine how the order of introducing each policy affects the performance. Similar to the previous study, we fix the targeted recovery ratio to 0.98, and keep allocating cache budget until the constructed cache hit the recovery ratio. For simplicity, make every examined order opt-in the C Special first, as it's typically the most important tokens and of super-low memory cost, suggested in Figure 1. We summarize the results in Table 5. Our current order (as Equation 2) achieves the highest win-rates and the highest pruned ratios. Meanwhile, using alternative orders leads to different trade-off between KV cache compression and generation quality. For example, using Cirequent + Clocal > Cpunct. leads to an improved KV cache compression ratio at the cost of generation quality.
A.2 SENSITIVITY STUDY.
We analyze the sensitivity of selecting different hyper-parametert for FastGen, illustrated in Fig- ure 6. We observe that altering these hyper-parameters does not have visible impact on the gen- eration quality, as the model maintains winrate over 45% in all situations. Meanwhile, it leads to a relative large change on the compression ratio. For example, changing the ratio for the frequency policy from 0.3 0.1 leads more KV cache. In our experiments, we the ratio to 0.3 for both and ry.
14


Reviews:
[{"review_id": "w8AOju6vTN", "reviewer": "Reviewer_XPHq", "summary": "This study introduces a lossy adaptive KV cache compression technique aimed at reducing the memory footprint of LLMs. The paper is guided by two key insights:\n\nDifferent attention heads typically exhibit distinct structures.\nThese attention head structures remain relatively consistent during inference.\nThe paper profiles the prompt encoding phase to identify the intrinsic structures of various attention heads and uses these structures to determine the optimal compression policy. This policy, determined during the prompt encoding phase, is then applied uniformly throughout all token generation iterations. The compression policies are combinations of the following four basic ones: special, punct, local, and frequent.\n\nThe results demonstrate that this approach yields improved model quality compared to fixed KV compression methods, with KV cache budgets ranging from 30% to 100%. The ablation study further reveals that frequency- and special-token-based compression policies have the most significant impact on compression ratio and win rate.", "questions": "-", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper introduces valuable insights drawn from LLMs: 1. Different structure in different attention 2. The same head structures persist. These insights are well-supported with empirical data and references to existing literature. \n\n- The authors leverage these insights to come up with an effective compression method that adapts to the structure of each attention head. The results show consistent compression rate and model quality improvement over prior SoTA fixed compression mechanisms.", "weaknesses": "- The paper could benefit from presenting actual GPU inference performance results using FastGen and comparing them with other compression methods. Additionally, providing a runtime breakdown would offer more insights into the overhead caused by the profiling, compression, and decompression processes.\n- It would be nice to look into the structure of KV in the multi-query attention design.", "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": "lfT0rmabIS", "reviewer": "Reviewer_XU71", "summary": "This paper discussed how to apply adaptive KV cache compression to improve the system efficiency, which conducts profiling to discern the intrinsic structure of attention modules. The proposed method can be deployed without resource-intensive fine-tuning or re-training. Solid empirical study was conducted to verify the efficiency and effectiveness of the proposed method.", "questions": "Would it be possible to address the minor issues I listed in the weakness section?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- This paper solves a critical research problem about efficient LLM inference with advanced algorithm design. The designed algorithm is straightforward and effective. \n\n- The presentation of the technical discussion is accurate and well-organized.\n\n- The organization of the evaluation sections is clear, and the presented results show the advance and efficiency of the proposed method.", "weaknesses": "- Based on my understanding, the proposed algorithm specializes in the most classic softmax-based attention. Is it possible to include a small section discussing the limitations of the proposed algorithm for more complicated attention mechanisms and some preliminary ideas about supporting those mechanisms in the future?\n\n- Given the scale of the benchmarked model (llama-70B fp16 on A100-80G), I guess there is a missing detail about the parallel strategies applied in the experiments.", "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": "TKcKski2aK", "reviewer": "Reviewer_BiHE", "summary": "The paper addresses the memory footprint reduction of LLMs during inference, in which the recent problem is the KV cache eviction/compression policies. The paper proposes an adaptive KV cache compression technique that operates in two stages, i) diagnose through profiling based on the attention heads and ii) applying an eviction strategy per each layer.", "questions": "Please refer to weaknesses section for questions.\n\n## Post rebuttal comments\n\nThe responses and the detailed analysis in the Tables1,2 address my concerns.\n\nHowever the authors seem to reserve one of the suggestions to the future works. Overall, very satisfied with the impressive work in the paper and raising the score to clear accept.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- Having an adaptive KV cache for each of the attention module type is a really interesting and exciting idea.\n- No fine-tuning costs of the proposed method is commendable. \n- The paper clearly positions within the body of existing literature, by distinguishing the proposed method as an adaptive and a diverse set of eviction strategies.\n- The paper is clearly written, the presentation is great, easy to follow along and digest the concepts.", "weaknesses": "- Although, the idea of adaptive KV cache compression sounds interesting, what is the overhead of book-keeping to support this adaptive and diverse ability based on the type of the attention? This is not discussed anywhere in the paper?\n  - That is, each layer id will be mapped to a eviction policy and is deployed with the model at hand. \n  - Next, what is the added computational complexity both asymptotically as well experimentally.\n- Table 3 shows an ablation on the policy order, why is this needed? Is the policy fixed per layer and the order will be dictated by the layer that needs a certain policy determined by the diagnosis step. Is it not true, clarify on this please.\n- Another interesting exploration/ablation to see is to experiment with long context tasks. What if the downstream task requires a long context window then what can be the best set of eviction strategies and the corresponding expected win rates?\n### Minor comments:\n- \"The resulting distribution is visualized in Figure As in Figure 3.\" can be rewritten as \" Figure 3 shows the resulting distribution\"\n- A minor nit, the paper has too much forward referencing, which disturbs the flow of reading and attention, general recommendation in research papers is to avoid such referencing..!\n- Better to define the new terms such as win-rate, KV cache budget, etc. when they were introduced for the first time. Similar applies to abbreviations when they are introduced first time, expand them, for the sack of saving readers time to search internet.", "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": "aNqNdIZp5T", "reviewer": "Reviewer_9wxD", "summary": "This paper proposes FastGen, an adaptive key-value (KV) cache compression method to reduce the memory footprint and accelerate inference for large language models (LLMs). The key ideas are: 1) Profiling attention modules to discern their intrinsic structures, such as primarily attending to local contexts or special tokens. 2) Constructing the KV cache adaptively based on the recognized structure to compress less useful contexts. 3) The lightweight attention profiling guides the KV cache compression without expensive fine-tuning.\n\nThe experiments are conducted on LLaMa models with sizes from 7B to 65B parameters on diverse generative tasks. Results show FastGen effectively compresses the KV cache to 40-50% smaller with negligible quality loss. It also outperforms non-adaptive baselines.", "questions": "Comments:\n-----------------\nC1:\tThe compression policies are combined in a simple naive way in FastGen, by just taking the union of multiple policies such as Cspecial + Cpunct + Cfrequent. This straightforward combination approach has several potential issues. First, is it possible that the union combination may introduce redundancy, as different policies could select overlapping important content, leading to suboptimal compression ratios? More intelligent strategies should consider the complementarity between modules to avoid duplicating the key contexts. Second, is it possible that existing policies may not be fully compatible? Some combinations could introduce conflicts and hurt generation quality. More systematic analysis should examine the compatibility between policies.\n\nC2:\tThe experiments in the paper are all conducted on the decoder-only LLaMa models, without validation on encoder-decoder models like BART and T5. These models are also widely used for generative tasks, so the efficacy of FastGen on them remains unclear. This is worth further investigation. \n\nC3:\tThe paper lacks sufficient analysis on the overhead and time cost of conducting attention profiling, which is important to judge the efficiency of FastGen in real deployment. Specifically, the time complexity of attention profiling needs analysis, and concrete profiler time under different model sizes should be provided or disscussed. Moreover, analyzing the extra memory or GPU memory required for the profiler and assessing its impact on deployment is necessary. In summary, quantitatively analyzing the resource overhead for profiling and demonstrating effective solutions to reduce it could strengthen the practicality of FastGen in real-world usage. Further experiments on optimized profiling and its cost-benefit trade-off with compression performance could provide more comprehensive insights into the efficacy of the approach.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Strengths:\n-----------------\n+ Adaptively compressing KV cache better aligns with model-specific attributes without retraining.\n+ Comprehensive experiments verify FastGen works for diverse models and tasks. Up to 50% compression on 65B LLaMa with little quality loss is remarkable.\n+ Ablation studies provide good insight into the design choices. The profiling method and compression policies are well motivated.", "weaknesses": "Weaknesses:\n-----------------\n- The compression policies are combined in a naive way. More advanced adaptive selection could be explored (see detailed in C1).\n- No experiment on encoder-decoder models. The efficacy on them is unclear (see detailed in C2). \n- More analysis on the overhead of profiling could be provided (see detailed in C3).", "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": "gc87DBhLPk", "reviewer": "Reviewer_rrTE", "summary": "This paper propose an adaptive KV cache compression technique to reduce the memory footprint of generative inferences of LLMs. The authors fist perform targeted profiling to indentify the intrinsic structure of attention modules, and then build an adaptive KV cache by evicting long-range contexts on attention heads emphasizing local contexts, removing non-special tokens on attention heads centered on special tokens, and using only the standard KV cache. The experimental results show the adaptive KV cache achieves large reduction on GPU memory consumption with trivial geneation quality loss.", "questions": "Please comment the two points in the weakness section.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper works on an important topic, i.e., reducing the memory footprint of GPU during generative inferneces of LLMs.\n2. The paper flows well.", "weaknesses": "1. The model profiling part is not clear. Did the authors do a profiling for each model on all datasets, or each model on a single dataset. \n2. The model profiling results have a huge impact on the final KV cache compression results. Although the authors show empirical data supporting the structure of the attention map is stable at different positions for all attention heads, the authors still need to discuss what if the structure of the attention map is not stable.", "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": "YN66lg0xUR", "reviewer": "Reviewer_XtqU", "summary": "Key-value cache takes the majority of GPU memory in LLM serving, and its extent is continuously growing along with the model size and context length. Therefore, if we can reduce the key-value cache memory while maintaining the generation quality, we can accelerate the LLM inference. \bThe authors propose FastGen, a framework for efficient generative inference by applying on-the-fly key-value compression. They analyzed the structural patterns of each attention head of layers and then categorized four policies. By adopting the optimal policy based on profiling, FastGen achieves a comparable generation quality to the full-cache (non-compression) inference.", "questions": "- What are the additional challenges for the models that use the grouped query attention technique?\n- In Figure 4, attention scores of special tokens always take more than half. Are there attention heads whose special token score is lower than half?\n- In Figure 5, compressing sometimes wins the full-cache strategy. How can we interpret such results?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The authors provide abundant experiments with varying model sizes and tasks\n- The authors provide informative ablation studies\n- Their work will motivate various related work, for example,\n    - efficient kernel which aware of compression\n    - as the model size grows, more redundant key-values exist where we have more room to optimize", "weaknesses": "- Since the inference time matters in practical serving, it would be helpful to understand more if the authors can provide corresponding results\n    - For example, how long does inference take compared to the full-cache strategy? I think it might become slower because the existing attention kernels may not efficiently deal with the sparsity\n    - How long does profiling take? Is it feasible for practical inference scenarios?\n- It seems that the StreamingLLM paper [1] is similar to this work. It sets sink tokens and performs local attention, where the sink tokens may correspond to $C_{special}$ (and maybe $C_{punct}$. Since the StreamingLLM paper has also recently been uploaded, it is unlikely to compare this paper with it. But it would be better if the differences in this paper were clarified.\n\n[1] Xiao, Guangxuan, et al. \"Efficient Streaming Language Models with Attention Sinks.\" arXiv preprint arXiv:2309.17453 (2023).", "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 (oral)