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

TL;DR

MixKVQ is a novel query-aware mixed-precision quantization method for KV caches in LLMs, effectively reducing memory use while maintaining high reasoning performance on complex tasks.

Contribution

It introduces a lightweight, query-aware algorithm that selectively preserves critical key channels for high precision, improving upon existing quantization techniques.

Findings

Outperforms existing low-bit quantization methods on reasoning tasks

Achieves near full-precision performance with significantly less memory

Effectively identifies and preserves critical key channels for reasoning

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.