KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache
About
Efficiently serving large language models (LLMs) requires batching of many requests to reduce the cost per request. Yet, with larger batch sizes and longer context lengths, the key-value (KV) cache, which stores attention keys and values to avoid re-computations, significantly increases memory demands and becomes the new bottleneck in speed and memory usage. Additionally, the loading of the KV cache causes the computational core to be idle, which limits the inference speed. A straightforward and effective solution to reduce KV cache size is quantization, which decreases the total bytes taken by KV cache. However, there is a lack of in-depth studies that explore the element distribution of KV cache to understand the hardness and limitation of KV cache quantization. To fill the gap, we conducted a comprehensive study on the element distribution in KV cache of popular LLMs. Our findings indicate that the key cache should be quantized per-channel, i.e., group elements along the channel dimension and quantize them together. In contrast, the value cache should be quantized per-token. From this analysis, we developed a tuning-free 2bit KV cache quantization algorithm named KIVI. With hardware-friendly implementation, KIVI can enable Llama, Falcon, and Mistral models to maintain almost the same quality while using $\mathbf{2.6\times}$ less peak memory (including model weight). This reduction in memory usage enables up to $\mathbf{4\times}$ larger batch size, bringing $\mathbf{2.35\times \sim 3.47\times}$ throughput on real LLM inference workload. The source code is available at https://github.com/jy-yuan/KIVI.
Key insights
What problem does KIVI solve and how?
KV cache memory bottlenecks batched LLM inference. KIVI reduces this by quantizing key cache per-channel and value cache per-token to 2-bit, using a sliding window of full-precision residual tokens to preserve accuracy.
How does KIVI differ from prior KV cache quantization?
KIVI quantizes keys and values differently. For Llama-2-13B, Table 3 reports a CoQA score of 66.23 with KIVI-2 versus 52.93 with the tested 2-bit per-token scheme. KIVI also retains a small full-precision residual cache.
What are the practical benefits of KIVI?
With Llama-2-7B on one NVIDIA A100 80GB GPU, KIVI supports up to four times the batch size and 2.35–3.47 times the throughput of the 16-bit baseline on synthetic ShareGPT workloads. Gains depend on the tested configuration.
What limitations and future work do the authors note?
Falcon-7B requires 4-bit KIVI to maintain accuracy due to its multi-query attention. The authors state future work will further optimize the implementation to reduce quantization overhead during prefill and decoding phases.