# A2ATS: Retrieval-Based KV Cache Reduction via Windowed Rotary Position Embedding and Query-Aware Vector Quantization

arXiv: https://arxiv.org/abs/2502.12665 · HTML: https://mlsanity.com/paper/2502.12665/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Placed in KV cache quantization, one of 12 directions in Efficient LLM inference and serving.
- Builds on 19 of its 29 references that are in the map.
- First posted on arXiv 18 Feb 2025.

## Placed in

- Efficient LLM inference and serving / [KV cache quantization](https://mlsanity.com/inference/kv-cache-quantization/)

## Builds on (most cited first, 19 of 19)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context](https://mlsanity.com/paper/2403.05530/) (2024)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding](https://mlsanity.com/paper/2308.14508/) (2023)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [SnapKV: LLM Knows What You are Looking for Before Generation](https://mlsanity.com/paper/2404.14469/) (2024)
- [KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache](https://mlsanity.com/paper/2402.02750/) (2024)
- [KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization](https://mlsanity.com/paper/2401.18079/) (2024)
- [Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference](https://mlsanity.com/paper/2406.10774/) (2024)
- [Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving](https://mlsanity.com/paper/2407.00079/) (2024)
- [PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference](https://mlsanity.com/paper/2405.12532/) (2024)
- [MagicPIG: LSH Sampling for Efficient LLM Generation](https://mlsanity.com/paper/2410.16179/) (2024)
- [PQCache: Product Quantization-based KVCache for Long Context LLM Inference](https://mlsanity.com/paper/2407.12820/) (2024)
- [Loki: Low-rank Keys for Efficient Sparse Attention](https://mlsanity.com/paper/2406.02542/) (2024)
- [Full Stack Optimization of Transformer Inference: a Survey](https://mlsanity.com/paper/2302.14017/) (2023)
- [ClusterKV: Manipulating LLM KV Cache in Semantic Space for Recallable Compression](https://mlsanity.com/paper/2412.03213/) (2024)
- [Transformer-VQ: Linear-Time Transformers via Vector Quantization](https://mlsanity.com/paper/2309.16354/) (2023)

## Built on by (most cited first, 7 of 7)

- [ThinKV: Thought-Adaptive KV Cache Compression for Efficient Reasoning Models](https://mlsanity.com/paper/2510.01290/) (2025)
- [Addressable Memory for Video World Models](https://mlsanity.com/paper/2608.07408/) (2026)
- [Multipole Attention for Efficient Long Context Reasoning](https://mlsanity.com/paper/2506.13059/) (2025)
- [Multipole Semantic Attention: A Fast Approximation of Softmax Attention for Pretraining](https://mlsanity.com/paper/2509.10406/) (2025)
- [ContiguousKV: Accelerating LLM Prefill with Granularity-Aligned KV Cache Management](https://mlsanity.com/paper/2601.13631/) (2026)
- [KV Cache Compression Through the Lens of Transform Coding](https://mlsanity.com/paper/2608.14191/) (2026)
- [A Target-Centric Survey of Quantization-Aware Training](https://mlsanity.com/paper/2608.29667/) (2026)

## Abstract (arXiv)

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 A2ATS, a novel retrieval-based KV cache reduction method. A2ATS 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 A2ATS 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 ×.