Accurate KV Cache Quantization with Outlier Tokens Tracing

TL;DR

This paper tackles the memory–accuracy trade-off in KV Cache quantization for large language models. It identifies a small subset of outlier tokens whose Keys in certain channels disrupt 2-bit quantization under existing channel-wise/token-wise schemes, and introduces Outlier Tokens Tracing (OTT) to dynamically trace and exclude these tokens from quantization during decoding. OTT maintains three KV caches (quantized, full-precision, and outlier) and uses an outlier pool that captures tokens based on Key magnitude, with group-based quantization occurring every $G$ tokens; decoding relies on a fused kernel to combine contributions efficiently. Across multiple LLMs and benchmarks, OTT yields significant improvements in accuracy over the KIVI baseline at 2-bit precision, achieving up to 6.4× memory reduction and 2.3× throughput gain, while maintaining close-to-FP16 performance on many tasks, thereby enabling faster and more memory-efficient inference for long-context scenarios.

Abstract

The impressive capabilities of Large Language Models (LLMs) come at the cost of substantial computational resources during deployment. While KV Cache can significantly reduce recomputation during inference, it also introduces additional memory overhead. KV Cache quantization presents a promising solution, striking a good balance between memory usage and accuracy. Previous research has shown that the Keys are distributed by channel, while the Values are distributed by token. Consequently, the common practice is to apply channel-wise quantization to the Keys and token-wise quantization to the Values. However, our further investigation reveals that a small subset of unusual tokens exhibit unique characteristics that deviate from this pattern, which can substantially impact quantization accuracy. To address this, we develop a simple yet effective method to identify these tokens accurately during the decoding process and exclude them from quantization as outlier tokens, significantly improving overall accuracy. Extensive experiments show that our method achieves significant accuracy improvements under 2-bit quantization and can deliver a 6.4 times reduction in memory usage and a 2.3 times increase in throughput.

Figures (10)

• Figure 1: Observations from preliminary experiments: (a) The Keys are distributed by channel and have some outlier channels. (b) The distribution of the Values does not exhibit any notable characteristics. (d) In certain outlier channels, a few tokens with low magnitude of Keys disrupt the originally uniform distribution within these channels. (e) Visualization of the sorted Keys in an outlier channel shows a rapid increase from a low value to very high values. (c) The L1 loss of attention output before and after quantization by retaining full-precision tokens based on different criteria. The best result is retaining full-precision tokens with the smallest magnitude of the Keys.
• Figure 2: Overview of OTT. Top: Decoding stage. Multiply the Query by each type of the Keys and concatenate the results to obtain the attention scores. Multiply the attention scores by each type of the Values and sum the results to get the attention output. Bottom: Quantization stage. Process the outlier tokens before quantization.
• Figure 3: Experiments on throughput and memory: (a) Comparison of throughput (tokens/s) for different methods across different batch sizes on NVIDIA A800 80G. (b) Peak memory usage (including model weights and other components) at different batch sizes on NVIDIA A800 80G. (c) Peak memory usage (including model weights and other components) at different sequence lengths when batch size = 1 on NVIDIA A100 40G. The results shows that OTT achieves a peak memory reduction of up to 6.4× and a throughput increase of 2.3×.
• Figure 4: The Keys in an outlier channel (up) and the magnitude of the Keys overall (down).
• Figure 5: Magnitude of the keys and Values for Llama-2-7B-chat-hf in head 17.