PaperlayerView exact revision on arXiv

Source-linked AI summary

Palu: Compressing KV-Cache with Low-Rank Projection

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

arXiv:2407.21118v2cs.AIcs.LG

TL;DR

KV-Cache compression methods such as quantization and token eviction do not exploit redundancy in hidden dimensions, despite the memory pressure of growing caches. Palu uses low-rank projection to cache compressed latent representations and reconstruct keys and values, with decomposition, rank-allocation, quantization, and kernel optimizations. Experiments report up to 2.91× RoPE-based attention speedup with 50% low-rank compression and 4-bit quantization, while preserving strong accuracy.

  • Problem

    Quantization and token eviction do not exploit redundancy in KV tensors’ hidden dimensions, while growing KV-Caches strain memory capacity and bandwidth.

    Source: §1 · pp. 1–2

  • Method

    Palu decomposes key and value projections into low-rank factors, caches latent representations, and reconstructs full states, with group-head decomposition, rank search, quantization compatibility, and fused kernels.

    Source: §1 · p. 3, §3.1 · p. 4

  • Results

    With 50% low-rank compression and 4-bit quantization, Palu achieves up to 2.91× RoPE-based attention speedup and up to 2.2× end-to-end speedup while preserving strong accuracy across benchmarks.

    Source: §6 · p. 11

  • Takeaways & Limitations

    Palu provides a low-rank compression dimension that can be combined with quantization to reduce KV-Cache memory while retaining strong accuracy and inference speed.

    Source: §1 · p. 2, §6 · p. 11

    Speedup decreases for sequence lengths beyond 16K as rising reconstruction costs shift online reconstruction from memory-bound to compute-bound.

    Source: §4.5.2 · p. 10

Abstract

from arXiv · show

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

PALU: KV-CACHE COMPRESSION WITH LOW-RANK PROJECTION

The supplied title-page excerpt lists the paper’s authors.

  • The listed authors are 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, and Kai-Chiang Wu.Source: PALU: KV-CACHE COMPRESSION WITH… · p. 1

1 INTRODUCTION

KV-Cache compression methods such as quantization and token eviction do not exploit redundancy in hidden dimensions. Palu applies low-rank projection to cache compressed latent states, with proposed decomposition, rank-allocation, quantization, and kernel optimizations; experiments report memory compression, accuracy, and speedup results.

  • Quantization and token eviction leave hidden-dimension redundancy in KV tensors unexploited, while Palu adds low-rank projection as an orthogonal compression dimension.Source: §1 · p. 2
  • Palu statically decomposes Key and Value projection matrices and caches low-rank latent representations, avoiding runtime decomposition of cached matrices.Figure 1 depicts down-projection to cached H and reconstruction of Y using B.Source: §1 · p. 2, Figure 1 · p. 2
  • Group-head low-rank decomposition balances accuracy and reconstruction efficiency between joint-head and per-head decompositions.Source: §1 · p. 2
  • A Fisher-information-based rank search assigns higher ranks to important matrices and lower ranks to less critical ones, improving accuracy at the same overall compression rate.Source: §1 · p. 2
  • Palu fuses transformation matrices into forward and backward matrices to mitigate low-rank-induced outliers without runtime efficiency impact.The paper presents this as a quantization-compatibility optimization.Source: §1 · p. 2
  • Over 91.25% compression (11.4× reduction) and 1.19 lower perplexity than KVQuant (87.5% compression) are reported for Palu with quantization.The introduction also reports strong accuracy and perplexity with up to 50% low-rank compression.Source: §1 · p. 2
  • At 50% KV-Cache compression without quantization, Palu reaches up to 1.89× and 2.2× speedup for RoPE-based and non-RoPE attention modules, respectively.With quantization, the reported maxima are 2.91× and 6.17× for RoPE-based and non-RoPE attention, respectively.Source: §1 · p. 2
  • Palu’s listed contributions also include customized kernels with operator fusion, alongside automatic rank allocation and quantization-compatibility optimization.Source: §1 · p. 3

2 BACKGROUND

The background introduces multi-head attention and singular value decomposition, then explains how truncating SVD yields low-rank factors for approximating a weight matrix.

2.1 MULTI-HEAD ATTENTION MECHANISM

  • MHA projects each input into per-head queries, keys, and values, computes attention outputs, then concatenates head outputs before the output projection.Source: §2.1 · p. 3

2.2 SINGULAR VALUE DECOMPOSITION (SVD)

  • Truncating SVD to the largest r singular values forms low-rank factors A and B that approximate W and reduce storage.SVD represents W using U, Σ, and V^T; the retained singular vectors correspond to the largest r singular values.Source: §2.2 · p. 3

3 THE PALU FRAMEWORK

Palu reduces cached KV state by projecting key and value states into low-dimensional latents and reconstructing them during decoding. It balances decomposition granularity and rank allocation, and incorporates quantization-compatible transforms.

3.1.1 INTEGRATION WITH THE ATTENTION MECHANISM AND OFFLINE MATRIX FUSION

  • Palu reconstructs keys on the fly but fuses the value reconstruction matrix into Wo, avoiding full value-vector reconstruction and reducing output-projection computation.The corresponding key reconstruction matrix can also be fused into the query projection matrix for attention-score computation.Source: Figure 2 · p. 4, §3.1.1 · p. 4

3.1 COMPRESSING THE KV-CACHE VIA LOW-RANK PROJECTION

  • Palu projects keys and values into low-dimensional latent representations for caching, then reconstructs the original states during decoding.The decomposition rewrites a linear projection as y = xAB, with A projecting into the latent space and B reconstructing the output.Source: §3.1 · p. 4

3.1.2 COMPABILITY WITH POSITIONAL EMBEDDING

  • RoPE's nonlinearity prevents attention-score matrix fusion, so Palu dynamically reconstructs keys; non-RoPE modules can use fusion and avoid reconstruction.A custom GPU kernel is used to improve reconstruction efficiency for RoPE-based attention.Source: §3.1.2 · pp. 4–5

3.2.1 MULTI-HEAD LOW-RANK DECOMPOSITION

  • Per-head M-LRD can lose accuracy, possibly because SVD fails to capture information shared across heads.Source: §3.2.1 · p. 5

3.2.2 JOINT-HEAD LOW-RANK DECOMPOSITION

  • J-LRD jointly decomposes all heads to preserve shared principal components, but reconstruction requires 4 times the multi-head FLOPs at equal latent size and costs n times M-LRD overall.The larger fused matrix also raises memory consumption relative to M-LRD.Source: §3.2.2 · p. 5, Figure 3 · p. 5

3.2.3 GROUP-HEAD LOW-RANK DECOMPOSITION

  • G-LRD groups heads to balance accuracy and reconstruction cost; at matched total rank size, it reduces J-LRD reconstruction cost and fused-matrix size by ng.Group-wise decomposition captures shared information within each group while limiting computational overhead.Source: §3.2.3 · p. 6

3.3 AUTOMATIC RANK ALLOCATION

  • Palu allocates ranks in proportion to each target matrix's Fisher information, giving more important layers higher ranks under a fixed compression rate.The method treats compression sensitivity as proportional to Fisher information.Source: §3.3 · p. 6

3.4 QUANTIZATION COMPATIBILITY

  • SVD-induced outliers in low-rank latents hinder quantization; folding WHT into decomposed weights eliminates them without runtime overhead and adds negligible perplexity, even at 2-bit or 3-bit.The outlier pattern stretches the latent-key distribution and hurts quantization accuracy.Source: §3.4 · pp. 6–7

4 EXPERIMENTS

The experiments evaluate Palu across LLM families, accuracy benchmarks, long-context tasks, and latency settings. They report accuracy–compression trade-offs across decomposition granularities, benefits from combining low-rank compression with quantization, and speedups that are strongest at longer contexts.

4.1 EXPERIMENTS SETUP

  • Evaluations cover Llama-2, Llama-3, Mistral, and LongChat using perplexity, zero-shot and LongBench accuracy; latency tests use Llama-2-7B on one RTX 4090 against FP16 and KIVI-4-bit baselines.Source: §4.1 · p. 7, §4.5 · pp. 8–9

4.2 RESULTS WITH DIFFERENT DECOMPOSITION GRANULARITY

  • At 50% compression on Llama-2-7B, J-LRD reaches 5.62 WikiText-2 perplexity and 0.23% average accuracy degradation; G-LRD records 6.01 perplexity and 2.4% degradation, versus M-LRD’s 7.26% drop.Source: Table 1 · p. 7, §4.2 · p. 8

4.3 RESULTS OF QUANTIZATION INTEGRATION

  • With 3-bit quantization, perplexity increases by 0.08 and 0.23 at 30% and 50% low-rank compression; at 2 bits, Palu reduces perplexity by 1.19 and 0.54 versus KVQuant while using 30% and 50% less memory.Source: Table 2 · p. 8, §4.3 · p. 8

4.4 EVALUATION ON LONG CONTEXT DATASETS

  • On LongBench, 30% compression causes less than 1% average accuracy drop for Mistral-7B and LongChat-7B; 3-bit latent-cache quantization adds less than 1% degradation, with 7.59× compression.Source: §4.4 · p. 8, Table 3 · p. 9

4.5.1 SPEEDUPS OF ATTENTION MODULE AND END-TO-END DECODING.

  • At 64K, Palu reaches 1.89× RoPE-attention speedup without quantization and 2.91× with 4-bit Value quantization; non-RoPE attention reaches 2.20× and 6.17×, respectively.Source: Figure 5 · p. 9, §4.5.1 · pp. 9–10
  • For end-to-end decoding, Palu achieves up to 1.71× and 2.05× speedups without quantization for RoPE and non-RoPE models, rising to 2.59× and 5.53× at 64K with 4-bit quantization.Source: Figure 5 · p. 9, §4.5.1 · p. 10

4.5.2 KERNEL FOR ROPE-BASED ATTENTION SCORE

  • For the RoPE attention-score kernel, G-LRD reaches up to 1.95× speedup, M-LRD 3.56× at 16K, while J-LRD slows across sequence lengths; speedup declines beyond 16K as reconstruction becomes compute-bound.Source: §4.5.2 · p. 10

5 RELATED WORK

The related work distinguishes Palu from weight-compression SVD methods, KV-cache quantization, and low-rank attention designs. Palu instead targets post-training KV-cache compression through low-rank projection and associated optimizations.

  • Earlier SVD methods compress model weights: FWSVD uses Fisher information [Hsu et al., 2022], ASVD considers activation outliers, and SVD-LLM reduces compression loss per singular value [Wang et al., 2024].Source: §5 · p. 10
  • A concurrent low-rank KV-cache method groups heads but requires LoRA finetuning; Palu decomposes projection weights without finetuning and adds rank search, quantization integration, and GPU kernels [Yu et al., 2024].Source: §5 · p. 10
  • KV-cache methods quantize values in different ways: KIVI uses per-channel and per-token quantization with group size 32 [Liu et al., 2024b], while KVQuant adds non-uniform quantization and sparse matrices [Hooper et al., 2024].Source: §5 · p. 10
  • Unlike MLA, which requires pre-training as a new attention mechanism, Palu adapts existing MHA or GQA models for post-training KV-cache compression.Source: §5 · p. 11

6 CONCLUSION

The conclusion presents Palu as a low-rank KV-cache compression framework with decomposition, rank allocation, quantization compatibility, and fused-kernel optimizations. The appendix excerpt describes quantization as representing high-precision values with discrete low-bit values.

6 CONCLUSION

  • Palu decomposes projection weights and caches compressed latent representations, combining group-head decomposition, automatic rank allocation, quantization compatibility, and fused GPU kernels.Source: §6 · p. 11

A QUANTIZATION BASICS

  • Quantization approximates floating-point tensors with discrete low-bit values; the excerpt identifies a zero-point and rounding operation in the asymmetric uniform scheme.Source: A QUANTIZATION BASICS · p. 18

B KERNEL IMPLEMENTATION DETAILS

Palu’s RoPE attention-score kernel reconstructs keys from cached low-rank states while computing attention scores. It tiles work along sequence length so each thread block performs reconstruction, positional embedding, and query multiplication on-chip.

C DISCUSSION REGARDING MEMORY USAGE

Palu's memory footprint depends on decomposition strategy and sequence length: M-LRD and G-LRD avoid the weight-storage increase possible with J-LRD, while low-rank compression reduces total memory at long contexts. LoRA integration improves some decomposition variants' accuracy results.

C DISCUSSION REGARDING MEMORY USAGE

  • With m=n and r=0.7n, a 30% KV-Cache compression rate can increase weight size by approximately 40%.The discussion cautions that KV-Cache compression rate alone does not capture total memory savings; with group-size-4 G-LRD, m=8n yields a 21.25% memory saving.Source: C DISCUSSION REGARDING MEMORY USAGE · p. 19
  • K and V projections are 2 of 7 transformer-block linear layers and 16% of Llama-2-7B parameters, limiting their overall memory impact.The discussion identifies M-LRD and G-LRD as alternatives that do not increase memory usage, in contrast to J-LRD overhead.Source: C DISCUSSION REGARDING MEMORY USAGE · p. 19

D OVERALL MEMORY REDUCTION

  • At 64k tokens, KV-Cache is 78% of total memory; Palu reduces total memory 1.7× with 50% low-rank compression and 4.6× with 2-bit quantization.Figure 8 covers total memory, including model weights and KV-Cache, across sequence lengths.Source: Figure 8 · p. 19, D OVERALL MEMORY REDUCTION · p. 19

E INTEGRATING Palu WITH LORA FINETUNE

  • LoRA adds trainable low-rank matrices A′ to capture task-specific nuances and recover information lost during compression.Palu's A matrices are fixed parameters derived from pretrained linear-layer weights; the reported fine-tuning setup uses 8k Alpaca samples, rank r′=32, and α=32.Source: E INTEGRATING Palu WITH LORA FINETUNE · p. 20
  • With LoRA, G-LRD differs in accuracy by 1.03% from J-LRD; average drops are 1.00% for J-LRD, 2.01% for G-LRD (gs=4), and 5.14% for M-LRD.The reported results compare LoRA-integrated variants with their non-LoRA counterparts and include J-LRD's average performance drop.Source: E INTEGRATING Palu WITH LORA FINETUNE · p. 20, Table 4 · p. 20

F MORE RESULTS ON ZERO-SHOT ACCURACY

Across these additional evaluations and ablations, Palu maintains competitive accuracy at 50% compression, while group size, Walsh-Hadamard Transform, and rank allocation affect performance. Rank allocation is non-uniform across key and value projections and layers.

F MORE RESULTS ON ZERO-SHOT ACCURACY

  • At 50% compression on Llama-2-13B, J-LRD, G-LRD, and M-LRD each show accuracy drops around 3% or less.The authors suggest M-LRD as a first choice for further efficiency optimization.Source: F MORE RESULTS ON ZERO-SHOT ACCURACY · p. 20

G.1 INFLUENCE OF DIFFERENT GROUP SIZE

  • Increasing G-LRD group size increases shared information and improves performance.The ablation evaluates perplexity on Wikitext-2 for Llama-2-7B at 50% compression.Source: §G.1 · p. 20, Table 6 · p. 21

G.2 INFLUENCE OF WALSH-HADAMARD TRANSFORM

  • At 2-bit quantization, the Walsh-Hadamard Transform yields a 4.17 perplexity enhancement; at 3-bit, its perplexity effect is slight.Offline preprocessing optimizes the transform so it adds no inference overhead.Source: §G.2 · p. 21, Table 7 · p. 21

G.3 AUTOMATIC RANK ALLOCATION VS. UNIFORM RANK ALLOCATION

  • At 50% compression, rank searching reduces perplexity by 2.18; allocation is non-uniform, generally giving values higher rank than keys and the first half of layers higher ranks.Figure 9 visualizes allocation across transformer blocks for key and value projections.Source: §G.3 · p. 21, Figure 9 · p. 22, Table 8 · p. 22

H.1 ZERO-SHOT EVALUATION DETAILS

The evaluation details specify six LM-eval zero-shot tasks and their accuracy metrics, along with eight LongBench tasks from four subgroups. LongBench evaluation uses a maximum sequence length of 31500 for Mistral and LongChat.

H.1 ZERO-SHOT EVALUATION DETAILS

  • Six LM-eval zero-shot tasks are evaluated using accuracy or accuracy normalized by sequence length, depending on the task.The tasks include OpenBookQA, HellaSwag, PIQA, ARC-Easy, ARC-Challenge, and WinoGrande.Source: §H.1 · p. 22

H.2 LONGBENCH EVALUATION DETAILS

  • LongBench evaluation covers eight tasks from four subgroups, with maximum sequence length set to 31500 for Mistral and LongChat.The listed tasks include QA, summarization, few-shot learning, and code-related evaluations.Source: §H.2 · pp. 22–23
Loading 2407.21118v2…