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

TL;DR

RateQuant introduces a novel rate-distortion theory-based method for mixed-precision quantization of KV caches in large language models, significantly reducing memory usage without performance loss.

Contribution

It proposes RateQuant, a calibration-based approach that accurately models distortion curves for mixed-precision quantization, overcoming previous mismatch issues.

Findings

Reduces perplexity from 49.3 to 14.9 at 2.5 bits.

Achieves 70% reduction in memory cost.

Calibration takes only 1.6 seconds on a single GPU.

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.