KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization

By Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W. Mahoney, Yakun Sophia Shao, Kurt Keutzer, A. Gholami. Neural Information Processing Systems, 2024

589 citations recorded by the source corpus.

LLMs are seeing growing use for applications which require large context windows, and with these large context windows KV cache activations surface as the dominant contributor to memory consumption during inference. Quantization is a promising approach for compressing KV cache activations; however, existing solutions fail to represent activations accurately in sub-4-bit precision. Our work, KVQuant, facilitates low precision KV cache quantization by incorporating several novel methods: (i) Per-Channel Key Quantization, where we adjust the dimension along which we quantize the Key activations to better match the distribution; (ii) Pre-RoPE Key Quantization, where we quantize Key activations before the rotary positional embedding to mitigate its impact on quantization; (iii) Non-Uniform KV Cache Quantization, where we derive per-layer sensitivity-weighted non-uniform datatypes that better represent the distributions; and (iv) Per-Vector Dense-and-Sparse Quantization, where we isolate outliers separately for each vector to minimize skews in quantization ranges. By applying our method to the LLaMA, Llama-2, Llama-3, and Mistral models, we achieve<0.1 perplexity degradation with 3-bit quantization on both Wikitext-2 and C4, outperforming existing approaches. Our method enables serving LLaMA-7B with a context length of up to 1 million on a single A100-80GB GPU and up to 10 million on an 8-GPU system. We develop custom CUDA kernels for KVQuant, showing that we can achieve up to ~1.7x speedups, compared to baseline fp16 matrix-vector multiplications, for the LLaMA-7B model.

Subfields: Artificial Intelligence, Information Systems

View Paper: https://www.semanticscholar.org/paper/b085968c4362fb286ad6c5ef71a5db9630da0498

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Damai Dai
2. Zhihong Shao
3. Huajian Xin
4. Zihan Wang
5. Daya Guo
6. Tri Dao

In Information Systems

1. Zhihong Shao
2. Daya Guo
3. Tri Dao

Citations in the corpus

• QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs (2024, unknown): “However, key vectors in the attention module are also known to suffer from outliers [Hooper et al., 2024, Liu et al., 2024].”
• KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache (2024, unknown): “A similar observation and approach has been independently discovered and developed in the concurrent work KVQuant (Hooper et al., 2024). vLLM (Kwon et al., 2023) and S3 (Jin et al., 2023) are system-level works, which include memory management through the use of PagedAttention or memory usage…”
• FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (2024, background): “Recent work suggests that for inference the KV cache is highly compressible down to 4, 3-, or even 2-bits [26, 33].”
• SGLang: Efficient Execution of Structured Language Model Programs (2023, unknown): “Additionally, SGLang is compatible with many other common inference optimizations [60, 39, 3, 23, 59, 10, 26, 15, 19, 32, 31, 11].”
• DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model (2024, methodology): “In addition, we also perform KV cache quantiza-15 tion (Hooper et al., 2024; Zhao et al., 2023) for DeepSeek-V2 to further compress each element in its KV cache into 6 bits on average.”