MixKVQ: Query-Aware Mixed-Precision KV Cache Quantization for Long-Context Reasoning

TL;DR

MixKVQ introduces a query-aware, mixed-precision quantization framework for KV caches in long-context LLM reasoning. By computing a Salience score that combines per-channel query relevance and quantization sensitivity, it selectively preserves critical key channels in BF16 while aggressively quantizing others, and quantizes the value cache per-token. The approach yields substantial memory and bandwidth savings, with performance on complex reasoning tasks remaining close to full-precision and significantly outperforming prior low-bit methods. Empirical results across mathematical and scientific reasoning benchmarks, as well as long-context generation tests, demonstrate strong accuracy retention and notable efficiency gains, making it a practical plug-and-play solution for memory-bound LLM inference.

Abstract

Long Chain-of-Thought (CoT) reasoning has significantly advanced the capabilities of Large Language Models (LLMs), but this progress is accompanied by substantial memory and latency overhead from the extensive Key-Value (KV) cache. Although KV cache quantization is a promising compression technique, existing low-bit quantization methods often exhibit severe performance degradation on complex reasoning tasks. Fixed-precision quantization struggles to handle outlier channels in the key cache, while current mixed-precision strategies fail to accurately identify components requiring high-precision representation. We find that an effective low-bit KV cache quantization strategy must consider two factors: a key channel's intrinsic quantization difficulty and its relevance to the query. Based on this insight, we propose MixKVQ, a novel plug-and-play method that introduces a lightweight, query-aware algorithm to identify and preserve critical key channels that need higher precision, while applying per-token quantization for value cache. Experiments on complex reasoning datasets demonstrate that our approach significantly outperforms existing low-bit methods, achieving performance comparable to a full-precision baseline at a substantially reduced memory footprint.

Figures (7)

• Figure 1: Complex reasoning performance of different 2-bit KV cache quantization. Reasoning score is the average accuracy of AIME 2024-2025, MATH 500, GPQA, and LiveCodeBench.
• Figure 2: Absolute quantization error of key and value cache for Qwen-2.5-14B-Instruct model.
• Figure 3: Analysis of Key channel properties on Qwen-2.5-14B-Instruct. Scatter plot of Query magnitude ($\mathcal{I}$) versus Key scale ($\mathcal{S}$) in Layer 0. Here, $\mathcal{I}$ denotes the average activation intensity of the Query vectors and reflects each channel’s contribution to the attention scores. Traditional methods assign high bit-widths to channels with high $\mathcal{S}$ but low $\mathcal{I}$ (shown as blue dots) due to their large $\mathcal{S}$ values; however, these channels are in fact non-crucial for preserving the fidelity of attention score computation. (b) Channel statistics for Head 4. In contrast to the Key scales ($\mathcal{S}$), which are densely clustered and offer limited discriminative capability, the combined salience score ($\mathcal{A}=\mathcal{I}\cdot\mathcal{S}$) effectively isolates critical channels that require high-precision retention. Bar colors denote the adaptive precision levels: green (BF16), orange (INT4), and grey (INT2).
• Figure 4: Workflow of MixKVQ. Keys and Values undergo per-channel and per-token quantization, respectively, managed by a full-precision buffer of size $R$. For Key cache, salience Score $\mathcal{A}_d$ are updated every $R$ tokens within the window, while Value cache utilize uniform 2-bit quantization.
• Figure 5: Memory usage and throughput comparison between MixKVQ and 16 bit baseline.