More Tokens, Lower Precision: Towards the Optimal Token-Precision Trade-off in KV Cache Compression

Jiebin Zhang, Dawei Zhu, Yifan Song, Wenhao Wu, Chuqiao Kuang +4 more
12/17/2024
cs.CL

Official code: zhzihao/QPruningKV (12 stars, MIT, last commit 2025-01-15).

Abstract

As large language models (LLMs) process increasing context windows, the memory usage of KV cache has become a critical bottleneck during inference. The mainstream KV compression methods, including KV pruning and KV quantization, primarily focus on either token or precision dimension separately. However, these works leaving the trade-off between these two orthogonal dimensions largely under-explored. In this paper, we comprehensively investigate the token-precision trade-off in KV cache compression.Experiments demonstrate that storing more tokens in the KV cache with lower precision,a strategy we term quantized pruning, can significantly enhance the long-context performance of LLMs. In-depth analysis of the token-precision trade-off across key aspects demonstrates that, quantized pruning achieves substantial improvements in retrieval-related tasks and consistently performs well across varying input lengths. Furthermore, quantized pruning demonstrates notable stability and effectiveness across different KV pruning methods, quantization strategies, and model scales. These findings offer valuable insights into optimizing KV cache compression through balanced token-precision trade-off strategies. Our code is available at https://github.com/zhzihao/QPruningKV.

View on arXivView PDF

Code Implementations(1)

zhzihao/QPruningKVOfficial100%
MIT

More Tokens, Lower Precision: Towards the Optimal Token-Precision Trade-off in KV Cache Compression(EMNLP 2025)

120Shell, PythonJan 5, 20251 years agoMIT

Similar Papers

Cite this paper

@article{zhang2024more,
  title  = {More Tokens, Lower Precision: Towards the Optimal Token-Precision Trade-off in KV Cache Compression},
  author = {Jiebin Zhang and Dawei Zhu and Yifan Song and Wenhao Wu and Chuqiao Kuang and Xiaoguang Li and Lifeng Shang and Qun Liu and Sujian Li},
  year   = {2024},
  eprint = {2412.12706},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2412.12706v2}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2412.12706.svg)](https://codeofpaper.com/paper/2412.12706)

Discussion