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

Thanks to the reviewer for the valuable comments.

Q1: Evaluating ZipCache on other generation or comprehending tasks.

As shown in Table C in the rebuttal PDF, we evaluate the performance of ZipCache on LongBench. The results show that ZipCache outperforms the previous state-of-the-art method, KIVI. Due to the limited time slot of the rebuttal, we will conduct experiments on more models and add the results to the revised version.

Q2: Using only the most recent token to determine token saliency.

As shown below, determining token saliency with only the most recent token leads to a 3.2% accuracy drop compared to ZipCache.

Table: The effect of different saliency metric on GSM8k with CoT prompts. Here, "H/L" denotes the bit-width for salient tokens (high-precision) and regular tokens (low-precision), respectively. The compression ratio is calculated with an average input length of $l=840$.

Q3: What does saliency mean in Figure 3c?

There are many attention layers and heads in a model. For each token, this value is derived by counting the number of times this token is selected as a salient token among all attention heads, then normalizing this count by the total number of attention heads.

Thanks to the reviewer for the valuable comments.

Q1: Evaluate ZipCache over LongBench.

As shown in Table C in the rebuttal PDF, we evaluate the performance of ZipCache on LongBench. The results show that ZipCache outperforms the previous state-of-the-art method, KIVI. Due to the limited time slot of the rebuttal, we will conduct experiments on more models and add the results to the revised version.

Q2: The font size in Figures 5 and 6 is too small.

Thanks for your valuable comment. We will revise it in the final version.

Thanks to the reviewer for the valuable comments.

Q1: The idea of channel-separable quantization is not novel.

Please refer to General Response Q1.

Q2: Performing the mean operation may cause a prefer to the latest tokens.

Indeed, there might be a prefer to the latest tokens since the values on the diagonal of the matrix are relatively large. However, the overall token saliency is adaptively determined in our method. As shown in Figure A in the rebuttal PDF, we present the token saliency with a 100-line retrieval sample as input. Besides the latest tokens, the system prompts in the front and related context in the middle are also assigned high saliency by our method. Moreover, Table B in the rebuttal PDF shows that our method achieves superior performance compared to using accumulated attention scores or a fixed prefer to the latest tokens.

Q3: Ablation studies are not provided.

As referred to Table 1 and Table 2 in the paper, we have conducted ablation experiments to demonstrate the effectiveness of channel-separable quantization scheme and the efficient approximation of the saliency metric, respectively. We further conduct ablation studies on our saliency metric, as shown in Table B in the rebuttal PDF. It demonstrates the superiority of our saliency metric over using accumulated attention scores or consistently prioritizing the latest tokens.

Q4: Remaining new tokens un-quantized can improve accuracy.

During the prefill phase, KIVI [i] maintains full-precision KV caches for recent tokens to ensure accuracy, while our approach adaptively quantizes all KV caches. During the decoding phase, we quantize all KV caches with mixed-precision every 100 tokens, while KIVI adopts a sliding window and always keeps KV caches for latest tokens (128 tokens by default) in full-precision. This implies that KIVI always retains more KV cache at full precision compared to our method. Moreover, the streaming strategy we adopted aligns with GEAR [ii] and is aimed at enhancing decoding speed. Therefore, the comparison with KIVI and GEAR fairly demonstrates the efficacy of our method.

Q5: The value of `H/L' should be carefully considered.

As referred to Q4, during the prefill phase, we quantize all KV caches, with the bit-width for salient tokens set to 4. This approach contrasts with KIVI's method of maintaining full-precision (FP16) for the recent tokens.

Q6: The equation (5) is not accurate.

Thanks for your valuable comment. Equation (5) has been revised as follows:

$\hat{\mathbf{x}}=\mathcal{Q}_U(\mathbf{x},k)=(\mathrm{clip}(\lfloor \frac{\mathbf{x}}{s}\rceil +z, 0, 2^{k}-1) - z) \cdot s.$

Q7: Comparisons between channel-separable quantization and static methods like WKVQuant.

Due to different settings, ZipCache and WKVQuant [iii] are not directly comparable. WKVQuant quantizes model weights and KV cache together, optimizing parameters through cross-block reconstruction regularization and a gradient descent algorithm. In contrast, our method focuses solely on KV cache compression and is entirely training-free.

Q8: Compare the channel-separable quantization with group quantization and static methods in terms of the inference latency.

Inference latency comparisons among channel-separable quantization, groupwise quantization, and static quantization methods are detailed in Table A of the rebuttal PDF. Using static quantization can slightly reduce latency (2556.03 ms vs. 2584.01 ms). Conversely, groupwise quantization increases inference speed (2664.05 ms vs. 2584.01 ms) and reduces the compression rate (3.81$\times$ vs. 4.43$\times$) due to massive quantization overhead.

[i] KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache.

[ii] GEAR: An Efficient KV Cache Compression Recipe for Near-Lossless Generative Inference of LLM.

[iii] WKVQuant: Quantizing Weight and Key/Value Cache for Large Language Models Gains More.

Dear Reviewer Nc8R,

We greatly appreciate the time and effort in reviewing our work. We have carefully considered your comments and suggestions and made significant revisions to address the concerns you raised. We are eager to ensure that our paper meets the high standards of our respected reviewers.

Please don’t hesitate to let us know if there is any additional feedback you might have at this stage.

Best regards,

Authors of #5969.

Dear Reviewer Nc8R,

Thank you for dedicating your time to reviewing our paper. As the discussion period deadline is approaching, we kindly invite any further comments or concerns you might have. Your feedback has been immensely valuable to us for refining the paper.

Best regards,

Authors of #5969.

Dear Reviewer Nc8R,

Thank you for your feedback. We truly appreciate your careful consideration of our responses and will carefully revise the paper based on your and other reviewers' suggestions.

Best regards,

Authors of #5969.

Thanks to the reviewer for the valuable comments.

Q1: The writing can be more native.

Thanks for your valuable comment. We will carefully proofread the paper in the final version.

Q2: The idea of group quantization (different for different K and V) is not new.

There might be some misunderstandings regarding groupwise quantization. Groupwise quantization is a scheme built upon tokenwise quantization, further processing outlier channels in distinct groups, as illustrated in Figure 2(c) of the paper. Additionally, applying different quantization schemes for K and V respectively (channelwise quantization for K and channel-separable tokenwise quantization for V) is not the key contribution of our paper. In terms of the quantization scheme, our main contribution lies in the introduction of channel-separable quantization to the KV cache compression, significantly reducing the overhead of quantization parameters without sacrificing performance, as detailed in General Response Q1.

Q3: The contribution is incremental.

Please refer to General Response Q1 for the detailed contribution of our channel-separable quantization scheme and to General Response Q2 for an elaboration on the overall contribution of our paper.

Q4: The speed comparison is not fair.

Thank you for your valuable comments. Firstly, we highlight that we are the pioneering work that enables adaptive KV cache compression to be compatible with FlashAttention, greatly enhancing the generation speed compared to previous approaches like H2O [i] and MiKV [ii]. Moreover, the reported speed results of GEAR [iii] and KIVI [iv] were obtained with their official implementations, which did not integrate FlashAttention at the time of our paper submission. To ensure a fair comparison, we implement GEAR and KIVI with FlashAttention integration, and the results are shown in Table A in the rebuttal PDF. Notably, the latency of ZipCache is lower than that of GEAR, which can be attributed to our efficient quantization scheme, whereas GEAR has a high overhead due to outlier extraction. Compared to KIVI, ZipCache's latency is slightly higher (2584.01 ms vs. 2482.26 ms), but ZipCache achieves a higher compression ratio and better performance. This difference is due to KIVI's fixed compression strategy, while we adaptively compress the KV cache based on the saliency. These results will be revised in the final version.

Q5: The details of system level implementation are missing.

The detailed processes of ZipCache for both prefill and decoding phases are summarized in Algorithms 2 and 3 in the Appendix of our paper. Currently, we implement ZipCache based on the Hugging Face Transformers library, with specific modifications to the KV cache module. Our method is also orthogonal to other system level frameworks such as vLLM [v]. It should be noted that we utilize FlashAttention to maximize computational efficiency for both the prefill and decoding phases, eliminating the need to customize additional matrix multiplication kernels. We will release the source code upon acceptance.

[i] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models.

[ii] No Token Left Behind: Reliable KV Cache Compression via Importance-Aware Mixed Precision Quantization.

[iii] GEAR: An Efficient KV Cache Compression Recipe for Near-Lossless Generative Inference of LLM.

[iv] KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache.

[v] Efficient Memory Management for Large Language Model Serving with PagedAttention.

Dear Reviewer ozSB,

We greatly appreciate the time and effort in reviewing our work. We have carefully considered your comments and suggestions and made significant revisions to address the concerns you raised. We are eager to ensure that our paper meets the high standards of our respected reviewers.

Please don’t hesitate to let us know if there is any additional feedback you might have at this stage.

Best regards,

Authors of #5969.

Dear Reviewer ozSB,

Thank you for dedicating your time to reviewing our paper. As the discussion period deadline is approaching, we kindly invite any further comments or concerns you might have. Your feedback has been immensely valuable to us for refining the paper.

Best regards,

Authors of #5969.

Q9: The probe tokens during decoding phase are kept at higher precision. This is not the case. All KV caches, including those for probe tokens, are quantized after every 100 new tokens are generated. As referred to Q5, during the decoding phase, the primary difference between probe tokens and other tokens is that we use the queries of probe tokens to compute attention scores explicitly, while other tokens are processed using FlashDecoding.

Q10: Keeping probe tokens at higher precision will impact accuracy as well as throughput. Please refer to Q9. Probe tokens are not kept in higher precision.

Q11: What is the computation and memory overhead for probe tokens? As shown below, computing attention scores with the queries of probe tokens introduces limited computation and memory overhead.

Table: Computation and memory overhead for probe tokens. Here, "ZipCache w/o Probe Tokens" denotes the token saliency is randomly generated rather than approximated with probe tokens. Data is collected by serving LLaMA3-8B on a NVIDIA A100 GPU with a batch size of 8 and sequence length of 3072.

[1] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models.

[2] No Token Left Behind: Reliable KV Cache Compression via Importance-Aware Mixed Precision Quantization.

[3] GEAR: An Efficient KV Cache Compression Recipe for Near-Lossless Generative Inference of LLM.

[4] KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache.

[5] On the Efficacy of Eviction Policy for Key-Value Constrained Generative Language Model Inference.

Thanks to the reviewer for the valuable comments.

Q1: The KV compression can be useful for longer context based evaluation. Firstly, in alignment with prior literature [1-4], we have evaluated our method on three challenging and widely-recognized benchmarks, including the long context Line Retrieval task. As highlighted in Table 3, our approach consistently surpasses previous state-of-the-art methods, underscoring its effectiveness. Secondly, as mentioned in lines 36-38 of our paper, during batch inference, the KV cache with an input length of approximately 4K is already a significant bottleneck in the storage system. For example, the KV cache can occupy 1.2TB of memory space with a batch size of 64 and an input length of 4096. Moreover, to further address your point, we have also evaluated the performance of ZipCache on LongBench using the longchat-7b-v1.5-32k model, as shown below. Due to the limited time slot of the rebuttal, we will conduct experiments on more long-context tasks and add the results to the revised version.

Table: Performance comparisons on LongBench with LongChat-v1.5-7B-32k. The saliency ratio is 60% for ZipCache.

Q2: Similar idea of normalized attention scores has been proposed. Thank you for sharing this literature [5] and we will include it in the references in the revised version. However, we would like to emphasize that our proposed saliency metric can be applied universally to all tokens without exception, whereas the approach in [5] excludes certain tokens from the eviction scope based on their standard deviation. Additionally, our work provides a comprehensive analysis of the limitations of using accumulated attention scores as a saliency metric, as discussed in lines 200-210 and illustrated in Figure 3 of the paper.

Q3: The generation efficiency is compared with only MiKV and the comparison is unfair. As referred to Figure 1 in the paper, Table A in the Appendix and Table A in the rebuttal PDF, we have compared the generation efficiency with H2O, GEAR and KIVI. Moreover, the previous adaptive KV cache compression methods [1-2] are not compatible with FlashAttention, which is a major drawback of them. By contrast, our method integrates seamlessly with FlashAttention, enhancing generation speed.

Q4: How do probe tokens help perform flashattention or flashdecoding efficiently? Firstly, as referred to lines 230-232 and Equation (9) of the paper, we select a small set of probe tokens and compute attention scores using their queries, which allows us to approximate the saliency for all tokens. Once we have determined the token saliency, we can efficiently compute the attention output using FlashAttention. For a more detailed explanation of our method, please refer to Algorithm 2 and Algorithm 3 in the Appendix.

Q5: Do you implement flashattention or flashdecoding during decoding phase? Yes. As mentioned in lines 267-268, 304-307, and detailed in Algorithm 3 in the Appendix, we implement FlashDecoding during the decoding phase. Similar to the prefill phase, we explicitly compute attention scores for a small set of probe tokens to approximate the saliency of all tokens, enabling the majority of tokens to be computed using FlashDecoding.

Q6: How does the selection of probe tokens help guide the bit width assignment of salient tokens if a token does not have any representative in the probe token set. As mentioned in lines 254-255 of the paper, our approach to selecting probe tokens involves a hybrid strategy, where 5% of the tokens are the latest ones, and 5% are randomly selected. All previous tokens are attended when computing attention scores for the latest token.

Q7: The algorithm of decode stage seems ambiguous. For example, what is the precision of the KV cache for new tokens. As referred to lines 267-268 of the paper, we implement a streaming strategy during the decoding phase, where the new KV cache is quantized every 100 new tokens generated. This approach is consistent with the strategy used in GEAR [3] and is designed to enhance decoding speed. Consequently, before reaching the 100-token threshold, the new KV cache is maintained in full precision. After 100 new tokens are generated, all KV cache will be quantized based on their estimated saliency.

Q8: Attention and softmax need to be computed at high precision to have numerical stability. Indeed. The KV cache will be dequantized to full-precision when calculating attention. This is consistent with previous KV cache quantization work [2-4].