ZipCache: Accurate and Efficient KV Cache Quantization with Salient Token Identification

Summary

The paper presents an adaptive, mixed-precision quantization method for compressing KV cache in LLMs. It proposes a channel-separable tokenwise quantization scheme to establish a robust quantization baseline, reducing the memory overhead of quantization parameters. A saliency metric, based on normalized attention scores, is employed to identify salient tokens accurately. This approach preserves essential information while aggressively quantizing less salient data, allowing for higher compression ratios with minimal impact on model accuracy. The method can be integrated with FlashAttention with an efficient approximation of the proposed metric. Experimental results validate the superiority of ZipCache on speed and accuracy over existing methods.

Strengths

1. The idea of preserving salient tokens while compressing less critical tokens is both intuitive and effective, with the proposed accurate saliency metric being crucial for achieving this. 2. The paper presents a novel quantization scheme tailored for channel-wise outliers in LLMs, providing a notable reduction in memory overhead compared to group-wise quantization schemes. 3. The paper demonstrates strong experimental results on GSM8k and HumanEval, convincingly showcasing the benefits of adaptive KV cache compression. 4. The integration with FlashAttention greatly enhances overall generation speed. 5. The paper is well-organized.

Weaknesses

1. Evaluating ZipCache on other generation or comprehending tasks will make the findings more robust.

Questions

1. Have you considered using only the attention scores from the most recent token to determine token saliency? What impact might this have on the overall performance? 2. What does saliency mean in Figure 3c? What is the probability of each token being selected as a salient token, as mentioned in the caption of Figure 3?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

Summary

The paper introduces ZipCache, an adaptive KV cache compression method for LLMs by accurately identifying salient tokens. It first presents a channel-separable tokenwise quantization scheme that reduces the overhead of quantization parameters. Next, it proposes a metric for identifying salient tokens based on normalized attention scores. By efficiently approximating the saliency metric, the method integrates seamlessly with fast attention mechanisms like FlashAttention. The efficacy of ZipCache is demonstrated through extensive experiments, showing superior performance in terms of compression ratio, speed, and minimal accuracy loss compared to existing methods.

Strengths

1.The introduction of normalized attention scores as a metric for token saliency is significant and promising for adaptive KV cache compression. It can accurately identify salient tokens, marking a substantial improvement over prior methods. 2.The approximated saliency metric can be integrated with FlashAttention, making it practical for real-world applications. 3.The idea of the channel-separable tokenwise quantization approach is novel and well-motivated, reducing memory overhead compared to groupwise counterparts. 4.The experimental results are promising, clearly demonstrating the efficacy of ZipCache in terms of compression ratio, speed, and minimal accuracy loss. 5.The paper is clearly written and easy to follow.

Weaknesses

1.The experiments should encompass a broader diversity of tasks to comprehensively assess the method's effectiveness, such as LongBench. 2.The font size in Figures 5 and 6 is too small, making them difficult to read. It is recommended to increase the font size for better clarity and accessibility.

Questions

See weaknesses

Rating

8

Confidence

5

Soundness

3

Presentation

3

Contribution

4

Limitations

See weaknesses

Summary

This paper proposes a post training quantization framework named ZipCache for quantizing the key and value cache of LLMs. The authors introduce a channel-separable token-wise quantization scheme, which consumes less memory than the group quantization in terms of the quantization parameters. They also select saliency tokens based on a normalized attention matrix to avoid the accumulated bias. Additionally, the authors introduce an approximation method to apply it to FlashAttention.

Strengths

- Overall, the writing of the paper is clear and easy to follow. Especially, the figures are drawn well. - By proposing the ZipCache, this work promotes the PTQ quality of LLMs on three chanllenging datasets. - The methods are well-integrated into popular framwork like FlashAttention.

Weaknesses

- The idea of channel-separable quantization is not novel. It is a common practice to smooth the outliers of activations in the channel dimension, e.g., SmoothQuant [1], OmniQuant [2]. - As shown in Table 3 (a), the values on the diagonal of the matrix are relatively large, while the values at other positions are relatively small. Based on this circumstance, if the mean operation is performed for none-zero positions, earlier tokens may have smaller $p_i$ for they have more tiny values counted. Therefore, it may cause a prefer to the latest tokens. I think this is a issue for the proposed technique in Section 4.2, which needs more analysis. - Ablation studies are not provided. - As stated in Line 268 in the paper, new tokens remain un-quantized unless they reach 100. As far as I know, keeping part of tokens in full precision can improve the accuracy [3, 4]. Therefore, it is better to distinguish the contributions of this implementation and the proposed techniques. - I think the value of 'H/L' should be carefully considered. This is because both KIVI and ZipCache keep new tokens in FP16 at first and then perform quantization. Since the 'H' value of KIVI is 16, why it becomes 4 for ZipCache? - The equation (5) is not accurate for not subtracting the zero-point 'z' in the de-quantization. [1]. SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models. [2]. Omniquant: Omnidirectionally calibrated quantization for large language models. [3]. KIVI: ATuning-Free Asymmetric 2bit Quantization for KV Cache. [4]. SKVQ: Sliding-window Key and Value Cache Quantization for Large Language Models.

Questions

- The authors have compared the channel-separable quantization with many dynamic methods, but lacks comparison of accuracy with static methods, e.g., WKVQuant [1]. It is better to further explore whether the $c_i$ in equation (6) can outperform well-trained parameters, since they consumes the same memory space. - How does the calculation of $c_i$ in equation (6) effect the inference latency? It is better to compare the channel-separable quantization with group quantization and static methods in terms of the inference latency. [1]. WKVQuant: Quantizing weight and key/value cache for large language models gains more.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

This paper contains the describtion of limitations. For suggestions, please refer to the Weakness and Questions.

Summary

The paper introduces a channel-separable quantization scheme that decouples the quantization along channel and token dimensions. This method significantly reduces the quantization overhead without compromising performance. To accurately recognize salient tokens, the paper introduces a new token saliency metric based on normalized attention scores, which alleviates the bias towards earlier tokens that accumulate more values. The authors have demonstrated on different LLM models with different tasks.

Strengths

The idea of channel separable token saliency quantization seems to work better than standard groupwise quantization. The results are strong and beats the SoTA! The demonstration with flashattention integration is useful

Weaknesses

The paper needs to be proof read by a native English speaker, to improve its readability. The contribution is incremental, and the idea of group quantization (different for different K and V) are not new as already noted by the authors. The idea of speeding up through flashattention and flashdecoding while the baselines not leveraging that is a bit of unfair comparison. Both GEAR and KIVI can exploit similar speed up benefit, thus Fig 1 comparison and later in table speed comparison is not fair. The details of system level implementation is missing, it will be good to help the reviewer provide how the speed up can be demonstrated on a single GPU inference system.

Questions

refer to weakness.

Rating

5

Confidence

5

Soundness

3

Presentation

2

Contribution

2

Limitations

The system benefits are not clear or thorough for the reviewer to clearly appreciate. Also more fair system eval is needed.

Review for rebuttal

Thanks for the efforts from the authors. My concerns are addressed in the rebuttal. I think this paper is well-motivated and has its merit to the community. It demonstrates strong performance in the experiments. Therefore, I raise my score after rebuttal.

Response.

Thanks to the authors for the new results with attached pdf. Particularly, I appreciate the fair pre-fill stage comparison with all quantization schemes leveraging flashattention. 1. While I appreciate the author's effort, I think further clarifications are needed. In specific, the KV compression can be useful for longer context based evaluation, the current manuscript including that in the one in rebuttal pdf. The LLaMA2 model has pretraining context length of ~4k which is not sufficient enough for the long context demonstration. Beyond long context evaluation the benefits of KV cache compression is not that usefully demonstrable. And the main manuscript demonstrates results with around 900 and 200 tokens on average as said by the authors. 2. The idea of normalized attention score is not new either, so I believe their is a bit of overclaim here. Please refer to this paper, that proposed something similar, namely MAS [1]. 3. Why the generation efficiency is compared with only MiKV, instead of other SoTA quantization methods like KIVI? It is understandable that MiKV falls short as it cant leverage flash-decoing/attention, this is again unfair comparison. 4. It is also not clear if the probe tokens are selected randomly how it will still help manage perform flashattention or flashdecoding efficiently. Also, during decode phase do you implement flashattention or flashdecoding? please provide more details. 5. Overall, I believe the paper is built on top of the main findings of probe tokens, however the details in that front is yet to be sufficient. Additionally, how the selection of probe tokens help guide the bit width assignment of salient tokens if a token does not have any representative in the randomly selected probe token set? 6. The algorithm of decode stage seems ambiguous, or not fully informative. For example, the new tokens getting added on top of previous KV, is it the hybrid H/L quantized KV or or some high precision KV? As we may need to compute attention and softmax etc at high precision to have numerical stability. This needs clarification. 7. It looks like the probe tokens during decode phase are kept at higher precision, did the author include that in their total compression scheme? Also, the probe token computation and memory overhead is not clearly discussed in the paper's results. This should impact accuracy as well as throughput compared to schemes like KIVI. [1] On the Efficacy of Eviction Policy for Key-Value Constrained Generative Language Model Inference, 2024.

Review for Rebuttal

Thanks to the authors for the detailed response. I understand the motivation of the channel-seperable token-wise quantization, as well as the effectiveness of the token saliency metric after the rebutall. The integration of this kind of methods to flash-attention is also a valuable contribution. Therefore, considering the answers from Q1 to Q5, I decide to raise my score to 5. Although, I think the description of the Section 4.1 can be improved, especially from the aspect of accuracy and latency. As claimed to be a more advanced baseline, it should process better generalization ability compared to group quantization and well-trained static parameters, rather than solely performing well on a single dataset. Additionally, I think the answers to Q6 and Q8 should be included into the main text.