Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

Palu: Compressing KV-Cache with Low-Rank Projection

About

Post-training KV-Cache compression methods typically either sample a subset of effectual tokens or quantize the data into lower numerical bit width. However, these methods cannot exploit redundancy in the hidden dimension of the KV tensors. This paper presents a hidden dimension compression approach called Palu, a KV-Cache compression framework that utilizes low-rank projection to reduce inference-time LLM memory usage. Palu decomposes the linear layers into low-rank matrices, caches compressed intermediate states, and reconstructs the full keys and values on the fly. To improve accuracy, compression rate, and efficiency, Palu further encompasses (1) a medium-grained low-rank decomposition scheme, (2) an efficient rank search algorithm, (3) low-rank-aware quantization compatibility enhancements, and (4) optimized GPU kernels with operators fusion. Extensive experiments with popular LLMs show that Palu compresses KV-Cache by 50% while maintaining strong accuracy and delivering up to 1.89x on the RoPE-based attention module. When combined with quantization, Palu's inherent quantization-friendly design yields small to negligible extra accuracy degradation while saving additional memory than quantization-only methods and achieving up to 2.91x speedup for the RoPE-based attention. Moreover, it maintains comparable or even better accuracy (up to 1.19 lower perplexity) compared to quantization-only methods. These results demonstrate Palu's superior capability to effectively address the efficiency and memory challenges of LLM inference posed by KV-Cache. Our code is publicly available at: https://github.com/shadowpa0327/Palu

Chi-Chih Chang, Wei-Cheng Lin, Chien-Yu Lin, Chong-Yan Chen, Yu-Fang Hu, Pei-Shuo Wang, Ning-Chi Huang, Luis Ceze, Mohamed S. Abdelfattah, Kai-Chiang Wu• 2024

Related benchmarks

TaskDatasetResultRank
Language ModelingWikiText2 v1 (test)
Perplexity45.4
383
Long-context Language UnderstandingLongBench
M-Avg47.47
292
Long-context language modelingLongBench
Average Score35.45
164
Question AnsweringARC Easy (test)
Accuracy28.16
54
Reasoning and Language ModelingReasoning Suite (ARC, HellaSwag, PIQA, WinoGrande, MMLU, OpenBookQA, Real-world QA) Zero-shot Llama-3.1-8B-Instruct with Alpaca calibration
PPL45.4
32
Long-context UnderstandingRULER
S1 Score100
20
Language Modeling EvaluationLM Evaluation Harness
ARC44.56
11
Showing 7 of 7 rows

Other info

Follow for update