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

2024-01-31

Neural Information Processing Systems

AI Summary

To address the KV cache memory explosion in large-context LLM inference, this work proposes a high-fidelity sub-4-bit KV quantization method. We introduce four novel techniques: channel-wise key quantization, RoPE-pre-quantization, layer-aware non-uniform quantization, and vector-level dense-sparse hybrid quantization—collectively overcoming the accuracy bottleneck of KV activations under ultra-low-bit (≤3-bit) quantization. Our custom CUDA kernels support LLaMA and Mistral model families. Experiments show that 3-bit quantization incurs <0.1 perplexity degradation; a single A100-80GB GPU sustains million-token context for LLaMA-7B, scaling to ten million tokens on an 8-GPU system, with ~1.7× inference speedup. This is the first practical, low-bit KV compression solution enabling efficient deployment of ultra-long-context LLMs.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Computer Vision: Large Vision Models
• Natural Language Processing: (Large) Language Models

Application Category

• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

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.

Problem

• Reduce KV cache memory consumption in large-context LLMs
• Enable accurate sub-4-bit KV cache quantization
• Achieve efficient million-token context inference

Innovation

• Per-Channel Key Quantization for better distribution matching
• Pre-RoPE Key Quantization to reduce embedding impact
• Non-Uniform KV Cache Quantization for sensitive layers