W4

Ada-KV[3] is a cache budget allocation method that tries to allocate different budgets across heads. EMS uses uniform cache budget for each head, which is orthogonal and compatible to the head-wise cache budget work.

The similar part between KVMerger [4] and EMS is that we both do KV Cache merging. However, they differs in merging mechanism and compression strategy. For the merging mechanism, EMS utilizes a cross-position strategy, where KVMerger is to merge adjacent tokens. For the compression strategy, EMS designs both eviction and unified merging operation, while KVMerger only do merging.

Nacl[5] is an eviction-based method that uses proxy tokens to accumulate the attention weights. EMS decouples the compression into 2 steps: identifying token importance and applying evict-then-merge strategy, and focuses more on the evict-then-merge strategy. Therefore, Nacl is generally orthogonal to EMS. Although they have also designed an accumulate pattern, which is a more intricate token importance indicator, their score can also be applied to EMS.

[1] Feng, Yuan, et al. "Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference."

[2] Wang, Zheng, et al. "Model tells you where to merge: Adaptive kv cache merging for llms on long-context tasks."

[3] Chen, Yilong, et al. "Nacl: A general and effective kv cache eviction framework for llms at inference time."

W5

From the perspective of identifying the token importance, both global and local score have certain advantages. As can be seen from Table 1, H2O and SnapKV perform well on different tasks. And our global-local score can adaptively combine them.

Although the current calculation of global scores is not compatible with FlashAttention, there is potential for optimization in this area:

• For instance, [1] can approximate global scores through sampling, which is orthogonal to and compatible with our approach.
• FlashAttention itself is evolving. We are currently working on improving its kernel to enable it to return accumulated attention scores, which will further accelerate inference. Nevertheless, this will not change the nature of our method. Furthermore, we find that the recent work [2] has implemented accumulated attention weights of FlashAttention by recomputation based on Paddle. So the dilemma of the global score can be solved.

[1] He, Yefei, et al. "ZipCache: Accurate and Efficient KV Cache Quantization with Salient Token Identification."

[2] Chen, Yilong, et al. "Nacl: A general and effective kv cache eviction framework for llms at inference time."

Thank you for taking the time to thoroughly read our article and for pointing out its shortcomings. Here is our response:

W1

We have supplemented the performance of different models using more budgets and their comparisons to baselines. We set $\gamma=4$, $\tau=0.6$ as ablated in Appendix B.

More results can be seen in Table 9. Through the results, we can see that EMS outperforms other methods on the 4 LLMs across different cache budgets.

W2

In the ablation part, the parameters are fixed to the best parameters obtained by the ablation ($\gamma=4$, $\tau=0.6$). In Table 1, we conducted multiple sets of tests on a subset of parameters ($\gamma=3:5$, $\tau=0.5:0.7:0.05$) around the best parameters and reported the best results. Here, we also provide the mean performance on this subset.

The experimental results show that even when using the mean results on the test subset, that is, without much parameter tuning, the performance of EMS still outperforms baselines.

W3

Thanks for pointing out this problem! This is not the problem of experiment pipeline. We checked the table and find that we incorrectly set the table header when organizing the table, which led to the biases in the reported scores. We have updated the table header and did not change any data. The results now can match.

Thanks for your comments on the design of Global-Local score. Here is our response for your concernings:

W1

The design of the compression strategy is essentially a trade-off among performance, memory, and time. The goal of EMS is to achieve the extreme compression of tokens, reduce memory overhead, and minimize the performance degradation as much as possible, which is of great significance for some long-context scenarios. Although EMS has some additional overhead, it can bring benefits in terms of performance and memory.

About EMS, perhaps you have some misunderstanding of our approach. StreamingLLM and H2O are two pioneering works in KV cache eviction. Unlike them, EMS decouples the compression process into two phases: identifying token importance and applying evict-then-merge strategy. In the first phase, we found the flaws of StreamingLLM and H2O and designed a better token indicator to alleviate the biases they exhibit. In the second phase, we achieved more extreme compression based on the token importance. Therefore, the first phase of EMS actually discovers the deficiencies of StreamingLLM and H2O and improves them.

W2

Mechanism of the combination of local and global attention: By mean-alignment, EMS dynamically integrates the former- and local-biased scores and designs the more balanced Global-Local score to indentify the token importance more effectively. We investigate more on the relative weights of global and local scores and found that changing the relative weight has the potential to achieve better results. And EMS can also achieve good performance when the relative weight is 1. This easy-to-implement mean-alignment can improve the performance. In the future, we will continue to explore better token indicator, such as task-aware ratios.

Thank you for your valuable feedback on our submission. We greatly appreciate your efforts and constructive comments. As a kind reminder, the discussion period is drawing close. Please let us know if there remains anything that we can further clarify to improve our work. We understand that this is a busy period and appreciate your time and consideration.

Thank you again for your contribution to the review process. With the clock ticking down on the discussion period, we kindly ask if you could revisit our submission, review our rebuttal, and share any additional remarks or conclusions you may have. Your continued engagement will ensure a more thorough and fair assessment of our work, which we believe holds significant potential contributions to the field.

W3

Motivation of expansion: We hope to conduct clustering on the KV tokens based on cosine similarity. In this way, we can use a small number of class center tokens to represent more tokens. Therefore, those similar tokens share the same KV entry and need to be expanded during computation.

Implementation of expansion: a position look-up-table is maintained and $N_{imp}$ tokens are expanded to $\gamma N_{imp}$ tokens at computation time of the certain layer.

To validate the effectiveness of expansion, here we provide the comparison between the cases with and without expansion across different budgets on Llama2. We can infer that the expansion can steadily bring performance gains, which demonstrates its necessity.

W4

Thanks for pointing out this problem! This is not the problem of experiment pipeline. We checked the table and find that we incorrectly set the table header when organizing the table, which led to the bias in the reported scores. We have updated the table header without changing any data. The scores now can match.

Besides, the implementation of EMS is partly based on the codebase of KIVI. Under our experiment pipeline, the performance of the model with full cache can match that of KIVI.

[1] KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache, ICML’24, https://arxiv.org/abs/2402.02750.

W5

We are using the haystack formed from a long corpus (Paul Graham Essays).

Thank you for your careful reading of this paper and for pointing out issues in the details. Here is our response:

W1

Thank you for pointing out our oversight. The case proposed by the reviewer is a minor one. We have retested the average performance on LongBench considering the absolute value of cosine similarity with a 256-cache budget. Although the improvement has brought a certain performance gain, this gain is relatively small, and there is even a slight performance drop in some case. Compared with the best baseline results, the main increase is still brought by our method. In the future, we will consider designing an adaptive method to make better use of similarity.

W2

Thank you for pointing out the shortcomings of the experiment. We have supplemented the performance results on LongBench across various cache budgets using different methods to make the evaluation of our approach more comprehensive.

Here we show the performance on Llama2.

More results can be seen in Table 9. Through the results, we can see that EMS outperforms other methods on the 4 LLMs across different cache budgets.

Thank you for your valuable feedback on our submission. We greatly appreciate your efforts and constructive comments. As a kind reminder, the discussion period is drawing close. Please let us know if there remains anything that we can further clarify to improve our work. We understand that this is a busy period and appreciate your time and consideration.

Thanks for your feedback. Here we conduct the statistics and performance test again.

Statistics. We have counted this minor case (key_sim < 0 and value_sim < 0 and key_sim * value_sim > tau) on four LLMs using Longbench samples and calculated its percentage among the cases we actually use (key_sim * value_sim > tau). The results show that this minor case scarcely occurs in Llama3 and Mistral, and it only accounts for 0.25% and 0.02% in Llama2 and LongChat respectively. Considering that taking this situation into account will instead add a mask operation, it is thus a negligible case.

Test code:

# -- and ++ cases
mm_pp = torch.count_nonzero(key_sim * value_sim > tau)
# ++ cases
pp = torch.count_nonzero((key_sim > 0) * keys_sim * values_sim > tau)

# accumulate all mm_pp and pp, then calculate:
rate = (mm_pp - pp) / pp

Performance test. We have supplemented the performance using ((key_sim > 0) * (key_sim * v_sim)) > tau. From the results, it can be seen that the impact of this minor case on the results is rather minimal and negligible.

Q1

Combine global and local scores for less biased indicator. We observe the intrinsic former and biases of $s_{Glo}$ and $s_{Loc}$, and hope to design a method that can preserve the advantages of both.

Mean-alignment to counteract the huge scale difference. Because $s_{Glo}$ accumulates all rows of attention weights, while $s_{Loc}$ only accumulates $L_{win}$ rows. For example, when the sequence length and local window length are 4096 and 32, we did a statistics on the mean values of $s_{Glo}$ and $s_{Loc}$ and they  can differ by a hundredfold.

Benefits of mean-alignment. As mentioned in W3, when the relative value of one score is too large, it will lead to performance degradation.

Q2

The difference lies in two aspects: the cosine similarity and the position encoding differences.

Merge with position information: with the position information, the cosine similarity between keys decays with the increase of relative distance, leading to less redundancy and merged chances. The advantage is that the precise position information of tokens is not changed, which is beneficial for tasks that require precise token localization.

Merge without position information: redo the positional encoding for each decoding as illustrated in [1]. Thus we can calculate cosine similarity of raw keys (without positional information), and more tokens can be merged potentially. Besides, re-apply the positional encoding for each decode is actually superimposing the encoding from 0 to the kv_len involved in the computation, which also enables streaming generation.

[1] Efficient Streaming Language Models with Attention Sinks, ICLR’24, https://arxiv.org/abs/2309.17453.

Q3

Sorry for the unclarity. Let me illustrate again.

We can first understand it from the inference stage. What we actually need to retain are $N_{imp}$ (also known as $N_{budger}$) tokens as class centers. They will be expanded into $\gamma N_{budget}$ tokens according to a look-up-table, and only these tokens will be used for computation.

From the perspective of compression, our goal is to compress the most important $\gamma N_{imp}$ tokens into $N_{imp}$ tokens, and evict all the remaining unimportant tokens. In this way, the three types of Key-Value (KV) cache is obvious.  The first type consists of $N_{imp}$ compressed class centers. The second type includes  $N_{tbm}$(equivalent to $(\gamma-1) N_{imp}$ ) tokens that are to be merged. The third type is the evicted irrelevant tokens.

Thanks for pointing out that the relative weight can affect the global-local score. In this paper, we only focus on the impact of global score, local score and global-local score. After some preliminary experiments, we found that changing the relative weight has the potential to achieve better results. This further illustrates the potential of the global-local score as a token importance indicator.

Generally, EMS can also achieve good performance when the relative weight is 1. We will take the dynamic weights setting as a direction for more exploration in the future. Thanks again for the enhancement of the token importance indicator!

Thank you so much for acknowledging this paper and for sharing your concerns about algorithm efficiency.

W1

For the calculation of the token importance score, H2O only calculates global scores, while SnapKV only calculates local scores and performs pooling. EMS adopts the Global-Local score, calculates both and combines them.

Here we provide the test results with different test sequence lengths using Llama2-7B-chat. All results are obtained by testing 5 times and then averaging the results. The format of the datas in the table is prefilling time (ms)/decoding time (ms).

It's worth mentioning that SnapKV only performs compression in the prefilling stage, so there is no decoding time in the table.

From the results, the calculation time of the Global-Local score is basically the sum of the calculation time of the global score and the local score, or slightly less than that. In addition, the longer the sequence length is, the smaller the additional overhead will be.

Regarding the overhead for calculating the score, EMS focuses on extreme compression while maintaining the performance. We have noticed that some recent works can be adapted to our work to accelerate the compuation .ZipCache [1] utilizes 10% probe tokens to approximate the accumulated scores. Besides, Nacl [2] integrates the calculation of the global score into the FlashAttention kernel. These works are all orthogonal to our work and can be applied to EMS to accelerate the calculation.

[1] ZipCache: Accurate and Efficient KV Cache Quantization with Salient Token Identification, NeurIPS’24, https://arxiv.org/abs/2405.14256.

[2] Nacl: A general and effective kv cache eviction framework for llms at inference time, ACL’24, https://arxiv.org/abs/2408.03675.

W2

We set the input length to 4096 and test the merging time under different cache budgets. The experiments are conducted on two RTX4090 GPUs.

Since the code is currently implemented in the initial version, we will conduct further optimizations to enhance the compression speed. Although the merging process takes up a certain amount of inference time, it reduces the memory overhead and also brings a significant performance gain. Here we supplement the performance of different models using more budgets and their comparisons to baselines.

More results can be seen in Table 9. Through the results, we can see that EMS outperforms other methods on the 4 LLMs across different cache budgets. Under an extremely low cache budget, while other baselines may break down, EMS can maintain its performance, which illustrates the robustness of our method.

Thank you for your valuable feedback on our submission. We greatly appreciate your efforts and constructive comments. As a kind reminder, the discussion period is drawing close. Please let us know if there remains anything that we can further clarify to improve our work. We understand that this is a busy period and appreciate your time and consideration.

Thank you again for your contribution to the review process. With the clock ticking down on the discussion period, we kindly ask if you could revisit our submission, review our rebuttal, and share any additional remarks or conclusions you may have. Your continued engagement will ensure a more thorough and fair assessment of our work, which we believe holds significant potential contributions to the field.

Thanks for your comment.

We set the cache budget 256 to show that even at very low budgets, EMS holds up well. In fact, the maximum context lengths for Llama-2-7B-Chat, Llama-3-8B-Instruct, LongChat-7B-v1.5-32k and Mistral-7B-Instruct-v0.2 are 4096, 8192, 32k and 32k. So the budget 256 is a very small budget and accounts for 6.25%, 3.125%, 0.8% and 0.8% for these four models. Even under this circumstance, our method can maintain 95.7%, 93.7%, 92.7%, 85.5% performance with a full cache.

To further validate the performance of EMS, we also tested the performance of different models using more budgets and their comparisons to baselines.

Here we show the performance on Llama2 using different budgets and the comparisons with baselines. More results can be seen in Table 9.

Through the results, we can see that EMS outperforms other methods on the 4 LLMs across different cache budgets.

Thanks for your comment.

We believe that the expansion will not cause this problem. The compressed tokens themselves exist in a continuous memory space. The way we obtain the expanded tokens is to first allocate a continuous memory space and then expand the compressed tokens according to the position look-up-table(LUT). The overhead brought by expansion is only for the currently computed layer, and other layers will still maintain compact storage. In fact, this is a trade-off among performance, memory, and latency. Although there is additional overhead, it leads to more compact memory and better performance.

Thanks for your comment.

Irregular pattern illustration: EMS set the same KV cache budget and expansion length for each head. The head-wise characteristic is utilized by expanding to the evicted zero or the merged values, which means that different heads have different expanded patterns.

The expansion length is indicated by the factor $\gamma$. EMS will first evict $N_{irr}$ tokens, and then perform a merge within the remaining $N_{imp}+N_{tbm}=\gamma N_{imp}$ tokens. The compressed $N_{imp}$ tokens are expanded to $\gamma N_{imp}$ tokens during the inference.

In Table 7, we have tested the impact of $\gamma$. Setting $\gamma=$'Inf' is equivalent to expanding shared entries to the full context length. And we found that expanding to full context length is sub-optimal, which might be due to the distraction of long context. So the the eviction in the first stage can not only reduce overheads but also improve performance.

Thank you very much for recognizing this paper and for raising concerns about memory. We believe there may be a slight misunderstanding regarding our method.

values such as $N_{tbm}$, $N_{irr}$ and $N_{imp}$ are non-uniform across all heads

Uniform head budgets: EMS does track $s_{Glo}$ and $s_{Loc}$ for each head. But $N_{tbm}$, $N_{irr}$ and $N_{imp}$ are same for each head, and the length of $s_{Glo}$ and $s_{Loc}$ are also the same and constant $\gamma N_{imp}$. So there won’t exist non-uniform problem.

support this irregular pattern and maintain data contiguity

Adapt irregular patterns by unified merging: we speculate that the irregular pattern you are concerned about lies in the possibility of different merge ratios in the unified

the KV cache would need to exceed the target budget

Scalarized score indicator is much smaller than vectorized KV tokens: for example, the $\gamma=4$ (ablated in Table 7), head_dim=128, the additional overhead of score indidator accounts for $(4*3*N_{imp})/(N_{imp}*128*2)=4.7$ of the retained KV cache, which is efficient.

Thank you for your valuable feedback on our submission. We greatly appreciate your efforts and constructive comments. As a kind reminder, the discussion period is drawing close. Please let us know if there remains anything that we can further clarify to improve our work. We understand that this is a busy period and appreciate your time and consideration.