RateQuant: Optimal Mixed-Precision KV Cache Quantization via Rate-Distortion Theory

2026-04-22

AI Summary

In large language model inference, the KV cache grows linearly with sequence length, becoming a primary memory bottleneck. Existing quantization methods apply uniform bit-widths across all attention heads, ignoring their varying importance; naively adopting mixed precision often degrades performance due to mismatched distortion models. This work identifies this issue and proposes RateQuant, a framework grounded in rate-distortion theory that fits per-head distortion models using minimal calibration data and computes the theoretically optimal bit allocation via a reverse water-filling algorithm. On Qwen3-8B, RateQuant achieves substantial compression—reducing KIVI’s perplexity from 49.3 to 14.9 (a 70% drop)—using only 2.5 average bits and 1.6 seconds of calibration time, outperforming QuaRot by 6.6 PPL without introducing any inference overhead.

Abstract

Large language models cache all previously computed key-value (KV) pairs during generation, and this KV cache grows linearly with sequence length, making it a primary memory bottleneck for serving. Quantizing the KV cache to fewer bits reduces this cost, yet all current quantizers assign the same bit-width to every attention head, ignoring the large variation in head importance. A natural idea is to allocate more bits to important heads and fewer to the rest. We show, however, that such mixed-precision allocation has a hidden pitfall: each quantizer follows a different distortion curve D(b)=alpha*beta^{-b}, and the decay rate beta varies from 3.6 to 5.3 across quantizer designs. Applying one quantizer's distortion model to another inverts the allocation order and makes performance worse than uniform quantization. We call this failure mode distortion model mismatch and propose RateQuant to resolve it. RateQuant fits a per-quantizer distortion model from a small calibration set, then solves the resulting bit-allocation problem in closed form via reverse waterfilling from rate-distortion theory. On Qwen3-8B at 2.5 average bits, calibrated RateQuant reduces KIVI's perplexity from 49.3 to 14.9 (70% reduction) and improves QuaRot by 6.6 PPL. The entire calibration takes 1.6 s on a single GPU and adds zero overhead at inference time.

Problem

• KV cache quantization
• mixed-precision allocation
• distortion model mismatch
• memory bottleneck
• attention heads

Innovation

• mixed-precision quantization
• KV cache
• rate-distortion theory
• distortion model mismatch
• reverse waterfilling