arXivDaily arXiv每日学术速递 周一至周五更新
​通知:联邦节假日调休:哥伦布日 / 原住民日(Columbus Day / Indigenous Peoples' Day),今日(13日)无更新。
arXiv 2505.18610cs.CL

PM-KVQ:用于长思维链语言模型的渐进式混合精度键值缓存量化

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

Tsinghua University(清华大学);Infinigence-AI;Columbia University(哥伦比亚大学);OPPO AI Center(OPPO人工智能中心);Shanghai Jiaotong University(上海交通大学)

Tengxuan Liu, Shiyao Li, Jiayi Yang, Tianchen Zhao, Feng Zhou, Xiaohui Song, Guohao Dai, Shengen Yan, Huazhong Yang, Yu Wang

AI总结:

针对长思维链语言模型因键值缓存内存开销大导致性能下降的问题,提出渐进式混合精度KV缓存量化(PM-KVQ),通过渐进量化策略和逐块内存分配减少累积误差,用带位置插值的校准策略增加校准长度,提升了推理性能和吞吐量。

摘要:

最近,通过长思维链(CoT)技术在开发具有推理能力的大语言模型(LLMs)方面取得了重大进展。然而,这种长CoT推理过程由于大的键值(KV)缓存内存开销而带来大量内存开销。训练后KV缓存量化已成为一种有前途的压缩技术,并在短上下文场景中得到广泛研究。但直接将现有方法应用于长CoT LLMs会导致显著性能下降,原因有二:一是大的累积误差,现有方法未充分利用可用内存,在每个解码步骤直接量化KV缓存,导致大的累积量化误差;二是短上下文校准,由于旋转位置嵌入(RoPE),在校准期间使用短上下文数据无法考虑键缓存中低频通道的分布,导致性能损失。我们提出用于长CoT LLMs的渐进式混合精度KV缓存量化(PM-KVQ)来解决上述问题,一方面通过设计渐进量化策略逐步降低每个块中KV缓存的位宽,并提出逐块内存分配为更敏感的Transformer块分配更高位宽来减少累积误差;另一方面通过提出带位置插值的新校准策略,利用带位置插值的短校准数据近似长上下文数据的分布来增加校准长度且无额外开销。在7B - 70B长CoT LLMs上的大量实验表明,在相同内存预算下,PM-KVQ比SOTA基线将推理基准性能提高了8%,并且比原始16位LLMs实现了2.73 - 5.18倍的吞吐量。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑