SYNTOLOGY HomeExplorerAtlasCodeMethodologyAboutDevelopersFeedPricing
Paper · 2409.10593 · NeurIPS · 2024

CSKV: Training-Efficient Channel Shrinking for KV Cache in Long-Context Scenarios

Xuefei Ning, Yu Wang, Zhihang Yuan, Shiyao Li, Guohao Dai, Shengen Yan, Luning Wang

arXiv · PDF · Open in the Atlas

Code that ran

We lifted 9 functions out of this paper's own repositories and ran 4 of them in a sandbox. "Ran" means the function executed on a synthesized input and returned a value. It is not a reproduction of the paper's results.

RepositoryRoleRan
wln20/CSKV canonical 4 of 9
FunctionStatusWhere it lives
forward_mistral Ran wln20/CSKV/cskv_src/modeling/modeling_mistral.py
pointer only (licence: NONE) · get_code("53c019a2f72f6891")
get_name Ran wln20/CSKV/cskv_src/utils/general_utils.py
pointer only (licence: NONE) · get_code("8717d7f11f169bd2")
init_svd Ran wln20/CSKV/cskv_src/utils/training_utils.py
pointer only (licence: NONE) · get_code("b11f68f08e771fbd")
init_svd_parallel_inference Ran wln20/CSKV/cskv_src/utils/training_utils.py
pointer only (licence: NONE) · get_code("fb2320a6cd6c2461")
get_dataset_for_calib Not yet run wln20/CSKV/cskv_src/utils/data_utils.py
pointer only (licence: NONE) · get_code("99c1c2ed2f9a3ae0")
get_dataset_for_trainer Not yet run wln20/CSKV/cskv_src/utils/data_utils.py
pointer only (licence: NONE) · get_code("6dd2377be358412b")
init_asvd Not yet run wln20/CSKV/cskv_src/utils/training_utils.py
pointer only (licence: NONE) · get_code("607710bbee79e7ac")
prepare_inputs_for_generation_llama Not yet run wln20/CSKV/cskv_src/modeling/modeling_llama.py
pointer only (licence: NONE) · get_code("90861e147e3c2c04")
prepare_inputs_for_generation_mistral Not yet run wln20/CSKV/cskv_src/modeling/modeling_mistral.py
pointer only (licence: NONE) · get_code("3472698614a9b449")

Repositories linked to this paper

Some links come from the archived Papers with Code dataset (CC BY-SA 4.0): attribution and licence.

Abstract

Large Language Models (LLMs) have been widely adopted to process long-context tasks. However, the large memory overhead of the key-value (KV) cache poses significant challenges in long-context scenarios. Existing training-free KV cache compression methods typically focus on quantization and token pruning, which have compression limits, and excessive sparsity can lead to severe performance degradation. Other methods design new architectures with le ss KV overhead but require significant training overhead. To address the above two drawbacks, we further explore the redundancy in the channel dimension and apply an architecturelevel design with minor training costs. Therefore, we introduce CSKV, a trainingefficient Channel Shrinking technique for KV cache compression: (1) We first analyze the singular value distribution of the KV cache, revealing significant redundancy and compression potential along the channel dimension. Based on this observation, we propose using low-rank decomposition for key and value layers and storing the low-dimension features. (2) To preserve model performance, we introduce a bi-branch KV cache, including a window-based full-precision KV cache and a low-precision compressed KV cache. (3) To reduce the training costs, we minimize the layer-wise reconstruction loss for the compressed KV cache instead of retraining the entire LLMs. Extensive experiments show that CSKV can reduce the memory overhead of the KV cache by 80% while maintaining the model's longcontext capability. Moreover, we show that our method can be seamlessly combined with quantization to further reduce the memory overhead, achieving a compression ratio of up to 95%. Code is available at https://github.com/wln20/CSKV.

For agents

The same record, over MCP at https://syntology.ai/mcp:

get_harvested_code_for_paper("2409.10593")
get_code_for_paper("2409.10593")
have("2409.10593")

Connect an agent — have() is free.