KVmix: Gradient-Based Layer Importance-Aware Mixed-Precision Quantization for KV Cache

The high memory demands of the Key-Value (KV) Cache during the inference of Large Language Models (LLMs) severely restrict their deployment in resource-constrained platforms. Quantization can effectively alleviate the memory pressure caused by KV Cache. However, existing methods either rely on static one-size-fits-all precision allocation or fail to dynamically prioritize critical KV in long-context tasks, forcing memory-accuracy-throughput tradeoffs. In this work, we propose a novel mixed-precision quantization method for KV Cache named KVmix. KVmix leverages gradient-based importance analysis to evaluate how individual Key and Value projection matrices affect the model loss, enabling layer-specific bit-width allocation for mix-precision quantization. It dynamically prioritizes higher precision for important layers while aggressively quantizing less influential ones, achieving a tunable balance between accuracy and efficiency. KVmix also introduces a dynamic long-context optimization strategy that adaptively keeps full-precision KV pairs for recent pivotal tokens and compresses older ones, achieving high-quality sequence generation with low memory usage. Additionally, KVmix provides efficient low-bit quantization and CUDA kernels to optimize computational overhead. On LLMs such as Llama and Mistral, KVmix achieves near-lossless inference performance with extremely low quantization configuration (Key 2.19bit Value 2.38bit), while delivering a remarkable 4.9x memory compression and a 5.3x speedup in inference throughput.

💡 Research Summary

Large language models (LLMs) rely on a key‑value (KV) cache during autoregressive decoding to avoid recomputing attention for previously generated tokens. While this cache dramatically reduces computation, its memory footprint grows linearly with sequence length, quickly exceeding the capacity of typical GPUs (e.g., a 70‑billion‑parameter model needs >50 GB for a 20 k‑token cache). Existing quantization approaches either apply a static, one‑size‑fits‑all bit‑width to every layer or perform dynamic quantization at high computational cost without distinguishing the relative importance of KV entries across layers and time steps. Consequently, practitioners must accept sub‑optimal trade‑offs among memory usage, model accuracy, and throughput.

KVmix addresses these limitations with three core contributions. First, it introduces a gradient‑based layer importance profiler. For each transformer layer i, the L2 norm of the gradients of the loss with respect to the key and value projection matrices (∥∇{W_k^i}L∥₂ and ∥∇{W_v^i}L∥₂) is computed over a small set of representative prompts. These norms serve as sensitivity scores s_k^i and s_v^i, indicating how much quantization error in that layer would affect the final loss. By averaging across prompts, KVmix obtains stable importance estimates and then allocates higher bit‑widths (e.g., 3‑4 bits) to the top‑20 % most sensitive layers while assigning aggressive low‑bit settings (e.g., 2 bits) to the remaining layers. The proportion is user‑configurable, allowing flexible memory‑accuracy balancing.

Second, KVmix adds a dynamic pivotal‑context selection mechanism. Leveraging the observation that recent tokens dominate attention (the “attention sink” phenomenon), the method keeps KV pairs for the most recent pivotal tokens in full precision (FP16) while compressing older KV entries according to the layer‑wise importance scores. This temporal compression adapts as decoding proceeds, ensuring that long‑context generation retains high quality without storing the entire history at full precision.

Third, the authors implement efficient asymmetric low‑bit quantization and custom CUDA kernels. Per‑channel scaling and offset are learned for each layer, enabling 3‑bit quantization with minimal bias. The kernels fuse quantization, dequantization, and attention matrix multiplication to keep the runtime overhead negligible; profiling is performed offline once, so inference incurs no extra cost.

Experiments span Llama‑2‑7B, Llama‑2‑13B, and Mistral‑7B models, evaluated on GSM8K, TruthfulQA, WikiText, and instruction‑following benchmarks. KVmix achieves average key and value bit‑widths of 2.19 bits and 2.38 bits respectively, yet the loss in accuracy relative to FP16 is under 0.1 % on all tasks. Memory consumption is reduced by a factor of 4.9×, and inference throughput improves by 5.3×. Compared with prior mixed‑precision methods such as QA‑Q and KVTuner, KVmix delivers 1‑2 % higher accuracy under the same memory budget and up to 1 × faster inference due to its lightweight profiling and optimized kernels.

The paper’s significance lies in unifying layer‑wise sensitivity analysis with temporal KV compression. By grounding importance in actual loss gradients, KVmix avoids heuristic or exhaustive search procedures, offering a fast, scalable way to determine per‑layer bit‑widths. The dynamic pivotal‑context strategy further exploits the natural decay of attention relevance over time, achieving near‑lossless generation even for very long sequences.

Future directions suggested include extending importance profiling to head‑level or token‑level granularity, combining the approach with sparsity‑based KV eviction, and exploring multi‑GPU or serving‑stack integration where per‑request KV budgets must be dynamically allocated. Overall, KVmix presents a practical, high‑impact solution for deploying LLMs on memory‑constrained hardware while preserving generation quality.