arXiv:2506.19505cs.CL2025-06被引 6

通过识别关键令牌,实现大模型缓存的高效低比特压缩。

AnTKV: Anchor Token-Aware Sub-Bit Vector Quantization for KV Cache in Large Language Models

  • 基于锚点令牌感知的向量量化,分两阶段优化压缩与精度。
  • 1%高敏感度令牌保留使1比特量化后困惑度降至6.32。
  • 兼容FlashAttention,支持840K上下文长度,吞吐提升3.5倍。

量化已成为降低大语言模型中键值(KV)缓存内存占用的有效轻量级方案。然而,如何在极低比特量化下最小化精度损失仍是重大挑战。尽管标量量化受限于1比特极限,向量量化利用向量内相关性,可实现亚比特压缩,更适合超低比特场景。我们发现注意力质量的退化在不同令牌间分布不均,提出锚点得分(Anchor Score)衡量各令牌对量化敏感度。分析与实验表明,保留最高锚点得分的1%令牌可显著缓解极端量化下的精度下降。为此,我们提出AnTKV:一种双阶段框架,结合离线令牌感知中心学习与在线锚点令牌选择,平衡压缩与准确率。为支持高效部署,设计了适配FlashAttention的在线锚点选择核函数。该方法使LLaMA3-8B在单张80GB A100上扩展至840K上下文长度,解码吞吐相较FP16基线最高提升3.5倍。实验表明,AnTKV在4比特下性能持平或优于现有方法,并在超低比特量化下显著降低困惑度:在Mistral-7B上1比特时困惑度达6.32,优于CQ的7.25与KVQuant的15.36。

原文摘要 · Abstract (English)

Quantization has emerged as an effective and lightweight solution to reduce the memory footprint of the KV cache in Large Language Models. Nevertheless, minimizing the accuracy degradation caused by ultra-low-bit KV cache quantization remains a significant challenge. While scalar quantization is constrained by 1-bit bound, vector quantization exploits intra-vector correlations and enables sub-bit regimes, making it more suitable for ultra-low-bit quantization. To further mitigate quantization-induced degradation, we reveal that the degradation is highly uneven across tokens in attention quality. To investigate this unevenness, we introduce anchor score to measure each token's sensitivity to quantization. Our analysis and experiments show that preserving a small subset (1\%) of tokens with the highest Anchor Score significantly mitigates accuracy loss under aggressive quantization. We propose AnTKV, a dual-stage framework that leverages anchor token-aware vector quantization to compress the KV cache. It combines offline token-aware centroids learning and online anchor token selection to balance compression and accuracy. To enable efficient deployment, we design an online anchor token selection kernel compatible with FlashAttention. It allows LLaMA3-8B to scale to 840K tokens on a single 80GB A100, while delivering up to $3.5\times$ higher decoding throughput over the FP16 baseline. Experiments demonstrate that AnTKV matches or surpasses prior methods at 4-bit, and significantly reduce perplexity under ultra-low-bit quantization, achieving 6.32 at 1-bit on Mistral-7B, compared to 7.25 for CQ and 15.36 for KVQuant.

KV缓存量化大模型压缩

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