Spend Bits Where Queries Look: KV Cache Vector Quantization with Attention-Preserving Transforms

2026-08-04

AI Summary

This work addresses the bandwidth bottleneck in long-context decoding of large language models caused by KV cache loading, aiming to compress the cache efficiently under a fixed bit rate while preserving attention accuracy. The authors formulate KV cache quantization as an attention-preserving transform coding problem and propose, for the first time, a non-orthogonal optimal linear transform derived from an attention distortion criterion. They establish a generalized Parseval relationship that reformulates attention fidelity as mean squared error minimization in the transform domain. By integrating calibration data-driven statistical modeling, high-resolution vector quantization, and equal-volume coefficient grouping, the method achieves rate-optimal compression under a fixed bitrate constraint. Experiments demonstrate that the proposed NOVA-KV approach substantially recovers the long-context retrieval accuracy degraded by scalar quantization at 2 bits per element, while maintaining high-throughput decoding performance.

Abstract

Long-context LLM decoding reads the key-value (KV) cache at every step. Loading it takes longer than computing attention over it, so throughput is bandwidth-bound. Hence, reducing the cache size can raise both decoding speed and serving capacity. The challenge is to reduce cache size while preserving the attention products, keeping reconstruction cheap, and using a fixed per-token bit count. At two bits per element, the most competitive methods rely on orthogonal transforms. However, existing techniques are either data-oblivious or use the query statistics without deriving the transform from a distortion criterion. Moreover, they rely on transforms built on top of random or Hadamard rotations, which equalize variances across entries rather than compacting energy, and fixed-width scalar quantizers, which are suboptimal at low rates. In this paper, we formulate KV cache quantization as a transform coding problem in which distortion is the error in the attention products. We derive closed-form optimal transforms for keys and values from calibration statistics, under a high-resolution model. We show that the optimal key transform is not orthogonal and satisfies a generalized Parseval relation: the attention-aware distortion becomes mean-squared error (MSE) in the transform domain. Thus, we can use MSE-optimal vector quantizers applied directly to the transformed key coefficients. To meet the fixed-width layout requirement, we show that grouping coefficients into equal-volume partitions makes equal-size codebooks attain the variable-rate optimum under the same high-resolution model. At two bits per element, our method, termed NOVA-KV, recovers most of the long-context retrieval accuracy lost by scalar quantization methods at comparable throughput.

Problem

• KV cache quantization
• attention preservation
• transform coding
• vector quantization
• long-context LLM

Innovation

• KV cache quantization
• attention-preserving transforms
• transform coding
• vector quantization
• long-context LLM