提出旋转量化方法,2比特压缩大模型缓存,精度损失小于0.3,内存减少近4倍。
RotateKV: Accurate and Robust 2-Bit KV Cache Quantization for LLMs via Outlier-Aware Adaptive Rotations
- 通过自适应旋转和通道重排,动态应对不同通道异常值分布。
- 在LLaMA-2-13B上实现2比特量化,WikiText-2 PPL下降不足0.3,长文本推理损耗低于1.7%。
- 适合高并发、长上下文场景,尤其对内存敏感的部署需求。
键值(KV)缓存通过避免重复计算历史键值,显著提升大语言模型推理效率。随着批量大小和上下文长度增加,庞大的KV缓存成为主要内存瓶颈,亟需高效压缩方案。现有量化方法依赖细粒度量化或保留大量高比特缓存,导致压缩率受限且在极低平均比特位宽下鲁棒性差。本文探索旋转技术在2比特KV量化中的潜力,提出RotateKV:(i) 异常值感知旋转,利用通道重排适配不同通道的异常值分布,不牺牲快速沃尔什-哈达玛变换(FWHT)的计算效率;(ii) RoPE预旋转分组头机制,缓解旋转位置编码(RoPE)对旋转策略的影响,并进一步平滑跨头异常值;(iii) 注意力汇聚点感知量化,利用大规模激活精准识别并保护注意力汇聚点。在使用LLaMA-2-13B模型时,于WikiText-2数据集上实现2比特量化,困惑度(PPL)下降小于0.3,保持强思维链推理与长上下文能力,GSM8K任务性能下降不足1.7%,优于现有方法。此外,峰值内存降低3.97倍,支持5.75倍更大批量,解码阶段速度提升2.32倍。
原文摘要 · Abstract (English)
Key-Value (KV) cache facilitates efficient large language models (LLMs) inference by avoiding recomputation of past KVs. As the batch size and context length increase, the oversized KV caches become a significant memory bottleneck, highlighting the need for efficient compression. Existing KV quantization rely on fine-grained quantization or the retention of a significant portion of high bit-widths caches, both of which compromise compression ratio and often fail to maintain robustness at extremely low average bit-widths. In this work, we explore the potential of rotation technique for 2-bit KV quantization and propose RotateKV, which achieves accurate and robust performance through the following innovations: (i) Outlier-Aware Rotation, which utilizes channel-reordering to adapt the rotations to varying channel-wise outlier distributions without sacrificing the computational efficiency of the fast Walsh-Hadamard transform (FWHT); (ii) Pre-RoPE Grouped-Head Rotation, which mitigates the impact of rotary position embedding (RoPE) on proposed outlier-aware rotation and further smooths outliers across heads; (iii) Attention-Sink-Aware Quantization, which leverages the massive activations to precisely identify and protect attention sinks. RotateKV achieves less than 0.3 perplexity (PPL) degradation with 2-bit quantization on WikiText-2 using LLaMA-2-13B, maintains strong CoT reasoning and long-context capabilities, with less than 1.7\% degradation on GSM8K, outperforming existing methods even at lower average bit-widths. RotateKV also showcases a 3.97x reduction in peak memory usage, supports 5.75x larger batch sizes, and achieves a 2.32x speedup in decoding stage.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。