RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval

## Abstract

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 latency 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 use approximate nearest neighbor search (ANNS) indexes for KV vectors in CPU memory and retrieves the most relevant ones with 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 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 shows that RetrievalAttention only needs to access 1–3% of data while maintaining high model accuracy. This leads to significant reduction in the inference cost of long-context LLMs with much lower GPU memory footprint. In particular, RetrievalAttention only needs a single NVIDIA RTX4090 (24GB) for serving 128K tokens in LLMs with 8B parameters, which is capable of generating one token in 0.188 seconds.

## 1 Introduction

Recent transformer-based Large Language Models [] have shown remarkable capabilities in processing long contexts. For instance, Gemini 1.5 Pro [] has supported the context window of up to 10 million tokens.
While this is promising for analyzing extensive data, supporting longer context windows also introduces challenges for inference efficiency due to the quadratic complexity of attention computation. To enhance efficiency, KV caching, a technique that retains past Key and Value vectors, has been
widely adopted to prevent redundant computations. However, KV caching-based systems face two primary issues: (a) substantial GPU memory requirements particularly for long contexts, e.g., the Llama-2-7B model requires approximately 500GB per million tokens in FP16 format; and (b) inference latency increases linearly to the context size,
primarily due to the time needed to access cached tokens — a common issue across various computing devices, including GPUs.
Therefore, reducing token access and storage costs is vital for enhancing inference efficiency.

The solution lies in leveraging the dynamic sparsity inherent in the attention mechanism []. This refers to the phenomenon where each query vector significantly interacts with only a limited subset of key and value vectors, with the selection of these tokens varying dynamically based on individual queries.
Prior work [, , , , , ] has proposed various techniques to capitalize on this observation to improve the efficiency of attention computation. However, most of these methods identify important tokens either statically [, ] or heuristically [, , ], leading to imprecise approximations that often result in a significant performance drop.
We observe that the Approximate Nearest Neighbor Search (ANNS) index is particularly effective in this context. When using the inner product as the similarity measurement, ANNS precisely aligns with the attention mechanism, represented as $\mathbf{q}\mathbf{K}^{T}$. It can directly identify the most critical key vectors with maximum inner product value, yielding a higher accuracy compared to previous static or heuristic methods (as illustrated in Figure 1).

**Figure 1: RetrievalAttention achieves similar task accuracy as full attention but exhibits extremely low decoding latency.**
| Prompt Length | 128K | 256K | 512K | 1M |
| --- | --- | --- | --- | --- |
| Total Latency (s) | 32.8 | 111 | 465 | 1,765 |
| FFN (s) | 7.6 | 15 | 31 | 70 |
| Attention (s) | 25.2 | 96 | 434 | 1,695 |
| GPU Memory<br>KV Cache (GB) | 62.5 | 125 | 250 | 500 |

Leveraging ANNS for attention mechanisms presents a unique challenge: the out-of-distribution (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 distribution. However, due to the different projection weights for query and key vectors in attention mechanism, this assumption often does not hold in this context. Unfortunately, the effectiveness of ANNS degrades significantly if the query vectors are OOD. In particular, our empirical analysis shows that maintaining an acceptable accuracy level for inference quality can require scanning 30% of all key vectors to identify the critical ones, which fails to exploit the inherent sparsity in attention mechanism. To the best of our knowledge, this is the first work that identifies the challenge of OOD in using ANNS for attention computation.

In this work, we present RetrievalAttention, an efficient method for accelerating long-context LLM generation. RetrievalAttention employs dynamic sparse attention during token generation, allowing the most critical tokens to emerge from the extensive context data. To address the challenge of OOD, RetrievalAttention proposes a vector index tailored for the attention mechanism, focusing on the distribution of queries rather than key similarities. This approach allows for traversal of only a small subset of key vectors (1% to 3%), effectively identifying the most relevant tokens to achieve accurate attention scores and inference results.
In addition, RetrievalAttention reduces GPU memory consumption by retaining a small number of KV vectors in the GPU memory following static patterns (e.g., similar to StreamingLLM []) and offloading the majority of KV vectors to CPU memory for index construction. During token generation, RetrievalAttention efficiently retrieves critical tokens using vector indexes on the CPU and merges the partial attention results from both the CPU and GPU. This strategy enables RetrievalAttention to perform attention computation with reduced latency and minimal GPU memory footprint.

We evaluate the accuracy and efficiency of RetrievalAttention on both commodity GPUs (4090) and high-end GPUs (A100) on three long-context LLMs across various long-context benchmarks like $\infty$-Bench [] and RULER [].
For the 128K context on the 4090 GPU, RetrievalAttention achieves $4.9\times$ and $1.98\times$ decoding-latency reduction compared to the retrieval method based on exact KNN and traditional ANNS indexing respectively, while maintaining the same accuracy as full attention.
To the best of our knowledge, RetrievalAttention is the first solution that supports running 8B-level models on a single 4090 GPU (24GB) with acceptable latency and almost no accuracy degradation.

## 2 Background and Motivation

### 2.1 LLM and Attention Operation

In the generation of the $t$-th token, the attention mechanism computes the dot product between the query vector $\mathbf{q}_{t}\in\mathbb{R}^{1\times d}$ (where $d$ is the hidden dimension) and the key vector of past tokens $\mathbf{k}_{i}\in\mathbb{R}^{1\times d}$ (for $i\leq t$). This product is scaled by $d^{-\frac{1}{2}}$ and normalized via a Softmax function to yield the attention score $a_{t,i}$. These scores then weight the values $\mathbf{v}_{i}$, resulting in the output $\mathbf{o}_{t}$.

$$ $z_{i}=\frac{\mathbf{q}_{t}\cdot\mathbf{k}^{T}_{i}}{\sqrt{d}},\quad a_{t,i}=\dfrac{e^{z_{i}}}{\sum_{j=1..t}{e^{z_{j}}}},\quad\mathbf{o}_{t}=\sum_{i=1..t}a_{t,i}\cdot\mathbf{v}_{i}$ (1) $$

LLM inference contains two stages: the prefill phase and decoding phase.
The prefill phase, which only happens once, computes the keys and values of the prompt with a time-complexity $O(n^{2})$.
In the decoding (token generation) phase, the newly generated token becomes the new query and computes attention scores with all past key vectors.
One common optimization to avoid repetitive calculation is to cache past KV states in the GPU memory, thereby reducing the complexity to $O(n)$.

### 2.2 Expensive Long-Context Serving

Due to the quadratic time complexity of attention operation, serving long-sequence input incurs extremely high cost. Table 1 shows the inference latency 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.
Although KV cache can reduce the decoding latency, it demands a huge amount of GPU memory for long contexts.
As shown in Table 1, 500 GB memory is necessary for storing the KV cache when the context length reaches 1 million tokens, which is far beyond the GPU memory capacity of a single A100 GPU (80 GB).
Offloading and reloading the KV cache between GPU and CPU memory is a potential solution but incurs excessive commutation overhead over PCIe [], degrading the inference performance especially on commodity GPUs.

### 2.3 Dynamic and Sparse Attention

Figure: (a)
Refer to caption: figures/ai_2.png

Despite the large size of the context, only a small proportion of tokens actually dominate the generation accuracy.
Figure 2(a) shows the distribution of $|a_{t,i}|$ in Equation 1 for a query vector from Llama-2-7B with a prompt of 64,000 tokens. We observe that the top 500 tokens dominate the values of $|a_{t,i}|$, while the remaining tokens contribute approximately zero.
A high attention score indicates that two vectors are close, as measured by the inner product.
Therefore, the sparsity of attention scores means that the relevant keys to the query are very few.
We measure the mean-square-error (MSE) of the attention output if we use the top-$k$ $|a_{t,i}|$ as an approximation. We find that it only needs 36 tokens to achieve a very low MSE ($<10^{-6}$) of the full attention, showing a high sparsity ratio ($>99.9\%$).

Moreover, we observe that as the LLM continues generating new tokens, critical key vectors change dynamically. This means that tokens considered important in previous queries may not be critical in subsequent queries, and vice versa.
The dynamic sparsity shows a promising path to approximately compute attention with greatly reduced cost and without sacrificing the model accuracy. For each query, if we can accurately determine the relevant key-value vectors with higher importance, minimum GPU memory and a much lower time complexity can be achieved for attention computation.

### 2.4 Challenges of Off-the-shelf Vector Search

Finding the most similar vectors using ANNS indexes is a widely studied problem [, ], which semantically aligns with the goal of attention to find the nearest key vectors to each query vector in the inner product space. However, we find that naively applying off-the-shelf vector indexes fails to provide good performance because of the OOD issue between query ($Q$) and key vectors ($K$).

In conventional vector databases, the distribution of vectors between content and query is often well-aligned because they are derived from the same embedding model.
However, naively using off-the-shelf vector indexes for attention computation suffers from an inherent distribution gap between queries and keys, which are projected by different model weights. Figure 2(b) (focus on $Q$ to $K$ for now) demonstrates the performance of widely-used vector indexes supported by Faiss [] using a query vector to retrieve the most similar key vectors.
It compares the percentage of keys scanned and the corresponding recall achieved (i.e., the overlapping ratio between the retrieved top-100 results and the ground truth).
IVF [] requires scanning $\sim 30\%$ data for a recall rate higher than 0.95, and HNSW[] falls into a local optimum. The results show that traditional vector indexes require scanning a large number of vectors to achieve high recall, highlighting the challenge of performing efficient vector searches for attention.

Fundamentally, the difficulty is due to the OOD between query and key vectors.
We quantify this using Mahanobis distance [], which measures the distance from a vector to a distribution.
We sample 5,000 vectors from $Q$ and $K$ respectively as the query set and compute the the Mahanobis distance from the query set to the
remaining vectors in $K$.
Figure 2(c) shows that the queries 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 to retrieve critical tokens as nearest neighbors from the query vectors’ viewpoint.
In contrast, Figure 2(b) 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.
Similarly, query vectors also follow the same distribution.
For efficient vector search, the index must consider the OOD characteristic of the attention computation by design.

## 3 RetrievalAttention Design

In this work, we focus on the acceleration of token generation and assume the prefill of the long-context prompt is done in advance, which is widely supported by existing LLM service providers (e.g., context caching [] or separation of prefill and decoding [, ]).

Figure: (a) Overall design of RetrievalAttention.
Refer to caption: x4.png

We propose RetrievalAttention that leverages attention-aware vector search to accurately approximate attention computation by CPU-GPU co-execution. 3(a) shows the overall design of RetrievalAttention. Based on our observation in §2.3, We derive an approximated attention by selectively retrieving relevant key-value vectors while discarding those that are negligible(§3.1).
To efficiently supports long context, we offload most KV vectors to the CPU memory, build vector indexes, and use attention-aware vector search to find critical tokens. (§3.2).
To better exploit the GPU devices, we leverage
the attention scores obtained in the prefill phase to select a proportion of KV cache that are consistently important during the decoding phase and persist them on GPU devices.
RetrievalAttention computes partial attention with dynamically retrieved from CPU memory and persistent key-value vectors in GPU memory in parallel, and finally combine them together(§3.3).

### 3.1 Approximated Attention

Based on the Equation 1, RetrievalAttention approximates the full attention output $\mathbf{o}_{t}$ by selectively utilizing the KV vectors associated with high attention scores (i.e., $a_{t,i}$).
Specifically, we define $\mathcal{I}_{t,\epsilon}$ as a subset of token indices for which the attention score surpasses $\epsilon$.
Consequently, a sparse attention mechanism, which only considers tokens located in $\mathcal{I}_{t,\epsilon}$, can be defined as follows:

$$ $\mathbf{o}_{t}=\sum_{i\in\mathcal{I}_{t,\epsilon}}{a_{t,i}\cdot\mathbf{v}_{i}}+\xcancel{\sum_{i\not\in\mathcal{I}_{t,\epsilon}}{a_{t,i}\cdot\mathbf{v}_{i}}}\approx\sum_{i\in\mathcal{I}_{t,\epsilon}}{\tilde{a}_{t,i}\cdot\mathbf{v}_{i}}\quad\text{where}\quad\tilde{a}_{t,i}=\frac{e^{z_{i}}}{\sum_{j\in\mathcal{I}_{t,\epsilon}}{e^{z_{j}}}}$ (2) $$

Based on the above approximation, we build RetrievalAttention to only consider important key-value vectors (i.e., $\mathcal{I}_{t,\epsilon}$) which are retrieved by vector indexes.

### 3.2 Attention-aware Vector Search

For each pair of key and value vectors, we first decide whether to hold them in CPU or GPU memory (decision method is elaborated in §3.3). The KV vectors offloaded to CPU memory will be indexed by $\mathbf{k}_{i}\in\mathcal{R}^{d}$ and queried by $\mathbf{q}_{t}$ to find for the most relevant ones.

To accelerate the vector search during token generation, RetrievalAttention leverages the existing query vectors in the prefill phase to guide the index building for key vectors, efficiently mitigating the distribution gap.
As shown in Figure 3(b),
RetrievalAttention explicitly establishes connections from the query vector to its nearest key vectors (i.e., exact $k$-nearest neighbors, or KNN).
The KNN results can be efficiently computed via GPU, forming a mapping from query vector distribution to key vector distribution.
Using this structure, the decoding query vector can firstly searches its nearest query vectors and then obtain the most relevant key vectors through the distribution mapping.

Therefore, the KNN connections from query vectors to key vectors serves as a bridge to reconcile their distribution differences.
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 [] to eliminate all query vectors.
Specifically, we project the KNN connections into key vectors by linking key vectors that are connected to the same query vectors, which efficiently streamlines the search.
This process connects key vectors that are perceived close from query vectors’ perspective, allowing efficient index traversal for future query vectors.

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 a high recall, significantly reducing the index search latency by 74% compared with 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 predictable ones on GPU and the dynamic ones on CPU, and then combine the partial attention outputs together.

We leverage the patterns observed in the prefill phase to predict KV vectors that are consistently activated during token generation.
Similar to StreamingLLM [], our current implementation uses fixed initial tokens and last sliding window of the context as the static pattern, and persist them in the GPU cache.
RetrievalAttention can be adapted to utilize more complex static patterns [, ], achieving the best trade-off between low inference cost and high accuracy.
To minimize data transfer over the slow PCIe interface, RetrievalAttention independently computes the attention results for the CPU and GPU components and then combines them, inspired by the FastAttention []. The detailed combination algorithm and overall execution flow of RetrievalAttention is illustrated in §B.

## 4 Evaluation

In this section, we compare the performance of RetrievalAttention in long-context LLM inference against full attention and other state-of-the-art methods. Through experiments, we mainly explore the following questions:
(1) How does RetrievalAttention affect the model’s inference accuracy? Specifically, we assess the generation accuracy of RetrievalAttention and other methods across various downstream tasks.
(§4.2)
(2) Can RetrievalAttention efficiently reduce the token generation latency of long-context LLM inference? We compare the end-to-end decoding latency of RetrievalAttention with that of other baselines. (§4.3).

### 4.1 Experimental Setup

Testbed, Models, and Configurations.
We conduct experiments on a server equipped
with one NVIDIA RTX 4090 GPU (24GB memory) and an Intel i9-10900X CPU with 20 cores and 128GB DRAM.
The experiment results using NVIDIA A100 GPU can be found in §A.4.
We implement RetrievalAttention on three state-of-the-art long-context LLMs,
including Llama-3-8B-Instruct-262k[], Yi-6B-200K[], and Yi-9B-200K[].
To show practical speedup of RetrievalAttention and ensure the CPU memory consumption in long contexts do not exceed the capacity of DRAM, we follow previous work [] to run the benchmark in real-world single-batch scenarios.

Baselines. We compare RetrievalAttention with the following training-free baselines.
(1) Full attention without KV cache as well as the version with KV cache using vLLM [].
(2) StreamingLLM[]: it retains initial tokens along with fixed-length recent tokens in the GPU memory and discards remaining tokens.
(3) SnapKV[]: it caches the active tokens observed from the last window of the prompts.
(4) InfLLM[]: it separates KV cache of continuous token sequences into blocks and select representative vectors for each block. In the decoding phase, the current query scans all representative vectors and retrieve top-$k$ blocks with the highest similarity.
We do not include Quest [] for evaluation because its implementation (^1^11https://github.com/mit-han-lab/quest) does not support the GQA model.

To better assess the effectiveness of our method, we introduce two additional baselines using traditional vector search methods from Faiss[].
Specifically, Flat is an exact KNN method that performs a linear scan of all key-value vectors, whereas IVF indexes key vectors through clustering.
By default, all indexing-based methods retrieve top-100 nearest key vectors to the current query vector.

Benchmarks. We adopt three representative long-context benchmarks for evaluation.

- •
$\infty$-Bench[]: this benchmark consists of 7 tasks, including three retrieval tasks (PassKey retrieval, Number retrieval, KV retrieval) and four realistic tasks (code debugging, dialogue and multiple-choices questions). The average context length of $\infty$-Bench is over 100K tokens.
- •
RULER[]: a comprehensive long-context benchmark consisting of 4 categories and 13 tasks, including retrieval, multi-hop tracing, aggregation, and QA tasks. It contains prompts of different lengths from 4K to 128K, allowing us to determine the actual context window size of models.
- •
Needle-in-a-haystack []: it challenges the models to accurately retrieve information (the "needle") hidden within a lengthy document (the "haystack").

**Table 2: Performance (%) of different methods and models on $\infty$-Bench. The size of the static pattern is consistently 640 (128 initial tokens + 512 tokens in the local window). All indexing based methods including Flat, IVF and RetrievalAttention 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 | Tokens | Retr.N | Retr.P | Retr.KV | Code.D | Math.F | En.QA | En.MC | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Llama-3-8B | FullAttention | 128K | 100.0 | 100.0 | 17.5 | 19.0 | 39.5 | 9.1 | 68.0 | 50.4 |
| StreamingLLM | 640 | 5.0 | 5.0 | 0.5 | 18.5 | 39.5 | 5.9 | 66.5 | 20.1 (-30.3) |  |
| SnapKV | 2K | 100.0 | 100.0 | 0.5 | 18.0 | 40.0 | 11.8 | 67.0 | 48.2 (-2.2) |  |
| InfLLM | 640+2K | 100.0 | 100.0 | 0.5 | 20.5 | 48.0 | 7.0 | 37.0 | 44.7 (-5.7) |  |
| Flat | 640+100/2K | 100.0 | 100.0 | 8.5/14.5 | 19.0 | 40.0 | 7.5 | 67.0 | 48.9 (-1.5) / 49.7 (-0.7) |  |
| 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 | 100.0 | 100.0 | 9.0/14.0 | 19.0 | 40.0 | 7.5 | 67.0 | 48.9 (-1.5) / 49.6 (-0.8) |
| Yi-9B | FullAttention | 128K | 100.0 | 100.0 | 30.5 | 25.5 | 36.5 | 9.8 | 67.0 | 52.8 |
| StreamingLLM | 640 | 5.0 | 5.0 | 0.5 | 23.5 | 30.5 | 6.3 | 69.5 | 20.0 (-32.8) |  |
| SnapKV | 2K | 63.0 | 100.0 | 0.5 | 23.0 | 33.0 | 10.3 | 68.5 | 42.6 (-10.2) |  |
| InfLLM | 640+2K | 100.0 | 100.0 | 0.5 | 20.5 | 43.0 | 9.4 | 44.0 | 45.3 (-7.5) |  |
| Flat | 640+100/2K | 100.0 | 100.0 | 21.0/30.0 | 23.0 | 35.0 | 10.8 | 68.5 | 51.2 (-1.6) / 52.5 (-0.3) |  |
| IVF | 640+100/2K | 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) |
| Yi-6B | FullAttention | 128K | 98.0 | 100.0 | 3.5 | 31.0 | 11.0 | 19.2 | 55.5 | 45.5 |
| StreamingLLM | 640 | 5.0 | 5.0 | 0.5 | 29.5 | 8.0 | 11.2 | 54.0 | 16.2 (-29.3) |  |
| SnapKV | 2K | 39.0 | 100.0 | 0.0 | 30.5 | 8.5 | 17.1 | 55.0 | 35.7 (-9.8) |  |
| InfLLM | 640+2K | 99.0 | 100.0 | 0.5 | 27.5 | 18.0 | 12.7 | 40.5 | 42.6 (-2.9) |  |
| 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 | 98.0 | 100.0 | 2.5/3.5 | 29.5 | 16.0 | 17.5 | 54.5 | 45.4 (-0.1) / 45.6 (+0.1) |  |
|  | RetrievalAttention | 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) |

### 4.2 Accuracy on Long Context Tasks

$\infty$-Bench.
As shown in Table 2, RetrievalAttention 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 generation, the critical tokens change dynamically according to the current query, invalidating the previously captured static patterns.
Although InfLLM supports dynamic retrieval of relevant blocks, it achieves nearly zero accuracy in complex tasks (i.e., KV retrieval) due to low accuracy of representative vectors.
Since RetrievalAttention can accurately identify the most relevant key vectors, it achieves the best accuracy in KV retrieval.
Moreover, by retrieving more tokens (i.e., top-2000 shown in the column of Retr.KV) in KV retrieval, RetrievalAttention achieves nearly the same accuracy as full attention, which demonstrating the effectiveness of our method in complex and dynamic tasks.

It is worth noting that Flat and IVF need to scan $100\%$ and $30\%$ of the past key vectors to achieve the same task accuracy as RetrievalAttention. In contrast, RetrievalAttention only requires scan $1-3\%$ vectors, resulting in much lower decoding latency.

RULER.
Table 3 demonstrates that models utilizing RetrievalAttention achieves nearly the same task accuracy as full attention in different context lengths.
In contrast, other training-free methods experience an noticeable reduction in accuracy, particularly for longer context sizes like 128K, as they fail to capture dynamically changed important tokens.

Needle-in-a-haystack.
As shown in Figure 4, RetrievalAttention can effectively focus on information at various positions across different context windows, ranging from 4K to 128K. In contrast, methods like StreamingLLM encounter difficulties when critical information lies beyond the range of the static patterns, whose results are shown in §A.2.

### 4.3 Latency Evaluation

As the context length increases, the decoding latency of full attention significantly increases due to its quadratic time comlexity.
Enabling the KV cache (vLLM) incurs out-of-memory (OOM) issues due to limited GPU memory.
The latency of StreamingLLM, SnapKV, and InfLLM remains relatively stable because of constant tokens involved in the attention computation but they suffer significant accuracy degradation as demonstrated in §4.2. Due to efficient attention-aware vector search, RetrievalAttention
achieves $4.9\times$ and $1.98\times$ latency reduction compared to Flat and IVF for the 128K context. The detailed latency breakdown can be seen in the Table 6 of §A.3.

**Table 3: Performance (%) of different methods and models on RULER.**
|  | Methods | Act. Tokens | Claimed | Effective | 4K | 8K | 16K | 32K | 64K | 128K | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Llama-3-8B | FullAttention | 128K | 262K | 32K | 93.13 | 90.49 | 89.27 | 85.11 | 82.51 | 78.74 | 86.54 |
| StreamingLLM | 640 | - | <4K | 28.74 | 16.49 | 15.46 | 13.34 | 11.30 | 9.45 | 15.80 (-70.75) |  |
| SnapKV | 2K | - | 4K | 91.51 | 80.70 | 75.53 | 70.84 | 65.44 | 58.68 | 73.78 (-12.76) |  |
| InfLLM | 640+2K | - | 4K | 85.20 | 52.86 | 38.29 | 32.44 | 27.94 | 25.71 | 43.74 (-42.81) |  |
| Flat | 640+100 | - | 16K | 92.71 | 87.93 | 87.01 | 84.97 | 80.99 | 74.34 | 84.66 (-1.89) |  |
| IVF | 640+100 | - | 16K | 92.73 | 87.86 | 87.22 | 84.74 | 78.46 | 68.21 | 83.20 (-3.34) |  |
|  | RetrievalAttention | 640+100 | - | 16K | 92.64 | 88.46 | 86.80 | 84.78 | 80.50 | 74.70 | 84.70(-1.85) |
| Yi-9B | FullAttention | 128K | 200K | 8K | 91.02 | 86.62 | 82.85 | 73.17 | 67.08 | 60.51 | 76.87 |
| StreamingLLM | 640 | - | <4K | 27.96 | 16.03 | 14.89 | 9.48 | 10.95 | 12.05 | 15.23 (-61.65) |  |
| SnapKV | 2K | - | 4K | 90.39 | 75.59 | 64.48 | 48.70 | 39.28 | 32.97 | 58.57 (-18.30) |  |
| InfLLM | 640+2K | - | <4K | 82.66 | 50.36 | 36.17 | 28.20 | 22.65 | 20.94 | 40.16 (-36.71) |  |
| Flat | 640+100 | - | 8K | 91.09 | 87.71 | 84.42 | 74.58 | 66.16 | 59.50 | 77.24 (+0.37) |  |
| IVF | 640+100 | - | 8K | 91.03 | 91.04 | 83.85 | 72.19 | 65.13 | 58.04 | 76.29 (-0.58) |  |
|  | RetrievalAttention | 640+100 | - | 8K | 90.78 | 86.32 | 82.95 | 73.73 | 65.67 | 59.15 | 76.43(-0.44) |
| Yi-6B | FullAttention | 128K | 200K | <4K | 84.52 | 77.77 | 69.12 | 61.64 | 58.36 | 55.77 | 67.86 |
| StreamingLLM | 640 | - | <4K | 24.77 | 12.88 | 11.21 | 7.33 | 8.76 | 9.86 | 12.47 (-55.40) |  |
| SnapKV | 2K | - | <4K | 80.94 | 59.55 | 45.36 | 36.11 | 33.43 | 29.53 | 47.49 (-20.38) |  |
| InfLLM | 640+2K | - | <4K | 76.42 | 44.38 | 34.11 | 27.11 | 25.28 | 25.33 | 38.77 (-29.09) |  |
| Flat | 640+100 | - | <4K | 83.69 | 77.26 | 67.28 | 60.58 | 57.27 | 50.63 | 66.12 (-1.75) |  |
| IVF | 640+100 | - | <4K | 83.25 | 76.90 | 67.00 | 58.94 | 55.99 | 50.31 | 63.33 (-2.53) |  |
|  | RetrievalAttention | 640+100 | - | <4K | 83.01 | 76.56 | 67.49 | 59.46 | 57.20 | 51.44 | 65.86(-2.00) |

## 5 Related Works

To accelerate the long-context LLM inference,
some works [, , , , , ] attempt to compress the size of 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.

**Figure 4: Performance of RetrievalAttention in Needle-in-a-haystack.**
| Methods | 4K | 8K | 16K | 32K | 64K | 128K |
| --- | --- | --- | --- | --- | --- | --- |
| Full (without cache) | 0.527 | 1.167 | 2.672 | 6.214 | 15.263 | 43.927 |
| vLLM | OOM | OOM | OOM | OOM | OOM | OOM |
| StreamingLLM | 0.029 | 0.030 | 0.029 | 0.030 | 0.030 | 0.029 |
| SnapKV | 0.029 | 0.028 | 0.028 | 0.029 | 0.029 | 0.028 |
| InfLLM | 0.058 | 0.063 | 0.063 | 0.065 | 0.067 | 0.069 |
| Flat | 0.140 | 0.178 | 0.226 | 0.328 | 0.522 | 0.922 |
| IVF | 0.128 | 0.140 | 0.162 | 0.201 | 0.253 | 0.373 |
| RetrievalAttention | 0.137 | 0.144 | 0.156 | 0.162 | 0.169 | 0.188 |

FlexGen [] and Lamina [] offload the KV cache to CPU memory, but they struggle with slow and costly full-attention computation.
By identifying the dynamic nature of important KV vectors for different queries, recent work chooses to retain all of the KV cache and dynamically attend to different parts of KV vectors based on the current query. Quest [] partitions the KV cache into blocks and selects a representative key vector for each block. For a given query, it scans all representative key vectors and attends top-$k$ blocks with the highest attention scores. InfLLM [] adopts a similar strategy as Quest but offloads most KV cache blocks to the CPU memory to support longer contexts. Due to block-based organization and retrieval, the accuracy of representative vectors significantly impacts the effectiveness of those methods for obtaining important tokens. SparQ [], InfiniGen [], and LoKi [] approximate the most relevant top-$k$ keys corresponding to a given query by reducing the head dimension. RetrievalAttention instead organizes the KV cache using ANNS indexes, allowing the retrieval of important tokens with high recalls and low cost. The concurrent work MagicPiG [] and PQCache [] employ locality-sensitive-hashing (LSH) and PQ centroids respectively to retrieve critical tokens.
However, they fail to address the OOD issue in attention, necessitating retrieving a large portion of KV cache (e.g., 20%) to maintain the model accuracy.

## 6 Conclusion

We propose RetrievalAttention, a method that offloads most KV vectors to CPU memory and leverages vector search for dynamic sparse attention to minimize inference cost.
RetrievalAttention identifies the different distributions of the query and key vectors and employs an attention-aware approach to efficiently find critical tokens for model generation.
Experimental results demonstrate that RetrievalAttention effectively achieves $4.9\times$ and $1.98\times$ decoding speedup than exact KNN and traditional ANNS methods, on a single RTX4090 GPU for a context of 128K tokens. RetrievalAttention is the first system that supports running 8B-level LLMs with 128K tokens on a single 4090 (24GB) GPU with an acceptable latency cost and without compromising model accuracy.

## Appendix A Additional Experimental Details and Results

### A.1 Model Architecture

**Table 5: Architecture overview of LLMs.**
| Model | Total Layer | Query Head | KV Head |
| --- | --- | --- | --- |
| Yi-6B | 32 | 32 | 4 |
| Yi-9B | 48 | 32 | 4 |
| Llama-3-8B | 32 | 32 | 8 |

### A.2 Additional Results on Needle-in-a-haystack

Figure 5 shows the results of other methods on Needle-in-a-haystack benchmark. StreamingLLM can only find the correct answer when the needle’s position is within the static pattern. InfLLM maintains high performance with shorter context lengths. However, as the length increases, its performance shows a significant decline. Although SnapKV, Flat, and IVF perform well on this benchmark, we have analyzed their disadvantages in accuracy and latency in the previous evaluation.

Figure: (a)
Refer to caption: figures/needle_test_result/full_output.png

**Table 6: Latency breakdown of retrieval attention-based algorithm on Llama-3-8B.**
| Methods | Retrieval | Attention | Others | Total |
| --- | --- | --- | --- | --- |
| Flat | 0.798 | 0.083 | 0.041 | 0.922 |
| IVF | 0.250 | 0.084 | 0.039 | 0.373 |
| RetrievalAttention | 0.064 | 0.081 | 0.043 | 0.188 |

### A.3 Latency Breakdown

Table 6 presents the breakdown of end-to-end latency for different retrieval attention-based algorithms on a single RTX 4090 server under the 128K context length.
Since RetrievalAttention effectively addresses the issue of out-of-distribution, it only requires 34.0% of the time for vector search. In contrast, Flat and IVF spends 86.6% and 67.0% of time, respectively. This is because RetrievalAttention scans less data, avoiding memory bandwidth contention when multiple heads are performing parallel retrieval in the CPU side. This advantage becomes more pronounced with longer context lengths.

### A.4 Decoding Latency on A100

We test the generality of RetrievalAttention by measuring its performance on a server with one A100 (80GB) and one AMD EPYC CPU with 24 cores and 220GB DRAM. We show the token-generation latency of different methods on three models in Table 8. Since the KV cache of full attention is disabled, all prompt tokens need to be recalculated during the decoding, incurring a very high decoding latency. By enabling the KV cache with the PageAttention optimization in vLLM, the decoding latency is significantly reduced. However, vLLM suffers from OOM issue with the increase of context length, which we elaborate further later.
Other KV cache dropping or block retrieval methods including StreamingLLM, SnapKV, and InfLLM achieve faster decoding speed, but this is at the expense of a significant drop in model accuracy.
In contrast, RetrievalAttention does not compromise generation accuracy while achieving much lower decoding latency than IVF and Flat because of the efficient mitigation of out-of-distribution 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 Table 8. To make sure there is enough CPU memory to hold the KV cache and indexes especially in the 1M context scenario, we use a powerful machine equipped with an AMD EPYC 7V12 CPU with 48 cores and 1.72 TB of memory. The machine is also equipped with the same 80GB A100 GPU.
The decoding latency of full attention with KV state re-computation increases quadratically with the context size. With the KV cache enabled in the the GPU memory, vLLM starts triggering OOM issue when the context size is larger than 200K. Static KV dropping methods such as StreamingLLM have no latency increase due to the constant KV cache involved for attention computation. Different from Flat and IVF whose latency numbers are sensitive to context size, RetrievalAttention only has a minor latency increase (8$\%$) when the context size increases 10$\times$ from 100K to 1M.

**Table 7: Per-token generation latency (s) of 128K context-length on A100.**
| Methods | Yi-6B | Yi-9B | Llama-3-8B |
| --- | --- | --- | --- |
| Full(without cache) | 31.61 | 47.51 | 33.38 |
| vLLM | 0.030 | 0.044 | 0.033 |
| StreamingLLM | 0.032 | 0.047 | 0.031 |
| SnapKV | 0.033 | 0.05 | 0.033 |
| InfLLM | 0.069 | 0.11 | 0.068 |
| Flat | 0.541 | 0.802 | 0.564 |
| IVF | 0.309 | 0.468 | 0.345 |
| RetrievalAttention | 0.150 | 0.227 | 0.155 |

## Appendix B RetrievalAttention Algorithm

### B.1 Formula of Combining Attention Results from the CPU and GPU Side

RetrievalAttention partitions the KV vectors for attention into two disjoint sets: predictable ones on GPU (denoted as $\mathcal{W}$) and dynamically retrieved ones on CPU (denoted as $\Omega$).

$$ $\mathcal{I}_{t,\epsilon}=\mathcal{W}\cup\Omega$ (3) $$

Attention operation is applied to the two sets of KV vectors separately on CPU and GPU, generating two partial attention outputs (denoted as $\mathbf{o}_{\mathcal{W}}$ and $\mathbf{o}_{\Omega}$, respectively). To guarantee the approximated attention output equals to the attention computation on $\mathcal{I}_{t,\epsilon}$, RetrievalAttention uses a similar idea of FlashAttention [] to combine $\mathbf{o}_{\mathcal{W}}$ and $\mathbf{o}_{\Omega}$ in the following equations:

$$ $\displaystyle\mathbf{o}_{\mathcal{W}}$ $\displaystyle=\text{Attn}(\mathbf{q}_{t},\mathbf{K}[\mathcal{W},:],\mathbf{V}[\mathcal{W},:])$ $\displaystyle=\frac{\sum_{i\in\mathcal{W}}{e^{z_{i}-\tilde{z}_{1}}\cdot v_{i}}}{\sum_{i\in\mathcal{W}}{e^{z_{i}-\tilde{z}_{1}}}}$ $\displaystyle\mathbf{o}_{\Omega}$ $\displaystyle=\text{Attn}(\mathbf{q}_{t},\mathbf{K}[\Omega,:],\mathbf{V}[\Omega,:])$ $\displaystyle=\frac{\sum_{i\in\Omega}{e^{z_{i}-\tilde{z}_{2}}\cdot v_{i}}}{\sum_{i\in\Omega}{e^{z_{i}-\tilde{z}_{2}}}}$ $\displaystyle\mathbf{o}_{t}$ $\displaystyle=\gamma_{1}\cdot\mathbf{o}_{\mathcal{W}}+\gamma_{2}\cdot\mathbf{o}_{\Omega}$ (4) $$

where $\tilde{z}_{1}=\max_{i\in\mathcal{W}}{z_{i}}$ and $\tilde{z}_{2}=\max_{i\in\Omega}{z_{i}}$ are the local maximum dot products in set $\mathcal{W}$ and $\Omega$ respectively. And $\gamma_{1}$ and $\gamma_{2}$ are re-scaling factors to guarantee the attention output is the same as that on $\mathcal{I}_{t,\epsilon}$, which are defined as follows:

$$ $\displaystyle\gamma_{1}$ $\displaystyle=\frac{e^{\tilde{z}_{1}-\tilde{z}}\cdot\sum_{i\in\mathcal{W}}{e^{z_{i}-\tilde{z}_{1}}}}{\sum_{i\in\mathcal{I}_{t,\epsilon}}{e^{z_{i}-\tilde{z}}}}$ (5) $\displaystyle\gamma_{2}$ $\displaystyle=\frac{e^{\tilde{z}_{2}-\tilde{z}}\cdot\sum_{i\in\Omega}{e^{z_{i}-\tilde{z}_{2}}}}{\sum_{i\in\mathcal{I}_{t,\epsilon}}{e^{z_{i}-\tilde{z}}}}$ $$

### B.2 Overall Execution Flow

Figure: Algorithm 1 RetrievalAttention

Algorithm 1 summarizes the above design of RetrievalAttention and elaborate the procedure in an algorithm. At the beginning of each token generation, RetrievalAttention predicts active KV vectors and move them to GPU memory, and compute partial attention using the FlashAttention [] kernel (1 - 1). In parallel with GPU computation, RetrievalAttention leverages the specially designed vector database to find the most relevant KV vectors to compute attention on CPU (1 - 1). Finally, RetrievalAttention combines the partial attention outputs on GPU and CPU using 4 and gets the approximated attention output (1).

## Appendix C Implementation

RetrievalAttention builds one individual vector index for the KV cache in one attention head.
RetrievalAttention did several optimizations to optimize the prompt prefill, accelerate the vector search, and reduce the CPU memory usage.

#### Optimization for the Prefill Phase.

During the prefill phase, the full attention computation of the long context is required for generating the output vector for the next level of the LLM.
Simultaneously, we move the KV vectors to the CPU side for the ANNS index building.
To accelerate the overall prefill process,
we overlap the cache movement to the CPU with the full attention computation on the GPU in a pipeline manner.
To minimize peak GPU memory usage during the prefill phase, attention computation is performed sequentially across multiple attention heads. This approach only slightly impacts the attention computation speed, as longer prompts can fully leverage GPU parallelism with FlashAttention.

#### Multi-head Parallelism on the CPU side.

To speed up the dynamic sparse attention computation on the CPU, we exploit the multi-thread parallelism in vector databases by leveraging the multi-core ability of modern CPU architecture.
Specifically, since the computation of different attention heads is independent, we launch multiple threads for parallel searching across different vector indexes to reduce the overall latency on the CPU side. For grouped query attention (GQA) [], although multiple query heads could share the same key-value vectors, we observe that the query vectors from different query heads in the same group exhibit different vector distributions.
Therefore, we build one vector index for each query head to leverage the specific query distribution of each head.

#### Minimize the CPU Memory Usage.

To reduce CPU memory consumption, the indexes in the same attention group share one copy of KV vectors by only storing the pointers to KV vectors in each index.
In the future, we plan to utilize scalar quantization to further compress the KV vectors, implementing an 8-bit representation in place of the original FP32 format. This compression is promising to reduce memory usage while preserving computational efficiency. Importantly, our initial results demonstrate that this quantization approach does not compromise the accuracy of the model outcomes, maintaining performance equivalent to the full-precision representation.

## Appendix 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 studied to design sparse transformers to reduce the computational and memory complexity of the self-attention mechanism.
Some works restrict the attention computation to predefined patterns including sliding windows [], dilated windows [], or mixture of different patterns [, ]. Some approaches use cluster-based sparsity based on hash value [] or KNN algorithms [, ].These solutions either require pre-training a model from scratch or target limited scenarios like CPU-only, which do not work for our target to out-of-box usage of LLMs on the GPU-CPU architecture. Although some approaches [, ] exploit the dynamic sparse nature of LLMs, they often use some estimation using low-rank hidden states or 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 commodity GPUs with limited GPU memory.

Additionally, some approaches accelerate the inference by employing dynamically sparse attention patterns [], separating the prefill and decoding stages [, ], and utilizing sequence parallelism [, ]. These methods are orthogonal to ours and can be in conjunction with our approach.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
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?
soundness: 3
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.
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.
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.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
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?
soundness: 3
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.
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.
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.

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
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.
soundness: 2
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.
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.
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.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
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
soundness: 2
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.
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.
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.

Meta-review: 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.

Decision: Reject