GEAR: An Efficient KV Cache Compression Recipe for Near-Lossless Generative Inference of LLM

Hao Kang; Qingru Zhang; Souvik Kundu; Geonhwa Jeong; Zaoxing Liu; Tushar Krishna; Tuo Zhao

Source-linked AI summary

GEAR: An Efficient KV Cache Compression Recipe for Near-Lossless Generative Inference of LLM

Hao Kang, Qingru Zhang, Souvik Kundu, Geonhwa Jeong, Zaoxing Liu, Tushar Krishna, Tuo Zhao

arXiv:2403.05527v4Mar 8, 2024cs.LGcs.AIcs.CL

TL;DR

Growing KV-cache memory use and accumulated approximation error constrain LLM generation efficiency and accuracy. GEAR combines quantization with low-rank residual approximation and sparse correction, achieving substantial accuracy, memory, and throughput improvements.

Problem
KV-cache memory grows rapidly with generation, while accumulated approximation error can degrade accuracy and divert model generations.
Method
GEAR augments KV-cache quantization with a low-rank matrix for quantization residuals and a sparse matrix for individual errors.
Results
GEAR improves average accuracy by 14.95% over the best-performing FP16 baseline, reduces peak memory by up to 2.39×, and delivers 2.10× ∼5.07× throughput.
Takeaways & Limitations
The reported results support GEAR as an efficient approach for reducing KV-cache approximation error while improving inference accuracy, memory use, and throughput.
Takeaways & Limitations
The three methods alone cannot achieve high compression ratios and cannot fully remediate the error efficiently.

Abstract

from arXiv · show

Key-value (KV) caching has become the de-facto to accelerate generation speed for large language models (LLMs) inference. However, the growing cache demand with increasing sequence length has transformed LLM inference to be a memory bound problem, significantly constraining the system throughput. Existing methods rely on dropping unimportant tokens or quantizing all entries uniformly. Such methods, however, often incur high approximation errors to represent the compressed matrices. The autoregressive decoding process further compounds the error of each step, resulting in critical deviation in model generation and deterioration of performance. To tackle this challenge, we propose GEAR, an efficient KV cache compression framework that achieves near-lossless high-ratio compression. GEAR first applies quantization to majority of entries of similar magnitudes to ultra-low precision. It then employs a low rank matrix to approximate the quantization error, and a sparse matrix to remedy individual errors from outlier entries. By adeptly integrating three techniques, GEAR is able to fully exploit their synergistic potentials. Our experiments demonstrate that compared to alternatives, GEAR achieves near-lossless 4-bit KV cache compression with up to 2.38x throughput improvement, while reducing peak-memory size up to 2.29x. Our code is publicly available at https://github.com/HaoKang-Timmy/GEAR.

1 Introduction

KV-cache memory grows with model size and sequence length, constraining throughput, while aggressive compression can introduce errors that accumulate during autoregressive generation. GEAR combines quantization, low-rank approximation, and sparse correction to reduce these errors and support efficient high-ratio compression.

Motivation: KV-cache memory grows rapidly with model size and sequence length, imposing significant constraints on inference throughput.A 30-billion-parameter model with length 1024 and batch size 128 can require up to 180 GB for its KV cache.
Motivation: Offloading KV caches to CPU memory or NVMe can introduce non-trivial overhead because of limited PCIe bandwidth.
Error accumulation: Approximation errors from compression can compound across autoregressive decoding steps, magnifying their effect on subsequent generation.The paper links this accumulation to degraded accuracy in GSM8k experiments.
GEAR: GEAR quantizes most similar-magnitude entries to ultra-low precision, uses a low-rank matrix for residual error, and uses a sparse matrix for outlier errors.The low-rank component captures coherent error structure, while the sparse component corrects incoherent individual outliers.
Results: 14.95% average accuracy improvement over the best-performing baseline is achieved at 2-bit KV quantization across various models and datasets.
Results: GEAR reduces peak memory up to 2.39× and delivers 2.10× ∼5.07× throughput compared with the FP16 baseline.

2 Background

Autoregressive decoding appends new Key and Value vectors to cached tensors and reuses the cache for attention. Background methods include group-wise quantization, while per-channel Key and per-token Value quantization reduce error but can add storage overhead.

Autoregressive decoding: During autoregressive decoding, each step appends the new Key and Value vectors to the existing KV cache before attention.
Group-wise Quantization: Group-wise quantization partitions tensor entries into groups and quantizes each group using its scale and zero-point.The group size controls how many consecutive entries are quantized together.
KV quantization: Key cache per-channel and Value cache per-token quantization targets large-magnitude Key channels and achieves state-of-the-art 2-bit compression.
KV quantization: Smaller quantization groups improve approximation accuracy but increase memory overhead from additional scaling factors and zero-points.
KCVT: KCVT uses coarse-grained per-vector grouping for per-channel Key and per-token Value quantization, reducing scaling and zero-point storage overhead.
GEAR motivation: GEAR is orthogonal to off-the-shelf quantization methods and can augment them to achieve near-lossless accuracy.Its motivation is supported by the observed rapid decay of the residual spectrum.

3 GEAR Framework

GEAR combines quantization, low-rank approximation, and sparse outlier correction to reduce KV-cache approximation error efficiently. A streaming buffer further limits decoding overhead.

GEAR Framework: GEAR decomposes each KV cache into a quantized backbone, a low-rank residual approximation, and a sparse matrix for individual outliers.The three components target entry-wise similarity, coherent residual structure, and sparse outlier information, respectively.
Motivation: High compression ratios make standalone quantization, token dropping, and sparse correction increasingly inaccurate or memory-intensive.Sparse storage can require two index vectors and one full-precision value vector, while iterative optimization introduces unacceptable latency.
Outlier-aware quantization: Outlier extraction stores the largest-magnitude entries sparsely before quantization, complementing the quantized backbone.The method filters per-vector outliers and stores them in full precision with a sparse matrix.
Low-rank approximation: Head-wise low-rank approximation captures coherent residual information because residual spectra drop rapidly at the beginning.GEAR uses a small rank; r = 4 is reported as sufficient for near-lossless high-ratio compression.
Design variants: GEAR recommends all three components for best 4-bit and 2-bit performance, while GEAR-L uses only low-rank approximation with lower memory overhead.GEAR and GEAR-L are described as orthogonal to existing quantization schemes and usable in a plug-and-play manner.
Streaming buffer: During decoding, a small full-precision buffer batches newly generated tokens before compression and applies low-rank approximation only to those tokens.The example buffer size is nb = 20, reducing repeated processing during long-sequence generation.

4 Experiments

Experiments evaluate GEAR as a plug-and-play KV-cache compressor across LLMs and generative tasks at 4-bit and 2-bit precision. It preserves or improves accuracy while reducing memory and increasing throughput.

Additional task results: 7.42% higher accuracy than KIVI is achieved by GEAR on 5-shot GSM8k with LLaMA3-8B at 2-bit compression.GEAR reaches 49.96% accuracy in this setting.
Efficiency: Throughput improves by up to 5.07× over FP16, while low-rank and sparse components introduce no unacceptable overhead.The primary complexity remains the model forward pass, and GEAR-L slightly exceeds KIVI throughput.

Appendix 11.

Ablations show that low-rank approximation is central to GEAR’s high-ratio accuracy, while small sparse and rank settings are generally sufficient. Applying error reduction broadly across tokens also matters.

Component ablation: Removing the low-rank component significantly degrades GEAR and GEAR-L performance, confirming its vital role in error reduction.The ablation identifies low-rank approximation as pivotal for near-lossless high-ratio compression.
Role of low-rank approximation: Outlier extraction alone improves quantized performance but does not reach GEAR’s near-lossless 2-bit performance.The comparison uses outlier-aware KIVI quantization under the same evaluation setting.
Token coverage: Applying low-rank error reduction to fewer tokens causes GEAR-L performance to deteriorate.The study varies coverage across input and generated tokens and reports degradation as fewer tokens receive correction.
Different compression ratios: GEAR and GEAR-L consistently outperform other quantization baselines across compression ratios on LLaMA3-8B GSM8k-CoT.The comparison reports near-lossless accuracy across the tested remaining-cache sizes.

5 Related Work

Prior KV-cache compression methods include weight-inspired quantization, activation/KV quantization, and attention-based token pruning. Their applicability is limited by KV-cache sensitivity, online changes, dense attention in reasoning tasks, and deployment overhead.

LLM weights compression: Weight-compression methods can reach 2–3 bits with acceptable accuracy loss but often require latency overhead and gradient information unsuitable for KV caches.KV caches have no trainable parameters and change at every generation stage, requiring lightweight online compression.
LLM KV cache compression: SmoothQuant achieves 8-bit activation and weight compression by adjusting scaling factors to reduce outlier error, with near-lossless results on simple generative tasks.
LLM KV cache compression: Atom compresses KV caches to 4 bits within 5% performance degradation on simple generative tasks using combined 4-bit and 8-bit channel-wise quantization.
LLM KV cache compression: H2O and FastGen reduce KV-cache size by dropping tokens according to attention scores, while SparQ also incorporates pruned value-cache error.
LLM KV cache compression: Token dropping is problematic for fine-tuned, chain-of-thought, and generative-reasoning settings because attention is denser, tokens carry important information, and scoring complicates FlashAttention deployment.
LLM KV cache compression: Attention sparsity also depends on the model’s non-linearity choice, making it vulnerable as a metric for KV compression.

6 Discussion and Conclusions

GEAR is presented as a plug-and-play error-reduction framework that augments KV-cache quantization with lightweight techniques for high-ratio compression. It reports improved accuracy and inference efficiency, while identifying fixed rank allocation as a limitation.

Discussion and Conclusions: GEAR augments any off-the-shelf KV-cache quantization scheme with two lightweight error-reduction techniques in a plug-and-play manner.
Discussion and Conclusions: GEAR demonstrates state-of-the-art performance on complex reasoning tasks, with a 14.95% average accuracy improvement at 2-bit KV quantization compared with alternatives.
Discussion and Conclusions: GEAR reduces peak memory compared with the FP16 baseline and brings up to approximately 5.07× throughput improvement.
Limitations and Broader Societal Impact: GEAR’s low-rank error approximation uses an identical rank for every Key/Value matrix, despite differing importance across layers and heads.
Limitations and Broader Societal Impact: Adaptive allocation of the low-rank approximation budget across Key and Value matrices can further improve GEAR performance.
Limitations and Broader Societal Impact: The authors associate GEAR with reduced inference energy and carbon footprint and more efficient deployment across hardware platforms.

8 Power Iteration Algorithm as SVDSolver

The section identifies the power iteration algorithm as the procedure presented in Algorithm 2.

Power Iteration Algorithm as SVDSolver: The power iteration algorithm is presented in Algorithm 2.

9 More Discussion on Related Works

Related work covers compression of weights, activations, and KV caches using quantization, outlier handling, token pruning, and low-rank error approximation. KV-cache methods must account for input sensitivity, changing cache contents, and the difficulty of preserving information in reasoning tasks.

LLM weights compression: Weight-compression methods can achieve 2–3 bits with acceptable accuracy loss but often require latency overhead and gradient information unsuitable for KV-cache compression.
LLM KV cache compression: KV caches change at every generation stage and therefore require efficient, lightweight methods for online compression.
LLM KV cache compression: Activation and KV-cache compression are harder than weight compression because they are more sensitive and related to model inputs.
LLM KV cache compression: SmoothQuant achieves 8-bit activation and weight compression by adjusting scaling factors to reduce outlier error, with near-lossless performance on simple generative tasks.
LLM KV cache compression: Atom compresses KV caches to 4 bits within 5% performance degradation on simple generative tasks through combined 4-bit and 8-bit channel-wise quantization.
LLM KV cache compression: H2O and FastGen prune KV caches by dropping tokens according to attention scores, while SparQ also incorporates pruned value-cache error.
LLM KV cache compression: Token-dropping methods face difficulty in fine-tuned, chain-of-thought, and generative-reasoning settings because attention is denser and tokens contain important information.
LLM KV cache compression: Attention-score token dropping is hard to deploy with FlashAttention, and attention sparsity varies with the model’s non-linearity choice.

10 Dataset Statistics

The paper reports dataset statistics covering input length during prefill, generation length, and the number of evaluation examples.

Dataset statistics include input length in the prefill phase and generation length.
These measurements characterize the datasets used for evaluation.
The statistics also report the number of evaluation examples.

11 More Inference Analysis Comparison

The paper evaluates GEAR’s inference efficiency, memory use, sequence-length capacity, and comparisons with alternative KV-cache compression approaches across several models and datasets. Results include up to 2.39× memory reduction, 2.10× throughput improvement, and roughly 2k longer maximum sequences in the reported settings.

Inference efficiency: 2.39× memory reduction is achieved by GEAR compared to the FP16 KV cache baseline.Peak memory is measured under the same batch size, with batch size 18 used for both systems.
Error reduction: GEAR-L degenerates when fewer tokens receive error reduction, supporting the importance of applying low-rank error reduction broadly.GEAR-L is the lite version that applies only low-rank approximation, and GEAR-L Prefill applies error reduction to prefill tokens.
Inference efficiency: 2.10× throughput improvement is achieved by GEAR-L Prefill, GEAR-L, and GEAR compared to the FP16 baseline.The comparison uses LLaMA2-7B on an RTX Titan GPU with 24GB memory.
KV-cache composition: GEAR and GEAR-L add overheads from sparsity and low-rank components, while quantization backbones also use scales, zero points, and residual-token buffers.KCVT has smaller streaming-buffer overheads, whereas KIVI has larger residual and scaling-factor overheads because of its smaller group size.
Maximum sequence length: GEAR increases maximum sequence length by around 2k, making previously impossible long-sequence generation feasible.The comparison uses LLaMA2-7B with batch size 1, 8-bit model weights, and FlashAttention.
Comparison with token dropping: Token dropping cannot effectively preserve performance or achieve a high compression ratio on GSM8K with LLaMA2-7B when dropping 50% of tokens.The paper links this limitation to reasoning tasks requiring close attention to most contextual information.