KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization

Thank you for taking the time to review our paper!

As a minor weakness, due to the extensive experimental results and content, there are many appendices linked in the main text and frequent parenthetical explanations, which can make the paper somewhat challenging to follow. There is room to improve readability by filtering out and emphasizing key points in the main text.

We appreciate the feedback on how to better structure our paper in order to represent the key contributions and takeaways from our analysis and results. In the final version, we will streamline the paper by focusing on key points in the main text (and ensuring that these are fully explained within the main text), while making sure the supplemental information is not necessary for understanding the core contributions of our work.

Can the distribution characteristics of Key and Value be generalized to widely used popularized LLMs? (e.g., Mistral, Gemma-2, Phi-3, …)

We have observed that these characteristics are consistent across different models in our work (eg. LLaMA, Llama-2/3, Mistral). Additionally, in the rebuttal PDF we have attached the analysis of the KV cache distributions for the Gemma-2-2B and Phi-3-Mini-128K-Instruct models for a randomly selected layer. These distributions are similar to what we observed in the analysis for LLaMA-7B included in Figure 2 of our paper (i.e. both models have outlier input channels with consistent average magnitudes pre-RoPE), indicating that pre-RoPE per-channel quantization would generalize to these networks.

Do you have any insights into the cause of the disparate characteristics in the distributions of keys and values?

With respect to the disparate distributions of keys and values, keys have specific outlier channels with significantly larger average magnitude than other channels, as highlighted in Figure 2. For values, we do not observe particular per-token or per-channel outliers. Given the consistent patterns that we observe across different model families and model sizes, we think that it is an interesting direction for further research to understand the root cause of the distributional characteristics of keys and values.

I am curious about the specific method used to measure PPL. I might have missed it, but when measuring generation PPL on wikitext and C4, did the authors use a teacher forcing style by providing all input tokens at once, applying quantization to the KV embeddings, and measuring loss with the final logits for all tokens? It would be helpful to have a detailed explanation of how PPL was measured.

This is a great question! We are indeed using teacher forcing by providing all input tokens in parallel and quantizing the Key and Value embeddings on-the-fly in the forward pass. Then we measure the perplexity using the output logits of all input tokens. We will clarify this point in the final version of the paper.

Thank you for the detailed feedback on our work!

The paper seems to be a combination of several existing quantization methods), e.g. choosing to use per-channel quantization for Key cache, using non-uniform quantization, and dense-and-sparse quantization to improve performance.

With respect to the novelty of our approach, we would like to highlight multiple contributions from our work:

• To the best of our knowledge, our work was the first to observe the impacts of RoPE on per-channel outliers in Keys and to propose pre-RoPE per-channel key quantization to mitigate this challenge.
• Our method leverages non-uniform layer-wise datatypes derived from calibration data, which facilitates more accurate quantization online for activations than prior works using fixed datatypes, without compromising on efficiency.
• Our work shows that dense-and-sparse quantization can be applied efficiently for activations, and demonstrates the necessity of per-vector outlier detection for applying dense-and-sparse quantization to activations with outlier channels (which was not explored in prior works on weight quantization).

The paper writing is not clear enough. The paper proposes too many techniques and details are not fully presented in the main text of the paper.

We appreciate the feedback. In the revised version, we will ensure to highlight all key takeaways in the main text such that they can be understood without referencing the appendix.

The experiments are only conducted on Wikitext-2 and C4 datasets and reported the PPL. Other challenging tasks such as reading comprehension, math problem solving, or code generation should be evaluated.

In our submission, we indeed included benchmarks other than PPL evaluation such as passkey retrieval results (Section 4.2, Table 2), as well as LongBench evaluation (Appendix Q, Table 18), which includes a range of downstream tasks such as QA, summarization, code generation, and in-context learning.

Below, as a part of our rebuttal, we include additional evaluation of KVQuant on various benchmarks including XSum (summarization), GSM8K (math reasoning), and RULER (retrieval, QA, aggregation, multi-hop tracing). Please check our response below for more details.

The evaluation section lacks detailed comparisons of the existing KV cache compression methods (like H2O, GEAR, KIVI).

We appreciate the suggestions for relevant comparisons for our work.

1. Comparison against KIVI

In our paper, we included comparisons with KIVI for passkey retrieval as well as LongBench evaluation with LLaMA-2-7B-32K. On passkey retrieval (Table 2), KIVI attains 68-76% with an average bit-width of 3, whereas KVQuant achieves an accuracy of >98% even with 2-bit quantization. On LongBench (Table 18), 3-bit KVQuant shows only 0.75 point degradation from the FP16 baseline, which is much smaller than the 1.92 point degradation of KIVI with a similar average bit-width (31.96 → 31.21 vs 30.04).

In the table below, we additionally compare KVQuant and KIVI on the RULER benchmark suite using LLaMA-2-7B-32K. Here, 3-bit KVQuant achieves ~14% better score against KIVI with similar average bit-width. Furthermore, our 2-bit KVQuant achieves similar accuracy to KIVI with 1.4x smaller bit-width. The results on passkey retrieval, LongBench, and RULER demonstrate how our method provides improved accuracy on long context-length tasks compared to KIVI.

LLaMA-2-7B-32K on RULER

• Accuracies for each subtask are in Table 1 of the rebuttal PDF

2. Comparison against H2O and GEAR

We have also included additional comparisons with H2O and GEAR to provide further insight into how our approach compares with other KV cache compression methods. The table below provides a comparison with H2O on XSum with 5-shot prompting for LLaMA-7B, using the evaluation script from the H2O repo. As demonstrated in the tables, our method achieves consistently better ROUGE scores (up to ~3 point) compared to H2O with the same compression rates. Furthermore, our method does not show performance degradation up to 5x compression rate.

LLaMA-2-7B on XSUM

Additionally, in the table below, we also included GSM8K comparisons with both GEAR and H2O using 8-shot CoT prompting using the environment from the GEAR repo. Here, our method achieves better GSM8K accuracy compared to both H2O and GEAR even with a bigger compression rate.

LLaMA-2-7B on GSM8K

What exactly is the word quantization signpost?

Quantization signposts refers to the discrete representable values when quantizing to reduced precision. For example, if the representable values with 2-bit quantization are [-0.7, -0.2, 0.2, 0.7], then these 4 values are the signposts.

Will calibration set lead to overfitting of the model? Will calibrating a quantized model on one dataset affect its performance on other datasets?

Please refer to our response to Reviewer 2 for our ablation study which demonstrates that our method is not overfitting to the calibration dataset. Furthermore, all experiments in our paper were calibrated using 16 examples from WikiText2 and evaluated on different downstream tasks.

Thank you for taking the time to review our paper.

In the per-vector dense-and-sparse quantization step, the numerical outliers are stored as high precision in a separate sparse matrix. It is not quite clear how to achieve this in practice in a hardware-friendly way without affecting inference speed.

It is correct that storing the numerical outliers separately in a sparse high-precision matrix can add considerable inference latency if the sparsity percentage is too high. However, we have observed that KV cache activations have a relatively small percentage of outlier elements (<1%) which need to be stored separately to alleviate quantization difficulty. By storing only this small percentage of outliers in a separate sparse matrix, we observe significant accuracy benefits; as highlighted in Figure 1, isolating 1% outliers improves the perplexity on Wikitext-2 for the LLaMA-7B model from 5.94 to 5.75, achieving <0.1 PPL degradation with 3-bit quantization. Since the sparsity percentage is extremely low, we find that the sparse matrix / dense vector computation doesn’t add significant runtime overhead, as shown in Appendix S. Our approach can therefore store outliers separately in higher precision to maintain high accuracy with low memory consumption, without adding significant latency overhead.

Additionally, the offline calibration step makes the implicit assumption on the access to data samples that are from a similar distribution as data from inference time.

In our experiments, we have only used a small amount of data from the Wikitext-2 training set for calibration (16 samples of length 2048), and then we do not recalibrate when running inference on different downstream tasks. We observe strong performance across a range of downstream tasks (such as LongBench in Table 18 and passkey retrieval in Table 2), demonstrating that the calibration process is not sensitive to the choice of calibration data.

To further support this claim, we have included an ablation in the table below showing the perplexity of LLaMA-7B on Wikitext-2 and C4 when we use each of the two datasets as our calibration dataset. These results demonstrate that the accuracy is relatively unchanged when running calibration using different training sets.

Thank you for taking the time to review our paper.

The paper does not sufficiently clarify the baseline methods and implementations used for comparison.

The baseline methods that we used to compare with KVQuant are (i) uniform quantization with and without grouping, (ii) non-uniform quantization (with and without grouping) using the fixed NormalFloat datatype introduced in QLora, and (ii) KIVI, a state-of-the-art KV cache quantization method. Particularly for uniform quantization, we followed the same configurations as FlexGen and Atom, which use group sizes of 64 and 128, respectively, and we used per-token quantization for both keys and values. We quantize keys post-RoPE for all baseline approaches since it has higher accuracy for per-token key quantization (as outlined in Appendix P). Please let us know if there are any other details about our baseline approaches which require clarification.

What is the FP16 baseline implementation when evaluating speedup of the KVQuant kernel?

The FP16 baseline implementations that we used to evaluate latency are the default kernels called by PyTorch for batched matmul operations. In particular, in Table 3 in the paper, the Key kernel baseline and the Value kernel baseline measure the latency for Query x Key and Attention Score x Value matmul operations, respectively. For the Key kernel, we additionally included the time for applying RoPE for a single key vector since the new key vector at each iteration needs to have RoPE applied before computing Query x Key. We will make this point clear in the revised version of our paper. Kernel latency is collected using the Pytorch CUDA profiling tool.

Compared to Post-RoPE quantization where post-RoPE Keys are stored, why does recalculating the RoPE in the KVQuant implementation can achieve faster speeds?

Recalculating RoPE in the KVQuant implementation is introduced to enhance the quantization performance, rather than to achieve faster speeds and efficiency. As outlined in Section 3.2 of our paper, pre-RoPE key quantization is necessary for accurate key quantization, improving the perplexity on Wikitext-2 for the LLaMA-7B model by 0.82 with 3-bit quantization. Our experiments demonstrate that we can quickly recompute RoPE on-the-fly, thereby attaining the significant accuracy benefits of pre-RoPE key quantization without compromising much efficiency.