KV Cache is 1 Bit Per Channel: Efficient Large Language Model Inference with Coupled Quantization

We sincerely thank the reviewer for the insightful review and invaluable feedback. We address the reviewer's concerns as follows.

[W1] Limited Evaluation Tasks: KV Cache compression is most relevant in long-text scenarios.

• We appreciate the reviewer's suggestion to test KV cache quantization in long-text scenarios. We present additional experiments in long-text settings below. Specifically, we test CQ and KVQuant using Llama-2-7b on GSM8K with chain-of-thought (CoT) and MMLU with CoT Fewshot (gsm8k_cot, mmlu_flan_cot_fewshot_humanities, mmlu_flan_cot_fewshot_stem, mmlu_flan_cot_fewshot_social_sciences, mmlu_flan_cot_fewshot_other from lm-evaluation-harness).
• In long-text settings, CQ mostly outperforms or performs similarly to KVQuant under the same bit width.

[W2] Incomplete Explanation of Channel Correlation: The paper doesn't fully explain the existence and variability of channel correlation across different layers, which may raise questions about the method's general applicability.

• The amount of correlation between channels does vary across layers. In the table below, we present the mean absolute correlation (MAC), excluding the diagonals, of different layers of LLaMA-7b on 262K tokens of WikiText-2.
• Although the key correlation of the first few layers are higher than the later layers, the key/value correlation of any layer is never very close to zero, meaning channel coupling can be effective for any layer.
• Correlation as a metric only captures the linear dependency between a pair of channels. In practice, we couple up to 8 channels together to leverage the higher order dependency between multiple channels.
• We thank the reviewer for the careful reading of our paper. We will include a discussion on the existence and variability of channel correlation across different layers in the final paper.

[Q1] Quantization Setting: Is the KV Cache quantized during the prefill stage, or only during the decoding stage (after full-precision prefill computation), similar to the KIVI approach?

• For the majority of the experiments presented in the paper, including Table 1,2,4,5,6, all tokens are quantized in both the prefill and the decoding stage. We have specified this on line 247 of the paper.
• For experiments with sliding window full-precision cache, including Table 3 and parts of Figure 1, we keep a constant number of the most recent tokens in full precision and quantize the rest of the tokens, and this holds true with respect to any token in the prefill and the decoding stage. This is slightly different from the KIVI implementation, which computes prefill in full precision and adopts a sliding window approach during decoding. We adopt this approach since perplexity testing and some tasks are log-likelihood-based and do not have a decoding stage.

[Q2] Channel Grouping Strategy: Would coupling non-adjacent channels with high mutual information potentially bring more benefits? --- Yes, very likely!

• Coupling channels with high mutual information will highly likely bring benefits. As shown in Figure 6 and 7, certain pairs of non-adjacent channels have higher correlations. By coupling them into the same channel group, we further reduce the joint entropy of channel groups, leading to better quantization accuracy.
• We thank the reviewer for pointing us in this direction, and leave this for future work due to the difficulty of system implementation and optimizations.

[Q3] CUDA Kernel Implementation Questions

• Yes, the efficiency may be improved by adjusting the thread block assignments and reduction strategies.
• First, it is important to note that the centroid size of each channel group is greater than 256 bytes. For example, for CQ-8c8b, the centroid size of each channel group is (num_centroids x num_channels x 2 bytes) $2^8 \times 8 \times 2=4096$ bytes. Assuming 100KB of shared memory per thread block, we can fit at most 25 groups of centroids into shared memory.
• Fitting more channel groups into the same thread block reduces the number of concurrent writes to the HBM, hence speeding up the computation.
• We thank the reviewer for this valuable suggestion. We will incorporate this kernel improvement into our implementation.

We sincerely thank the reviewer for the support of our paper and the insightful suggestions. We address the reviewer's concerns as follows.

[W1, L1] Include latency experiments.

• In the table below, we present additional latency measurements of CQ with comparison to FP16 cache and KIVI. We use Llama-2-7b with a batch size of 1 and a prompt of 2000 tokens to decode 100 tokens.
• We observe that CQ achieves comparable efficiency with KIVI. We plan to further optimize CQ at system level to reduce latency and improve throughput.

[Q1] Combining with other weight/activation quantization methods.

• Our KV cache quantization method can be combined with other weight/activation quantization methods. In the standard CQ calibration process, the centroids are learned using the KV cache of the full-precision model based on a calibration dataset.
• To combine CQ with other weight/activation quantization methods, we use the KV cache produced by the quantized model for centroid learning to minimize the distortions introduced by CQ.

[Q2] Does this mean that the estimation needs to be done for each dataset and cannot be reused? Which dataset was used for the centroid learning in Table 8? --- No, centroid learning only needs to be done once.

• Centroids of CQ are learned once during the calibration phase, and can be used for different downstream tasks.
• For Table 8, the centroids of CQ are learned on a set of 16 sequences from WikiText-2, each with 2048 tokens. We have specified this on line 243 in the paper. The centroids are learned only once and used for all perplexity and accuracy experiments in the paper.
• We present additional experiments below on the accuracy of CQ with different calibration datasets. We use 16 sequences of 2048 tokens from WikiText-2 and C4 as calibration set and evaluate CQ on 4 downstream tasks.
• Despite using different calibration datasets, CQ performs similarly in various downstream tasks.

[Q3] Ablation study on the overhead of centroid loading.

• We perform an additional ablation study to understand the overhead of centroid loading. We profile the KQ multiplication kernel for CQ-2c8b with a single query and 4K or 16K keys using Llama-2-7b hidden dimensions. We enable and disable loading centroids from global memory to shared memory. The results shown in the table below are average over 1000 kernel runs.
• As shown in the table, centroid loading does not significantly contribute to the overall latency. The latency primarily comes from reading quantized KV cache and queries from global memory, and writing the results to global memory. Hence reusing the centroids will not significantly reduce the latency.

[Q4] Integration with vllm or other serving frameworks. From Figure 4, it seems all CQ methods perform worse than the fp16 version with small batch sizes. Can you explain why? Is it due to the overhead of (de)quantization?

• We anticipate compatibility between our proposed CQ and PagedAttention in vLLM, and believe their integration is a promising avenue for future exploration. Given the implementation difficulty, we leave this to future investigations.
• CQ has higher latency than FP16 in small batch sizes due to the overhead of (de)quantization. As suggested by Reviewer rmbd, the efficiency of our kernel can be further optimized by loading centroids of more channel groups at once into shared memory, and adjusting the thread block assignment and reduction strategies. We will incorporate these improvements after the rebuttal period.

[Q5] How does the sliding window affect inference latency?

• Please see the first table for a latency comparison between sliding window and no sliding window. Sliding window full-precision cache does not significantly contribute to the overall latency in the prefill or the decoding stage.

[Q6] How do you determine the best sliding window size to achieve the optimal tradeoff? Also, what happens if the number of coupled channels exceeds 8?

• The sliding window size needs to be larger for lower bit widths to compensate for the precision loss. For CQ-2c8b (4-bit quantization), a small window size of 16 tokens or no window may suffice, but for CQ-8c8b (1-bit quantization), a larger window of 128 tokens may be necessary for preserving quality.
• The number of coupled channels can exceed 8. We present additional experimental results below using LLaMA-7b with CQ-16c12b (12 bits per 16 coupled channels, averaging 0.81 bits per activation).

[L2] It would be more helpful to show the actual memory usage in GB instead of the parameter count in Table 8.

• We will include the actual memory usage in the final paper. The memory overhead of centroids can be calculated as (parameter count x 2) bytes.

We express our sincere gratitude to the reviewer for their thoughtful comments and suggestions. We address the reviewer's concerns as follows.

[W1] The code is not available.

• We will open source our code during the camera-ready phase. To provide additional context, we have included expanded experimental results and implementation profiling details below.

[W2] Section 3 should be revised for clarity.

• We will revise and polish section 3 to describe the random variables and entropy in mathematical notations. We appreciate the reviewer's careful reading of our paper and their valuable feedback.

[W3] Extension to LLAMA3.

• We present additional experimental results comparing CQ and KVQuant on LLaMA-3-8b in the table below. CQ mostly outperforms KVQuant, especially in lower bit widths. We would like to gently remind that LLaMA-3 was released only one month before the NeurIPS deadline.

[W4] Comparison with QJL and KIVI.

• We tried our best to compare with QJL in the limited timeframe of the rebuttal. However, we ran into some issues due to incompatibilities of the QJL codebase with our hardware. Specifically, we ran into the following error: File ".../QJL/models/llama3_utils_qjl.py", line 138, in build_sketch self.key_states_norm = torch.norm(key_states, dim=-1) RuntimeError: CUDA error: limit is not supported on this architecture. This could be caused by incompatibilities of our V100 GPUs with the code. We will try to resolve this issue during the discussion period. We would like to kindly remind that QJL was first published online in June, which is after the NeurIPS deadline.
• We present additional experiments in the table below comparing CQ with KIVI on accuracy of LongBench with Llama-2-7b. For KIVI, we use 2-bit quantization, a full-precision sliding window (residual length) of 32 tokens, and a group size of 32. For CQ, we use 2-bit quantization and a sliding window size of 32. CQ mostly outperforms KIVI across different tasks.

• In the table below, we present additional latency measurements of CQ with comparison to KIVI and FP16 cache. We use Llama-2-7b with a batch size of 1 and a prompt of 2000 tokens to decode 100 tokens.

[W5] Longer context datasets such as LongBench.

• Please see the table above for CQ's results on LongBench with comparison to KIVI.
• We present additional experiments with longer context datasets comparing CQ and KVQuant using Llama-2-7b on GSM8K with chain-of-thought (CoT) and MMLU with CoT Fewshot (gsm8k_cot, mmlu_flan_cot_fewshot_humanities, mmlu_flan_cot_fewshot_stem, mmlu_flan_cot_fewshot_social_sciences, mmlu_flan_cot_fewshot_other from lm-evaluation-harness). In long-context settings, CQ mostly outperforms or performs similarly to KVQuant under the same bit width.

[Q1, L1] Overhead regarding storing the centroids. I would like you to mention this overhead and explain how you set those values.

• We kindly refer the reviewer to Section A.4 for a discussion on the overhead of storing and learning the centroids. We have reported the number of centroid parameters and the learning time for each CQ configuration and model, and described how the overhead is calculated in detail.
• CQ only has two hyperparameters: the number of channels and the number of bits in a code. These parameters are explicitly reported in our experimental results. There is no need of extensive hyperparameter tuning for our method.

We thank the reviewer again for carefully reviewing our paper and providing constructive feedback. We have conducted additional experiments using the official QJL codebase (https://github.com/amirzandieh/QJL) with an A100 40GB GPU. We evaluate QJL and CQ with Llama-2-7b on LongBench. For QJL, we used a sliding window of size 32 and a group size of 32 (buffer_size=32,group_size=32), and set other hyper-parameters following the codebase (key_quantization_bits=256,key_quantization_bits_initial_layers=512,initial_layers_count=15,outlier_count_general=8,outlier_count_initial_layers=8,value_quantization_bits=2). For CQ, we use the 4c8b (2-bit) configuration and a sliding window of size 32. The results are presented in the table below.

We also encountered some challenges when attempting to conduct additional experiments.

1. For QJL on TriviaQA, we ran into the following error: File ".../QJL/models/llama2_utils_qjl.py", line 143, in _update_outliers self.outlier_indices = torch.cat([self.outlier_indices, outlier_indices], dim=2).contiguous() TypeError: expected Tensor as element 0 in argument 0, but got NoneType.

2. We tried to directly compare CQ with QJL by following the experimental settings (longchat-7b-v1.5-32k on LongBench) in Table 1 of the QJL paper. However, we ran into out-of-memory issues with CQ-4c8b due to GPU memory constraints (Nvidia A100 40G).

Given the time constraint of the discussion period, we have made our best effort to provide a fair comparison between CQ and QJL. We are open to further investigation and would welcome specific suggestions from the reviewer on how to resolve the issues above. We are also happy to provide additional clarification on any follow-up questions. We respectfully request that the reviewer reconsider our paper in light of these responses.

Dear Reviewer KwWk,

We would like to express our sincere gratitude for your time and effort in reviewing our paper. Your feedback has been invaluable to us.

As the discussion period is drawing to a close, we kindly request that you review our previous responses to your review. If you have any additional questions or concerns, we would be happy to address them promptly.

Thank you again for your valuable contributions.

We sincerely thank the reviewer for their careful review of our paper and the insightful suggestions. We address your concerns as follows.

[W1] Focused on the Llama family of models. -- We have added Mistral model results!

• We would like to draw the reviewer's attention to the results on Mistral in Tables 1, 2, and 6.

[W2] The novelty of the paper is limited.

Although novelty is a multifaceted concept in research, we believe it can be viewed roughly from two aspects: empirical novelty, which involves the discovery of properties unknown to the community (e.g. the lottery ticket hypothesis [1]), and technical novelty, which refers to the development of new solutions (e.g. Attention [2]). We argue that our work exhibits both forms of novelty, as detailed below.

• Empirically Novel Observation that KV Channels are Highly Interdependent: To the best of our knowledge, we are the first to highlight the phenomenon that channels of KV cache exhibit high amounts of mutual information or correlation. This discovery opens up new avenues for KV cache compression.
• Technically Novel Approach that Enables Extreme KV Cache Compression: Based on our novel observation and concepts from Information Theory, we propose a new approach for KV cache quantization by coupling multiple KV channels. Existing approaches such as KVQuant and KIVI quantize KV cache channel-wise or token-wise independently, which cannot take advantage of the high mutual information shared between channels. By exploiting the interdependency among channels, we enable KV cache compression rates previously difficult or impossible to achieve, i.e., 1-bit quantization of KV cache.

[Q1] Does CQ affect the training process or is it purely an inference-time optimization? --- It does not!

• CQ is a post-training quantization (PTQ) approach that does not affect the training process of the models. CQ can improve the inference efficiency by saving memory and increasing batch size.

[Q2] Do you envision a way of using non-contiguous channels in the coupling, perhaps with some limited search of finding groups of channels with higher interdependency? --- Yes!

• Coupling non-contiguous, highly interdependent channels will likely further improve the quantization accuracy. One potential way of achieving that is to place channels with the highest correlation into the same channel group through a greedy search. We leave this for future work due to the difficulty in system implementation and optimizations.
• We thank the reviewer for pointing us in this direction.

[Q3] Can you clarify the meaning of bold and underline in Tables? --- Sure!

• The bolded number is the best perplexity/accuracy achieved under the same bit-width, while the underlined number is the second best. We will clarify this in the final paper.

References

[1] Frankle, Jonathan, and Michael Carbin. "The lottery ticket hypothesis: Finding sparse, trainable neural networks." arXiv preprint arXiv:1803.03635 (2018).

[2] Vaswani, Ashish, et al. "Attention is all you need." Advances in neural information processing systems 30 (2017).

We sincerely thank the reviewer for their careful consideration of our work and for providing valuable feedback. We have addressed their comments in detail below.

[W1] Table 2 shows that CQ only outperforms the baselines in the 1bit regime, which already suffers great accuracy loss.

• It is essential to consider the bits per activation metric when comparing quantization methods in Table 2. Under identical whole-number bit widths (4.00 bits, 2.00 bits, and 1.00 bit), CQ consistently surpasses KVQuant in average accuracy, particularly in the 2-bit and 1-bit regimes where the improvement exceeds 10%. When considering similar bit widths, CQ outperforms or matches the performance of sparsity-based KVQuant.
• We highlight that, in Table 2, CQ-8c8b (with a bit width of 1) outperforms KVQuant-2b (with a bit width of 2) in average accuracy despite using half the memory, demonstrating the efficacy of our proposed approach.
• Additionally, it's important to note that KVQuant+1% sparse is a sparsity-based method that relies on storing outlier activations in a sparse matrix. This approach can potentially introduce latency overhead and deployment complexities. In contrast, our CQ method maintains a fully dense representation, avoiding these potential issues.

[W2] Evaluations with long generation.

• We present additional experimental results below with long-context datasets comparing CQ and KVQuant using Llama-2-7b on GSM8K with chain-of-thought (CoT) and MMLU with CoT Fewshot (gsm8k_cot, mmlu_flan_cot_fewshot_humanities, mmlu_flan_cot_fewshot_stem, mmlu_flan_cot_fewshot_social_sciences, mmlu_flan_cot_fewshot_other from lm-evaluation-harness). In long-context settings, CQ mostly outperforms or is comparable to KVQuant under the same bit width.

[W3] Baselines are absent in Table 3.

• To add more baselines to Table 3, we present additional experiments on LongBench below comparing CQ and KIVI with sliding window full-precision cache using Llama-2-7b. For KIVI, we use 2-bit quantization, a full-precision sliding window (residual length) of 32 tokens, and a group size of 32. For CQ, we use 2-bit (4c8b) quantization and a sliding window size of 32. CQ mostly outperforms KIVI across different tasks.

[W4] Testing the proposed method on long context benchmarks or multi-hop retrieval tasks.

• Please see the first table above for evaluations on long-context benchmarks including GSM8k with CoT and MMLU with CoT Fewshot.
• We tried our best to evaluate our method on the RULER dataset in the limited timeframe of the rebuttal. However, we could not run RULER on our servers due to compatibility issues with docker. We will keep trying our method on RULER during the discussion period.

[W5] The quality of the centroids may not hold under distribution shifts.

• We have added an ablation study as follows to suggest that calibration on language modeling datasets provides transferable performance on downstream tasks. We use 16 sequences of 2048 tokens from WikiText-2 and C4 as the calibration set and evaluate CQ on 4 downstream tasks.
• Despite using different calibration datasets, CQ performs similarly in various downstream tasks.

[Q1] How are shared memory bank conflicts handled during centroid lookup?

• Due to a relatively high number of centroids (256 centroids for 8-bit codes and 1024 centroids for 10-bit codes), we did not notice shared memory bank conflicts to significantly impact system performance. We thank the reviewer for raising this insightful question and will continue to explore potential optimizations in this area.

We appreciate the reviewer for carefully reviewing our paper and offering thoughtful feedback. We have conducted additional experiments on passkey retrieval, following the setup in [1], with Llama-2-7b at its maximum context length of 4096. The passkey retrieval task is similar to the needle retrieval task in RULER. We still have trouble running RULER due to compatibility issues with Docker, and we will continue to work on it. As shown in the table below, CQ consistently outperforms KVQuant at various bit widths on passkey retrieval. We are also happy to provide additional clarification on any follow-up questions. We respectfully request that the reviewer reconsider our paper in light of these responses.

References

[1] Zhu, Dawei, et al. "Pose: Efficient context window extension of llms via positional skip-wise training." arXiv preprint arXiv:2309.10400 (2023).