CommVQ: Commutative Vector Quantization for KV Cache Compression

We thank the reviewer for the positive feedback and thoughtful suggestions.

1. Full Longbench and Needle-in-a-Haystack results

The full LongBench consists of 21 tasks in total. Following prior works such as KIVI and KVQuant, we report results on the same eight representative tasks for fair comparison and consistency. To address your comment, we now provide results on the full LongBench benchmark, including all 21 tasks, using the LLaMA-3.1 8B model.

The results below show that CommVQ-2 maintains accuracy across nearly all subtasks, while CommVQ-1 achieves competitive performance even under extreme 1-bit compression.

These results further validate CommVQ’s generalization across a wide range of long-context tasks and domains.

We also provide the results for the Needle-in-a-Haystack test using the LLaMA-3.1 8B model. The NIAH result figures are included in this link: https://github.com/commvq/CommVQ/blob/main/fig1.png

We can see that CommVQ-2 preserves the full retrieval capability of the FP16 baseline, while our 1-bit quantized version, CommVQ-1, significantly outperforms its counterpart, KIVI-1.

2. Triton kernel implementation details

Our Triton kernels implement the techniques described in Section 4 of the main paper. Specifically, they include:

1. A kernel that fuses RoPE application with commutative codebook decoding, reducing intermediate memory operations;
2. A mixed-precision batched matrix multiplication for efficient computation;
3. Loading low-bit representations on the fly, avoiding the need to upcast them to higher precision as required in native PyTorch implementations.

Together, these optimizations reduce memory access overhead and improve compute utilization by fully leveraging Triton’s capabilities. This implementation is under active development and will continue to be optimized. We will release the Triton kernels, along with code and implementation details, upon acceptance.

We thank the reviewer for the detailed and constructive feedback. Below, we address the key concerns raised.

1. Latency comparison

Please see our response to Reviewer 86AS (2. Latency comparison with the FP16 baseline and prior methods such as KIVI) for a detailed latency comparison. In summary, CommVQ-1 is consistently faster than KIVI-1, especially as context length increases, as measured on the LLaMA-3.1 8B model using a single NVIDIA H100 80G GPU.

2. Compatibility with other compression methods

CommVQ is compatible with other KV cache compression techniques like token eviction and quantization.

For token eviction or retrieval-based methods, which retain only key tokens during decoding, CommVQ integrates naturally by quantizing just those essential tokens. This allows further compression as only selected tokens are encoded and decoded using the learned codebook.

CommVQ can also benefit from codebook quantization, reducing storage and speeding up decoding via low-bit matrix operations.

In short, CommVQ is orthogonal to other methods and can be combined with them for greater compression. Exploring such combinations is a promising direction for future work.

3. Robustness under domain shift

Thanks for pointing this out. In practice, we find that our codebooks and encoder trained on general pre-training datasets (e.g., FineWeb-Edu) transfer reasonably well across tasks and domains. To validate this point, we conducted an ablation study on the LLaMA-3.1 8B model to show CommVQ-2's perplexity changes compared to the FP16 baseline on 4 datasets:

1. FineWeb-Edu: a general dataset.
2. GSM-8K: a math benchmark.
3. Repobench-p: a code retrieval and completion benchmark.
4. KV_Retrieval in InfiniteBench: a synthetic UUID key-value retrieval benchmark.

The first dataset represents in-domain evaluation, while the last three represent evaluations with domain shifts—i.e., the codebooks and encoder are trained on general text and tested on math, code, and synthetic UUID data.

We find no significant increase in perplexity (PPL) due to domain shifts when compared to in-domain evaluations. This suggests that our method performs consistently well across domains that differ from the calibration data, including synthetic UUID data, which is unlikely to appear in the calibration set. Overall, we conclude that our method is robust and generalizable under domain shifts.

Finally, if a significant domain shift is encountered, we recommend further fine-tuning the encoder and codebook on domain-specific data—similar to best practices in other calibration-based quantization methods.

4. Additional benchmark (e.g., LongBench v2)

Thank you for the suggestion. Due to the limited time available during the rebuttal phase, we conducted additional evaluations on LongBench v2 using the LLaMA-3.1 8B model and compared them against KIVI, KVQuant, and VQLLM. KIVI and KVQuant fail to produce meaningful output with 1-bit quantization, so their results are omitted from the table. As shown below, CommVQ continues to outperform the baseline methods at comparable average quantization bit levels.

5. Explanation on how to incorporate the commutative constraints during codebook training

Thank you for pointing this out. The commutativity constraint is enforced by restricting each 2D subvector in the codebook to the form [[x, y], [-y, x]], which ensures commutativity with RoPE rotations. In other words, for each 2D subvector, we only learn two scalars—x and y—and use them to construct the subvector for computation. We will provide a more detailed explanation in our future revision.

6. Code and Implementation

Upon acceptance, we will open-source our training pipeline, model weights, and optimized Triton kernels to facilitate reproducibility and further research. Please refer to the response to Reviewer 9BoD (2. Triton kernel implementation details) for explanations of the Triton kernel implementation.

We thank the reviewer for the thoughtful feedback and positive assessment of our direction and formulation. Below, we address the key concerns raised.

1. Experiments on larger models (e.g., 70B)

We focused on 8B models (LLaMA-2, Mistral, and LLaMA-3.1) due to their popularity and our resource constraints. However, to address your suggestion, we have now evaluated CommVQ on the LLaMA-3.1 70B model under 1-bit quantization using LongBench. Results are shown below:

These results show that CommVQ generalizes well to the 70B model, achieving strong accuracy using a 1-bit KV cache. This confirms CommVQ’s scalability and practical utility in larger models.

2. Actual practical benefits for 1-bit quantization

Compressing the KV cache to 1-bit offers several practical benefits. The most significant benefit is a 16× memory reduction compared to the standard FP16 KV cache. These memory savings are applicable regardless of the model's precision, meaning the model doesn't need to operate at 1-bit to take advantage of 1-bit KV cache quantization.

Such a substantial memory reduction not only allows a single GPU to handle longer context lengths and larger batch sizes (as shown in Figure 2 of our main paper), but it also significantly reduces data transfer time when offloading the KV cache to non-GPU memory. This benefit is especially notable in offloading scenarios, which are common when serving large models or deploying on edge devices. In these cases, limited PCIe bandwidth often makes KV cache transfer a major bottleneck. Reducing the size of the KV cache helps mitigate this issue, leading to faster overall inference.

3. Clarification on "87.5% reduction"

We clarify that the 87.5% reduction refers to the reduction from FP16 (16-bit) KV cache to 2-bit, i.e., (1 - 2 / 16) = 87.5%. This applies to CommVQ-2. For CommVQ-1 (1-bit), the reduction reaches 93.75%, enabling even more aggressive compression.

4. More insights behind the codebook, encoder and EM algorithm

Thank you for your suggestions. We chose to use a simple neural network as the encoder, as it proved to be simple yet effective in our preliminary experiments. We selected the EM algorithm because it offers two advantages:

1. The EM algorithm converges very quickly, requiring far fewer iterations than gradient-based approaches.
2. The EM algorithm incorporates principled techniques—such as soft assignment and annealing—to prevent mode collapse, where some centroids become obsolete after a few iterations. This has been one of the major challenges in quantization algorithms.

Regarding the codebook configurations, we included the ablation study in the Appendix primarily due to page limitations. As suggested, we will provide more insights in the main paper in future revisions.

We thank the reviewer for the thoughtful feedback and helpful suggestions. Below, we address your concerns.

1. Applicability to retrieval tasks such as RULER

We appreciate the suggestion to evaluate additional retrieval-specific tasks. We have included results on the RULER benchmark using the LLaMA-3.1 8B model. We set the context length to 128K. As shown below, CommVQ-2 achieves the highest average score among the methods that can achieve an average quantization bit of 2. CommVQ-1 also retains competitive retrieval ability under the extreme 1-bit quantization.

These results demonstrate that CommVQ can preserve retrieval capability even under aggressive compression while performing better than other methods under the same compression rate.

Apart from RULER, we also conducted experiments on the Needle-in-a-Haystack test, which also focuses on retrieval; please refer to our response to Reviewer 9BoD (1. Full LongBench and Needle-in-a-Haystack results) for details. The NIAH result figures are presented in this link.

In summary, CommVQ under 2-bit quantization can preserve the full retrieval capability of the FP16 model, while our 1-bit quantization version performs better on the NIAH test than KIVI's 1-bit quantization version. Moreover, from Table 2 in our main paper, CommVQ achieves good performance on retrieval tasks (namely R.PK, R.Num, and R.KV), especially in the extreme 1-bit quantization case where CommVQ significantly outperforms the baselines. Our comprehensive results demonstrate our method's effectiveness on retrieval tasks.

2. Comparison with low-rank decomposition

Thanks for the suggestion, and we chose to compare CommVQ with Palu, a state-of-the-art low-rank decomposition KV cache compression method published at ICLR 2025. We conducted experiments on LongBench using the LLaMA-3.1 8B model. We set the context length to 128K. As shown below, CommVQ consistently outperforms Palu at both 2-bit and 1-bit quantization levels. This shows that our method is more effective than low-rank decomposition methods under various compression rates.

We appreciate your positive feedback and insightful suggestions. Below, we address your concerns.

1. CommVQ applied to latest models such as Qwen-2.5

To evaluate our method's generalization to the latest models, we applied CommVQ to the Qwen-2.5 7B model and evaluated it on LongBench. Due to the limited time for rebuttal and compatibility issues (e.g., KIVI and KVQuant do not officially support Qwen in their open-sourced code), we used KV-4, the built-in 4-bit KV cache quantization method (HQQ) in HuggingFace Transformers v4.46.2, as our baseline.

CommVQ achieves strong performance while compressing the KV cache to an average of 2 bits, significantly outperforming the default 4-bit quantization baseline. This indicates that our method can be effectively applied to the latest models. We also demonstrate that CommVQ can be applied to much larger models, such as the LLaMA-3.1 70B model, in our response to Reviewer yREq (1. Experiments on larger models). All these results show the broad applicability and effectiveness of CommVQ.

2. Latency comparison with the FP16 baseline and prior methods such as KIVI

We report latency per generated token (in seconds) using the LLaMA-3.1 8B model on a single NVIDIA H100 80G GPU. Both KIVI and our method are optimized with Triton kernels. We set the batch size to 1 and vary the context length. As shown below, CommVQ-1 is consistently faster than KIVI-1, especially as the context length increases.

While both CommVQ and KIVI appear slower than the FP16 baseline in our current measurements, we believe this is primarily due to practical factors such as the usage of flash attention (the FP16 model uses flash attention by default, while KIVI and our method currently do not support flash attention during the generation stage), Triton kernel launch overhead, and the hardware and model we chose. Importantly, as the context length grows, CommVQ achieves better efficiency than KIVI. We are actively optimizing our Triton implementation and expect further latency improvements in future versions.

3. Lack of some detailed proofs, such as why RoPE embedding is commutative in Property 1

We thank the reviewer for pointing this out. Due to the page limit, we have omitted some proofs in the main paper. We will add detailed proofs to the main paper in our future revisions.

As for why RoPE embedding is commutative in Property 1, since

R_m^iC = \\left(\\begin{array}{cc} \\cos m \\theta_i & -\\sin m \\theta_i \\\\ \\sin m \\theta_i & \\cos m \\theta_i \\\\ \\end{array}\\right) \\begin{pmatrix} x & y \\\\ -y & x \\\\ \\end{pmatrix}= \\begin{pmatrix} x\\cos m\\theta_i + y\\sin m\\theta_i & y\\cos m\\theta_i - x\\sin m\\theta_i \\\\ x\\sin m\\theta_i - y\\cos m\\theta_i & y\\sin m\\theta_i + x\\cos m\\theta_i \\end{pmatrix}

and

CR_m^i = \\begin{pmatrix} x & y \\\\ -y & x \\\\ \\end{pmatrix} \\left(\\begin{array}{cc} \\cos m \\theta_i & -\\sin m \\theta_i \\\\ \\sin m \\theta_i & \\cos m \\theta_i \\\\ \\end{array}\\right)=\\begin{pmatrix} x\\cos m\\theta_i + y\\sin m\\theta_i & y\\cos m\\theta_i - x\\sin m\\theta_i \\\\ x\\sin m\\theta_i - y\\cos m\\theta_i & y\\sin m\\theta_i + x\\cos m\\theta_i \\end{pmatrix}

We can see that they are equivalent, so we can confirm Property 1.