Source-linked AI summary
ZipCache: Accurate and Efficient KV Cache Quantization with Salient Token Identification
Yefei He, Luoming Zhang, Weijia Wu, Jing Liu, Hong Zhou, Bohan Zhuang
TL;DR
Long-context KV caches create a memory bottleneck, and existing adaptive approaches can misidentify salient tokens while requiring costly full attention matrices. ZipCache combines channel-separable quantization with normalized-attention saliency and an efficient approximation. Experiments report state-of-the-art compression ratio, accuracy, and generation efficiency, while the saliency ratio remains manually specified.
Problem
KV caches consume substantial memory as batch size and context length grow, while accumulated-attention saliency can be inaccurate and requires full attention matrices.
Source: §1 · p. 2Method
ZipCache combines channel-separable quantization with normalized-attention saliency and an efficient approximation compatible with fast attention implementations.
Source: §1 · pp. 2–3, §4.3 · p. 6Results
Experiments show ZipCache achieves state-of-the-art performance in compression ratio, accuracy, and generation efficiency.
Source: §1 · p. 3, §6 · p. 10Takeaways & Limitations
The paper concludes that adaptive quantization based on token saliency can improve compression ratios without sacrificing model performance.
Source: §6 · p. 10The saliency ratio is manually specified before evaluation and cannot be automatically adjusted based on task datasets.
Source: §6 · p. 10
Abstract
from arXiv · showhide
KV cache stores key and value states from previous tokens to avoid re-computation, yet it demands substantial storage space, especially for long sequences. Adaptive KV cache compression seeks to discern the saliency of tokens, preserving vital information while aggressively compressing those of less importance. However, previous methods of this approach exhibit significant performance degradation at high compression ratios due to inaccuracies in identifying salient tokens. In this paper, we present ZipCache, an accurate and efficient KV cache quantization method for LLMs. First, we construct a strong baseline for quantizing KV cache. Through the proposed channel-separable tokenwise quantization scheme, the memory overhead of quantization parameters are substantially reduced compared to fine-grained groupwise quantization. To enhance the compression ratio, we propose normalized attention score as an effective metric for identifying salient tokens by considering the lower triangle characteristics of the attention matrix. Moreover, we develop an efficient approximation method that decouples the saliency metric from full attention scores, enabling compatibility with fast attention implementations like FlashAttention. Extensive experiments demonstrate that ZipCache achieves superior compression ratios, fast generation speed and minimal performance losses compared with previous KV cache compression methods. For instance, when evaluating Mistral-7B model on GSM8k dataset, ZipCache is capable of compressing the KV cache by $4.98\times$, with only a $0.38\%$ drop in accuracy. In terms of efficiency, ZipCache also showcases a $37.3\%$ reduction in prefill-phase latency, a $56.9\%$ reduction in decoding-phase latency, and a $19.8\%$ reduction in GPU memory usage when evaluating LLaMA3-8B model with a input length of $4096$.
1 Introduction
Growing KV caches create a substantial memory bottleneck, while existing adaptive methods can misidentify salient tokens and require costly full attention matrices. ZipCache combines lower-overhead quantization with normalized-attention saliency and an efficient approximation compatible with fast attention.
1 Introduction
- For a 175B LLM at batch size 64 and context length 4096, KV cache can occupy 1.2TB, compared with 350GB for model weights.KV-cache size grows with batch size and input context length.Source: §1 · p. 2
- Accumulated attention scores can bias saliency toward earlier tokens and require storing full attention matrices, using O(l^2) memory versus O(l) for fast attention.The lower-triangular attention matrix gives earlier tokens more scores to accumulate.Source: §1 · p. 2
- ZipCache combines a channel-separable quantization baseline, normalized-attention saliency, and an efficient probe-token approximation compatible with fast attention.Its contributions aim to reduce quantization overhead, improve salient-token identification, and enhance generation speed.Source: §1 · pp. 2–3
- On LLaMA3-8B Line Retrieval, ZipCache achieves the highest accuracy, generation speed, and compression ratio among the compared methods.Figure 1 presents the comparison; details are in the supplementary material.Source: Figure 1 · p. 2
2 Related Work
The related work describes quantization granularity and token importance as central trade-offs in KV-cache compression. Prior approaches either discard tokens or quantize them, while fine-grained quantization and accumulated-score saliency have limitations.
2.1 Model Quantization
- Finer quantization granularity can improve performance but adds quantization parameters and memory overhead; mixed precision assigns bit-widths according to component sensitivity.The paper identifies differing token importance as motivation for mixed-precision KV-cache quantization.Source: §2.1 · p. 3
2.2 KV Cache Compression
- KV-cache compression includes token dropping and quantization; dropping erases information and can be suboptimal for retrieval, while outlier handling in quantization adds overhead.The paper proposes channel-separable quantization to reduce overhead and a more accurate saliency metric than accumulated attention scores.Source: §2.2 · p. 3
- Prior methods commonly use accumulated attention scores for token importance, which the authors find inaccurate and capable of degrading performance at low bit-widths.ZipCache instead uses a more accurate metric to identify salient tokens.Source: §2.2 · p. 3
3 Preliminary
The preliminary section introduces the attention and quantization setup used in the paper. It describes prefill and decoding, KV-cache updates, and uniform quantization with full-precision parameters.
3.1 Attention Block in LLMs
- During prefill, the model computes and stores KV cache for prompt tokens; during decoding, it generates tokens and updates the cached key and value matrices.The setup defines batch size b, attention heads h, sequence length l, and head dimension d.Source: §3.1 · p. 4
- The attention output is computed from attention scores, and the key dimension and attention scores are denoted d_k and A, respectively.The output projection is omitted in this formulation.Source: §3.1 · p. 4
3.2 Model Quantization
- The paper uses uniform quantization of floating-point vectors to k bits, while full-precision quantization parameters can create substantial overhead when quantization is fine-grained.This motivates reducing the number of stored quantization parameters.Source: §3.2 · p. 4
4 Method
ZipCache combines a reduced-overhead quantization baseline with normalized-attention saliency and a probe-based approximation compatible with fast attention implementations.
4.1 A Strong Baseline for KV Cache Quantization
- Channel-separable tokenwise quantization disentangles channel and token dimensions, reducing quantization-parameter overhead while balancing channel outliers and token representations.It normalizes each channel before tokenwise quantization and restores its magnitude afterward; the scheme uses hd + 2bl quantization parameters.Source: §4.1 · pp. 4–5
- Channelwise key quantization plus channel-separable tokenwise value quantization outperforms groupwise quantization with reduced overhead, establishing a strong baseline.The comparison uses 4-bit KV cache quantization on LLaMA3-8B and GSM8k.Source: §4.1 · p. 5, Table 1 · p. 5
4.2 Accurate Salient Token Identification
- Normalized attention scores mitigate initial-row bias in accumulated scores; in a GSM8k example, they identify question tokens as salient where accumulated scores mark them low-saliency.The example uses LLaMA3-8B with chain-of-thought prompting.Source: §4.2 · p. 6, Figure 3 · p. 6
4.3 Efficient Approximation of Saliency Metric
- Probe-token attention scores approximate saliency for all tokens, avoiding explicit score computation for non-probes and enabling fast attention implementations such as FlashAttention.Figure 4 contrasts probe-based computation with full attention scores and tiled FlashAttention computation.Source: §4.3 · p. 7, Figure 4 · p. 7
- Among four probe strategies, the recent-plus-random hybrid is most effective; ZipCache uses 5% recent and 5% random probes by default.The strategy comparison uses a 10% probe-token proportion, with salient tokens quantized to 4-bit and remaining tokens to 2-bit.Source: §4.3 · p. 7, Table 2 · p. 7
5 Experiment
ZipCache is evaluated on math, code, and retrieval tasks, with mixed-precision KV quantization and streaming compression. The reported results show improved accuracy over MiKV on GSM8k and Line Retrieval, alongside lower latency and memory use.
5.1 Implementation Details
- Experiments use Mistral, LLaMA2, and LLaMA3 on GSM8k, HumanEval, and Line Retrieval, with salient tokens at 4-bit and others at 2-bit.Keys use channelwise quantization, values use channel-separable tokenwise quantization, and compression repeats every 100 generated tokens; HumanEval results are in the supplementary material.Source: §5.1 · pp. 7–8, §5.2.2 · p. 9
5.2.1 Evaluation on GSM8k
- 18.27% accuracy enhancement over MiKV on LLaMA3-8B is reported on GSM8k with normalized attention scores.The reported comparison also notes a 20.4% accuracy decrease for MiKV under high compression.Source: Table 3 · p. 8, §5.2.1 · p. 8
5.2.2 Evaluation on Line Retrival
- 42% accuracy improvement over MiKV on Mistral-7b at 200 lines demonstrates ZipCache’s stronger Line Retrieval performance.The evaluation compares retrieval accuracy across different numbers of lines.Source: Figure 5 · p. 9, §5.2.2 · p. 9
5.3 Generation Efficiency
- At an input length of 4096, ZipCache reduces prefill latency by 37.3%, decoding latency by 56.9%, and GPU memory usage by 19.8% versus MiKV.The comparison serves LLaMA3-8B on an Nvidia A100; ZipCache computes attention scores for 10% of tokens and uses FlashAttention or FlashDecoding for the rest.Source: Figure 6 · p. 9, §5.3 · p. 9
6 Conclusion and Future Work
ZipCache combines lower-overhead KV cache quantization, normalized-attention-based token saliency, and an efficient approximation compatible with fast attention implementations. The paper reports strong compression performance, accuracy, and generation speed, while noting that its saliency ratio is manually specified.
6 Conclusion and Future Work
- ZipCache combines channel-separable quantization, normalized-attention saliency, and an efficient approximation compatible with FlashAttention and FlashDecoding.The paper reports state-of-the-art compression performance in compression ratio, accuracy, and generation speed.Source: §6 · p. 10
- The saliency ratio is manually specified before evaluation and cannot be automatically adjusted for task datasets.The paper also notes potential use of ZipCache to generate malicious content.Source: §6 · p. 10
A Calculation of Overhead for Different Quantization Schemes
Under the stated assumptions, the appendix calculates compression ratios for groupwise, tokenwise, and proposed-baseline quantization. The proposed baseline’s calculated ratio is marginally higher than the tokenwise ratio.
A Calculation of Overhead for Different Quantization Schemes
- 3.200 is the calculated compression ratio for groupwise quantization with group size n = 32.The calculation assumes b = 8 and hd = l = 4096, with the KV cache quantized to 4-bit.Source: A Calculation of Overhead for… · p. 13
- 3.992 is the calculated compression ratio for tokenwise quantization under the stated assumptions.Those assumptions are b = 8, hd = l = 4096, and a 4-bit KV cache.Source: A Calculation of Overhead for… · p. 13
- 3.995 is the calculated compression ratio for the proposed quantization baseline under the stated assumptions.The appendix calculates this ratio with b = 8, hd = l = 4096, and a 4-bit KV cache.Source: A Calculation of Overhead for… · p. 13
B Implementation Details of ZipCache
The implementation appendix describes channel-separable tokenwise quantization and ZipCache’s prefill and decoding procedures. It also explains how the attention computation avoids splitting and merging outputs.
B Implementation Details of ZipCache
- During prefill and decoding, FlashAttention computes outputs for all tokens while attention scores are calculated for probe tokens.This avoids matrix-splitting and merging memory accesses, which the paper says enhances generation speed.Source: B Implementation Details of ZipCache · p. 13
- CSTQuant normalizes each channel of the input, then applies tokenwise quantization using the target bit-width.The algorithm is named Channel-separable Tokenwise Quantization.Source: B Implementation Details of ZipCache · p. 13
C.1 Accuracy and Efficiency Comparisons of various KV cache compression methods
On LLaMA3-8B retrieval, ZipCache reports the highest accuracy, compression ratio, and generation speed among the compared methods. Its prefill and decoding procedures use probe attention scores to identify salient tokens and adaptively quantize the cache.
C.1 Accuracy and Efficiency Comparisons of various KV cache compression methods
- ZipCache improves accuracy by 10.0% and decreases prefill latency by 38.0% versus MiKV on LLaMA3-8B retrieval.The paper attributes the accuracy improvement to more accurate salient-token identification and the latency reduction to FlashAttention integration.Source: §C.1 · p. 13
- During prefill, ZipCache estimates token saliency from probe attention scores, splits keys and values into salient and regular sets, then quantizes them at different bit-widths.The salient and regular groups receive kh and kl bit-widths, respectively.Source: §C.1 · p. 14
- During decoding, ZipCache periodically recompresses cached tokens using accumulated probe scores; probes include recent and randomly selected tokens.The algorithm re-compresses every 100 tokens and uses probe tokens consisting of 5% recent and 5% random tokens.Source: §C.1 · p. 14
C.2 Evaluation on HumanEval
On HumanEval, ZipCache retains performance at high compression on Mistral-7B and outperforms KIVI-2 on LLaMA3-8B. The comparison also notes that KIVI’s full-precision recent-token reserve reduces compression on short inputs.
C.2 Evaluation on HumanEval
- ZipCache outperforms KIVI-2 by 7.32% on LLaMA3-8B, with compression ratios of 4.39× versus 2.55×.Source: §C.2 · p. 14
- With an average input length of 119 tokens, KIVI’s practice of retaining 32 recent tokens in full precision considerably diminishes its compression ratio.Source: §C.2 · p. 14
- 29.27% accuracy at 4.94× compression matches FP16 on Mistral-7B HumanEval.The table reports 29.27 Acc.(%) for both ZipCache and FP16.Source: §C.2 · pp. 14–15