Running Huge Context Windows On Tiny GPUs

Thank you for your insightful comments and questions. We answer them below.

1. Why should we use this method rather than other more proven methods?

In long context modeling, quality can drop as context length grows, while simultaneously attention approximations must be made to fit contexts into GPU memory. Our solution, Top-k, removes the quality tradeoff to replace it with a vector database.

Attention based on true top-k score selection is of vastly higher quality than approximate attention mechanisms. Attention Sink have known failure cases for Needle In A Haystack, as can be shown in figure 8 on page 9 of our newest draft. Attention Sink is incapable of retrieving information outside of its local window or it's early sink tokens, whereas Top-k is able to do so out to over 1 million tokens with k=10.

Additionally, there are stark differences between true Top-k and approximate Top-k as reported in a recent paper, Loki[1]:

Top-k attention maintains 95% quality with a 1% context budget, while H2O and Loki require over 30% context for the same threshold. Our performant ratio also holds at extremely long contexts.

2. Can you provide a table or graph showcasing the potential compute/memory savings?

The primary challenge with long context inference is that the KV cache quickly becomes too large to fit into GPU memory. Top-k attention solves this by offloading the KV cache to a vector database, leaving a negligible memory footprint on the GPU. This is what enables 1-million token context lengths on a 16 GB GPU with an 8 billion parameter model.

The peak GPU memory consumption Flash Attention and Top-k attention can help illustrate this:

$N$: Context length $D$: Head dimension $H_q$: Number of Query heads $H_{kv}$: Number of Key-Value heads $L$: Number of layers

With Top-k attention, only a small subset of the cache is moved to the GPU for each attention computation and the memory requirements scale with $k$ rather than with the context length ($N$).

For a concrete comparison of memory requirements between Top-k and full attention, where $k$ is chosen as 1% of the total context length:

Top-k attention presents a theoretical compute advantage because the vector database supports sublinear search time for a given query.

3. How does this method scale to bigger models?

Top-k attention's GPU memory advantage scales differently depending on how the model is made larger. If the model is made deeper, Top-k attention becomes even more advantageous because it's GPU memory requirements do not include layers as a scaling factor, as shown in the table above. If the model is made wider, the scale of efficiency advantage will remain the same since both Top-k and competing attention methods include the model dimension equally in their theoretical memory requirements.

4. Why is there no comparison to RAG?

Given that Top-k attention can operate at high fidelity over ultra-long document stores, we have occasionally wondered how long context models in general compare to RAG methods. While there are some trivial ways to show failure cases in RAG models (use a query that exceeds the context length of the retrieval model), a proper comparison between RAG and long context models deserves its own paper.

Thank you for your thoughtful suggestions and the effort you put in to improving the quality of this work. We put significant effort into making improvements and running additional experiments based on your comments. Also, we were pleased that noted the ease with which this could be dropped into existing frameworks. Given our responses and updates, we would appreciate it if you would consider raising your score. We would love to continue the discussion and get this paper improved to the point that the results can be shared with and analyzed by the broader community. [1]: Loki: Low-rank Keys for Efficient Sparse Attention

Thank you for running the additional experiments, I have adjusted the score accordingly. I think it would be also worthwhile to include these experiments and additional discussion about implications regarding compute/memory complexity in the paper.

1. How does this method compares with existing KV cache sparsity approaches?

It turns out that attention based on true top-k score selection is vastly superior to cache sparsity methods like H2O, which can be thought of as approximate top-k methods. In the recent Loki paper[1] they compare their method and H20 to true top-k and the results are stark:

The same stark differences can be seen in simple Needle In A Haystack evaluations for methods like Attention Sink. We have added a figure to our newest draft -- figure 8 on page 9 -- demonstrating this. The red cells show that Attention Sink (StreamingLLM) is incapable of retrieving information outside of its local window or it's early sink tokens, whereas Top-k Attention (IVF Index) is able to do so out to 1 million tokens of context and beyond:

Two of the main contributions of this paper are demonstrating that the high fidelity of Top-k attention shown in the table above hold for contexts at the million-token scale, and also quantifying exactly how sparse you can go while still maintaining some desired threshold of quality. As you mentioned, we found that Top-k attention can maintain 95% of full attention quality with a 1% context budget, while methods like H2O and Loki require more than 30% to reach this same threshold. The fact that this ratio holds true at ultra-long contexts is what makes true Top-k attention so promising.

2. Can you provide latency or throughput measurements?

Here are our current inference time results. These were benchmarked on a 32 core AMD EPYC-7313 CPU with a 16GB NVIDIA RTXA4000 GPU using a Llama-3-8B model. Times are reported in seconds/token:

3. What is the method's performance on reasoning tasks such as GSM8K and MATH?

The favorable quality to sparsity ratio that we see on other benchmarks also holds for reasoning-focused benchmarks such as GSM8K. For a 2% context budget ($k$=56), Top-k attention achieves 90% of full attention quality. Here are results for increasing $k$, evaluating GSM8K with 5-shot in-context examples:

We also see a similar trend in the MATH dataset. Here we downloaded the Llama 3.1 8B model from huggingface, and ran the LM Eval harness on the minerva_math task with varying k. This task is a four-shot variant of the MATH Benchmark. Our results are here:

We recover the performance of the full attention for our model with much smaller k. The context lengths of these examples vary in size but are around 1024. We note that our absolute performance matches that reported by the OpenLLM Leaderboard

Thank you for taking time to make this a better paper. It is encouraging that you find this method to be effective and practical for implementation. Based on your comments we put significant effort into revising and running additional experiments. We would appreciate it if you would consider raising your score in light of our response. We would love to continue the discussion and get this paper improved to the point that the results can be shared with and analyzed by the broader community. [1]: Loki: Low-rank Keys for Efficient Sparse Attention

Thank you for taking the time to read and review our paper. We found your questions and comments insightful and answer them below.

1. What is the GPU memory consumption after applying this method?

Here is the theoretical peak memory consumption of Top-k attention. For reference we compare with Flash Attention because it is the default attention implementation and has $O(N)$ memory complexity.

$N$: Context length $D$: Head dimension $H_q$: Number of Query heads $H_{kv}$: Number of Key-Value heads $L$: Number of layers

The peak consumption in popular implementations like Flash Attention or VLLM comes from maintaining the KV cache on GPU, whereas with Top-k attention the KV cache is offloaded to a vector database. Only a small subset of the cache is moved to the GPU for each attention computation, so the memory requirements scale with $k$ rather than with the context length ($N$).

Here is a concrete comparison of memory requirements between Top-k and full attention, where $k$ is chosen as 1% of the total context length:

2. How does the CPU memory vector database handle efficient searching?

There is a large body of work on efficient searching over vector databases and specialized compute instances optimized for this type of work. Our implementation uses the FAISS library which we think is a sensible choice but was not optimized for this workload. We benchmarked our inference times on a 32 core AMD EPYC-7313 CPU with a 16GB NVIDIA RTXA4000 GPU using a Llama-3-8B model. Times are reported below in seconds/token and Flash Attention times are reported for reference:

These runtimes could certainly be further optimized, but our method enables generation at long contexts on limited hardware where other methods degrade.

3. What is the expected context length for a given model with this method?

Here is the formula for maximum context length given a specific GPU and a specific model configuation. This is an inversion of the formula listed in question 1 above. Here we choose $k_{budget}$ to be 1% of the total context length:

As shown above, there is no practical limit to the context length based on GPU memory for Top-k attention. In this regime the maximum expected context length becomes limited by the storage capacity of the vector database.

Again, we'd like to thank you for the time you put in to helping sharpen and improve this publication. Your comments pointing out the formatting issues in Algorithm 2 were particularly helpful. We made a significant effort to revise our paper and run additional experiments in light of your feedback. We are pleased that you think this method of extending context length could have substantial impact, and would appreciate it if you would consider raising your score in light of our response. We would be happy to continue the discussion and get this paper improved to the point that the results can be shared with and analyzed by the broader community.

Thank you for providing additional benchmarking results! As noted, the majority of the time is spent on nearest neighbor search. I agree with the authors that the overall design could benefit from improved algorithms or novel search methods on the CPU side to reduce latency. However, I would also like to highlight that the proposed method may not yet be ready for production due to the potential for high latency in scenarios with long context lengths. I will raise my score to a 5 with more confidence.