KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache
Liu, Yuan, Liu et al. · 2024 · ICML 2024
Proposes asymmetric quantization of the KV cache to 2-bit precision, reducing memory usage by 2.6x with negligible quality loss. Key insight: keys have per-channel outliers (quantize per-channel), while values have per-token outliers (quantize per-token).
Key Idea
During LLM inference, the KV cache grows linearly with sequence length and batch size, often becoming the primary memory bottleneck. KIVI observes that keys and values have fundamentally different outlier distributions: keys have large-magnitude outlier channels (columns) that persist across all tokens, while values have large-magnitude outlier tokens (rows) that affect all channels. This asymmetry motivates different quantization strategies for K and V.
How It Works
- Key quantization: Per-channel (across the token dimension). Each channel gets its own scale/zero-point, catching the persistent outlier channels. Quantized to 2-bit with group size.
- Value quantization: Per-token (across the channel dimension). Each token gets its own scale/zero-point, catching the sporadic outlier tokens. Quantized to 2-bit.
- Residual length: The most recent few KV pairs (e.g., 128 tokens) are kept in FP16 to preserve quality for recent attention, since quantization noise matters more for tokens currently being attended to.
- Tuning-free: No retraining or calibration needed, just apply the quantization scheme at inference time.
Why It Matters
- 2.6x KV cache compression at 2-bit with minimal quality loss (< 0.1 perplexity degradation on LLaMA-2-7B)
- Enables longer contexts: A 128K context window that would require 32GB of KV cache at FP16 fits in ~12GB at 2-bit
- Orthogonal to model quantization: KIVI compresses the cache, not the model weights, it stacks with weight quantization (GPTQ, AWQ) for compounding savings
- Production-ready: vLLM, TensorRT-LLM, and other serving frameworks have adopted KV cache quantization inspired by this work
- Practical impact: Directly increases the max batch size and sequence length you can serve on a given GPU
Key Takeaways for Interviews
- Why KV cache is the bottleneck: for a 7B model serving 4K context at batch 32, KV cache alone is ~16GB (larger than model weights in INT4)
- Asymmetric insight: K has per-channel outliers → quantize per-channel; V has per-token outliers → quantize per-token. Getting this backwards destroys quality.
- Residual tokens: keep recent ~128 tokens in full precision, attention to recent context is more sensitive to quantization noise
- Pairs with PagedAttention: KIVI reduces per-block memory, PagedAttention reduces fragmentation, both are needed for production LLM serving
- In system design: cite when asked about serving LLMs at scale, long-context support, or reducing per-request memory cost