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

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 ^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 .

## introduction large language models (llms) with long context windows @xcite @xcite @xcite are driving advancements in ai applications. however, these models pose significant challenges for efficient serving. their transformer-based @xcite architecture generates and maintains a key-value (kv) cache during inference to store intermediate results and avoid re-computation. as the context length increases, the size of the kv cache grows proportionally, leading to severe overheads. first, the size of kv cache accessed when generating each token increases, resulting in a gpu memory bandwidth bottleneck. moreover, the large kv cache size of each request limits the maximum feasible batch size, resulting in suboptimal gpu utilization. various methods were proposed to address these challenges from different perspectives. quantization-based methods @xcite compress kv cache by using lower bit-width representations for kv cache elements. eviction-based methods @xcite @xcite reduce the kv cache size by directly evicting unimportant tokens from memory. retrievalbased methods @xcite @xcite @xcite offload the complete kv cache to cpu memory and retrieve necessary tokens on demand during inference. however, these methods still face challenges of limited compression ratio, unsatisfactory accuracy degradation, or extra retrieval overhead. to address the above 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 vector quantization technique to key states, thereby enabling efficient and precise retrieval of the top-k tokens. in order to achieve this goal, we face two main challenges. first, the position-dependent nature of key states after applying position embedding hinders the direct application of shared codebooks across varying inputs. second, directly utilizing the conventional vector quantization fails to guarantee an accurate approximation of attention scores. to overcome these challenges, the main contributions in this paper are as follows: ## method in a 2 ats, there are two stages. during the offline pre-processing stage (1st stage), a 2 ats constructs a shared codebook on a representative dataset for each attention head of each layer. during the inference stage (2nd stage), a 2 ats applies quantization functions to key states to map them to the nearest codewords. at each autoregressive decoding step, a 2 ats first utilizes codebooks and codeword indices to approximate attention scores, then retrieves the top-k tokens with the highest attention scores for computation, thereby mitigating the memory overhead of accessing the entire kv cache. ## conclusion in this paper, we propose a 2 ats, a novel retrievalbased kv cache reduction method. first, we propose windowed rotary position embedding to decouple the positional dependency from query and key states after position embedding. then, we propose query-aware vector quantization to achieve an accurate attention score approximation. next, we introduce the heterogeneous inference design for kv cache offloading, which increases available batch size. experimental results demonstrate that a 2 ats achieves lower performance degradation with comparable or lower overhead compared to existing methods, thereby boosting long context serving throughput by up to 2.7×.