PM-KVQ: Progressive Mixed-precision KV Cache Quantization for Long-CoT LLMs
Journal:
arXiv
Published Date:
May 24, 2025
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. Our
code is available at https://github.com/thu-nics/PM-KVQ.
