RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval

W1. "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."

Q1. "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?"

Response: Thank you for your detailed comments. In our current implementation, we do not update the index and, consequently, the nearest-keys set for prefilled queries during the decoding phase. Instead, we maintain the newly generated tokens in the GPU memory as a static pattern and include all of them in the attention computation for subsequent generation steps. This design choice is based on the observation that newly generated tokens are typically much fewer in number compared to the long context.

The maximum generation length of long-context benchmark Ruler and InfiniteBench is 128 tokens. Existing studies [13][14] have also demonstrated that although modern LLMs can handle extremely long contexts (e.g., 128K to 1M tokens), the maximum generation length usually does not exceed 2K tokens. Therefore, newly generated tokens can either remain in the GPU memory as static patterns or be offloaded to the CPU memory and efficiently retrieved using a brute-force KNN method. Since we do not update the original index, the index quality for searching critical tokens within the prompt context remains unchanged as the decoding queries follow the same distribution as prefilled queries. To confirm this, we add Figure 9 in our revised paper, which illustrates that the Mahalanobis distances from the newly generated queries to the prefill queries are nearly identical to the distances among the prefill queries themselves.

Although we have adopted above design choices, our ANNS can accept incremental inserts without updating the nearest-keys of prefilled queries. Specifically, we can utilize the update strategy from RoarGraph [15] to achieve this. For each new key vector ($v$), we find the closest prefilled query ($q$), get the previous nearest-key set ($S$), and use it to select neighbors for $v$ without modifying the pre-existing set. This insert strategy is shown to be highly efficient while maintaining the indexing quality [15].

We have clarified these details in the revised paper.

Q2. "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?"

Response: Thank you for your comment. As stated in the introduction of our paper, we focus on scenarios involving long contexts, where the prefilled queries should be sufficiently large to represent the distribution of queries. In cases like the one you mentioned, since the generation limit of all LLMs is typically small (<2K tokens), it is feasible to perform full attention over the newly generated tokens or offload them to the CPU memory for retrieval using a brute-force KNN method, which remains efficient for a small number of tokens.

[13] WildChat: 1M ChatGPT Interaction Logs in the Wild, ICLR 2024

[14] LongWriter: Unleashing 10,000+ Word Generation From Long Context LLMs, 2024

[15] RoarGraph: A Projected Bipartite Graph for Efficient Cross-Modal Approximate Nearest Neighbor Search, VLDB 2024

W2. "Certain baselines are missing, such as QUEST."

Q4. "Although does not support GQA, I believe it is still worthwhile to include QUEST as one of the baselines for comparison with vanilla attention."

Response: Thanks for your suggestion. We have included Quest as a new baseline. We have also added another baseline InfiniGen, as requested by another reviewer. We evaluated their performance on InfiniteBench and RULER, with the results presented in RTable 3 and RTable 4, respectively.

InfiniGen exhibits a noticeable drop in model accuracy compared to full attention due to inaccurate speculation of important tokens from previous layers. Although Quest performs better than InfiniGen, it achieves zero accuracy in the highly dynamic KV retrieval task (Retr.KV in the InfiniteBench) due to low representativeness of centroid vectors. Moreover, both Quest and InfiniGen have severe performance degradation on the RULER benchmark. In contrast, RetrievalAttention performs best and achieves nearly the same accuracy as full attention across two benchmarks.

Regarding the decoding latency, Quest encounters OOM issues on the RTX 4090 due to its fully GPU-resident design, making it unsuitable for supporting 128K contexts in consumer-grade GPUs. Since the current decoding optimizations of InfiniGen does not support GQA for LLMs [8], we will implement InfiniGen for the GQA model and test its decoding latency in the future version of our paper.

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

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

W3. "I couldn't find the code implementation of the proposed method, which makes it difficult to assess its performance in practice."

Response: We are in the legal process to make our code public. Will release once it is ready.

Q3."For CPU-GPU co-execution, it is unclear how to separate the two disjoint sets of KV cache vectors between CPU and GPU..."

Response: In the design of RetrievalAttention, during the prefilling phase, we physically separate consistently important tokens and dynamic tokens by placing them in GPU memory and CPU memory, respectively. The ANNS index is built only for the remaining dynamic tokens in the CPU memory. Therefore, there is no overlap between static tokens in the GPU cache and dynamic critical tokens retrieved from ANNS indexes. The accurate identification of consistently important tokens is well-studied and is orthogonal to the design of RetrievalAttention. Currently, we adopt StreamingLLM to keep initial and last window tokens in the GPU cache, but we are adaptable to support more complex static patterns. We'll include this in the next version of the paper.

Thank you again for reviewing our work! Please let us know if we have sufficiently addressed your questions and concerns.

Authors

Thank you for your valuable feedback. We apologize for any confusion and would like to address your concerns as follows:

• Clarification of Attention Computation: We realize the description of "full attention over the newly generated tokens" in our previous reply was unclear and may have led to misunderstandings. Our current implementation incorporates newly generated tokens into the local window of the static pattern. However, the attention computation over the prompt context always relies on ANNS for dynamic sparse attention.

• Design Choice Rationale: The design choice mentioned above was made with a focus on long-context scenarios where the generation lengths are generally small (fewer than 2K tokens), compared to input tokens (e.g., 128K). This is a common setting in existing long-context benchmarks and real-world scenarios, like repo-level code debuging, long-document QA. The strong performance on downstream tasks, as demonstrated in our paper, supports the effectiveness of our method for these scenarios. Moreover, as the generation length increases, the index quality does not decrease even without updating the index. We demonstrate this by adding Figure 10 in our revised paper, where we generate 1200 tokens from the 128K context using a summary task. Figure 10 shows the index recalls on prefill tokens consistently remain at a high level (>0.95) throughout the decoding steps from the first token to the 1200th token. This results is corresponding to Figure 9, which we added previously.

• Ability to Handle Long Generation: As the number of generated tokens increases, we can insert the newly generated tokens into the index by adopting the update strategy from RoarGraph [15] instead of keeping them in the GPU cache. To show this, we have implemented the insert operation and added Figure 11 in our revised paper. In this experiment, we proactively insert the key vector of every newly generated token into the index at each decoding step, and recall rates were evaluated among all existing tokens. The variation of inserted index recall at each decoding step can be observed in Figure 11, the index recalls remain high (>0.95) across the whole 1200 decoding steps with a slight increase (6%) of per-token generation latency. This demonstrates that handling newly generated tokens is NOT a fundamental limitation of our approach.

• Future Work on Extremely Long Generation: We acknowledge that future LLMs may extend their generation lengths beyond 2K, particularly in models like O1 utilizing self-play reasoning. Therefore, we consider the study of extremely long generation as an important area for future work.

• Primary Contribution: Finally, we emphasize the primary contribution of our work: introducing vector retrieval methods into attention computation for long-context scenarios (Reviewer #MYWT, #Afco, #boJw). We designed a CPU-GPU framework utilizing OOD-aware vector indexes to address distribution discrepancies between queries and keys. While our current design choices impose limitations for certain long-generation scenarios, this is not the primary focus of our paper. These limitations can be addressed through incremental inserts to the index, as supported by our experiments (Figure 11 in our revised paper).

We hope that above points are helpful to clarify your concerns.

Thanks for your reply. While the scenario of very short input lengths with long outputs is not the primary focus of our paper, we have nonetheless conducted additional experiments, as suggested by the reviewer. It is important to note that our previous demonstrations and the revised paper have already established the index's capability to handle long-generation tasks within a 128K long-context. We further address the reviewer's concern in the following two aspects.

• Insertion latency. By conducting incremental inserts, it takes a 6% latency increase in per token generation, making 0.188s per token to 0.199s per token.

• Short Input Length and Long Generation. To assess the performance of the index on tasks involving short inputs and relatively long generations, we conducted experiments using the code generation task from LongBench. Specifically, we varied the input lengths between 151 and 719 tokens and allowed the model to generate up to 2,000 tokens. We conducted the incremental inserts to the index and measured the index recall rates at each step of the decoding process. As shown in the newly added Figure 12 in our revised paper, although there is a slight decline in recall rates as the generation progresses, the index maintains a high recall rate (> 0.95) even in long-generation scenarios where the output length significantly exceeds the input by multiple times ($2.8\times$ to $13\times$). This underscores the index's ability to effectively handle long generation tasks while preserving high recall, even as output length significantly grows.

Regarding two core concerns from the reviewer, (1) The nearest-keys set for prefilled queries is NOT updated during token generation, instead, we choose the lightweight update strategy from RoarGraph, (2) The above experiment demonstrates that the prefilled queries are representative enough for newly generated queries, even in the case that prefill context is very short.

Thank you to the authors for their detailed response. I would like to kindly reiterate my main concern: when the initial prompt length is short and the expected token generation length is long (such as code generation and multi-round dialogue), the proposed approach may encounter issues due to the two concerns outlined in my original review:

1. The nearest-keys set for prefilled queries (indexed during the prefill phase) is not updated during token generation.
2. The prefilled tokens may not remain representative enough for the query vectors of the newly generated tokens.

Unfortunately, the newest response from the authors does not fully address these two points.

Specifically, while the authors note that:

the attention computation over the prompt context always relies on ANNS for dynamic sparse attention.

This explanation still leaves my original concerns unresolved. Additionally, I would like to clarify that I disagreed that "full attention over the newly generated tokens" in the authors' first round comment is unclear.

Regarding the example provided:

Where we generate 1200 tokens from the 128K context using a summary task.

I believe a context length of 128K is too long in light of my concerns.

On the third point regarding the improved method (insert newly generated tokens into the index), I appreciate the authors’ solution, as it is a promising approach to address my concern. However, could the authors also provide details about the latency associated with this solution?

Please feel free to point any possible error in my response out.

Thank you for your detailed review. However, we respectfully disagree with the characterization of long-generation scenarios as a weakness of our work. Below, we clarify our reasons:

1. "Out of Scope": We want to reiterate that our work focuses on long-context scenarios (i.e., long inputs with short outputs). This is consistent with the scope of related works[1-4] and contemporary studies[5-7] in the field. For example, SnapKV mentions in Sec. 1[3]: "In practical applications like chatbots and agents, ..., prompts are often much larger than generated responses..." While long-generation scenarios are undoubtedly important, they remain largely unexplored in the current community, with no established benchmarks, designs, or experiments addressing this issue.

2. "Feasible Solution Proposal": We appreciate the reviewer’s suggestion. During the rebuttal period, we supplemented our work with a proposed solution for long-generation scenarios, utilizing an incremental insertion strategy. We included an analysis of attention recall for varying input sizes when generating up to 2k tokens. As suggested by the reviewer, now we have rerun the experiments using incremental inserts on KV retrieval, which is one of the most challenging task on long context, to further validate the approach shown in RTable 5. Due to time limit, we will include more experiments in future versions of the paper.

RTable 5. Performance on KV retrieval task in InfiniteBench using Llama-3-8B-262K in RTX 4090. We cannot get the decoding latency of Quest because it incurs OOM issue in RTX 4090.

Additionally, our paper comprehensively demonstrates the rationale of our insights and the effectiveness of our proposed method. Compared to baselines (e.g., Quest, InfiniGen), our approach achieves higher accuracy and lower decoding latency across various datasets and long-context LLMs.

Regarding the code release, we are committed to releasing it promptly after the review process. It is important to note that releasing the code during the submission phase is not a mandatory requirement for ICLR, and therefore, we do not consider this a weakness of our paper. Additionally, our paper provides a comprehensive explanation of the implementation and experimental setup in Sections 3.3, 4.1, and Appendix C, ensuring the reproducibility of our results.

Finally, we trust that the reviewer will fairly reassess the paper and adjust the score upon addressing your concerns.

[1] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models, NeurIPS 2023.
[2] Efficient Streaming Language Models with Attention Sinks, ICLR 2024.
[3] SnapKV: LLM Knows What You are Looking for Before Generation, NeurIPS 2024.
[4] Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference, ICML 2024.
[5] MagicPIG: LSH Sampling for Efficient LLM Generation, 2024.
[6] PQCache: Product Quantization-based KVCache for Long Context LLM Inference, 2024.
[7] Squeezed Attention: Accelerating Long Context Length LLM Inference, 2024.

Due to the approaching rebuttal deadline, we want to ensure that we have addressed your concerns one by one.

In the First Round of Rebuttal

The reviewer is satisfied with our additional experiments:

1. Successfully clarifying CPU-GPU co-execution
2. Comparing with Quest (accuracy: 60.5%, 2K tokens) and InfiniGen (acquired by R3, accuracy: 43.1%, 2K tokens). In contrast, we only activate 100 + 640 tokens, achieving 74.7% accuracy, which is the closest to full attention on RULER benchmark (RTable 4, Table 9).

In the Subsequent Comments

First, the reviewer is concerned about the robustness of our index when facing long-generation tasks.

We provided experiments with 128K inputs and 1.2K outputs (where 1.2K represents a specific long-generation case in modern benchmarks) on standard datasets. These experiments show that the data distribution during the prefill and decoding phases is similar (Figure 9), and our index is robust to insertions (Figures 10 and 11). These results suggest that our index performs effectively even in long-generation tasks.

Second, The reviewer is concerned about the input context being too long and requests a scenario with very short input and long output.

Although we cited relevant works showing that such cases are rarely evaluated, and no well-established benchmarks exist for them, we conducted additional experiments on a code generation task from LongBench. In this task, the model generates outputs of up to 2K tokens from inputs ranging from 155 to 719 tokens. This scenario, where the output length is significantly longer than the input (ranging from 2.8X to 13X), aligns closely with the reviewer’s "extremely short input and long-generation" concern. We inserted all newly generated key vector into our indexes at each decoding step and evaluate the index recall. As shown in Figure 12 of our revised paper, our index maintains a high recall rate (> 0.95) even in the long-generation scenarios.

Subsequently, the reviewer expressed concern that most of our experiments in original submissiopn "only" focus on long-context.

We presented detailed evaluation results for context lengths ranging from 4K to 128K in the initial version of our paper (Table 3, Table 4, Figure 7), which covers a broad range of contexts used by most of related works. Besides, the title of the paper clearly indicates that the focus is on long-context scenarios, which are the primary objective of our study.

After providing additional experiments, the reviewer still raised concerns about the "potential risk" of performance in the reviewer's constrainted secenarios -- extremely short input and long-generation.

Despite of presenting additional evaluations, we clearly stated that it is an "out-of-scope" and subjective opinion since it is inconsistent with all related works.

And then, the reviewer asked to rerun experiment with insertion for downstream tasks.

According to our design and additional experiments (Figures 9, 10, 11), the task accuracy of RetrievalAttention remains robust. In response to the reviewer’s request, we reran a challenging long-context task, KV retrieval, involving insertions and observed no accuracy drop in this scenario (see RTable 5), with slight increased latency.

Thanks for your detailed and insightful comments, which have been very helpful for our work.

W1. "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..."

Response: As stated in the beginning of Section 3, this paper focuses on the decoding phase by assuming that prefill and index construction are performed in advance (please refer to the response to W3 below for more details). Therefore, for simplicity, our current implementation does not integrate the KNN computation from Q to K into FlashAttention. Instead, we are using a brute-force method that relies on Faiss Flat [7] for KNN computation on the GPU.

However, the KNN computation can be fused into FlashAttention because: (1) FlashAttention inherently includes the step of calculating the inner product of querys and keys, and (2) the top-K algorithm only consumes registers that are redundant in flash-attention kernel for modern GPUs. We can initialize a min-heap for each query in registers and update it with the matrix multiplication results as the thread block moves through key blocks. Heap operations are performed in CUDA cores and may take up to 2-3$\times$ the prefill latency, which is acceptable. More details are provided in the Algorithm 2 in our revised paper. Exisitng KV cache compression work such as H2O [8] also requires KNN computation during the prefilling phase. Therefore we regard optimizing the KNN computation as important future work to benefit broader research works.

W2. "The article should adopt a more appropriate baseline...Quest...InfiniGen..."

Response: Thanks for your suggestion. In the revised paper, we have added Quest and InfiniGen as new baselines. We evaluated their performance on InfiniteBench and RULER, with the results presented in RTable 3 and RTable 4, respectively.

InfiniGen exhibits a noticeable drop in model accuracy compared to full attention due to inaccurate speculation of important tokens from previous layers. Although Quest shows better performance than InfiniGen, it achieves zero accuracy in the highly dynamic KV retrieval task (Retr.KV in the InfiniteBench) due to low representativeness of centroid vectors. Moreover, both Quest and InfiniGen have severe performance degradation on the RULER benchmark. In contrast, RetrievalAttention performs best and achieves nearly the same accuracy as full attention across two benchmarks.

Regarding the decoding latency, Quest encounters OOM issues on the RTX 4090 due to its fully GPU-resident design, making it unsuitable for supporting 128K contexts in consumer GPUs. Since the current decoding optimizations of InfiniGen does not support GQA for LLMs [9], we will implement InfiniGen for the GQA model and test its decoding latency in the future version of our paper.

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

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

[7] https://github.com/facebookresearch/faiss

[8] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models, NeurIPS 2023

[9] https://github.com/snu-comparch/InfiniGen/tree/main/speedup

Thank you again for reviewing our work! Please let us know if we have sufficiently addressed your questions and concerns.

Authors

1. Regarding the issues of prefix caching and consumer-grade GPUs, we believe the reviewer have completely misunderstood our point. First of all, prefix caching is indeed supported by major LLM providers, which does not conflict with using consumer-grade GPUs for decoding[1,2]. After applying our method, long-context decoding can be done on consumer-grade GPUs, making it possible for major LLM providers to use these more affordable devices for inference, thereby offering cheaper services. Additionally, we mentioned using relatively high-end GPUs to address the reviewer’s concerns. With the currently widely-used architecture separating prefilling and decoding, high-end GPUs can hide the index construction overhead of our method in prefilling. Although we focus on prefix caching scenarios and do not concern ourselves with prefilling overhead, we provided this information in the previous response because the reviewer raised concerns. Moreover, this also does not conflict with using consumer-grade GPU for decoding.

2. We have already explained that our focus is on decoding scenarios and that prefix caching has covered prefilling and index building, which was also stated in original submission. Therefore, we do not have the responsibility to report the experiment outside our scope. Regarding Squeezed Attention, we are deeply disappointed that the reviewer have misunderstood both our point and the referenced paper. The paper also targets scenarios where the index for frequently-used prefixes/fixed prompts has been pre-built, so its prefilling latency does not include index building time. Instead, its prefilling time includes sparse attention on fixed-context (with pre-built index) and very brief full attention on questions, hence its prefilling latency is lower than full attention. If the reviewer acknowledges this scenario, which we have been consistently emphasizing, we can similarly report such prefilling latency with pre-built index. We hope the reviewer can carefully review our response, rather than simply seeking incorrect evidence to counter our claims while disregarding content that supports our perspective.

3. The reviewer frequently challenges our paper's practical applicability while ignoring our core contributions. We question whether a research paper must be practically applicable to be considered valuable and worthy of acceptance by ML conferences such as ICLR. Our experiments indicate that many previously proposed KV cache compression methods significantly degrade the original model performance. Without guaranteed model performance, the practical use of these methods by LLM providers, for instance, is limited. However, we believe these methods still have high research value as they enhance our understanding of KV cache. Similarly, our proposed method advances the field by offering a novel approach to long-context inference through vector retrieval, which achieves the best model performance compared to all prior methods. This underscores the research merit and potential impact of our work. Furthermore, we have detailed potential scenarios, such as prefix caching, which illustrate the practical applicability of our solution.

In conclusion, we are disappointed with the misunderstandings of our points, our referenced papers, and the unreasonable standards applied by the reviewer. We believe that a good research paper is not the final and definite answer to a problem. We would like to reiterate our paper's core contribution: we empirically validate the alignment of vector retrieval and sparse attention in long-context scenarios through extensive evaluations, such as RULER; and we identify the critical issue of out-of-distribution (OOD) discrepancies between queries and keys, proposing a feasible solution to address it. We believe our work opens up the possibility for more in-depth studies and practical applications from the vector retrieval fields, benefiting the field of the long-context inference.

We hope the reviewer can recognize the importance of this point rather than getting constrained by out-of-scope details. We believe that asking for perfect solutions can stifle innovation within this niche area and the broader community, preventing innovative methods from gaining the visibility they deserve.

[1] SGLang: Efficient Execution of Structured Language Model Programs, NeurIPS 2024.
[2] Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023.

Thank you for the response. I think most of my questions have been addressed. I will increase my score to 8.

Thank you for your thorough review. We apologize for any confusion caused.

W1. "...powerful CPUs (in terms of compute speed and memory) are necessary to ensure this method achieves good performance...I think this method distributes the large cost(computation and memory) between the CPU and GPU to reduce the strain on the GPU...Additionally, the communication between the CPU and GPU could also negatively impact the speed..."

Response: Thanks for your comments. While higher CPU power may indeed enhance overall decoding speed, we argue that RetrievalAttention does not inherently require a powerful CPU by design.

• Clarification on CPU Specification: We have clarified the setting of Intel i9-10900X CPU in our revised paper. The Intel i9-10900X CPU is a desktop-level CPU with 10 physical cores, not 20, as the latter number refers to logical cores when hyperthreading is enabled. And this CPU is often paired with the RTX 4090 as a standard configuration. Hence, our experiments were conducted on a mainstream consumer-level CPU rather than on high-end CPUs.
• Efficiency of Attention-Aware ANNS Algorithm: Our highly efficient attention-aware ANNS algorithm requires scanning much fewer vectors (1-3% of all data), allowing it to achieve low search latency even with less powerful CPUs.

Regarding cost concerns, since GPUs must be plugged into a CPU server to operate, we can fully utilize existing idle CPU servers, without additional investment in CPUs and memory. Moreover, the total cost of ownership for CPUs and DRAM is significantly lower than for GPUs. Specifically, the combined cost of an i9-10900X CPU and DRAM is approximately 850 dollars (CPU: 650 dollars + DRAM: 200 dollars), compared to a single RTX 4090 GPU, which costs around 2600 dollars. Thus, it is a reasonable trade-off to shift some computational costs to the CPU.

During the decoding phase, we minimize communication between the CPU and GPU by only transferring the query vector from the GPU to the CPU and transferring the partial attention result instead of top-k vectors back to the GPU. This design significantly reduces GPU-CPU communication overhead.

Q1. "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?"

Response:

• As observed in previous static compression methods [5][6], some tokens, such as initial tokens (also known as attention sinks), are consistently important for subsequent generation.
• If relying solely on static tokens, the performance is significantly lower, as demonstrated by the results of StreamingLLM and SnapKV in the Table 2 and 3 of our paper.
• In the design of RetrievalAttention, during the prefilling phase, we physically separate consistently important tokens and dynamic tokens by placing them in GPU memory and CPU memory, respectively. Therefore, there is no overlap between static tokens in the GPU cache and dynamic critical tokens retrieved from ANNS indexes. The accurate identification of consistently important tokens is well-studied [5][6] and is orthogonal to the design of RetrievalAttention. Currently, we adopt StreamingLLM [5] to keep initial and last window tokens in the GPU cache, but we are adaptable to support more complex static patterns.

Q2. "Does the vocabulary size influence efficiency?"

Response: We adopted the vocabulary used in original LLMs for consistency. This vocabulary choice does not affect the index efficiency because the index construction and retrieval are performed solely on the KV cache.

RetrievalAttention benefits models with both small and large vocabulary sizes, as long as the attention computation over the tokens in the context becomes the bottleneck.

[5] Efficient Streaming Language Models with Attention Sinks, ICLR 2024

[6] SnapKV: LLM Knows What You are Looking for Before Generation, NeurIPS 2024

W1."...lacks a comparison with other OOD methods..."

Response: Thanks for your comments. We have included a new baseline, RobustVamana[1], a well-known solution for handling OOD queries. RobustVamana enhances graph connectivity by adding edges based on query vectors. In Figure 6 of our revised paper, we compare the number of vectors that need to be scanned to achieve different recall rates.

As shown in this figure, to reach a recall rate of 0.8, RobustVamana requires scanning 25% and 75% vectors in the model Yi-9B and Llama-3-8B respectively. Even worse, it falls into a local optimum on Yi-6B model. In contrast, RetrievalAttention only requires scanning a very limited subset of vectors (1-3%) to achieve a high recall rate (0.95), which is consistent across all models. Consequently, RetrievalAttention demonstrates superior search efficiency and accuracy compared to alternative OOD solutions.

W2."..larger models(30B, 70B)..."

Response: We appreciate this comment and agree that evaluating larger models will help demonstrate the generalizability of our methods. Therefore, we have evaluated our method on Llama-3-70B-262k. To accomodate the 70B model, we use a server with eight 40GB A100 GPUs and partition the model by layers across GPUs. Due to time limit, we choose the most complex task KV retrieval in InfiniteBench to stress test the efficiency of RetrievalAttention and other baselines, as shown in RTable 1.

RTable 1. Performance and latency of different methods on Ret.KV and Llama-3-70B-262K.

Specifically, RetrievalAttention achieves nearly the same task accuracy as the exact KNN method Flat, and outperforms Quest by 80%. The decoding speed of RetrievalAttention is $3.5 \times$ faster than Flat as it effectively reduces the vectors to scan by mitigating the distribution gap between queries and keys.

Q1."...layer-wise or head-wise... improve performance?"

Response: Thanks for your comments. Exploring different sparsity degrees across layers and heads is an intriguing research direction and is complementary to our work. Recent research work such as PyramidKV[2] shows that applying appropriate memory budgets in different layers based on the spartiy degree can enhance the model performance while reducing overall GPU memory usage.

We compare the performance of the original RetrievalAttention w/ and w/o PyramidKV budget allocation strategy on the InfiniteBench, as shown in RTable 2. Specifically, for the original RetrievalAttention, we set a fixed budget of 2000 tokens for all heads in all layers. In contrast, PyramidKV dynamically adjusts the retrieval size across different layers, allocating more in lower layers and less in higher ones.

The results in RTable 2 show that the PyramidKV allocation strategy achieves better performance in Retr.KV tasks, though it slightly decreases performance in the En.QA task. On average, the accuracy slightly surpasses that of the original RetrievalAttention. This indicates that dynamic budget allocation is promising but may require task-specific allocation strategies.

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

Q2."...other ANNS index, product quantization or hierarchical clustering"

Response: We have considered various ANNS indexes and found that the graph-based indexes are the most suitable for our scenario (i.e., low-latency index search on CPU) for the following reasons:

• Graph-based indexes generally require scanning fewer vectors to achieve a traget recall compared to clustering-based ANNS, thereby saving more CPU bandwidth. This is critical for high index search efficiency and improving overall decoding speed.
• Quantization methods, such as product quantization or scalar quantization, can be applied to both graph-based[3] and clustering-based indexes[4]. Although we did not apply quantization, we tried scalar quantization (from 16 to 8bits) to our attention-aware graph-based ANNS and found that it reduce the decoding latency by 20% without compromising the model performance. We intend to further explore efficient combinations of quantization methods with attention-aware ANNS indexes to improve decoding speed in future work.

[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

[3] DiskANN: Fast Accurate Billion-point Nearest Neighbor Search on a Single Node, NeurIPS 2019

[4] Product Quantizer Aware Inverted Index for Scalable Nearest Neighbor Search, ICCV 2021