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
Primary area: foundation or frontier models, including LLMs
Keywords: Large Language Model; Efficient Inference; Long Context; Vector Retrieval

Abstract:
Transformer-based Large Language Models (LLMs) have become increasingly important. However, due to the quadratic time complexity of attention computation, scaling LLMs to longer contexts incurs extremely slow inference speed and high GPU memory consumption for caching key-value (KV) vectors. This paper proposes RetrievalAttention, a training-free approach to both accelerate attention computation and reduce GPU memory consumption. By leveraging the dynamic sparsity of attention mechanism, RetrievalAttention proposes to build approximate nearest neighbour search (ANNS) indexes for KV vectors in CPU memory and retrieve the most relevant ones through vector search during generation. Unfortunately, we observe that the off-the-shelf ANNS indexes are often ineffective for such retrieval tasks due to the out-of-distribution (OOD) between query vectors and key vectors in the attention mechanism. RetrievalAttention addresses the OOD challenge by designing an attention-aware vector search algorithm that can adapt to the distribution of query vectors. Our evaluation demonstrates that RetrievalAttention achieves near full attention accuracy while only requiring access to 1–3% of the data. This leads to a significant reduction in the inference cost of long-context LLMs, with a much lower GPU memory footprint. In particular, RetrievalAttention only needs a single NVIDIA RTX4090 (24GB) to serve 128K tokens for LLMs with 8B parameters, which is capable of generating one token in 0.188 seconds.

Decision: reject
Reviewer scores: 8, 8, 3, 3

Reviews and discussion:
<thread>
creation date:
1729465243779

thread_type:
official_review

<thread_object>
creator: reviewer_bxgi

<content>
summary:
This paper introduces a new KV-cache optimization method called RetrievalAttention, which accelerates attention computation and reduces GPU memory consumption. It leverages the dynamic sparsity of the attention mechanism. It uses an attention-aware vector search algorithm that can adapt to the distribution of query vectors, to build indexes for KV vectors in CPU memory and retrieve the most important ones through search during the token generation. The authors compare the performance of RetrievalAttention in long-context LLM inference against full attention and other state-of-the-art methods. The comparisons mainly focus on the accuracy and latency of the proposed approach.

strengths:
1. The task this paper tries to solve is crucial. A longer context window for LLMs allows the model to (1) enhance the consistency and coherence flow of information across a larger span of text; (2) improve understanding of complex text; (3) handle multi-step reasoning.
2. The idea of CPU-GPU offloading for LLM KV-cache optimization is interesting. Although there are several works already contributing in similar directions, they mainly focus on different aspects such as channel-dimension reduction and block-based organization.
3. Identify the challenge of OOD in using the ANNS index for attention computation.

weaknesses:
1. Some details in the method section are unclear and could benefit from further elaboration, particularly regarding the ANNS algorithm and CPU-GPU co-execution. Please refer to Question (1) for more specific inquiries.
2. Certain baselines are missing, such as QUEST[1].
3. I couldn't find the code implementation of the proposed method, which makes it difficult to assess its performance in practice.

questions:
1. Will the nearest-keys set for prefilled queries (indexed during the prefill phase) be updated during token generation? If not, will the retrieval recall degrade as more new tokens are generated?
2. Even if the nearest-keys set is updated, in extreme cases like code generation, where there are very few prefilled tokens and many newly generated tokens, how can the method ensure that the prefilled tokens remain representative enough for the query vectors of the newly generated tokens to effectively use them as a bridge to find the nearest key vectors?
3. For CPU-GPU co-execution, it is unclear how to separate the two disjoint sets of KV cache vectors between CPU and GPU. From my understanding, it seems challenging to ensure that the significant keys selected by ANNS are not included in the set of fixed initial tokens or within the last sliding window.
4. Although does not support GQA, I believe it is still worthwhile to include QUEST[1] as one of the baselines for comparison with vanilla attention.

I am willing to raise my score if the authors can address the limitations mentioned above. 

References:
[1] https://github.com/mit-han-lab/Quest/tree/main
</content>
</thread_object>

</thread>
<thread>
creation date:
1730340848293

thread_type:
official_review

<thread_object>
creator: reviewer_bojw

<content>
summary:
The authors propose a method that offloads KV Cache to the CPU combined with an ANNS index for sparse attention computation, aim for addressing GPU memory burdens during inference and accelerating decoding speed.

strengths:
1. The approach of using Approximate Vector Search is intuitive and appears reasonable in this offloading attention computation.
2. The authors propose a simple but effective ANNS index that addresses potential OOD issues in this scenario.
3. This method outperforms baseline methods in multiple benchmarks.

weaknesses:
1. It seems that this work relies heavily on obtaining global attention weights during the prefilling process to build the ANNS index. This may be incompatible with the existing FlashAttention technique, which is essential for efficient long-sequence prefilling computation. FlashAttention reduces prefilling overhead by using computation fusion, which avoids explicitly computing global attention weights, yet provides final outputs directly—making it a commonly used approach in long-sequence inference. Empirically, this issue will significantly increase the computing time and memory usage in the prefilling phase. In evaluation, the authors only provided decoding latency and did not consider potential additional overhead caused by the incompatibility with FlashAttention. Since other baselines generally do not encounter this issue, the authors should include experimental results to address these concerns.
2. The article should adopt a more appropriate baseline. Current baselines, such as SnapKV and StreamingLLM, discard a substantial portion of the KV Cache compared to the fully stored KV Cache method used in this paper, predictably leading to lower accuracy. For this work, a more suitable baseline for comparison with the fully stored Cache and offloading method would be:
(1.) Quest[1]: Quest also fully stores the KV Cache and retrieves several important cache entries during decoding. Although the authors note that Quest’s implementation does not support GQA and thus was excluded, since SnapKV can be integrated with GQA, the authors should consider implementing Quest with GQA support or explain why such an implementation would be challenging. (2.) InfiniGen[2]: InfiniGen is a comparable approach that offloads the Cache to the CPU, retrieving only a subset of the Cache during decoding to achieve approximate decoding, resulting in GPU memory savings and accelerated performance. This would likely be the closest baseline to the method in this paper.
3. It appears that this paper does not account for the overhead of index construction, which also directly impacts overall inference time. The authors should provide results showing the time and memory required for index construction in different input lengths.

[1] Tang, Jiaming, et al. "Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference." arXiv preprint arXiv:2406.10774 (2024).

[2] Lee, Wonbeom, et al. "{InfiniGen}: Efficient generative inference of large language models with dynamic {KV} cache management." 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24). 2024.

questions:
Does sequence length during prefilling impact the effectiveness of the index? Intuitively, longer sequences would result in larger indexes, suggesting that variations in sequence length could affect index performance. Testing the index's effectiveness in needle in a Haystack by locating the 'needle' cache across different sequence lengths could help clarify this effect.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730498335948

thread_type:
official_review

<thread_object>
creator: reviewer_afco

<content>
summary:
The paper targets the slow inference speed and high memory consumption issue in long context generation when conducting KV caching. They propose RetrievalAttention. The method is training free as it directly leverage the dynamic sparsity of attention mechanism to retrieve the most relevant KV vectors. One challenge discussed in the paper is the out-of-distribution (OOD) between query vectors and key vectors. To mitigate this problem, the authors have designed an attention-aware vector search algorithm. The paper evaluates the accuracy and efficiency of RetrievalAttention across various long-context benchmarks and both commodity and high-end GPUs.

strengths:
The problem addressed in this paper is highly meaningful. RetrievalAttention enables the execution of 8-billion-parameter models on a single RTX 4090 GPU, demonstrating the significance and practical impact of this research.

The observation regarding the critical key vectors could inspire further research and is beneficial for advancing the entire field.

The experiments are comprehensive and effectively support the main points.

The paper is well-written and easy to follow.

weaknesses:
My primary concern is that powerful CPUs (in terms of compute speed and memory) are necessary to ensure this method achieves good performance. Otherwise, it could become a bottleneck. For instance, most KV vectors will be offloaded to the CPU memory. The experiment was conducted using Intel i9-10900X CPU with 20 cores and 128GB of DRAM, which is quite powerful. Therefore, I think this method distributes the large cost(computation and memory) between the CPU and GPU to reduce the strain on the GPU. However, the total cost reduction is less significant than just comparing the GPU cost. Additionally, the communication between the CPU and GPU could also negatively impact the speed. The evaluation results may vary across different server architectures.

questions:
As the recovery rate is still 71%, does this mean that some tokens are consistently important? Is there any idea about the overlap between the dynamic and static selected critical tokens? 

Does the vocabulary size influence efficiency?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730687851369

thread_type:
official_review

<thread_object>
creator: reviewer_mywt

<content>
summary:
The paper introduces RetrievalAttention to improve inference speed and reduce memory usage for LLMs with long contexts by using Approximate Nearest Neighbor Search (ANNS) to retrieve only the most relevant KV vectors during token generation. By offloading most KV vectors to CPU memory and leveraging dynamic sparsity in attention, RetrievalAttention efficiently retrieves necessary vectors from a small subset, substantially lowering GPU memory requirements. It addresses the OOD problem between query and key vectors by designing an attention-aware search that adapts to query distribution, enabling retrieval with only 1–3% of data and minimal accuracy loss.

strengths:
- The paper introduces a training-free approach to accelerate long-context LLM inference by leveraging approximate nearest neighbor search (ANNS) for KV vector retrieval.
- The paper presents a well-thought-out and rigorously tested approach, with substantial empirical evidence demonstrating the performance gains of RetrievalAttention.
- RetrievalAttention addresses a significant limitation in the scalability of LLMs, especially as context lengths continue to grow in practical applications. By enabling long-context inference on a single commodity GPU with minimal accuracy loss, RetrievalAttention makes long-context LLMs more accessible and cost-effective, thus broadening their usability.

weaknesses:
- While the paper presents a unique approach to address the OOD challenge between query and key vectors, it lacks a comparison with alternative methods for handling OOD issues in high-dimensional vector retrieval.
- The evaluation is limited to relatively small models, so it’s unclear how well RetrievalAttention scales for significantly larger models (e.g., 30B, 70B, or more). Larger models would place greater demands on both memory and compute resources, potentially exposing limitations in RetrievalAttention's design, especially concerning the OOD handling mechanism and ANNS index efficiency at scale.

questions:
- RetrievalAttention applies the same retrieval sparsity across all layers, but it’s known that different layers and heads in transformers contribute differently to attention computation. Could layer-wise or head-wise retrieval strategies improve performance? 
- Have other ANNS index structures, such as product quantization or hierarchical clustering, been considered?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732204608492

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
Thanks to all the reviewers and the ACs for the detailed and insightful reviews. We have revised our paper and highlighted the modifications in brick-red color.
 
Before delving into the detailed response to each review comment, we would like to first clarify the core contribution of our work. We empirically validate the alignment of vector retrieval and sparse attention in long-context scenarios through extensive evaluations, such as RULER. Additionally, we identify the critical issue of out-of-distribution (OOD) discrepancies between queries and keys and propose a feasible solution to address it.
 
We acknowledge the importance of the points raised by the reviewers, including the overhead of index construction and long-generation scenarios, which deserves attention and further exploration. However, it is important to note that the current version of RetrievalAttention still has extensive practical use cases for cached contexts supported by mainstream LLM providers, allowing indexes to be constructed offline without impacting overall inference latency. We believe our work opens up the possibility for more in-depth studies and practical applications from the vector retrieval fields. Some of these aspects are already under consideration in our ongoing research, while others will benefit from collaborative efforts across the research community. 

Lastly, we are grateful for all reviewers' recognition of the potential impact of our work. We have supplemented additional experiments to further support the effectiveness of our method, as outlined below:

1. **Comparison with Other OOD Indexes.** As shown in Figure 6 of our revised paper, to reach a recall rate of 0.8, the well-known OOD index RobustVamana [1] requires scanning 25% and 75% vectors in the model Yi-9B and Llama-3-8B respectively. In contrast, RetrievalAttention only requires scanning a very limited subset of vectors (e.g., 1--3%) to achieve a high recall rate (0.95), which is consistent across all models.

2. **Evaluation on Large LLMs.** We evaluated various methods on the KV retrieval task using Llama-3-70B-262K, as shown in RTable 1. The results demonstrate that our method achieves both high efficiency and effectiveness, not only with 8B LLMs but also with larger models like 70B.

|Metric|Full|StreamingLLM|Quest|Flat|RetrievalAttention|
|-|-|-|-|-|-|
|Task accuracy(%)|35.0|0.0|13.0|24.0|23.5|
|Decoding latency(s)|248|0.14|1.36|5.68|1.62|

RTable 1. The task performance and latency of different methods on KV retrieval using Llama-3-70B-262K.

3. **Allocate Different Retrieval Budgets across Layers.** We have investigated the impact of adjusting the retrieval budget according to the sparsity degree across layers, by adopting the budget allocation policy from PyramidKV [2]. The results shown in RTable 2 indicates that dynamic budget allocation is promising for improving the task accuracy but may require task-specific allocation strategies.

|Methods|Retr.N|Retr.P|Retr.KV|Code.D|Math.F|En.QA|En.MC|Avg.|
|-|-|-|-|-|-|-|-|-|
|RetrievalAttention|100.0|100.0|14.5|18.5|40.0|8.7|67.5|49.9|
|RetrievalAttention w/ PyramidKV Budget Allocation|100.0|100.0|16.0|18.5|40.0|8.5|67.5|50.1|

RTable 2. Performance of different retrieval budget allocations on InfiniteBench using the Llama-3-8B-262K. 


4. **Add New Baselines: InfiniGen and Quest**. We have added new baselines including InfiniGen and Quest. RetrievalAttention performs better than the two new baselines on both InfiniteBench and RULER benchmark. Particularly, for the highly dynamic task KV Retrieval (Retr.KV in RTable 3), InfiniGen and Quest achieve zero accuracy due to inaccurate identification of critical tokens, while RetrievalAttention performs best and achieves nearly the same accuracy as full attention when retrieving top-2000 tokens.  

|Methods|Act. Tokens|Retr.N|Retr.P|Retr.KV|Code.D|Math.F|En.QA|En.MC|Avg.|
|-|-|-|-|-|-|-|-|-|-|
|FullAttention|128K|100.0|100.0|17.5|19.0|39.5|9.1|68.0|50.4|
|InfiniGen|2K|99.5|100.0|0|17.5|39.0|7.3|57.5|45.8|
|Quest|2K|100.0|100.0|0|18.0|40.0|8.2|67.0|47.6|
|Ours|640+100/2K|100.0|100.0|9.0/14.0|19.0|40.0|7.5|67.0|48.9/49.6|

RTable 3. Performance (%) of different methods in 128K context length on InfiniteBench using the Llama-3-8B-262K.


|Methods|Act. Tokens|S1|S2|S3|M1|M2|M3|MQ|MV|VT|CW|FW|Q1|Q2|Avg.|
|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|
|FullAttention|128K|100.0|100.0|100.0|98.0|98.0|73.0|94.5|97.0|87.0|1.0|72.2|58.5|44.5|78.7|
|InfiniGen|2K|99.0|91.5|24.5|82.5|25.0|0.0|30.3|27.8|67.3|1.2|45.4|33.0|32.5|43.1|
|Quest|2K|100.0|100.0|98.5|98.5|36.5|0.0|48.9|64.3|89.4|1.0|64.5|45.0|39.5|60.5|
|Ours|640+100|100.0|100.0|100.0|99.0|98.0|45.0|92.8|93.0|88.0|1.1|49.3|60.5|44.5|74.7|

RTable 4. Performance (%) of different methods in 128K context length on RULER using the Llama-3-8B-262K.

[1] OOD-DiskANN: Efficient and Scalable Graph ANNS for Out-of-Distribution Queries, 2022 

[2] PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling, 2024
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces RetrievalAttention, a novel method that leverages Approximate Nearest Neighbor Search (ANNS) to accelerate attention computation and reduce GPU memory consumption in large language models (LLMs). Notably, RetrievalAttention employs an attention-aware vector search algorithm to effectively address the challenge of out-of-distribution (OOD) data. This contribution is particularly significant for LLM applications requiring long context processing.

Extensive experiments demonstrate a considerable improvement in inference speed achieved by the proposed algorithm. While reviewers acknowledged the value of additional experiments conducted during the rebuttal phase, such as ablation studies, evaluations on larger LLMs, and comparisons with more baselines, concerns remain regarding the time required for offline index construction.