arXiv:2505.18610cs.CL2025-05中稿 · ICLR被引 15

提升长思维链大模型推理效率,降低显存占用且不损失性能。

PM-KVQ: Progressive Mixed-precision KV Cache Quantization for Long-CoT LLMs

  • 分步渐进式量化+块级内存分配,减少累积误差。
  • 利用位置插值扩展校准长度,适配长上下文分布。
  • 适合需要高效长推理的AI研究者与部署工程师。

近期,通过长思维链(Long-CoT)技术显著提升了大语言模型的推理能力,但该过程带来巨大的键值缓存(KV Cache)内存开销。后训练阶段的KV缓存量化作为压缩技术备受关注,但在短上下文场景外表现不佳。主要原因在于:(1) 累积误差大——现有方法未充分利用内存,在每一步解码中直接量化,导致误差累积;(2) 校准上下文过短——受旋转位置编码(RoPE)影响,使用短上下文数据校准时无法覆盖键缓存中低频通道的分布,造成性能下降。本文提出渐进式混合精度KV缓存量化(PM-KVQ),从两方面解决上述问题:(1) 设计渐进式量化策略,逐步降低每块的位宽,并采用块级内存分配,为敏感模块分配更高精度;(2) 提出基于位置插值的新校准策略,仅用短数据即可逼近长上下文的数据分布。在7B至70B规模的长CoT LLM上实验表明,PM-KVQ在相同内存预算下比当前最优基线提升推理性能最高达8%,推理吞吐量达到原始16位模型的2.73–5.18倍。

原文摘要 · Abstract (English)

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.18x throughput over the original 16-bit LLMs.

大模型推理量化长上下文KV缓存

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。