Paper: Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs (ICLR-2024)

Reviewer comment (Evaluation):
Since the inference time matters in practical serving, it would be helpful to understand more if the authors can provide corresponding results. 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.

Author response:
To address reviewers’ concerns on the end-to-end speedup of FastGen, we implement a sparsity kernel for KV-cache pruning and present the end-to-end latency improvement in Table 1. Please refer to general response #1 for detailed settings and analysis.

As shown in Table 1, 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 a decent 16.04% latency improvement over the full-cache baseline on a short generation length of 512. In the best cases, we can achieve up to 55.0% latency reduction with Fastgen at a generation length of 16k. 
We can also observe a clear and consistent tendency of larger relative speed-up as generation length becomes longer. For example, given batch_size = 1, FastGen’s relative speed-up rises from 16.04% to 55.0%, as the generation length grows from 512 to 16384. The phenomenon can also be observed in other batch settings. 

This analysis confirms that FastGen can achieve major speed-up in real development, especially in long-generation settings. Meanwhile, the efficiency of the customized kernels can be further improved. We leave this unique research and engineering challenge to future works.

Reviewer comment (Evaluation):
How long does profiling take? Is it feasible for practical inference scenarios?

Author response:
To better understand the overhead of the profiling step, we compare the profiling time with the total generation time across different generation lengths. We present the result in Table 2. Please refer to general response #2 for detailed settings and analysis. 

Table 2 shows the profiling time of the LLaMA65b in different generation length settings. We can observe that the profiling time only accounts for a very small percentage of the total generation 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 strategies, C_frequent, 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 the KV cache is (batch_size, num_of_head, sequence_len, hidden_dimension), while the dimension of extra memory for the cumulative attention scores is (batch_size, num_of_head, sequence_len). Considering hidden_dimension=128 for all model sizes, the memory overhead is 1/128=0.78% compared to storing KV cache only, which is a negligible cost.

In conclusion, the overhead introduced by the profiling step is nearly negligible in both time and memory, which confirms FastGen’s potential for real-world deployment.

Reviewer comment (Novelty):
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.

Author response:
Thanks for mentioning this concurrent work! We were not aware of it at the time of submission given it was posted after the ICLR deadline. After reading the paper, we think it is a great parallel work on long-context LLM, and we are happy to provide a comparison between it and FastGen! Generally, the two differ in two aspects:
1) Goal and Setting: StreamingLLM aims to extend the context length of LLM, while FastGen aims to improve the efficiency of general LLM inference (normal+long contexts). As a result, FastGen focuses specifically on generation tasks, which is different from the general task settings in StreamingLLM.
2) Method: StreamingLLM uses a fixed attention pruning strategy for all attention heads, while FastGen focuses on adaptively choosing different compression strategies according to the attention structure of each head.

Reviewer comment (Theory):
What are the additional challenges for the models that use the grouped query attention technique?

Author response:
Thanks for mentioning grouped query attention (GQA). We are working on accommodating FastGen to GQA. In GQA, heads within each group share the same KV vectors. Instead of head-wise pruning, we could modify FastGen to perform group-wise pruning. Specifically, we could individually evaluate each query by calculating the recovery ratio of its attention map (Q*K), and then average ratios from all heads within the same group, using the averaged ratio as the criteria to find the optimal strategy for each group. As analyzed, adapting FastGen to GQA requires minimal methodology change and imposes little additional cost. We are working on obtaining concrete experimental results to show the effectiveness of GQA.

Reviewer comment (Evaluation):
In Figure 5, compressing sometimes wins the full-cache strategy. How can we interpret such results?

Author response:
Thank you for mentioning this, we think the performance instability could be attributed to two aspects:
1): Figure 5 presents win-rate changes, which are calculated from GPT-4’s pair-wise voting and thus could fluctuate due to the uncertainty in GPT-4 generation sampling. 
2): Small perturbations in KV cache may add some uncertainty to model performance, which could sometimes improve robustness and lead to slightly higher results.

Reviewer comment (Experiments):
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.

Author response:
**W1: 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.**

A1: The profiling is run on every data instance (each sentence). For the same model, the profiling would be different if the input is different. So, we do the profiling on the fly during deployment. Such fine-grained flexible adaptation allows FastGen to reduce more memory footprint while preserving the model quality. We also provide more analysis on the overhead of profiling in Table 2 of **General Response**, which shows the cost is nearly negligible.

Reviewer comment (Evaluation):
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.

Author response:
**W2: the authors still need to discuss what if the structure of the attention map is not stable.**

A2 We are not quite clear what the reviewers mean by "the structure of the attention map is not stable". We hypothesize that the reviewer was referring to the results in Figure 4, which shows the accumulated attention scores remain consistent/stable during the entire generation phase. The reviewer might be wondering if this observation can be generalized to all situations (e.g., different datasets and models). If that's the case, we think the reviewer raises a valid concern about FastGen, as some strategies such as static eviction policy (e.g., punctuation tokens) may no longer provide an accurate prediction, and more adaptive policies are needed.  To remedy this, we propose to repeat the diagnosis step multiple times across the generation process of one sentence. To be more specific, we choose a new pruning policy once the number of decoded tokens reaches $k$, where $k \in$[1,sentence_len] is a predefined hyperparameter. For each sentence, the diagnosis is repeated for sentence_len//$k$ times. From the analysis in the second part of General Response, we could infer that the profiling overhead is relatively small even if it is repeated several times. On the other hand, it would be interesting to study the theoretic explanation of our observation, which we would like to explore in the future.

Reviewer comment (Experiments):
The compression policies are combined in a naive way. More advanced adaptive selection could be explored. In FastGen, the compression policies are combined by just taking the union of multiple policies such as Cspecial + Cpunct + Cfrequent. This straightforward combination approach has several potential issues. First, 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, 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.

Author response:
We agree that different compression strategies have overlapping tokens. For example, C_frequent is sometimes overlapped with C_special and C_local, as special tokens and local contexts usually accumulate higher attention scores. However, FastGen is designed to progressively evict tokens. When there is an overlapped policy, we only consider the complementary tokens it contains. It is guaranteed that unioning the policy monotonously will only bring in newly evicted tokens.
We further confirm that different policies are compatible in the ablation study (section 5.3). We study (1) how removing one strategy and (2) how changing relative policy order affects the overall performance.

We can draw an empirical conclusion that little-to-no conflicts exist between different policies. In fact, we observe complementary effects between existing strategies. That is to say, the performance of standalone strategies can usually be boosted significantly by introducing other strategies when necessary.

Some combinations can introduce inferior/superior performance than others. In our experiment, over all datasets, we find that following the order of  C_special, C_punct, C_frequen, C_local consistently gets the best performance. We leave investigating the intrinsic mechanism to future work.

Reviewer comment (Experiments):
No experiment on encoder-decoder models. The efficacy on them is unclear. The 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.

Author response:
Thanks for the suggestion, FastGen could be easily adapted to encoder-decoder models by pruning the KV cache in their decoder. We will elaborate on this in the introduction and add several related works in the future version. In encoder-decoder models, KV cache is still used in the decoder to save computation. During generation, it is a standard practice to fix the encoder inputs as existing prompts or instructions, e.g., as in BART [1] and FlanT5 [2]. In such scenarios, the decoder part still works autoregressively to output newly generated tokens, and it is flexible to directly apply FastGen to their KV cache. In this paper, we focus on decoder-only models as most of the prevailing LLMs are decoder-only models, e.g., LLaMa, OPT, and GPT. We agree that additional discussions and experiments should be added, and we will revise accordingly.

Reviewer comment (Experiments):
More analysis on the overhead of profiling could be provided. The 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 discussed. Moreover, analyzing the extra memory or GPU memory required for the profiler and assessing its impact on deployment is necessary. 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.

Author response:
Thanks for the valuable advice. We provide extra experimental results on book-keeping overhead in Table 2. Please refer to the general response #2 for a detailed time and memory analysis of profiling cost. 

In short, Table 2 shows the profiling time of the LLaMA65b in different generation length settings. We can observe that the profiling time only accounts for a very small percentage of the total generation 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 strategies, C_frequent, 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 the KV cache is (batch_size, num_of_head, sequence_len, hidden_dimension), while the dimension of extra memory for the cumulative attention scores is (batch_size, num_of_head, sequence_len). Considering hidden_dimension=128 for all model sizes, the memory overhead is 1/128=0.78% compared to storing KV cache only, which is a negligible cost.

In conclusion, the overhead introduced by the profiling step is nearly negligible in both time and memory, which confirms Fastgen’s potential for real-world deployment. We additionally provide end-to-end system latency improvement in Table 1. It shows that FastGen can achieve major speed-up in various generation settings. Please refer to General Response #1 for more analysis.

Reviewer comment (Theory):
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. That is, each layer id will be mapped to an eviction policy and is deployed with the model at hand. Next, what is the added computational complexity both asymptotically as well as experimentally.

Author response:
**W1: What is the overhead of book-keeping to support this adaptive and diverse ability based on the type of the attention? What is the added computational complexity both asymptotically as well experimentally?**

A1: Thanks for the valuable advice. We provide extra experimental results on book-keeping overhead in **Table 2 of General Response**. Please refer to the second part of General Response for a detailed time and memory analysis of profiling cost. 

In short, Table 2 shows the profiling time of the LLaMA65b in different generation length settings. We can observe that the profiling time only accounts for a very small percentage of the total generation 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 strategies, C_frequent, 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 the KV cache is (batch_size, num_of_head, sequence_len, hidden_dimension), while the dimension of extra memory for the cumulative attention scores is (batch_size, num_of_head, sequence_len). Considering hidden_dimension=128 for all model sizes, the memory overhead is 1/128=0.78% compared to storing KV cache only, which is a negligible cost.

In conclusion, the overhead introduced by the profiling step is nearly negligible in both time and memory, which confirms Fastgen’s potential for real-world deployment. We additionally provide end-to-end system latency improvement in Table 1. It shows that FastGen can achieve major speed-up in various generation settings. Please refer to the first part of General Response for more analysis.

Reviewer comment (Experiments):
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.

Author response:
**W2: Table 3 shows an ablation on the policy order, why is this needed?  Is the policy fixed per layer and determined by the diagnosis step?**

A2: The policy is determined in the diagnosis step, and it is fixed per head in each layer. As introduced in section 3.4 “Hybrid Policies”, we search for the optimal hybrid policy according to a predefined order. The order is greedily designed to prioritize cache policy with smaller memory costs, e.g., C_special. Once the optimal policy is determined, it will stay fixed in the generation process.

In Table 3, the order ablation study aims to show that FastGen is agnostic to small changes in searching order. By shuffling the relative order of C_punct and C_local, we observe a different trade-off between KV cache compression and generation quality. Overall, our current order (as in Equation 2) achieves the highest win-rates.

Reviewer comment (Writing):
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.

Author response:
We really appreciate the reviewer’s suggestions on our writing. We will reduce the forward referencing and improve the term definition in the future version.

Reviewer comment (Writing):
Better to define the new terms such as win-rate, KV cache budget, etc. when they are introduced for the first time. Similar applies to abbreviations when they are introduced the first time — expand them, for the sake of saving readers time to search the internet.

Author response:
We really appreciate the reviewer’s suggestions on our writing. We will reduce the forward referencing and improve the term definition in the future version.

Reviewer comment (Reproducibility):
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.

Author response:
We will add more details about our implementation in the future version. During inference, we perform model parallel by equally sharding the model weights to different GPUs within the same node. Attention heads are evenly distributed across GPU for parallel attention computation. During inference, the model_parallel_size is 8 for 70B, 4 for 30B, 2 for 13B, and 1 for 7B. During finetuning, we use 32 A-100 80G GPUs with model_parallel_size=4, data_parallel_size=8,  and batch_size_per_GPU=4 for all model sizes.

Reviewer comment (Experiments):
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.

Author response:
Thanks for the valuable advice. We provide extra experimental results in Tables 1 and 2 of **General Response** section. Please refer to the general response for a detailed analysis of actual end-to-end latency and profiling cost. 

In short, we can observe from Table 1 that FastGen achieves significant end-to-end speed-up across all the generation settings. For example, given batch_size = 1, FastGen’s relative speed-up rises from 16.04% to 55.0%, as the generation length grows from 512 to 16k. The phenomenon can also be observed in other batch settings. This analysis confirms that FastGen can achieve major speed-up in real development, especially in long-generation settings. 

Additionally, Table 2 shows the profiling time of the LLaMA65b in different generation length settings. We can observe that the profiling time only accounts for a very small percentage of the total generation 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 conclusion, the overhead introduced by the profiling step is nearly negligible, which confirms FastGen’s potential for real-world deployment.