Title: PM-KVQ: Progressive Mixed-precision KV Cache Quantization for Long-CoT LLMs
Primary area: infrastructure, software libraries, hardware, systems, etc.
Keywords: KV Cache Quantization

Abstract: Recently, significant progress has been made in developing reasoning-capable Large Language Models (LLMs) through long Chain-of-Thought (CoT) techniques.
However, this long-CoT reasoning process imposes substantial memory overhead due to the large Key-Value (KV) Cache memory overhead.
Post-training KV Cache quantization has emerged as a promising compression technique and has been extensively studied in short-context scenarios.
However, directly applying existing methods to long-CoT LLMs causes significant performance degradation due to the following two reasons: 
(1) Large cumulative error: Existing methods fail to adequately leverage available memory, and they directly quantize the KV Cache during each decoding step, leading to large cumulative quantization error.
(2) Short-context calibration: Due to Rotary Positional Embedding (RoPE), the use of short-context data during calibration fails to account for the distribution of less frequent channels in the Key Cache, resulting in performance loss.
We propose Progressive Mixed-Precision KV Cache Quantization (PM-KVQ) for long-CoT LLMs to address the above issues in two folds:
(1) To reduce cumulative error, we design a progressive quantization strategy to gradually lower the bit-width of KV Cache in each block. Then, we propose block-wise memory allocation to assign a higher bit-width to more sensitive transformer blocks. 
(2) To increase the calibration length without additional overhead, we propose a new calibration strategy with positional interpolation that leverages short calibration data with positional interpolation to approximate the data distribution of long-context data.
Extensive experiments on 7B–70B long-CoT LLMs show that PM-KVQ improves reasoning benchmark performance by up to 8% over SOTA baselines under the same memory budget and achieves 2.73–5.18$\times$ throughput over the original 16-bit LLMs.
Our code will be released soon.

TL;DR: We propose Progressive Mixed-Precision KV Cache Quantization (PM-KVQ) for long-CoT LLMs.