Loading...
A2ATS: Retrieval-Based KV Cache Reduction via Windowed Rotary Position Embedding and Query-Aware Vector Quantization
He, Junhui Xing, Junna Wang, Nan Xu, Rui Wu, Shangyu Zhou, Peng Liu, Qiang Xue, Chun Jason Li, Qingan
He, Junhui
Xing, Junna
Wang, Nan
Xu, Rui
Wu, Shangyu
Zhou, Peng
Liu, Qiang
Xue, Chun Jason
Li, Qingan
Files
Supervisor
Department
Machine Learning
Embargo End Date
Type
Conference proceeding
Date
License
Language
Collections
Research Projects
Organizational Units
Journal Issue
Abstract
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<sup>2</sup>ATS, a novel retrieval-based KV cache reduction method. A<sup>2</sup>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<sup>2</sup>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×.
Citation
J. He, J. Xing, N. Wang, R. Xu, S. Wu, P. Zhou , et al., "A2ATS: Retrieval-Based KV Cache Reduction via Windowed Rotary Position Embedding and Query-Aware Vector Quantization," 2025, pp. 12451-12463.
Source
Findings of the Association for Computational Linguistics: ACL 2025
Conference
Findings of the Association for Computational Linguistics: ACL 2025
Keywords
46 Information and Computing Sciences, 4605 Data Management and Data Science
Subjects
Source
Findings of the Association for Computational Linguistics: ACL 2025
Publisher
Association for Computational Linguistics