Palu: Compressing KV-Cache with Low-Rank Projection

TL;DR

Palu tackles the memory bottleneck of KV-Cache in long-context LLMs by introducing a post-training, low-rank KV-Cache compression framework. It decomposes Key and Value projection weights with SVD, caches a compact latent representation, and reconstructs full KV states on the fly, augmented by group-head low-rank decomposition, Fisher-information-driven rank allocation, and a quantization-friendly design with offline fusion and Hadamard-based outlier mitigation. The approach yields substantial memory reductions (up to 11.4× with quantization) and significant latency improvements (up to 2.91× RoPE-attention, 5.53× end-to-end in long sequences) while preserving perplexity and zero-shot accuracy. These results demonstrate Palu’s practical potential for efficient, scalable LLM inference, especially in memory- and bandwidth-constrained deployments.

Abstract

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

Figures (9)

• Figure 1: Palu’s low-rank projection method for KV-Cache reduction. A weight matrix W of linear projection is decomposed into two low-rank matrices. Input X is down-projected to latent representation H, which is cached. Y can be reconstructed from H using the up-projection matrix B.
• Figure 2: Palu uses low-rank decomposition $(\mathbf{W} \approx \mathbf{A}\mathbf{B})$ to project the key (or value) to a lower-dimensional latent representation $(\mathbf{h})$, thereby reducing the size of the KV-Cache. The original key $(\mathbf{K}_t)$ is reconstructed on-the-fly with $\mathbf{B}^k$, and $\mathbf{B}^v$ is fused into $\mathbf{W}^o$ to avoid reconstruction overhead. The fusion also reduces the computational burden for output projection.
• Figure 3: Performing decomposition at different granularities. Jointly decomposing multiple heads can achieve higher accuracy. Assuming the same total size of the latent representations (i.e., $4\cdot r_i=2\cdot r_g=r_{\text{joint}}$), the FLOPs for reconstruction overhead in joint-head decomposition schemes are 4 times larger than those in multi-head ones.
• Figure 4: Activation distribution of the low-rank key caches at the $4^{th}$ Llama-2 attention layer.
• Figure 5: Normalized speedup for both the attention module and end-to-end model decoding. Solid lines represent exact measurements, while dashed lines indicate the FP16 baselines are out of memory, and the speedpus are compared the estimated baseline's latency.