arXiv:2502.12665cs.CL2025-02ACL被引 7

通过量化与位置编码优化,实现长文本生成的高效低损推理。

A$^2$ATS: Retrieval-Based KV Cache Reduction via Windowed Rotary Position Embedding and Query-Aware Vector Quantization

  • 用分窗旋转位置编码解耦位置依赖,提升检索精度。
  • 查询感知量化直接优化注意力得分,减少误差。
  • 支持更大批量处理,长上下文推理速度最高快2.7倍。

长上下文大语言模型在服务时面临巨大的内存占用和访问开销问题,主要源于键值缓存(KV cache)的规模。基于检索的KV缓存压缩方法通常将完整缓存移至CPU,并在推理时按需检索必要标记。然而,此类方法仍存在精度下降和额外检索开销的问题。本文提出A²ATS,一种新型基于检索的KV缓存压缩方法。A²ATS通过向量量化技术对键状态进行处理,以获得注意力得分的精确近似,从而实现高效且精准的Top-K标记检索。首先,提出分窗旋转位置编码(Windowed Rotary Position Embedding),在位置嵌入后解耦查询与键状态的位置依赖关系;其次,设计查询感知向量量化(Query-Aware Vector Quantization),直接优化注意力得分逼近目标;最后,构建异构推理架构用于缓存卸载,支持更大批量的长上下文服务。实验表明,A²ATS在相似或更低开销下实现更小性能退化,使长上下文服务吞吐量最高提升2.7倍。

原文摘要 · Abstract (English)

Long context large language models (LLMs) pose significant challenges for efficient serving due to the large memory footprint and high access overhead of KV cache. Retrieval-based KV cache reduction methods can mitigate these challenges, typically by offloading the complete KV cache to CPU and retrieving necessary tokens on demand during inference. However, these methods still suffer from unsatisfactory accuracy degradation and extra retrieval overhead. To address these limitations, this paper proposes A$^2$ATS, a novel retrieval-based KV cache reduction method. A$^2$ATS aims to obtain an accurate approximation of attention scores by applying the vector quantization technique to key states, thereby enabling efficient and precise retrieval of the top-K tokens. First, we propose Windowed Rotary Position Embedding, which decouples the positional dependency from query and key states after position embedding. Then, we propose query-aware vector quantization that optimizes the objective of attention score approximation directly. Finally, we design the heterogeneous inference architecture for KV cache offloading, enabling long context serving with larger batch sizes. Experimental results demonstrate that A$^2$ATS can achieve a lower performance degradation with similar or lower overhead compared to existing methods, thereby increasing long context serving throughput by up to $2.7 \times$.

KV缓存长上下文向量量化推理优化

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