RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval
Authors: Di Liu
Meng Chen
Baotong Lu
Huiqiang Jiang
Zhenhua Han
Qianxi Zhang
Qi Chen
Chengruidong Zhang
Bailu Ding
Kai Zhang
Chen Chen
Fan Yang
Yuqing Yang
Lili Qiu
Venue: NeurIPS 2025 poster ()
Keywords: Large Language Model
Efficient Inference
Long Context
Vector Retrieval
Primary area: deep_learning
TL;DR: RetrievalAttention improves decoding speed and reduces GPU memory usage in Transformer-based LLMs by using pre-built, attention-aware KV vector indexes stored in CPU memory, achieving significant efficiency gains without compromising accuracy.

Abstract: Transformer-based Large Language Models (LLMs) have become increasingly important. However, scaling LLMs to longer contexts incurs slow inference speed and high GPU memory consumption for caching key-value (KV) vectors. This paper presents RetrievalAttention, a training-free approach to both accelerate the decoding phase and reduce GPU memory consumption by pre-building KV vector indexes for fixed contexts and maintaining them in CPU memory for efficient retrieval. Unlike conventional KV cache methods, RetrievalAttention integrate approximate nearest neighbor search (ANNS) indexes into attention computation. We observe that off-the-shelf ANNS techniques often fail due to the out-of-distribution (OOD) nature of query and key vectors in attention mechanisms. RetrievalAttention overcomes this with an attention-aware vector index. Our evaluation shows RetrievalAttention achieves near full attention accuracy while accessing only 1-3\% of the data, significantly reducing inference costs. Remarkably, RetrievalAttention enables LLMs with 8B parameters to handle 128K tokens on a single NVIDIA RTX4090 (24GB), achieving a decoding speed of 0.107 seconds per token.

RetrievalA Accelerating Long-Context LLM Inference via Vector Retrieval
Di Liu 2 Meng Chen 1,3 Baotong Lu Huiqiang Jiang', Zhenhua Han', Qianxi Zhang' Qi Chen', Chengruidong Zhang', Bailu Ding', Kai Zhang', Chen Chen' Fan Yang'
Yuqing Yang', Lili Qiu'
Microsoft Research 2 Shanghai Jiao Tong University Fuuan University
(baotonglu, hjiang, qiazh, cheqi, fanyang, yugyang)@mice wwwcoooootttcoooooommooom com 2 (liu-di, chen-chen)@sjtu. edu. cn "mengchen220m fudan. fudan.edu.ccm edu.cn, en, zhangk&fudan. edu.cnn
Abstract
Transformer-base Large Language Models (LLMs) have become increasingly important. However, scaling LLMs to longer contexts incurs slow inference speed and high GPU memory consumption for caching key-value (KV) vectors. This paper presents RetrievalAAtention, training-free approach both accelerate the decoding phase and reduce GPU memory consumption by pre-building KV vector indexes for fixed contexts and maintaining them in CPU memory for efficient retrieval. Unlike conventional KV cache methods, Retrieval integrate ap- proximate nearest neighbor search (ANNS) indexes into attention computation. We observe that off- she-the--heff ANNS techniques often fail due to the (OOD) nature of query and key vectors in attention mechanisms. RetrievalAtten- tion overcomes this with an attention- aware vector index. Our evaluation shows RetrievalA achieves near full attention accuracy while accessing only 1-3% of the data, significantly reducing inference costs. Remarkably, Retrieval enables I.LMS with 8B parameters to handle 128K tokens on a single NVIDIA RTX4090 (24GB), achieving a decoding speed of O. 107 seconds per token.
Introduction
The advancements of transformer-base large language models (LLMs) [ have shown exceptional capabilities in processing long contexts, with support up to 10 million tokens per context 2 3 This capability of analyzing extensive data facilitates numerous downstream applications, including document-based question answering (QA) 4 repository-level code analysis and generation 5 and personalized AI agents (6). However, the use of long contexts poses challenges in inference efficiency, as both the size of the key-value (KV) cache and inference latency for KV accesses increase linearly with the context length. Consequently, how to reduce storage costs and token accesses crucial for enhancing long-context inference efficiency, particularly in resource-constrained environments such local deployments 7 edge devices
Recent studies 10 reveal that many long-contex, applications, such as document question answering (QA) and in-context learning, significant portion of the context remains constant and is reused across multiple requests. For instance, when conducting QA on a corpus, these documents are introduced to the LLM as prefix and then re-utilized to answer different questions (i.e., user queries). As another example, some fixed contexts containing domain-speciie knowledge are used to
Work performed during the intersship Microsoft Research. 8 Corresponding authors: Baotong Lu and Chen Chen.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

RetrievalAttention
Legomputa FlexGen
ÁKNN
Traditional
accuracy (the higher the better)
Model
AQest
Full attention by KV recompute offloading
Total Latency (8)
Prompi Length
128K
32.8
256K
111
512K
465
1,765
1M
FFN (s)
7.6
31
Attention (8)
25.2
434
1,695
SnepKV InfLLMM
Heuristic token selection
KV Cache (GB)
GPU Memory
15.6
31.2
62.5
125
Static KV compression
Decoding latency (the lower the better) Figure I: RetrievalAttention achieves similar task accuracy as full attention but exhibits extremely low decoding latency.
Table I: Decoding latency and memory re-
quired for KV cache of Llama-3-8B across different context lengths on one A100 GPU.
supplement system user prompts, improving generation quality. These observations have motivated existing research [9] [0] to pre-process fixed contexts for generating KV caches and reusing them in subsequent requests. This approach avoids redundant computation and thus enhances the efficiency of long-context inference, particularly by reducing the time to generate the first token.
However, the generation of new tokens during the decoding phase still suffers from inherent limitations of the conventional KV cache both in GPU memory consumption and inference latency, as all tokens from the cached context still participate in attention computation in each decoding step. A promising solution lies in leveraging the dynamic sparsity inherent in the attention mechanism [12), which subset of tokens that critical to the current query vector.
allows for selectively attending However, existing studies typically rely on static 13 14] or heuristic 15 16 token selection
methods, which often result in inaccurate identification and significant drop in model accuracy.
We propose integrating Approximate Nearest Neighbor Search (ANNS) [18] into attention compu- tation for the fixed contexts to enhance the inference efficiency. Our key insight stems from the observation of the mathematical equivalence between sparse attention computation and maximum inner product search" Specifically, searching the ANNS index built for key vectors can return those with the highest inner product values sublinear time. The retrieved key vectors are precisely aligned with the critical tokens required by the sparse attention mechanism, resulting in higher approximation accuracy compared to previous identification methods, as illustrated in Figure [ The fixed contexts allow for building the indexes offline and reusing these indexes across requests when user queries come. Moreover, most ANNS algorithms are compatible with CPU implementations, enabling strategic allocation of GPU and CPU memory, which facilitates long-context inference on devices with limited GPU memory.
Leveraging ANNS for attention mechanism presents unique challenge: the (OOD) problem between query and key vectors. Most ANNS engines operate under the assumption that both query and key vectors are drawn from the same data distribution. However, this assumption does not hold in this context due to the different projection weights for query and key vectors in attention mechanism. Our measurement study shows that query vectors deviate more than 10x farther from key vectors compared to that between in-distribution data. Unfortunately. the effectiveness of ANNS degrades significantly under the OOD problem. In particular, our empirical analysis indicates that maintaining high inference accuracy requires conventional ANNS scanning more than 30% of all key vectors identify the critical ones. This fails fully leverage the inherent sparsity of attention and nullifies the ANNS efficiency benefits. To the best of our knowledge, we are the first identify the challenge of OOD in using ANNS indexes to retrieve important tokens for sparse attention computation a factor that is crucial for inference efficiency and accuracy.
In this work, present RetrievalAAtention, long-context inference method that leverages ANNS indexes for efficient dynamic sparse attention in scenarios or applications where the fixed context takes large proportion of prompt. To address the challenge of OOD, RetrievalAAtteniion proposes vector index tailored for the attention mechanism to mitigate the distribution gap between query and
'Maximum inner product search can be viewed as similarity search and efficiently solved by ANNS
indexes [88

---

key vectors. This approach allows for the traversal of only small subset of key vectors (1-3%) to identify the most relevant tokens in attention, resulting in sub-linear time complexity with the context length while guaranting high end-to-end task accuracy. Retrieval Attention also reduces GPU memory consumption by offloading the index and KV vectors to CPU memory.
We thoroughly evaluate the accuracy and efficiency of IetrievalAttentiee across three long-context LLMs, using well-known long-context benchmarks like xc-eench [19] and RULER (20). RetrievalAt- tention maintains the nearly same accuracy as full attention on different tasks and context lengths. Regarding the inference efficiency, for the 128K context on the RTX4090 GPU, RetrievalAAtennion achieves x and x decoding latency reduction compared to the retrieval method based on exact KNN and traditional ANNS index, respectively.
2 Background and Motivation
2.1 LLM and Long-Context Serving
Attention is the core component of the transformer-biie based LLM. In the generation process of the t-th token, the attention operation computes the dot product between the query vector qe E l xd (where d is the hidden dimension) and the key vectors of all preceding tokens k, 1 (for i S t). This product is scaled by d' and normalized via Softmax function to yield the attention weights oz.. Then it adds up values v, with different weights, resulting in the output Or.
K. Vd
Σ ae i=
(I)
LLM inference contains two stages: the prefill phase and decoding phase. The prefill phase, which only happens once, takes all tokens of the prompt as input and performs attention with a time- complexity On))). In the decoding (token generation) phase, the newly generated token added to the input and computes attention scores with same complexity.
Due to the quadratic time complexity of attention computation, serving long- long-sequence input incurs extremely high costs. Table I shows the inference latency of Llama-3-8B without KV cache. When the prompt length reaches 1 million tokens, generating every token requires 1,765 seconds with over 96% of latency spent on attention operations. Since many long-context applications such as document QA and in-context learning involve a large proportion of fixed contexts, existing works [10 pre-process these contexts offline by generating their KV cache and reusing them in successive prompts from users. The fixed contexts include documents, source codes, as well as system prompts or user prompts with domain-specific knowledge. Reusing the KV cache of these contexts across successive prompts effectively avoids redundant computation and reduce the time first token.
However, the KV cache comes with a significant increase in GPU memory consumption and the attention complexity is still linear to the context length. As shown in Table T 125 GB memory is necessary for storing the KV cache when the context length reaches million tokens, which is far beyond the GPU memory capacity of commodity GPUs such as the RTX4090 (24GB) or even high-end GPUs like A100 (40GB or 80GB). This necessitates either scaling to more GPUs to accommodate the large KV cache (1 repetitively offloading and reloading the entire KV cache between CPU and GPU memory over PCIe (8). resulting in excessive communication overhead. Neither approach provides an efficient and cost-effective solution for long-context inference.
2.2 Dynamic and Sparse Attention
Corroborating recent work [3] [4], observe that attention computation in LLMs exhibits sig- nificant sparsity. Despite the large context length, only a small fraction of tokens with the highest attention weights, also known critical tokens, contribute significantly to the output. Figure 2a demonstrates the inherent sparsity, as measured by the mean squared error (MSE) of attention output between full attention and the sparse attention using only the top-k critical tokens in the context of 128K tokens. The MSE drops sharply when is even smaller than 200 tokens, indicating the high sparsity in long-context processing.
Furthermore, we observe that as the model continues to generate new tokens, the positions of critical key vectors in the context dynamically change, heavily depending on the current query vector.
3

---

X10
Step
Step
Step
Attention weights
10-1
10
MSE 1o
10-2
0
10-3
80.0K
120.0K
2a
2
2
0K
40K
80K
120K
Value of tap-k
Token Position
(a) Attention sparsity.
(h) Sparsity is dynamic.
Figure 2: (a) Mean squared error (MSE) between full attention outputs and approximations based on the top-k tokens in the 128K context. (b) Attention weights distribution across decoding steps.
Figure 2b illustrates the attention-weight distribution of the 10p-100 critical tokens across three decoding steps of single inference process. The positions of those important tokens vary significantly. with only 30% overlap of these three steps. This highlights the dynamic and unpredictable nature of attention sparsity, necessitating an efficient and effective method for identifying critical tokens.
The dynamic sparsity shows a promising path to approximately compute attention with greatly reduced cost while maintaining model accuracy. For each query, if we can accurately identify the relevant key-value vectors with higher importance, minimum GPU memory and a much lower time complexity can be achieved for attention computation.
2.3 Challenges of Off-the-shelf Vector Search
Based on Equation D one key vector is critical for a query vector they have a large inner product. With the inner product as a similarity function, performing searches ANNS indexes aligns well with the goal of the attention mechanism to efficiently find critical key vectors to a query.
Traditional ANNS indexes generally cluster similar vectors and select the representative vector for each cluster (22 or directly build connections between similar vectors to form a proximity graph 23 Both methods typically require scanning a limited subset of all vectors (e.g.. 1%) to accurately identify the most similar vectors to the query, achieving high search efficiency and accuracy. Moreover, these indexes can be constructed offline based on fixed contexts (e.g., documents) and reused when serving different user queries. However, we find that naively applying off-the-she vector indexes fails provide good performance because of the OOD issue between query (Q) and key vectors (K) in attention mechanism.
In conventional vector databases, the distribution of vectors between content and query is aligned (e.g.. image search). Naively using traditional vector indexes for attention computation suffers from an inherent distribution gap between queries and keys, which are projected by different weights. Figure 3b (focus on Q K for now) demonstrates the performance of widely-used vector indexes supported by Faiss [44 using a query vector to retrieve the most similar key vectors. It compares the percentage of keys scanned and the corresponding recall (i.e., the overlapping ratio between the retrieved top-100 results and the ground truth). Cluster-based IVF [22 requires scanning ~30-50% data for recall higher than 0.95, and graph-based HNSW 25 falls into a local optimum. The results highlight the challenge of performing efficient vector searches attention data.
Fundamentally, the difficulty due the OOD between query and key vectors. To quantify this, we adopt the Mahalanobis distance [26]. standard metric widely used in vector search literature (e.g., OOD-DisKANN (27) measure the distance from a vector to distribution. We sample 5,000 vectors from Q and K respectively as the query set and compute the Mahanobis distance from the vectors of Q to the remaining vectors in K. Figure 3a shows that the vectors from Q are significantly distant from the K vectors (OOD) while K themselves are very close. Therefore. traditional index building based solely on the closeness between key vectors does not align with the attention mechanism, which requires retrieving critical tokens nearest neighbors from the query vectors' viewpoint. contrast, Figure 3b shows that using sampled K as the queries (K to K) can easily achieve a high recall by only scanning 1-5% vectors because they are in the same distribution.
In this context, "similar" and "close" to indicate vectors with larger inner product interchangeably.

---

YI-98-200K
Llama-3--8 3-t.-------nnnttt-ooo
Yi-98-20KK
Jama-3-8B-instrinn
1.0
1.30
1500
ta K
Q
to
Recall
0.8
0.75
1000
0.6
0.50
1000
0.4
HBSW:
0.25
0.1
HRO to
HEW
200
0.50
0.75
0.00
0.5
1.00
Mahalanobis Distance
(Pencentage)
(a) Different distribution.
(b) ANNS index performance.
Figure 3: (a) Query vectors deviate from key vectors, while key vectors themselves are close. (b) Query vectors (Q) and key vectors (K) are dumped from Yi-9B and Llama-3-80 with prompt length of 128,00 28,000 tokens. Off-the-shelf ANNS indexes perform poorly on Q to K searches while working well for K to K searches.
Similarly, query vectors in each attention head also follow the same distribution they are projected by the same model weight. Therefore, for efficient vector search, the index must consider the OOD characteristic of the attention computation by design.
RetrievalAttention Design
In this work, we focus accelerating the decoding phase of long contexts in scenarios where the fixed context takes a large proportion of the whole context (e.g., document QA). We propose RetrievalAttention that leverages attention-aware vector search to approximate attention computation by CPU-GPU co-execution. RetrievalAAtention build vector indexes offline for the KV cache of fixed contexts (e.g., documents) and reuse them for successive prompts from users.
Figure shows the overall design of RetrievalAAtenttoon Based on our observation in Section 2.2 we derive an approximated attention by selectively retrieving relevant KV vectors while discarding those that are negligible ($3.1). To efficiently support long context, we store the pre-built vector indexes and KV vectors of fixed contexts in CPU memory, and use attention-aware vector search to find critical tokens during each decoding step ($3.2).
To better exploit the GPU devices, we leverage the attention pattern obtained in the offline index building phase to select small subset of KV cache that consistently important during the decoding phase and persist them (640 tokens of 128K context) on GPU devices. Retrieval/Ateenion computes partial attention in parallel using dynamically retrieved KV vectors from CPU memory with persistent KV vectors stored in GPU memory and combines them together to get the final output ($3.3).
3.1 Approximated Attention
Based the Equation I Retrieval approximates the full attention output o, by selectively utilizing the KV vectors associated with high attention weight (i.e., ar,;). Specifically, we define Ite subset of token indices for which the attention weight surpasses Consequently, sparse attention mechanism, which only considers tokens located in Le can be defined as follows:
e
O, Σ a V. E Oit V. SEI.
āe. Vi where ãt, Ejet E iEI,
Σ
Based on the above approximation, we build RetrievalAAttenioon to only consider important key-value vectors (i.c., Ir,)) that are persistent in GPU cache and dynamically retrieved by vector indexes.
3.2 Attention-aware Vector Index
RetrievalAtteenion processes fixed contexts offline to generate the KV vectors and decides whether to build indexes for them (the decision method is elaborated Section (33). All KV vectors with indexes are offloaded to CPU memory and used only when user queries comes.
To accelerate the vector search during decoding phase, Retrieval diverges from traditional indexes that only consider the closeness between key vectors for index building. Instead, it leverages
5

---

GPU-side Attention
GPU
KV
Combine
Query vectors
Predictable
KV Vectors
Officed Most KV Vectors to CPU Vactor DB
Attention
Query Vector
Output
KNN
Retrieval Search
CPU-side Attention
CPU ANNS Index
[indexed by key vectors)
Nearest KV Vectors
Key vectors
O (dynamically retrieved)
(b) Key building procedure of OOD-aware index.
(a) Overall design of retrievalAtteenionn
Figure 4: (a) RetrievalAttention offloads most KV tokens to vector databases in CPU, which are retrieved during the decoding phase to find the most relevant KV tokens with queries. (b) During the index construction, we link each query to its exact top-k nearest key vectors (KNN).
the existing query vectors to guide the index building for key vectors, efficiently mitigating the distribution gap. As shown in Figure 4bb during the index construction, RetrievalAtttenton explicitly establishes connections from the query vector to its nearest key vectors (i.e., exact k-nearest neighbors, or KNN), forming a mapping from query to key distribution. The KNN connections effectively serve as bridge reconcile their distribution differences. Using this structure, the decoding query vector can first search its nearest query vectors and then obtain the most relevant key vectors.
However, this structure still has imperfections in both memory overhead and search efficiency because we need to store and access query vectors besides key vectors. To address this problem, we leverage the projection technique from the state-of-the-art cross-modal ANNS index RoarGraph [28] to eliminate all query vectors. Specifically, we first project the KNN connections into key vectors by linking key vectors that are connected to the same query vectors, which efficiently streamlines the search. To enhance the connectivity of the graph structure, we then follow the conventional graph index building strategy perform nearest neighbor search in the graph. If key vector does not have enough neighbors, additional connections are added. The projection process connects key vectors that are perceived close from the query vectors' perspective, allowing efficient index traversal for future decoding queries.
Our evaluation shows that, by effectively modeling the proximity relationship between the query and key vectors, the vector database only requires scanning 1-3% key vectors to reach high recall, significantly reducing the index search latency by 79% compared to the IVF indexes.
3.3 CPU-GPU Co-Execution
To exploit GPU parallelism and accelerate attention computation, RetrievalAttention decomposes the attention computation into two disjoint sets of KV cache vectors, the persistent ones the GPU and the dynamic ones on the CPU, and combines the partial attention outputs during online inference.
We leverage the patterns observed in the offline phase to predict KV vectors that are consistently activated during token generation (i.e., static tokens). Similarly to StreamingLLM [3], our current implementation uses fixed initial tokens and the last sliding window of the context as the static pattern and persists them in the GPU cache. All remaining tokens require index building for retrieval.
During L.M inference, we first load KV vectors of the fixed context from CPU memory to GPU memory layer by layer to combine with the user query for context prefilling using full attention. When proceeding to the next layer, the KV cache of the previous layer, except for static tokens, is discarded to reduce GPU memory consumption. The time to first token is minimized because the attention computation for the fixed context is avoided by reusing their KV cache. The decoding phase leverages vector indexes for accelerating the token generation. query vector is sent to the CPU side for index traversal and attention computation by leveraging the compute capability of CPU, effectively reducing the data transfer over the PCle. The partial attention output from CPU is returned to GPU and merged with the partial attention output on static tokens computed by GPU, inspired by the FlashAttention description can be found in Appendix B.2
each decoding step. only a single
29 More detailed execution and algorithm

---

4 Evaluation
In this section, we compare the performance of RetrieealAAtention in long-contexi LLM inference against full attention and other state-of-the-ar methods. Through experiments, we mainly explore the following questions: (I) How does RetrievalAttention affect the model's inference accuracy? Specifically, we assess the generation accuracy of Retrieval and other methods across various downstream tasks ($4.2). (2) Can petrievalAttentiin efficiently reduce the token generation latency of long-context LLM inference? We compare the end-to-end latency of RetrievalAttention with that of other baselines ($4.3)
4.1 Experimental Setup
Testbed, Models, and Configurations. We conduct experiments on server equipped with one NVIDIA RTX4090 GPU (24GB memory) and an Intel i9- 19-10900X CPU with 10 physical cores and 128GB DRAM. The experiment results using NVIDIA A100 GPU can be found in Appendix A.55 We
implement RetrievalAttention on three long-context LLMs, including Llama-3-8B Instruct-2222 30). Yi-9B-200K 31 and Yi-6B-200K (32].
Baselines. We compare Retrieval/Attention with the following training-fres free baselines. (1) Full attention: it offloads the KV cache to CPU memory and selectively loads each layer's KV cache onto the GPU when needed. (2) StreamingLLM [[3] it retains initial tokens along with recent tokens and discards remaining tokens. (3) SnapKV [4] it only caches the critical tokens observed from
the last window of the prompt. (4) InfLLM 15: it separates the KV cache into blocks and selects representative vectors for each block. The query scans all representative vectors and retrieves top-k related blocks. (5) Quest [17]: it keeps track of the minimal and maximal values of KV cache pages and computes the criticality of a page using the query vector. (6) InfiniGen [33: it prefetches the essential KV cache by speculating attention output for subsequent layers. (7) MagicPIG [44] it
uses local sensitive hashing to partition key vectors and hash collision for queries to get related keys.
Detailed configurations of baselines can be found in Appendix A.1
To better assess the effectiveness of attention-aware index, we introduce two additional baselines using traditional vector search methods from Faiss (24). Specifically, Flat is an exact KNN method that performs linear scan of all key-value vectors, whereas IVF indexes key vectors through clustering. While HNSW is widely used in vector search, it always fails to retrieve the most relevant key vectors in attention scenarios across various parameter settings (Figure 3p. Therefore, we exclude HNSW as a baseline due to its unsatisfactory end-to-end accuracy. By default, all indexing- based methods, including RetrievalAAteention, retrieve the top-100 nearest key vectors and use 640 tokens (128 initial tokens 512 local window tokens) on GPU as the static pattern. All evaluated methods are applied across all model layers for sparse attention during the decoding phase.
Benchmarks. We adopt three representative long-context benchmarks for evaluation. (I) RULER [00]: comprehensive and widely used long-context benchmark consisting of retrieval, aggregation, QA tasks, and so on. (2) co-Bench
[19] this benchmark consists of retrieval tasks
and realistic tasks. The average context length of co-Bench is over 100K tokens. (3) Needle-in- a-haystack [35]: it challenges the models to accurately retrieve information (the "needle") hidden within a lengthy document (the "haystack").
4.2 Accuracy on Long Context Tasks
RULER. Table 2 demonstrates that Retrieval Attention establishes state-of-the-arr task accuracy across all evaluation scales, achieving nearly the same task accuracy as full attention in different context lengths. With the Llama-3-8B model, RetrievalAAtention achieves 81.7% average accuracy (vs. full attention's 83.91%) with only a 2.21% % performance drop. The loss of accuracy is only 0.52% on Yi-9B. Evaluations in Appendix A.2 also show similar superiority of RetrievalAAtention on Yi-6B. In contrast, other training- faiei methods experience a noticeable reduction in accuracy, particularly for longer context sizes like 128K, they fail to capture dynamically changed important tokens.
co- Bench. As shown in Tabel 3 Retrieval achieves an average accuracy of 49.6% across all tasks in oc-Bench for Llama-3-8B, comparable to full attention (50.4%). Retrieval Attention also achieves high accuracy on other models, demonstrating improvement over existing token selection methods. While Flat and IVF sometimes achieve marginally higher accuracy, they incur
7

---

Table Performance (%) of different methods and models on RULER.
Methods
Act. Tokens
Claimed
16K
32K
64K
128K
Avg.
Full attention
128K
262K
89.27
85.11
82.51
78.74
83,91
Llama 3-8B
StreamingLLM
2K
20.99
16.36
12.52
11.34
15.30
(-68.61)
SnapKV
2K
75.53
70.84
65.44
58.68
67.62 (-16.29)
Inf7.IM
2K
38.29
32.44
27.94
25.71
31,10 (-52.81)
InfiniGen
2K
87.16
81.60
72.52
64.10
76.35 (-7.56)
Quest
2K
82. 12.12
76.33
67.43
60.08
71.49 (-12.42)
MagicPIG
86.93
81.61
71.48
59.84
74.97 (-8.94)
Flat
640+100
87.01
84.97
80.99
74.34
81.83 (-2.08)
IVF
87.22
84.74
78.46
68.21
79.66 (-4.25)
RetrievalAttention
640+100
84.78
80.50
74.70
81.70(-2.21)
Full attention
128K
200K
82.85
73.17
67.08
60.51
70.90
StreamingL
2K
19.08
13.48
12.53
12.81
14.48 (-56.42)
Yi-98
SnapKV
2K
64.48
48.70
39.28
32.97
46.36 (-24.54)
InfLLM
36.17
28.20
22.65
20.94
26.99 (-43.91)
InfiniGen
72.87
62.10
56.76
68.55 (-2.35)
Quest
77.99
65.92
56.14
48.85
(-8.67)
MagicPIG
72.22
62.04
52.57
66.78 (-4.12)
Flat
640+100
84.42
74.58
66.16
59.50
71.17 (+4.27)
IVF
640+100
83.85
72.19
65.13
58.04
69.80 (-1.10)
RetrievalAttention
640+100
82.95
73.73
65.67
59.15
70.38 (-0.52)
significantly higher search costs (scanning 100% and 30% of all vectors, respectively) compared to RetrievalAttention (1-3%). Therefore, RetrievalAAtention achieves competitive accuracy while drastically reducing search overhead, as reflected in the latency evaluation presented Section 4.3 Detailed comparisons of different tasks provided in Appendix A.3
shown in Figure RetrievalAAtention can effectively focus on information at various positions across different context windows, ranging from 4K to 128K. Evaluation under extremely (e.g., 1,000K) long-context shows that RetrievalAAteenionn is capable of maintaining robust performance. The result is presented in Appendix A.4
Needle-in-a-haystacl
Llama-3-80 RetrievalAttention
Methods
Llama-3-8B
Yi-9B
Yi-6B
1.0
4K
Full attention
Streamingl.IMM L.M
20.2 (-30.2)
50.4
20.9 (-31.99
52.8
16.5 (-29.0)
45.5
Context Length
8K-
0.8
SnapKV
48.2 (-2.2)
42.6 (-10.2)
35.7 (-9.8)
16K
0.6
InfLLM
44.7 (-5.7)
(-7.5)
4-269)
InfiniGen
47.5 (-2.9)
48.8
44.9
32K-
-0.4
Quest
(-2.8)
(-.9
44.5
64K
MagicPIG
(-1.0)
(-4.8)
41.1 (44.))
0.2
128K-
Flat
49..
(-0.3)
(+0.2)
0.0
IVF
48.8 (-1.6)
(-0.5)
45.6 (+0.1)
0 25
50
75
100
RetrievalAAttentoon
49.6 (-0.8)
52.2(-0.6)
45.0 (-0.5)
Depth Percent (%)
Table Performance (%) of different methods and models on
Figure 5: Performance of Llama-3- 8B on Needle-in-a-hhasstcc.
oo-Benc.
4.3 Inference Latency
In this section, we study the inference efficiency of RetrievalAttention using the combination of different fixed context lengths and user input lengths. Specifically, we first measure the decoding latency for an user input of 100 tokens while varying the length of the fixed context from 16K to 128K. Then we study the prefill and decoding latency with varying user input lengths (from 10 to 2K tokens), using a fixed context length of 128K tokens.
Figure 6 presents the decoding latency per token of AttettoalAtteenion and other baselines. As context length increases, the decoding latency of full attention scales linearly, as loading the whole KV cache from CPU to GPU saturates the PCIe bandwidth. Because of fewer tokens involved in the attention computation, the latency of other basclines remains relatively low, but they suffer

---

Full attention
MagicPIG
IVF
Prefill latency
Decoding latency 0.20
InfLLM
Flat
RetrievalAttention
Decoding latency [s]
Decoding latency (5)
2.0
Prefill latency
0.15
1.0
2-
0.10
0.5
0.0
0.05
16K
32K
64K
128K
10
100
1K
2K
Context Length (log scale)
Input Length
Figure 6: Decoding latency (s) context length
Figure Prefill and decoding latency (s) of varying input length on Llama-3-8B.
varies from 16K to 128K on Llama-3-8B.
significant drop in model accuracy. Due to efficient attention- vector search, ReterivalAtteenion achieves 7.93 x and 2.80 latency reduction compared to Flat and IVF for the 128K context, and the latency shows a slow growth with context length due to sub-linear time complexity of the index traversal. Benefiting from high recall with minimal data scanning, RetrievalAAteenioon avoids memory bandwidth and computation contention caused by parallel retrieval across multiple heads on the CPU side. We present detailed latency breakdown in Appendix A.6 Compared with Flat and IVF, Attention effectively reduces the index search latency by 94% and 79%, respectively. As shown in Appendix A.5] this advantage becomes more pronounced with longer context lengths.
We compare both prefill and decoding latency as the user input ranges from 10 to 2K tokens under fixed context of 128K tokens and show the results in Figure 7 During the prefill stage, Retrieval loads the KV cache of the fixed contexts into the GPU memory to compute and generate the first token. The prefill latency increases slightly with the length of user input but remains low because the attention computation for the fixed context are avoided. Moreover, the decoding latency remains nearly constant, as short user inputs can be efficiently computed on the GPU together with the static pattern. Moreover, thanks to the parallel CPU-GPU design of RetrievalAAteenionn, attention computation on the GPU could be effectively overlapped with CPU computation.
4.4 Sensitive Study and Micro Analysis
Scalability. Figure presents the scalability of Retrieval/Attention on Llama-3-8B model under 128K context length. The throughput improves by 1.6x as the batch sizes increases. This improvement arises because RetrievalAttention offloads most of the KV cache to high-capacity CPU memory and transfers computation results of small size to the GPU. Full attention reaches throughput of only 0.45 tokens per second, since the PCIe transfer of all the KV cache becomes bottleneck even with a batch size of 1. The throughput of RetrievalAttennion eventually bound by the peak processing capacity of CPU.
Throughput (token/s)
20
RetrievalAttention
16
Full attention
12
Batch Size
Figure 8: Throughput of different batch sizes on Llama-3-8B
Index building Time. Table 4 presents the index building costs of AetrievalAtteenion for con- text of 100K length. The costs mainly includes KNN search, projection and graph nodes linking. RetrievalAttention currently focus on scenarios with largely fixed prefix and the index building done offline, thereby not affecting prefill latency during online serving. We believe more advanced acceleration [36] deserves exploration minimize the building cost for more dynamic use cases.
Projection. We study the effect of the projection technique by comparing RetrievalAttention's 's vector index with a vector index without projection. Specifically, the index without projection includes query vectors with connections to key vectors that identified as critical by query vectors. Searching on such graph structure requires routing between query and key vectors, resulting in excessive accesses. Table 5 presents the percentage of scanned vectors required by the two indexes for achieving a recall @ 100 of 0.95. The results clearly show that indexes without projection are inefficient.

---

Table Index building time (s)
Table The effect of projection in vector index quality.
of 100K context length.
Model
Recall 100
Scanned vectors
Model
Time (s)
w/o projection ours
w/o projection ours
Llama-3-8B
788.2
Llama-3-8B
0,951 / 0.954
13.2% 1.7 %
Yi-9B
1121.2
Yi-9B
0,956 / 0.954
23.4% 1.6%
Yi-6B 746.1
Yi-6B
0.954 / 0.958
27.8% 1.7 %
Static pattern size. presents RetrievalAAtteniions 's accuracy on RULER under variable static pattern sizes. The model used Llama-3-8B and the context length is 128K. As the size of static pattern increases, the accuracy gradually improves and tends to be stable. This is consistent with previous works that the number of sink tokens is relatively small [3]. The decoding latency remains unchanged as the static pattern is computed efficiently on the GPU. Moreover, our method also allows setting tokens at any position a static pattern.
Table 6: Performance (%) of Retrieval/Attenioon with different static pattern sizes on RULER.
Sizes
Accuracy (%)
16+64
73.93
32+128
74.15
64+256
74.52
128+512
74.70
5 Related Works
Scaling Context Windows of LLMs. Recent studies have explored extending the context window of pre-trained LLMs, thereby enhancing their capability to handle more complex applications. Mem- orizing Transformer (37 and Focused Transformer [X incorporate external memory mechanisms for context storage and leverage kNN-based retrieval to access relevant past information. However, they require additional model training to mitigate distraction issues. In contrast, work builds upon existing long-com long-context LLMs and enhances long-context inference efficiency without requiring any model retraining or architecture modification.
KV Cache Compression. As the context length increases, the enormous KV cache not only puts a strain on GPU memory but also slows down inference speed. making it one of the most significant challenges in the long- long-context LLM inference. Some works [39 40 13 41 42 [4] attempt to prune the KV cache by leveraging the sparsity of attention. However, these methods often suffer from significant model accuracy drops due to the dynamic nature of attention sparsity. Some works [43 44 quantize the KV cache size by approximating high-precision floating points with discrete low-bit values. These approaches are orthogonal to our work.
Dynamic Sparse Attention. By identifving the dynamic nature of important KV vectors for different queries, recent works 17
45 retain all KV cache and dynamically attend to different parts of them. RetrievalAttennion instead organizes the KV cache using attention- ANNS indexes, allowing the retrieval of important tokens with high recalls and low cost. Concurrent work PQCache [46) employs product quantization to compress KV vectors and retrieve critical tokens. Similarly, GPU-only AnteeezedAttentiion [ employs K-means clustering strategy, similar to IVF. However, these methods fail to address the -uttt--distribution (OOD) issue in attention mechanisms, often requiring retrieval of a large portion of the KV cache (e.g.. 20% in PQCache) maintain high model accuracy. Recent efforts [8) further build on our observations to develop vector storage systems for end-to-end long-context accelcration.
15
Conclusion
We propose RetrievalAttennion, method that offloads KV vectors CPU memory and leverages vector search for dynamic sparse attention to minimize inference cost. Retrieval/Atention identifies the different distributions of the query and key vectors and employs an attention -wwrre approach to efficiently retrieve critical tokens for model generation. Experimental results demonstrate that AtteetoalAAteentoon effectively achieves 7.93x and 2.80 decoding speedup than exact KNN and tra- ditional ANNS methods, on single RTX4090 GPU for context of 128K tokens. RetrievalAttention the first system that supports long-contex inference for 8B-level models on a single RTX4090 GPU (24GB) with acceptable latency and without compromising model accuracy.
10

---

References
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 5998-6008, 2017.
[2] Gemini Team. Gemini 1.5: Unlocking multimodal understanding across millions tokens of context, 2024.
[3] Meta. The Ilama 4 herd: The beginning of new era of natively multimodal ai innovation. https://ai. meta.
2025.
[4] Zukang Yang, Zixuan Zhu, and Jennifer Zhu. Curiousllm: Elevating multi-document ques- tion answering with Ilm-enhanced knowledge graph reasoning. In Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 3: Industry Track), pages 274-286, 2025.
[5] Qinyu Luo, Yining Ye, Shihao Liang, Zhong Zhang, Yujia Qin, Yaxi Lu, Yesai Wu, Xin Cong. Yankai Lin, Yingli Zhang, et al. Repoagent: An llm-powered open-source framework for repository-level code documentation generation. arXiv preprint #rXiv2242.166677 2024.
[6] Hao Li, Chenghao Yang, An Zhang, Yang Deng, Xiang Wang, and Tat-Seng Chua. Hello again! llm-powered personalized agent for long-term dialogue. arXiv preprint arXiv:2406.05925 2024.
[7] KVCache.Ar AI. Ktransformers, Accessed: 2025-05-15.
http:///wcccccceaaaaaaaaaaaarrrr 2024.
[8] Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Ré, lon Stoica, and Ce Zhang. Flexgen: high-throun high-thrroughut generative inference of large language models with single gpu. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JoLR.org, 2023.
[9] Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh. Monishwaran Maheswaran, June Paik, Michael W Mahoney, Kurt Keutzer, and Amir Gholami. Squeezed attention: Accelerating long context length Ilm inference. arXiv preprint a1.1v2411..996 2024.
[10] LMCache, Lmcache.
http://gttttubboooo/////aeeaaaaa 2024. Accessed: 2025-05-
15.
[11] Yuhan Liu, Hanchen Li, Kuntai Du, Jiayi Yao, Yihua Cheng, Yuyang Huang, Shan La, Michael Maire, Henry Hoffmann, Ari Holtzman, et al. Cachegen: Fast context loading for language model applications. ArXiv preprint, abs/2310.07244 2023.
[12] Yichuan Deng, Zhao Song, and Chiwun Yang. Attention is naturally sparse with gaussian distributed input, 2024.
[13] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, 2024.
[14] Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. SnapKV: LLM knows what you are looking for before generation. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[15] Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, and Maosong Sun. InfLLM: Training-free long-context extrapolation for LLMs with an efficient context memory. The Thirtt-eighth Annual Conference on Neural Information Processing Systems, 2024.

---

[16] Luka Ribar, Ivan Chelombiev, Luke Hudlass- Hudlass-Galley, Charlie Blake, Carlo Luschi, and Douglas Orr. Sparg attention: Bandwidth-efficient Ilm inference. In Forty-firss International Conference on Machine Learning, 2024.
[17] Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. QUEST: Query-aware sparsity for efficient long-context LLM inference. Forty-first International Conference on Machine Learning. 2024.
[18] Stanislav Morozov and Artem Babenko. Non-metric similarity graphs for maximum inner product search. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Desa-Bianchi, and Roman Garnett, editors, NeurIPS 2018, pages 4735-4735, 2018.
[19] Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Hao, Xu Han, Zhen Thai, Shuo Wang, Zhiyuan Liu, and Maosong Sun. ocBench: Extending long context evaluation beyond 100K tokens. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pages 15262-15277, Bangkok, Thailand, 2024. Association for Computational Linguistics.
[20] Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jía, and Boris Ginsburg. RULER: What's s the real context size of your long-context language models? In First Conference on Language Modeling, 2024.
[21] Hao Liu, Matei Zaharia, and Pieter Abbeel. Ringattention with blockwise transformers for near-infinite context. In The Twelfth International Conference on Learning Representations, 2024.
[22] Sivic and Zisserman. Video google: A text retrieval approach object matching in videos. In Proceedings ninth IEEE international conference on computer vision, pages 1470-1477. IEEE, 2003.
[23] Mengzhao Wang, Xiaoliang Xu, Qiang Yue, and Yuxiang Wang. comprehensive survey and experimental comparison of graph-based approximate nearest neighbor search. seed, 4(2):1.
[24] Mathijs Douze, Alexandr Guzhva, Chengqi Deng, Jeff Johnson, Gergely Szilvasy, Pierre- Emmanuel Mazaré, Maria Lomeli, Lucas Hosseini, and Hervé Jégou. The faiss library. 2024.
[25] Yu A Malkov and Dmitry Yashunin. Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs. IEEE transactions on pattern analysis and machine intelligence, 42(4):824 836, 2018.
[26] Prasanta Chandra Mahalanobis. On the generalized distance in statistics. Sankhyā: The Indian Journal of Statistics, Series A (2008-), 80:S1-S7, 2018.
[27] Shikhar Jaiswal, Ravishankar Krishnaswamy, Ankit Garg, Harsha Vardhan Simhadri, and Sheshansh Agrawal. Ood-diskann: Efficient and scalable graph ANNS for -iitt--diit queries. ArXiv preprint, abs/2211.12885 2022.
[28] Meng Chen, Kai Zhang. Zhenying He, Yinan Jing, and X. Sean Wang. Roargraph: A projected bipartite graph for efficient cross cross-modas modal approximate nearest neighbor search. Proc. VLDB Endow., (7(11):2735-274) 2024.
[29] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Ftestont:ntion: Fast and memory-efficient exact attention with 10-awareness. In Advances in Neural Information Processing Systems, 2022.
[30] Gradient AI. Llamma3-Sb-innsrret-2666 Llama- 3- 8B- Iww.-nnssttttt 262k 2024. Accessed: 2024-07-01.
http:://hggiingfaccccooooraaaaan
[31] 01-ai. Yi-9b-200k. http:////wwuggiige.....////////. 2024. Accessed: 2024- 07-01.
[32] 01-ai. Yi-6b-200k. httpss//uuginnafccccc//// 07-01.
2024. Accessed: 2024-
12

---

[33] Wonbeom Lee, Jungi Lee, Junghwan Seo, and Jaewoong Sim. InfiniGen: Efficient generative inference of large language models with dynamic KV cache management. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), pages 155-172, Santa Clara, CA, 2024. USENIX Association.
[34] Zhuoming Chen, Ranajoy Sadhukhan, Zihao Ye, Yang Zhou, Jianyu Zhang, Niklas Nolte, Yuandong Tian, Matthijs Douze, Leon Bottou, Zhihao Jia, et al. Magicpig: Lsh sampling for efficient I1m generation. ArXiv preprint, abs/2410. 16179, 2024.
[35] Greg Kamradt. Needle in haystack pressure testing llms. https://githun com/ 2023. Accessed: 2024-08-12.
Gkannratt/LLMTest_Meelll
[36] Yuxiang Yang. Shiwen Chen, Yangshen Deng, and Bo Tang. Paragraph: Accelerating graph indexing through gpu-cpu parallel processing for efficient cross-modal anns. In Proceedings of the 21st International Workshop Data Management New Hardware, DaMoN '25, New York, NY, USA, 2025. Association for Computing Machinery.
[37] Yuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, and Christian Szegedy, Memorizing transformers. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022, 2022.
[38] Szymon Tworkowski, Konrad Staniszewski, Mikolaj Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr Milos. Focused transformer: Contrastive training for context scaling. In Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023, 2023.
[39] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark W. Barrett, Zhangyang Wang, and Beidi Chen. H2O: heavy-hitter oracle for efficient generative inference of large language models. In Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
[40] Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of impor- tance hypothesis for IIm kv cache compression at test time. Advances in Neural Information Processing Systems, 36, 2024.
[41] Chi Han, Qifan Wang, Hao Peng, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. LM-infinite: Zero-shot extreme length generalization for large language models. In Kevin Duh, Helena Gomez, and Steven Bethard, editors, Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies (Volume 1; Long Papers), pages 3991-4008, Mexico City, Mexico, 2024. Associ- ation for Computational Linguisties.
[42] Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive KV cache compression for LLMs. In The Twelfth International Conference Learning Representations, 2024.
[43] Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen (Henry) Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. KIVI: tuning mieeg-free asymmetric 2bit quantization for KV cache. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024, 2024.
[44] Coleman Hooper, Schoon Kim, Hiva Mohammadzadch, Michael W. Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length L.M inference with KV cache quantization. In Advances in Neural Information Processing Systems Annual Conference on Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
38:
[45] Prajwal Singhania, Siddharth Singh, Shwai He, Soheil Feizi, and Abhinav Bhatele. Loki: Low- rank keys for efficient sparse attention. The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
13

---

[46] Hailin Zhang, Xiaodong Ji. Yilin Chen, Fangcheng Fu, Xupeng Miao, Xiaonan Nie, Weipeng Chen, and Bin Cui. Pqeache: Product quantization-bastt kveache for long context Ilm inference. ArXiv preprint, abs/2407.1285 112200, 2024.
[47]
Yangshen Deng, Zhengxin You, Long Xiang, Qilong Li, Peiqi Yuan, Zhaoyang Hong, Yitao Zheng, Wanting Li, Runzhong Li, Haotian Liu, Kyriakos Mouratidis, Man Lung Yiu, Huan Li, Qiaomu Shen, Rui Mao, and Bo Tang. Alayadb: The data foundation for efficient and effective long-context LLM inference. In Volker Markl, Joseph M. Hellerstein, and Azza Abouzied, editors, Companion of the 2025 International Conference Management of Data, SIGMOD/PODS 2025, Berlin, Germany, June 22-27, 2025, pages 364-377. ACM, 2025.
[48] Yaoqi Chen, Jinkai Zhang, Baotong Lu, Qianxi Zhang, Chengruidong Zhang. Jingjia Luo, Di Liu, Huiqiang Jiang. Qi Chen, Jing Liu, Bailu Ding, Xiao Yan, Jiawei Jiang, Chen Chen, Mingxing Zhang, Yuqing Yang, Fan Yang, and Mao Yang. Retroinfer: A vector-storage approach for scalable long-context LLM inference. CoRR, abs/2505.02922 2025.
[49] Shikhar Jaiswal, Ravishankar Krishnaswamy, Ankit Garg, Harsha Vardhan Simhadri, and Sheshansh Agrawal. Ood-diskann: Efficient and scalable graph anns for out-of-distribution queries. ArXiv preprint, abs/2211. 2022.
[50] Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, and Wen Xiao. Pyramidkv: Dynamic KV cache compression based on pyramidal information funneling. CoRR, abs/2406.02069, 2024.
[51] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. ArXiv preprint, abs/1904.10000 2019.
[52] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. ArXiv preprint, abs/2004.05150, 2020.
[53] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philipp Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Essun-TTie Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
[54] Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. ETC: Encoding long and structured inputs in transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 268 288, Online, 2020, Association for Computational Linguistics.
[55] Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir Abdi, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. MInference 1.0: Accelerating pre-filling for long-context LLMs via dynamic sparse attention. In The Thirth-eighth Annual Conference on Neural Information Processing Systems, 2024.
[56] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: 8th International Conference Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
The efficient transformer. In
[57] Amanda Bertsch, Uri Alon, Graham Neubig, and Matthew Gormley. Unlimiformer: Long-range transformers with unlimited length input. Advances in Neural Information Processing Systems, 36, 2024.
[58] Yuzhen Mao, Martin Ester, and Ke Li. Iceformer: Accelerated inference with long-sequence transformers on CPUs. In The Twelfth International Conference Learning Representations. 2024.
14

---

A Additional Experimental Details and Results
A.1 Detailed Configurations of the Baselines
Detailed experimental configurations are as follows: (1) StreamingLLLM: it uses 512 initial tokens and 1536 recent tokens for attention computation. (2) SnapK it employs an observation window of 32 tokens during the prefill phase and retains the top-2048 tokens. (3) InfLLM: it uses block size of 128, 4 representative tokens per block, and 16 blocks in total. (4) Quest: it uses a page size of 64, 2 representative tokens per page, and 32 pages in total. (5) InfiniGen: it fetches the top-20488 tokens after speculative attention. (6) MagicPIG: it follows the original paper settings for K and I in
different context lengths: 16K (7, 72), 32K (8, 95), 64K (9, 120), 128K (10, 150)
A.2 Additional Results RULER
Table demonstrates the advantage of RetrievalAttention on Yi-6B, with only 2.32% drop in accuracy. In contrast, other training-free methods suffer from larger accuracy degradation.
Table 7: Performance (%) of different methods on RULER of Yi-6B.
Methods
Act. Tokens
Claimed
16K
32K
64K
128K
Avg.
Full attention
128K
200K
69.12
61.64
58.36
55.77
61.22
StreamingLLM
2K
15.82
9.70
9.77
11.54
11.71
(-49.51)
Yi-6B
SnapKY
2K
45.36
36.11
33.43
29.53
36.11 (-25.11)
InfLLM
2K
34.11
27.11
25.28
25.33
27.96 (-33.26)
InfiniGen
2K
68.04
57.12
53.73
45.88
56.19 (-5.03)
Quest
2K
63.18
56.05
48.32
43,20
52.69 (-8.53)
MagicPIG
2K
66.61
57.69
51.10
44.04
54.86 (-6.36)
Flat
640+100
67.28
60.58
57.27
50.63
58.94 (-2.28)
IVF
640+100
67.00
58.94
55.99
50.31
58.06 (-3.16)
RetrievalAttention
640+100
67.49
59.46
57.20
51.44
58.90 (-2.32)
A.3 Additional Results on oo-Bench
As shown in Table Retrieval Attention achieves comparable accuracy to the full attention, benefiting from its efficient dynamic retrieval of important tokens. Static methods, such as StreamingLLM and SnapKV, lack this capability and, therefore, achieve sub-optimal accuracy. During token genera- tion phase, the critical tokens change dynamically according to the current query, invalidating the previously captured static patterns. InfiniGen exhibits a drop in model accuracy compared to full attention due inaccurate speculation of important tokens from previous layers. Although InfLLM and Quest support dynamic retrieval of relevant blocks, they achieve nearly zero accuracy in complex tasks (e.g., ... KV retrieval) due to the low accuracy of representative vectors. MagicPIG achieves good task accuracy in oc-Bench by employing LSH for critical token retrieval. However, the retrieval accuracy degrades in more challenging benchmarks such as RULER, as demonstrated in Table 2 Since Retrieval Attention can accurately identify the most relevant key vectors, it achieves the best accuracy in KV retrieval. Moreover, by retrieving more tokens (c.g.. top-2000 shown in the column of Retr.KVV KV) ) in KV retrieval, RetrievalAAtention achieves nearly the same accuracy as full attention, which demonstrates the effectiveness of our method in complex and dynamic tasks.
is worth noting that Flat and IVF need to scan 100% and 30% of the key vectors to achieve the same task accuracy Retrieval/ttenttio.. In contrast, RetrievalAttention only requires scan 1-3% vectors, resulting in much lower decoding latency.
A.4 Performance of Extremely Long-context Inference
Figure 9 shows the evaluation results of Attentin for extremely long contexts using the --ama-3-8B-1000 1048K model. AteetionlAttentinn still passes all test cases when ranging the context length from 250K to 1,000K, which demonstrates the robustness of our attention-aware indexes.
15

---

Table 8: Performance (%) of different methods and models co-Bench. The size of the static pattern is consistently 640 (128 initial tokens 512 tokens in the local window). All indexing- methods, including Flat, IVF, and RetrievalAAttetioon retrieve top-100 key vectors by default. In the relatively complicated task KV Retrieval, we include the results of retrieving top-2000 key vectors.
Methods
Act. Tokens
Retr.N
Retr.P
Retr.VV
Code.D Math.F
En.QA
En.MC
Avg.
Full attention
128K
100.0
100.0
17.5
19.0
39.5
9.1
68.0
50.4
Llana- 1oa----83 3-88
LtreamingLLM
2K
5.0
5.0
1.0
18.5
40.0
6.0
66.0
(-30.2)
SnapKV
2K
100.0
100.0
0.5
18.0
40.0
11.8
67.0
(-5.7) (-2.2)
InfLLM
2K
100.0
100.0
0.5
20.5
48.0
7.0
37.0
InfiniGen
2K
100.0
100.0
0.5
19.0
40.5
8.1
64.5
47.5 (-2.9)
Quest
2K
100.0
100.0
0.0
18.0
40.0
8.2
67.0
47.6 (-2.8)
MagiePIG
2K
9.5
19.5
8.7
67.5
49.4 (-1.0)
Flat
640+100V2K
100.0
100.0
8,5714.5
19.0
40.0
7.5
67.0
48.9
(115)-15576050 (-07)-
IVF
640+100/2K
94.0
100.0
9.5/14.0
19.0
40.0
7.8
67.0
48.2 (-2.2) / 48.8 (-1.6)
RetrievalAttention
640+100/2K
9.0/14.0
67.0
48.9 (-1.5) 49.6
Full attention
128K
100.0
30.5
25.5
36.5
9.8
67.0
LLreamingI-LM
2K
5.0
5.0
0.5
24.0
33.5
6.4
72.0
(-39
Yi.08
SnapKV
2K
63.0
100.0
0.5
23.0
33.0
10.3
68.5
42.6 (-1.2)
IafLLM
2K
100.0
100.0
0.5
20.5
43.0
9.4
44.0
(-7.5)
InfiniGen
2K
100.0
8.0
24.0
35.5
10.3
65.5
48.8
Quest
2K
99.5
100.0
0.0
24.0
32.0
11.0
69.0
47.9 (-4.9)
MagiePIG
2K
75.0
99.0
18.5
28.5
37.0
10.9
66.5
48.0 (-.8)
Flat
640+100/2K
100.0
100.0
21.0/30.0
23.0
35.0
68.5
512 (-16))
IVF
6402K
99.0
100.0
19.5/29.5
23.0
35.0
10.7
69.0
50.9 (-1.9) 52.3 (-0.5)
RetrievalAttention
640+100/2K
99.5
100.0
20.0/30.0
23.0
35.0
9.5
68.5
50.8 (-2.0) / 52.2(-0.6)
Full attention
128K
98.0
100.0
3.5
31.0
11.0
19.2
55.5
45.5
StreamingLLM
2K
5.0
0.5
27.5
11.0
12.2
54.0
16.5 (-29.0)
Y--B
SnapKV
2K
39.0
100.0
0.0
30.5
8.5
17.1
55.0
(-9.8)
InfLLM
2K
99.0
100.0
0,5
27.5
18.0
12.7
40.5
42.6 (-2.9)
InfiniGen
2K
99.0
100.0
1.5
30.0
11.5
18.1
54.0
44.9 (40.6)
Quest
2K
98.5
100.0
0.0
31.0
10.0
17.8
54.5
44.5 (-1.0)
MagicPIG
2K
69.5
99.5
1.5
31.5
12.0
19.3
54.5
41.1 (-4.3)
Flat
640+100/2K
98.5
100.0
2.5/3.0
30.5
16.0
17.7
54.5
45.7 (+0.2) 45.7 (+0.2)
IVF
640+100/2K
2.5/3.5
17.5
34.5
45.4 (-0.1) 45.6 (+0.1)
RetrievalAttentin
640+-100/2K
95.0
99.0
3.0/3.0
30.0
16.0
17.6
54.5
45.0 (-0.5) 45.0 (-0.5)
Haystack Llama-3-83-10488 RetriavaiAtterrtoo,
Degth Percent (%)
25
50-
84
250K
150R Context Length
1000K
Figure 9: Performance of RetrievalAAtention on Needle-in-a- in-a-haystack test of 1,000K context.
A.5 Decoding Latency A100 GPU
We test the generality of RetrievalAttention by measuring its performance on server with one A100 GPU (40GB) and one AMD EPYC 7V13 CPU with 12 physical cores and 220GB DRAM. We show the token generation latency of different methods on three models in Figure 1O Full attention with KV cache offloading suffers from the PCIe transmission bottleneck. Infl.LM achieve faster decoding speed, but this is at the expense a significant drop in model accuracy. contrast, RetrievalAttention does not compromise generation accuracy while achieving much lower decoding latency than Flat and IVF because of the efficient mitigation of the OOD problem.
We also evaluate how the decoding latency changes when the context length varies from 100K to 1M tokens on Llama-3-8B model and the results can be found in Figure [T] To make sure there is enough CPU memory to hold the KV cache and indexes, especially in the TM context scenario, we use powerful machine equipped with an AMD EPYC 7V12 CPU with 48 cores and 1.72 of memory. The machine is also equipped with the same 40GB A100 GPU. Different from Flat, IVE, and MagicPIG, whose latency numbers are sensitive to context size, RetrievalAttention only has minor latency increase when the context size increases 10x from 100K to IM.
16

---

Full attention
MagicPIG
Full attention
MagicPIG
IVF
InfLLM
Flat
RetrievalAttention
InfLLM
Flat
RetrievalAttention
Decoding lateney (5)
0.8
Decading latency
0.6
0.4
0.2
0.0
YI-6B
Yi-9B Model
Llama-3-88
100K
200K
500K
1000K
Context length
Figure 10: Decoding latency (s) of 128K context length on A100 GPU.
Figure 11: Decoding lateney (s) of Llama-3-8B as context length varies from 100K to 1M.
A.6 Detailed Latency Breakdown
Table Table 10 show detailed breakdown of decoding latency the RTX4090 GPU. Thanks to its decoupled CPU-GPU co-execution architecture, RetrievalAttennioo performs partial attention computation both the CPU and GPU, and transfers the CPU-side results to the GPU for final combination. The data transfer volume is extremely low. As shown in Table IO] the PCIe transfer overhead accounts for only 1.8% of the end-to-end latency with 128K context. Moreover, the computations of CPU and GPU could be fully overlapped.
Table 9: Decoding latency breakdown (s) on Llama-3-8B.
Methods
Retrieval
CPU Atm
Others
Total
Flat
0.791
0,016
0.041
0.848
IVF
0.243
0.017
0.040
0.300
RetrievalAttention
0.051
0.016
0.040
0.107
Table 10: Detailed breakdown (s) of each component.
Retrieval
CPU Attn
GPU Attn
PCIe Comm.
Result Combination
Other Layers
Total
Latency
0.051
0.016
0.006
0.002
0.006
0.026
0.107
A.7
Scalability with CPU Cores
Figure 12 demonstrates the scalability of RetricvalAttention. The model we used Llama-3-8B and the context length is 128K. The decoding latency decreases with more CPU cores. This stems from the ability to perform concurrent vector searches across attention heads.
0.3
RetrievalAttention
Decoding latency
0.2
0.1
0.0
Number of CPU cores
Figure 12: Decoding latency (s) under different CPU cores.
A.8 Index Recall vs. Scanning Vectors
We analyze the efficiency attention-aware vector index by examining the relationship between recall and the number of scanned key vectors. The number of key vectors scanned to achieve target
17

---

Q
RobustVamana, RetrievalAttention,
IVF,
RobustVamana,
RetrievalAttention,
YI-6B-200K
YI-9--00KK
1ama-3-8B-Instruct-v
Recall@:10
0.75
0.75
1.0
1.0
0.25
0.25
0.00
0.a5
0.00
0.05
0.00
0.05
0.0
0.5
0.0
0.5
0.0
0.5
Scanned Vectors (Percentage)
Figure 13: Recall scanning key vectors when using Q K the query.
recall serves as indicator of search efficiency (upper left is better). Figure 13 demonstrates that for the Q K search (q in attention), RetrievalAttention requires scanning only a very limited number of key vectors (1-3%) to reach recall @ 100 higher than 0.95, whereas traditional indexes necessitate retrieving a significantly higher number of keys. We also included an OOD-optimized solution, Robust [49]. for comparison. However, it performs poorly attention vectors. The efficiency of Retrieval Attentionn arises because effectively mitigates the OOD issue between query and key vectors in attention. In contrast, for the in-distribution K K search, all indexes exhibit. good performance.
A.9 Dynamic Retrieval Budget Allocation
We investigate the impact of adjusting the retrieval budget based on the sparsity degree across layers, using the budget allocation policy from PyramidKV [50]. We compare the performance of RetrievalAAtennioon with and without the PyramidK benchmark. For original detrievalAttention all layers. contrast, PyramidKV V dynamically adjusts the retrieval size across different layers, allocating more in lower layers and less in higher ones.
budget allocation strategy on the oc-Bench we set fixed budget of 2000 tokens for all heads in
The results in Table 11 show that Py allocation strategy achieves better performance in Retr.KV tasks, though it slightly decreases performance in the En.QA task. On average, the accuracy of dynamic budget allocation strategy slightly surpasses the original AetriivalAttentionn 8B. This indicates that dynamic budget allocation is promising, though it may require adjustment based on specific tasks.
Llama-3-
Table 11: Performance (%) of RetrievalAttention and RetrievalAAtention w/ PyramidKy V on oc-Bench.
Methods
Retr.N
Retr.P
Retr.V
Code.D
Math.F
En.QA
En.MC
Avg.
Full Attention
100.0
100.0
17.5
19.0
39.5
9.1
68.0
50.4
Ours
100.0
100.0
14.0
18.5
40.0
8.7
67.5
49.8
Ours w/ PyramidKV
100.0
100.0
16.0
18.5
40.0
8.5
67.5
50.1
A.10
Performance on the Larger Models
To demonstrate the generalizability AtteiivalAttention on larger models, we evaluate on Llama-3-70B-2613 using server equipped with eight 40GB A100 GPUs. We choose the most complex task, KV retrieval in oo-Bench, to stress test the efficiency of RetrievalAtttenion.
The results in Table 12 show that Retrieval achieves nearly the same task accuracy as the exact KNN method, Flat, and outperforms Quest. The decoding speed of AetrievalAAteetion faster than Flat as it effectively reduces the number of vectors to be scanned.
A.11 Index Quality
The index quality and construction time trade-off by adjusting parameters such as graph node degree and node linking traversal depth during index building. We conduct experiment of "the
18

---

Performance (%) and decoding latency (s) of different methods with Llama-3-700 on the Table 12:
complex task KV retrieval in co- Bo-Bench.
Full attention
StreamingI
Quest
Flat
RetrievalAAttnntonn
Accuracy (%)
35.0
0.0
13.0
24.0
23.5
Decoding latency (s)
14.8
0.14
1.36
5.48
1.52
number of accessed tokens to achieve recall @ 100=1.0 .. during search" versus "construction time" presented Table T3 All numbers are scaled based on index used our paper (i.e., the construction time and the number of accessed tokens for the index used in our paper are represented as 1). We observe that when index construction accounts for 51% the total time, search efficiency drops by 2.21 more tokens are involved in the computation. This reflects trade-off between search efficiency and index construction overhead the current implementation.
Table 13: Trade-off between construction time and required scanned tokens (normalized results).
Construction Time
1,00
0.84
0.73
0.63
0.51
0.48
0.45
Scanned Tokens
1.00
1.52
1.72
2.09
2.21
3.50
5.17
B RetrievalAAtention Algorithm
B.1 Formula of Combining Attention Results from the CPU and GPU Side
RetrievalAttennion partitions the KV vectors for attention into two disjoint sets: persistent ones on GPU (denoted as W) and dynamically retrieved ones on CPU (denoted as Q2).
It, W UΩ
Attention operation is applied to the two sets of KV vectors separately on CPU and GPU, generating two partial attention outputs (denoted as ow and og, respectively). To guarantee the approximated attention output equals to the attention computation on Iz,,, RetrievalAttention uses similar idea of FlashAttenticn 29
combine ow and on the following equations:
Ow Autn(q: K[W:,] v[W,:])
e
on Aun(q: K[S,;), v(2,:)
ΣiEn 8: 621-22 -2.
Σien e#:
71 Ow + on
where ž maxiew z; and 822 = maxien =; are the local maximum dot products set W and Ω respectively. And 71 and 72 are re-scaling factors to guarantee the attention output is the same as that on Ie,,, which are defined as follows:
e -1- Eiew ?1
EiEI,, Efi-8
(5)
72 EiErne E31-3
B.2 Overall Execution Flow
T summarrzess D the design of RetetivalAAteenioon and elaborates the procedure in an algorithm. (L) Offline vector index construction. RetrievalAttention pre-computes the KV vectors and builds the
19

---

attention-awee aware vector index offline for fixed contexts. The index and KV vectors are stored in CPU memory once the construction is completed. (2) Prefill phase. RetrievalAttention loads the cached KV vectors from CPU to GPU and computes full attention for the relatively short new user queries (#6). Since the number of these new tokens relatively small, we store the corresponding KV cache on the GPU with static patterns (+7). (3) Decoding phase. RetrievalAttentin performs vector search over the index for each query vector on the CPU, transfers the results to the GPU, and combines them with the GPU-computed results to obtain the final output. In detail, RetrievalA
computes
partial attention using the FlashAttention [2) kernel (+B) for static pattern KV vectors that are always stored in the GPU. In parallel with GPU computation, RetrievalAttention leverages the specially designed attention-awarr vector index to find the most relevant KV vectors and computes attention on CPU (#9) #10). Finally, CetrievalAtttntioon combines the partial attention outputs GPU and CPU using Equation and gets the approximated attention output 11
Algorithm 1: RetrievalAAtentioon
Input: Query vector q: E R
Data: KV Cache in GPU Kw, Vw E Rili
Data: CPU-based Vector Database H
Output:
Attention output o, E R
// Offline: Find the predictable KV vectors and build index
w <
PredictSatiicctiivTToooe...]
for fili H UW] do
// Will be computed on GPU H.remove();; Wiinser((i); for {ili & W" A H) do
// Will retrieved from CPU W.remove(i); H.insert(i);
// Online prefill user request
UK, Uy t Peeuuet W Winsert((; u/k e UK. Uv):
// Online Decoding
// Attention on GPU
Ow + FlashhAtenntionqq. Kw. Vw)
// Search in vector database to retrieve most
relevant vectors to compute attention on CPU
Ω < VectorSearch(qu))
10 og & AttentionCPU(());
// Combine partial attention outputs
11 0, = 71 Ow + 72 on: // Equation 4.5
C Theoretical Proofs of Retrieval Accuracy or Complexity Bounds
Existing widely-used vector retrieval methods, including graph and cluster- indexes, typically do not guarantee retrieval accuracy because their designs are heuristic-based on vector proximity. We follow common practice by stress-testing retrieval accuracy through extensive empirical evaluations. The results indicate that Atteeto provides better and more robust retrieval accuracy than other methods.
Index building time complexity: (1) KNN search for matching query vectors to key vectors takes O(n d) for the context length of and the vector dimensionality of d. (2) Projection node linking on CPU have a complexity of O(n C') + (n log n), where the constant C related to the in KNN and node linking requires traversing all nodes. The O(n log n) part refers connectivity enhancement following traditional graph-based indexing methods.
Index search time complexity: Assume the projection technique addresses the OOD problem (empirically verified), the time complexity of the index search close O n), similar to other graph-based methods.
20

---

D Additional Related Work
Sparse Transformers. Since the quadratic complexity of attention has become the bottleneck of LLM efficiency for long context applications, numerous works have been studied to design sparse transformers to reduce the computational and memory complexity the self-attention mechanism.
Some works restrict the attention computation to predefined patterns, including sliding windows [ST],
55 Some approaches use cluster- These solutions either require pre-training model from scratch or target limited scenarios like CPU-only, which do not work for our target out-of-box usage of LLMs the GPU-CPU architecture. Although some approaches [15 [6] exploit the dynamic sparse nature of LLMs, they often use some estimation using low-rank hidden states post-statistical approaches, which incurs high overhead but with low accuracy. Moreover, all these approaches have to maintain full KV vectors on GPU with only accelerated inference by reduced memory movement, which does not solve the challenge of limited GPU memory.
dilated windows [52], or a mixture of different patterns 153
54
based sparsity based on hash value [56] or KNN algorithms
57
58
E
Broader Impacts
RetrievalAttenticn bridges the fields of efficient attention and vector retrieval. It significantly accel- erates the decoding speed and reduces GPU memory consumption for long-context LLM serving, facilitating their deployment and practical application. By reducing latency and memory usage, Re- trievalAttention lowers the costs of long context serving, and thus contributes to the democratization of advanced AI. Moreover, it encourages further research and innovation in related areas.
F Limitations
AeteiivalAttention is designed to accelerate decoding and reduce GPU memory consumption in long- context inference. It focuses scenarios where fixed contexts are reused across multiple generations, allowing offline index construction. The prefill phase is not optimized in the current design. Exploring efficient index construction strategies to support more general long-context inference remains important direction for future work.
21

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We clearly and accurately present the contributions and scope in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of RetrievalAAtention in
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
22

---

Justification:
The methods proposed in this paper are based on empirical observations and do not involve theoretical proof.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide detailed descriptions of our experiments and implementation in and I
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
23

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide detailed descriptions of our experiments and implementation in and A Additionally, we submit our code in supplementary materials.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.com
aablic/guudees/Cddeebbbmsss www.oeeeSuuuiissiooooooooooo for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ wwtp:/cccppppgguuuiiaeeeae
Stteee///ooSSSeesssS.Si.a www.ssssiicgcyo for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide detailed descriptions of our experiments and implementation in and $A.I
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [NA]
Justification: We follow previous works by conducting tests on public benchmarks, without including error bars similar information.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
24

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide detailed descriptions of our experiments and implementation including computation resources in and AA.5
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics Ethics/https:
Answer: [Yes]
Justification:
We follow the guidelines of the NeurlPS Code Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. if there special consid- eration due to laws or regulations in their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the broader impacts of RetrievalAAtentioon in
Guidelines:
The answer NA means that there is no societal impact of the work performed.
25

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work does not involve the release of new models or data.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
The models, datasets, and codebase used in work comply with open-source licenses and can be used for scientific research.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
26

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/dataseta
datasets.papersstt
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [NA]
Justification: We provide new algorithm that can be used on many existing models, and do not introduce any more assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing experiments or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification:
This paper does not involve research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
27

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development this work does not involve L.MSs as any important, original, non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy httt::://earrpppcc/oooffeeereeee /c//////oeeeeeece//22////.. for what should or should not be described.
28


Reviews:
[{"review_id": "Bfb377j5FO", "reviewer": "Reviewer_efRu", "summary": "This paper introduces RetrievalAttention, a framework that accelerates long-context inference by pairing a retrieval‐based attention mechanism with a hybrid CPU–GPU KV-cache design. By building an attention-aware index over a large, mostly fixed prefix and fetching only the most relevant keys at run time, the approach reduces GPU memory usage and latency while preserving model quality. Experiments across several models and benchmarks show substantial speed-ups with minimal accuracy loss.", "questions": "1. Please report wall-clock build time, peak RAM, and index size for 128 K / 1 M tokens. How would you update the index if new tokens are appended or removed?\n\n2. Can RetrievalAttention operate when all tokens arrive online (no fixed prefix)? What is the end-to-end latency if the index must grow on the fly?\n\n3. Some attention-free approaches also claim large-context support. Could you add empirical results or explain why comparison is infeasible?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strenghts**:\n\n1. This paper systematically quantify how ANN search degrades when query and key distributions diverge, then designs an index that is query-aware by construction.\n\n2. The CPU–GPU split plus attention-aware index enables commodity GPUs to serve 100 K-scale contexts without quantization, delivering tangible memory and latency gains.\n\n**Weakness**:\n\n1. It seems the paper omits concrete numbers for build time, peak RAM, and index size at 128 K–1 M tokens, making deployability hard to judge.\n\n2. The evaluation parts did not include attention-free compression methods into comparison.\n\n3. The method assumes a largely fixed prefix. For interactive or rapidly changing inputs (chat, code completion) the indexing cost and latency overhead are unclear and untested.", "was_revised": true, "final_justification": "The authors clarified several design choices and provided further comparisons, which addressed my main initial concerns.\n\nThe rebuttal improved the presentation of the method, especially regarding how retrieval integrates with attention. Additional experiments strengthened the empirical evaluation, showing efficiency gains in long-context settings. While the improvements are not dramatic in all cases, they are consistent and relevant.\n\nGiven the clearer positioning and solid empirical validation, I find the contribution sufficient for acceptance.", "rebuttal": "We sincerely appreciate the reviewer's constructive feedback. As recognized by the reviewer, our key contribution lies in the novel application of attention-aware vector indexes to long-context LLM inference, coupled with systems designs for both high accuracy and efficiency. We will take the review comments to further improve the paper. We address each review comment below. \n\nW1. ***\"It seems the paper omits concrete numbers for build time, peak RAM, and index size at 128 K–1 M tokens, making deployability hard to judge.\"***\n\nQ1. ***\"Please report wall-clock build time, peak RAM, and index size for 128 K / 1 M tokens.\"***\n\nThank you for this comment. We provide the requested statistics for index build time, peak RAM consumption, and index size when constructing indexes on vector data derived from Llama-3-8B. The index is built layer by layer, from the initial to the final layer. Therefore, peak memory occurs during the index building of the final model layer, comprising the KV vectors, the cumulative index size from all preceding layers, and intermediate states of the single-layer index build.\n\nAs detailed in Table 1, all three metrics increase with context length due to the growing size of the KV vectors. The index building time is acceptable, given that this process is performed offline on CPUs. Regarding peak RAM and index size, the observed memory footprint remains manageable on commodity hardware.\n\nIndex building and memory footprint can be optimized with more advanced techniques. To further reduce build time, we conducted preliminary experiments using a recent GPU-based index building method [1] and observed a significant speedup of approximately 3$\\times$. To optimize memory consumption, we identify vector quantization and offloading data to storage devices as promising directions for future research.\n\n\n**Table 1:** Statistics of the index building.\n| Context length | 128K     | 256K     | 512K      | 1M        |\n| -------------- | -------- | -------- | --------- | --------- |\n| Build time     | 844 s    | 1568 s   | 3812 s    | 9174 s    |\n| Peak RAM       | 28.8 GiB | 56.6 GiB | 109.5 GiB | 208.8 GiB |\n| Index size     | 9.7 GiB  | 19.4 GiB | 39.1 GiB  | 77.9 GiB  |\n\n\nW2. ***\"The evaluation parts did not include attention-free compression methods into comparison.\"***\n\nQ3. ***\"Some attention-free approaches also claim large-context support. Could you add empirical results or explain why comparison is infeasible?\"***\n\nWe appreciate this comment. Since the reviewer did not specify which “attention-free methods” they referred to, we assume the term denotes models that avoid the use of attention mechanisms, such as Mamba [2] and RWKV [3].\n\nInitially, our focus was on comparing RetrievalAttention with other attention-based long-context methods, as our approach directly enhances the transformer's attention efficiency by leveraging sparsity. However, we acknowledge the growing interest in attention-free architectures.\n\nWhile these attention-free models offer compelling theoretical advantages, recent comprehensive long-context benchmarks, such as RULER [4], have indicated that their accuracy on diverse long-context tasks still lags significantly behind attention-based transformer architectures. Given that RetrievalAttention aims to achieve accuracy comparable to full attention models while significantly reducing computational overhead, it is expected to outperform these attention-free methods on long-context tasks. \n\nTo empirically validate this, we conducted an additional comparison with a popular attention-free model, Falcon-Mamba-7B-Instruct, on the RULER benchmark. As shown in Table 2, even at relatively modest context lengths (e.g., 16K), Mamba-7B's accuracy is substantially lower than that of both attention-based models (represented by Llama-3-8B) and RetrievalAttention. The result supports our initial reasoning and highlights that while attention-free methods are an active area of research for long-context modeling, they have not yet demonstrated competitive performance for practical deployment. \n\nWe plan to include a discussion of this aspect in the revised version of the paper, to clarify the competitive landscape and positioning of RetrievalAttention.\n\n**Table 2:** Accuracy of different models and methods on the RULER benchmark.\n|Context length|Falcon-Mamba-7B|Llama-3-8B|RetrievalAttention|\n|-|-|-|-| \n|16K|21.79%|89.27%|86.80%|\n|32K|5.66%|85.11%|84.78%|\n\n\nW3. ***\"The method assumes a largely fixed prefix. For interactive or rapidly changing inputs (chat, code completion) the indexing cost and latency overhead are unclear and untested.\"***\n\nQ2. ***\"Can RetrievalAttention operate when all tokens arrive online (no fixed prefix)? What is the end-to-end latency if the index must grow on the fly?\"***\n\n\nWe thank the reviewer for this comment and would like to clarify the scope of RetrievalAttention and potential extensions to fully dynamic scenarios. \n\nWe acknowledge that RetrievalAttention currently focus on scenarios with a largely fixed prefix. As explicitly stated in our paper, this design choice addresses a critical need in many practical use cases, such as document question answering, where the context (e.g., a document) remains consistent while queries change.\n\nWe agree with the reviewer that extending RetrievalAttention to handle rapidly changing inputs is an important direction. When all tokens arrive online, RetrievalAttention requires building indexes on the fly during the prefill phase. This dynamic index construction adds overhead and can impact the Time-to-First-Token (TTFT). For a 128K context, as shown in the Table 1 above, the TTFT would be approximately 884 seconds without specialized acceleration. GPU-based acceleration [1] can be used to reduce such overhead to approximately 300 seconds. We believe more advanced acceleration deserves exploration to minimize the indexing cost for more dynamic use cases. \n\nOnce the index is built, the subsequent token-by-token decoding latency of RetrievalAttention remains highly efficient, consistent with the results reported in our paper for fixed-prefix scenarios.\n\n\nQ1. ***\" How would you update the index if new tokens are appended or removed?\"***\n\nThanks for this comment. Our current design prioritizes index stability to enable index reuse across user requests. Therefore, for newly generated tokens, we choose to append them to the GPU memory as a static pattern, avoiding incremental index updates. This is a deliberate design choice to not change the index of the shared prefix (e.g., documents). Otherwise, indexes require handling deletion or reconstruction, impacting the serving latency. For example, if a long document is used as a shared prefix, inserting tokens at runtime would change the index, making it incompatible with other user requests using the same prefix context.\n\nWhile our design favors static appending for shared prefixes, we clarify that RetrievalAttention is capable of supporting incremental index updates, by adopting the update strategy in [5]. To validate this, we conducted an experiment using a 128K summary generation task with Llama3-8B, generating approximately 1.2K tokens. We inserted all generated tokens into the index and evaluated the recall performance. Our findings demonstrate that RetrievalAttention's index robustly maintains high recall quality even when new key vectors are continuously inserted. Moreover, by conducting incremental inserts, it takes a 10% latency increase in per token generation, making 0.107s per token to 0.118s per token.\n\n[1] ParaGraph: Accelerating Graph Indexing through GPU-CPU Parallel Processing for Efficient Cross-modal ANNS, DaMoN@SIGMDO 25\n\n[2] Mamba: Linear-Time Sequence Modeling with Selective State Spaces, 2023.\n\n[3] RWKV: Reinventing RNNs for the Transformer Era, 2023\n\n[4] RULER: What’s the Real Context Size of Your Long-Context Language Models? COLM 2024\n\n[5] RoarGraph: A Projected Bipartite Graph for Efficient Cross-Modal Approximate Nearest Neighbor Search，VLDB 2024"}, {"review_id": "eRnK44u53S", "reviewer": "Reviewer_b7VJ", "summary": "The paper introduces RetrievalAttention - a training-free approach for scaling long-context inference that off-loads most key/value vectors to CPU memory and serves them through an attention-aware ANN index, while keeping only a small “always-useful” window on the GPU. By retrieving roughly 1-3 % of tokens per query and merging the partial results with a FlashAttention-style kernel, it maintains accuracy within two percentage points of full attention but delivers up to an 8× speed-up and enables 128 K–1 M-token contexts on a single 24 GB GPU.", "questions": "1. In Appendix A.14, index building takes significant time—have you experimented with GPU-accelerated variants to make it more practical?\n2. How long does the offline KNN + projection index construction take for a 1 M-token corpus, and how often must it be rebuilt if the fixed context drifts? Providing wall-clock numbers would clarify practical deployability.\n3. How sensitive is performance to the choice of static patterns (e.g., varying the 640-token window)?", "limitations": "Yes - it discusses that the prefill side is not optimized in this work. It would be helpful to include CPU overhead in the limitations as well.", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strenghts:\n* Strong empirical results of decoding speedup at 128K context length, implying practical usability.\n* Novel framing of sparse attention as a cross-distribution MIPS problem and corresponding OOD-aware index design.\t\n* Thorough empirical study on open-source LLMs and three long-context benchmarks and ablations on index recall.\n\n\nWeaknesses:\n* The CPU cost of building the index need more transparency in the main part of the paper\n* Related-work section does not include prior work on kNN attention such as Memorizing Transformer (Wu et al., 2022) and Focused Transformer (Tworkowski et al., 2023)", "was_revised": true, "final_justification": "The authors clarified the points, I decide to keep my original positive score.", "rebuttal": "We sincerely appreciate the reviewer's constructive feedback and will take the review comments to further improve the paper. We address each review comment below. \n\nW1. ***\"The CPU cost of building the index need more transparency in the main part of the paper.\"***\n\nWe thank the reviewer for this comment and agree that the building cost should be presented in the main paper. We will do this in the revised paper. \n\nQ1. **\"In Appendix A.14, index building takes significant time—have you experimented with GPU-accelerated variants to make it more practical?\"**\n\nYes, we have tried to integrate a recent GPU-based index building method [1] into RetrievalAttention. Our preliminary experiments shows that it significantly reduces the building time by approximately 3$\\times$.\n\nQ2. ***\"How long does the offline KNN + projection index construction take for a 1 M-token corpus, and how often must it be rebuilt if the fixed context drifts? Providing wall-clock numbers would clarify practical deployment.\"***\n\nThanks for this comment. We evaluated the index building time using CPU for all layers on longer contexts, following the same setting in Appendix A.14. \n\nAs detailed in Table 1, index building time increases with context length, because of the increasing size of KV vectors. For the 1M tokens, it takes 9174 seconds. As noted in the response to the first question, GPU-based acceleration can be used to reduce this cost.  \n\n**Table 1:** Index building time with different context lengths.\n| Context length | 128K     | 256K     | 512K      | 1M        |\n| -------------- | -------- | -------- | --------- | --------- |\n| Build time     | 844 s    | 1568 s   | 3812 s    | 9174 s    |\n\n\nRegarding the case of fixed context drift, we categorize it into two scenarios. \n1. If the original context remains unchanged and new content is only appended to the end of the original one, RetrievalAttention is capable of supporting incremental index updates with relatively low cost, by following the update strategy in [2]. \n2. If the content of the fixed context changes, the KV cache would change so we need to rebuild the indexes on the updated KV cache.\n\n\nW2. ***\"Related-work section does not include prior work on kNN attention such as Memorizing Transformer (Wu et al., 2022) and Focused Transformer (Tworkowski et al., 2023)\"***\n\nWe appreciate this comment and agree that a discussion about KNN attention is important to clarify the landscape and positioning of RetrievalAttention.\n\nWhile Memorizing Transformer and Focused Transformer also leverage kNN-based retrieval, RetrievalAttention differs from them in the design goals and application scenarios. Memorizing Transformer and Focused Transformer focus on extending the context window and sometimes requires model training to mitigate the distraction issue. In comparison, RetrievalAttention builds upon existing models with long-context capacity and aims to improve the long-context inference efficiency.\n\nWe will incorporate these related works into the revised revision of our paper.\n\n    \nQ3. ***\"How sensitive is performance to the choice of static patterns (e.g., varying the 640-token window)?\"***\n\nThanks for this comment. We conducted a sensitivity analysis by varying the size of the static pattern using Llama-3-8B on the RULER benchmark (128K). \n\nAs shown in the Table 2, as the size of the static pattern increases, the accuracy gradually improves and tends to be stable. This is consistent findings in previous works that the number of sink tokens is relatively small [3]. We also observe that the decoding latency of RetrievalAttention remains unchanged, because the static pattern is computed efficiently on the GPU and fully overlapped with index operations on the CPU.\n\nWe will include above experiments and analysis in the revised version of paper.\n\n**Table 2: RetrievalAttention's accuracy with different static pattern sizes.**\n|Static pattern size|16+64|32+128|64+256|128+512|\n|-|-|-|-|-|\n|Accuracy on RULER|73.93%|74.15%|74.52%|74.70%|\n\n\n[1] ParaGraph: Accelerating Graph Indexing through GPU-CPU Parallel Processing for Efficient Cross-modal ANNS, DaMoN@SIGMOD 2025.\n\n[2] RoarGraph: A Projected Bipartite Graph for Efficient Cross-Modal Approximate Nearest Neighbor Search，VLDB 2024\n\n[3] Efficient Streaming Language Models with Attention Sinks. ICLR 2024"}, {"review_id": "h3ukytrpEH", "reviewer": "Reviewer_6mEN", "summary": "This paper presents RetrievalAttention, a system that accelerates long-context inference by offloading the KV cache to a CPU-based vector index. The work provides a strong analysis showing that standard ANNS methods fail due to an OOD problem between query and key vectors. The core contribution is an attention-aware index that adapts a projection technique from prior work to solve this issue, achieving significant speedups with minimal accuracy loss.", "questions": "- why is a simple IVF often comparable to RetrievalAttention? (e.g. Table 3, Yi-*B)", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**S1.** Provides a simple and effective system design for efficient long-context inference.\n\n**S2.** The paper is easy to follow, with a clear and compelling motivating analysis of the OOD problem.\n\n**W1.** The paper's core technique is heavily adapted from RoarGraph [28] but isn't explained in detail, making the work not self-contained and its novelty appear limited.\n\n**W2.** Key design choices, such as the vector projection technique and the static token heuristic, lack comparative ablation studies to justify their selection over other alternatives.", "was_revised": true, "final_justification": "authors address all the concerns so keeping the rating", "rebuttal": "We sincerely appreciate the reviewer's constructive feedback. As recognized by the reviewer, our key contribution lies in the novel application of attention-aware vector indexes to long-context LLM inference, coupled with systems designs for both high accuracy and efficiency. We will take the review comments to further improve the paper. We address each review comment below. \n    \nW1. ***\"The paper's core technique is heavily adapted from RoarGraph [28] but isn't explained in detail, making the work not self-contained and its novelty appear limited.\"***\n\nThanks for your valuable feedback. We will detail the techniques adapted from RoarGraph to make the paper self-contained in a revised manuscript. Below we explain how we adapt the projection technique to RetrievalAttention. \n\nThe projection technique aims to construct a graph index by retaining only key vectors, while ensuring that relevant key vectors identified by the query vectors remain closely connected in the graph structure. Therefore, benefiting from the elimitation of query vectors, it reduces the memory footprint and accelerates the search process. \n\nDuring index construction, we first use query vectors to identify the top-$m$ important key vectors to query vectors with highest inner product, as demonstrated in Figure 4(b) in our paper. After that, it comes to the projection phase to form a graph index. \n\n- The first phase is to connect $m$ key vectors, which are recognized as relevant by a query vector, in the graph structure. This process iterates for each query vector with top-$m$ key vectors. After this process, we can obtain a graph data structure that only has key vectors. However, the connectivity of this graph may not be good enough because some key vectors do not have enough neighbors. \n- The second phase is to enhance the connectivity of the graph structure. For each key vector, we follow the traditional graph index building to perform the nearest neighbor search in the graph to add more connections if one key vector does not have enough neighbors. \n- One of the most important adaption of RetrievalAttention is that, we use Euclidean distance instead of inner product score to evaluate distance between key vectors during the projection. We made this design choice because we found that, in the Q/K vectors from attention mechanism, some key vectors may dominate the neighboring space and have too large inner product score. Other key vectors tend to connect to these vectors, which will have large inner product score, within a limited node degree. This makes other key vectors hard to connect with each other and incurs low connectivity. However, we found that changing it to the pure geometry space metric (Eudclidean distance) can effectively modeling key vectors' relationship, helping establish edges between key vectors successfully.\n\nAlthough projection technique is introduced by RoarGraph, the technique details regarding the adaptation to RetrievalAttention would enhance the novelty of our paper. We will take this review comment to add more details and further improve the paper. \n\n\nW2. ***\"Key design choices, such as the vector projection technique and the static token heuristic, lack comparative ablation studies to justify their selection over other alternatives.\"***\n\nThank you for your suggestion. We conducted ablation experiments of the two techniques and make the following analysis.\n\n1. We first study the effect of the projection technique by comparing RetrievalAttention's vector index with a vector index without projection. Specifically, the index without projection includes query vectors with connections to key vectors that identified as critical by query vectors. Searching on such bi-directional graph structure will require routing between query and key vectors, resulting in excessive accesses. \n\nFor the Q/K vectors dumped from different models, Table 1 list the percentage of number of accessed vectors required by the graph index without and with projection for achieving recall@100$\\ge$ 0.95. The results show that indexes without projection are inefficient to achieve high accuracy.\n\n**Table 1: The effect of projection in vector index quality.**\n|Model|Recall@100 (w/o projection) | Scanned vectors (%) (w/o projection) | Recall@100 (RetrievalAttention, w/ projection)| Scanned vectors (%) (RetrievalAttention, w/ projection)|\n|---|---|---|---|---|\n|Llama3-8B-Instruct-262k| 0.951 | 13.2% |**0.954**|**1.7%**|\n|Yi-6B-200K| 0.954 | 27.8% |**0.958**|**1.7%**|\n|Yi-9B-200K| 0.956 | 23.4% |**0.954**|**1.6%**|\n\n2.  We conducted a sensitivity analysis by varying the size of the static pattern using Llama-3-8B on the RULER benchmark (128K). As shown in the Table 2, as the size of the static pattern increases, the accuracy gradually improves and tends to be stable. This is consistent findings in previous works that the number of sink tokens is relatively small [1]. We also observe that the decoding latency of RetrievalAttention remains unchanged, because the static pattern is computed efficiently on the GPU and fully overlapped with index operations on the CPU.\n\nWe will include above experiments and analysis in the revised version of paper.\n\n**Table 2: RetrievalAttention's accuracy with different static pattern sizes.**\n|Static pattern size|16+64|32+128|64+256|128+512|\n|-|-|-|-|-|\n|Accuracy on RULER|73.93%|74.15%|74.52%|74.70%|\n\n\nQ1.***\"why is a simple IVF often comparable to RetrievalAttention? (e.g. Table 3, Yi-*B)\"***\n\nThank you for the comment. You have pinpointed a crucial trade-off that is important to our work's motivation: the balance between retrieval accuracy and the computational overhead required to achieve it. \n\nA simple IVF can sometimes achieve comparable accuracy because it incurs a significantly higher search cost (scan 30% of all vectors) than RetrievalAttention (scan 1--3% of vectors). Our work aims to achieve strong accuracy while dramatically reducing search overhead. As shown in the Table 7 of Appendix A.8 in the original submission, the index retrieval latency of IVF is approximately 5$\\times$ higher than RetrievalAttention, and its end-to-end decoding latency is 2.8$\\times$ higher.  \n\nWe will clarify this in the revised version of this paper. \n\n[1] Efficient Streaming Language Models with Attention Sinks. ICLR 2024"}, {"review_id": "xBUJS4mJif", "reviewer": "Reviewer_krY7", "summary": "This paper presents RetrivalAttention, which dynamically retrieves the KV cache to the GPU during inference to reduce the GPU memory usage and speed up the inference.\n\nUnlike traditional methods that do static or dynamic KV cache eviction, this paper offloads the KV cache to the CPU and builds a vector-search index that helps nearest neighbor search and retrieval during the decoding process.\n\nEvaluation results show that the proposed solution can help run very long context inference on a consumer-grade GPU with limited memory, without causing too much quality drop.", "questions": "1. What's the detailed setup of \"full attention\" in Figure 6? Why does the decoding latency grow superlinearly as the context length increases? Shouldn't the time complexity be O(N)?\n\n2. I'm wondering if the proposed solution can increase the maximum batch size during the LLM inference, and thus improve the overall LLM inference throughput?", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "# Strengths\n\n1. It's innovative to use the vector-search methods for the KV cache retrieval during the LLM inference.\n2. The evaluation is solid, including both quality and system efficiency evaluation.\n\n# Weaknesses\n\n1. Clarity issue about the evaluation setup: the \"full attention\" baseline is not clearly explained in the evaluation setup.", "was_revised": true, "final_justification": "I think this is a good paper and should qualify for NeurIPS, so the final score is 5", "rebuttal": "We sincerely appreciate your constructive feedback, which is invaluable for further improving our paper. We address each review comment below. \n\nW1.***\"Clarity issue about the evaluation setup: the \"full attention\" baseline is not clearly explained in the evaluation setup.\"***\n\nQ1.***\"What's the detailed setup of \"full attention\" in Figure 6? Why does the decoding latency grow superlinearly as the context length increases? Shouldn't the time complexity be O(N)?\"***\n\n\nWe apologize for the confusion of the \"full attention\" baseline. We will clarify this in the revised version of the paper.\n\nFor the long-context inference on memory-constrained GPU devices, there are two common strategies for \"full attention\" implementation while avoiding Out-of-Memory (OOM) errors: KV recomputation or KV cache offloading to CPU memory. For the KV recomputation, the KV vectors are recomputed for all previous tokens at each decoding step, leading to quadratic computation complexity and high decoding latency. For example, decoding a single token with a 128K context length takes 43 seconds. This makes it impractical for real-time inference in long-context scenarios. \n\nTo mitigate the high latency of recomputation, we implemented \"full attention\" using KV cache offloading. Specifically, we offload the entire KV cache to the CPU memory and selectively load the KV cache for one model layer onto the GPU on demand during inference. This strategy drastically reduces the decoding latency to approximately 2.3 seconds per token for a 128K context, making it a more feasible baseline for comparison in long-context settings.\n\nRegarding the observed superlinear growth in decoding latency for \"full attention\" in Figure 6, we would like to clarify that the x-axis (Context Length) in the figure is plotted on a log scale for better clarity across a wide range of context lengths. While the plot visually suggests a non-linear trend of full attention, the actual decoding latency of \"full attention\" has an O(N) time complexity with respect to the context length N. This is because the KV cache size grows linearly with the context length, and consequently, both the data transfer latency over PCIe and the attention computation also scale linearly with N. \n\nTo address this, we will explicitly state in the caption of Figure 6 that the x-axis uses a log scale. We are also open to re-plotting the figure with a linear x-axis.\n \n\nQ2.***\"I'm wondering if the proposed solution can increase the maximum batch size during the LLM inference, and thus improve the overall LLM inference throughput?\"***\n\nThank you for this insightful question. Yes, our proposed solution indeed enables a larger maximum batch size during LLM inference, thereby improving overall LLM inference throughput. This is because RetrievalAttention can offload most KV cache to the CPU memory with high capacity. \n\nWe conducted experiments to demonstrate this capability, and results are presented in Appendix A.9 of our original submission. It shows that as the batch size increases, the decoding throughput of our method increases by approximately 1.6$\\times$. The throughput eventually becomes constrained by the CPU's compute power. \n\nWe will highlight this more prominently in the revised manuscript."}]

Author rebuttal:

Meta-review: This paper proposes an approach (RetrivalAttention) which offloads the KV cache to the CPU and builds a vector-search index that helps nearest neighbor search and retrieval during the decoding process. Empirical results show that the proposed solution can run long context inference on a consumer-grade GPU with limited memory with a small quality drop. Overall, reviewers pointed out as strengths the innovative use of vector-search methods for KV cache retrieval and the solid evaluation, and as weaknesses the lack of detail about the overheads of creating the index and a lack of clarity about the "full attention" baseline. The rebuttal clarified the second point and promised to provide details on the index creation. Overall this seems a solid paper. I urge the authors to follow the suggestions from the reviewers.
Decision: Accept (poster)