Elamur Research⌕
☰
home / papers / 2024 / RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval

RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval

Paper · 2024 · Di Liu, Meng Chen, Baotong Lu et al.

Summary

Serving LLMs with very long contexts is expensive. Attention cost grows with sequence length, and the KV cache for Llama-3-8B reaches roughly 125GB at one million tokens, which no consumer GPU can hold. RetrievalAttention is a training-free method that exploits the fact that each decoding step depends strongly on only a small, changing set of tokens. Most KV vectors are moved to CPU memory and organised into approximate nearest neighbor search (ANNS) indexes based on inner product. At each step, the current query retrieves the relevant keys from these indexes. A few tokens chosen by fixed patterns, similar to StreamingLLM, stay on the GPU, and the partial attention results from both devices are combined.

The authors found that standard ANNS indexes perform poorly here because query and key vectors come from noticeably different distributions. They therefore built an index that adapts to the distribution of queries. Evaluated on ∞-Bench and RULER with several long-context models, the method nearly matches full attention while reading only a small fraction of the keys. This makes 128K-token inference with 8B models feasible on a single RTX4090.

Key points

  • Attention is sparse and dynamic. In Llama-3-8B, the top-1000 tokens recover about 89% of attention mass on average. Reusing the tokens selected for the first generated token drops this to around 71%.
  • Query–key distribution gap. Measured by Mahalanobis distance, queries lie more than 10× farther from keys than in-distribution data would. As a result, off-the-shelf ANNS must scan 30–50% of keys, whereas the attention-aware index needs only 1–3%.
  • Speed on a single RTX4090 (24GB). With an 8B model at 128K context, decoding takes 0.188 s per token. This is 4.9× faster than exact KNN retrieval and 1.98× faster than a conventional ANNS index.
  • RULER accuracy at 128K (Llama-3). Average scores are:
    • RetrievalAttention: 74.7%
    • Full attention: 78.7%
    • Quest: 60.5%
    • InfiniGen: 43.1%
  • Scaling to Llama-3-70B-262k (8× A100 40GB):
    • Accuracy roughly matches exact Flat search (23.5 vs 24.0).
    • Decoding is 3.5× faster than Flat.
    • Under grouped query attention (GQA), a separate index is built for each query head.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🏢MicrosoftRedmond · United States378 works · 161,630 citations🎓Fudan UniversityShanghai · China202 works · 22,509 citations🎓Shanghai Jiao Tong UniversityShanghai · China212 works · 19,544 citations🏢Microsoft Research AsiaBeijing · China82 works · 61,484 citations🏢Adobe SystemsSan Jose · United States40 works · 4,248 citations🎓University of Central FloridaOrlando · United States15 works · 39,739 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 worksComputer Vision and Pattern Recognition · 3,486 works

🧭 Topics 2Research subjects and languages the work is about.Language models & pretraining · 3,894 worksInformation retrieval & search · 552 works

Cite

Cite as Di Liu, Meng Chen, Baotong Lu, Huiqiang Jiang, Zhenhua Han, Qianxi Zhang, Qi Chen, Chengruidong Zhang, Bailu Ding, Kai Zhang, Chen Chen, Fan Yang, Yuqing Yang and Lili Qiu. 2024. RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval. arXiv:2409.10516.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2409.10516 · arXiv 2409.10516

@misc{liu2024retrievalattention,
  title = {{RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval}},
  author = {Di Liu, Meng Chen, Baotong Lu, Huiqiang Jiang, Zhenhua Han, Qianxi Zhang, Qi Chen, Chengruidong Zhang, Bailu Ding, Kai Zhang, Chen Chen, Fan Yang, Yuqing Yang and Lili Qiu},
  year = {2024},
  doi = {10.48550/arxiv.2409.10516},
  eprint = {2409.10516},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2409.10516}
}

References

57 references, 26 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. 01-ai. Yi-6b-200k. https://huggingface.co/01-ai/Yi-6B-200K, 2024a. Accessed: 2024-07-01. 01-ai. Yi-9b-200k. https://huggingface.co/01-ai/Yi-9B-200K, 2024b. Accessed: 2024-07-01. ↗
  2. Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. ETC: Encoding long and structured inputs in transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EM…
  3. Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. In The 2023 Conference on Empirical Methods in Natural Language Processing, 2023.
  4. Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. ArXiv preprint, abs/2004.05150, 2020. URL https://arxiv.org/abs/2004.05150. ↗
  5. Amanda Bertsch, Uri Alon, Graham Neubig, and Matthew Gormley. Unlimiformer: Long-range transformers with unlimited length input. Advances in Neural Information Processing Systems, 36, 2024.
  6. Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, and Wen Xiao. Pyramidkv: Dynamic KV cache compression based on pyramidal information funneling. CoRR, abs/2406.02069, 2024.
  7. Meng Chen, Kai Zhang, Zhenying He, Yinan Jing, and X. Sean Wang. Roargraph: A projected bipartite graph for efficient cross-modal approximate nearest neighbor search. Proc. VLDB Endow., 17(11):2735–2749, 2024a. ISSN 2150-8097. doi: 10.14778/3681954.3681959. URL https://doi.org/10.14778/3681954.36819… — [doi:10.14778/3681954.3681959](https://doi.org/10.14778/3681954.3681959) ↗
  8. Shaoyuan Chen, Yutong Lin, Mingxing Zhang, and Yongwei Wu. Efficient and economic large language model inference with attention offloading. ArXiv preprint, abs/2405.01814, 2024b. URL https://arxiv.org/abs/2405.01814. — [arXiv:2405.01814](https://arxiv.org/abs/2405.01814) ↗
  9. Zhuoming Chen. Magicpig: sparse inference engine for llm. Infini-AI-Lab/MagicPiG/, 2024. Accessed: 2024-08-01. https://github.com/ ↗
  10. Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. ArXiv preprint, abs/1904.10509, 2019. URL https://arxiv.org/ abs/1904.10509. ↗
  11. Scott Cost and Steven Salzberg. A weighted nearest neighbor algorithm for learning with symbolic features. Machine learning, 10:57–78, 1993. find
  12. Paul Covington, Jay Adams, and Emre Sargin. Deep neural networks for youtube recommendations. In Shilad Sen, Werner Geyer, Jill Freyne, and Pablo Castells (eds.), Proceedings of the 10th ACM Conference on Recommender Systems, Boston, MA, USA, September 15-19, 2016, pp. 191–198. find
  13. ACM, 2016. doi: 10.1145/2959100.2959190. URL https://doi.org/10.1145/ 2959100.2959190. — [doi:10.1145/2959100.2959190](https://doi.org/10.1145/2959100.2959190) find
  14. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems, 2022. find
  15. Yichuan Deng, Zhao Song, and Chiwun Yang. Attention is naturally sparse with gaussian distributed input, 2024. find
  16. Matthijs Douze, Alexandr Guzhva, Chengqi Deng, Jeff Johnson, Gergely Szilvasy, Pierre-Emmanuel Mazaré, Maria Lomeli, Lucas Hosseini, and Hervé Jégou. The faiss library. 2024.
  17. Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive KV cache compression for LLMs. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum? id=uNrFpDPMyo. ↗
  18. Google Cloud. Context caching overview. https://cloud.google.com/vertex-ai/ generative-ai/docs/context-cache/context-cache-overview, 2024. Accessed: 2024-07-01. ↗
  19. Gradient AI. Llama-3-8b-instruct-262k. https://huggingface.co/gradientai/ Llama-3-8B-Instruct-262k, 2024. Accessed: 2024-07-01. ↗
  20. Greg Kamradt. Needle in a haystack - pressure testing llms. https://github.com/ gkamradt/LLMTest_NeedleInAHaystack, 2023. Accessed: 2024-08-12. ↗
  21. Chi Han, Qifan Wang, Hao Peng, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. LM-infinite: Zero-shot extreme length generalization for large language models. In Kevin Duh, Helena Gomez, and Steven Bethard (eds.), Proceedings of the 2024 Conference of the North American Chapter of the Association f…
  22. Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. Ruler: What’s the real context size of your long-context language models? ArXiv preprint, abs/2404.06654, 2024. URL https://arxiv.org/abs/2404.06654. ↗
  23. Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Leon Song, Samyam Rajbhandari, and Yuxiong He. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models. ArXiv preprint, abs/2309.14509, 2023. URL https://arxiv.org/abs/2309.14509. ↗
  24. Shikhar Jaiswal, Ravishankar Krishnaswamy, Ankit Garg, Harsha Vardhan Simhadri, and Sheshansh Agrawal. Ood-diskann: Efficient and scalable graph anns for out-of-distribution queries. arXiv preprint arXiv:2211.12850, 2022. — [arXiv:2211.12850](https://arxiv.org/abs/2211.12850) find
  25. Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin-Yew Lin, et al. Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention. ArXiv preprint, abs/2407.02490, 2024. URL https://arxiv.org/abs… ↗
  26. Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020.
  27. OpenReview.net, 2020. URL https://openreview.net/forum?id= rkgNKkHtvB. ↗
  28. Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023.
  29. Wonbeom Lee, Jungi Lee, Junghwan Seo, and Jaewoong Sim. InfiniGen: Efficient generative inference of large language models with dynamic KV cache management. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), pp. 155–172, Santa Clara, CA, 2024. USENIX Association. ISBN… find
  30. Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Llm knows what you are looking for before generation. ArXiv preprint, abs/2404.14469, 2024. URL https://arxiv.org/abs/2404.14469. ↗
  31. Hao Liu, Matei Zaharia, and Pieter Abbeel. Ringattention with blockwise transformers for near-infinite context. In The Twelfth International Conference on Learning Representations, 2024a. URL https://openreview.net/forum?id=WsRHpHH4s0. ↗
  32. Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time. Advances in Neural Information Processing Systems, 36, 2024b. find
  33. Prasanta Chandra Mahalanobis. On the generalized distance in statistics. Sankhyā: The Indian Journal of Statistics, Series A (2008-), 80:S1–S7, 2018. find
  34. Yu A Malkov and Dmitry A Yashunin. Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs. IEEE transactions on pattern analysis and machine intelligence, 42(4):824–836, 2018.
  35. Yuzhen Mao, Martin Ester, and Ke Li. Iceformer: Accelerated inference with long-sequence transformers on CPUs. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=6RR3wU4mSZ. ↗
  36. Stanislav Morozov and Artem Babenko. Non-metric similarity graphs for maximum inner product search. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural In… find
  37. Aditya Pal, Chantat Eksombatchai, Yitong Zhou, Bo Zhao, Charles Rosenberg, and Jure Leskovec. Pinnersage: Multi-modal user embedding framework for recommendations at pinterest. In Rajesh Gupta, Yan Liu, Jiliang Tang, and B. Aditya Prakash (eds.), KDD ’20: The 26th ACM SIGKDD Conference on Knowledge… find
  38. ACM, 2020. URL https://dl.acm.org/doi/10. 1145/3394486.3403280. ↗
  39. Pratyush Patel, Esha Choukse, Chaojie Zhang, Aashaka Shah, Íñigo Goiri, Saeed Maleki, and Ricardo Bianchini. Splitwise: Efficient generative llm inference using phase splitting. In 2024 ACM/IEEE 51st Annual International Symposium on Computer Architecture (ISCA), pp. 118–132.
  40. Ruoyu Qin, Zheming Li, Weiran He, Mingxing Zhang, Yongwei Wu, Weimin Zheng, and Xinran Xu. Mooncake: Kimi’s kvcache-centric architecture for llm serving. ArXiv preprint, abs/2407.00079, 2024. URL https://arxiv.org/abs/2407.00079. ↗
  41. Luka Ribar, Ivan Chelombiev, Luke Hudlass-Galley, Charlie Blake, Carlo Luschi, and Douglas Orr. Sparq attention: Bandwidth-efficient llm inference. In Forty-first International Conference on Machine Learning, 2024. URL https://openreview.net/forum?id=OS5dqxmmtl. ↗
  42. Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Ré, Ion Stoica, and Ce Zhang. Flexgen: high-throughput generative inference of large language models with a single gpu. In Proceedings of the 40th International Conference on Machine Learning, IC…
  43. Prajwal Singhania, Siddharth Singh, Shwai He, Soheil Feizi, and Abhinav Bhatele. Loki: Low-rank keys for efficient sparse attention. ArXiv preprint, abs/2406.02542, 2024. URL https://arxiv.org/abs/2406.02542. — [arXiv:2406.02542](https://arxiv.org/abs/2406.02542) ↗
  44. Sivic and Zisserman. Video google: A text retrieval approach to object matching in videos. In Proceedings ninth IEEE international conference on computer vision, pp. 1470–1477. IEEE, 2003. find
  45. Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. QUEST: Query-aware sparsity for efficient long-context LLM inference. In Forty-first International Conference on Machine Learning, 2024. URL https://openreview.net/forum?id=KzACYw0MTV. ↗
  46. Gemini Team. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context, 2024.
  47. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neur…
  48. Mengzhao Wang, Xiaoliang Xu, Qiang Yue, and Yuxiang Wang. A comprehensive survey and experimental comparison of graph-based approximate nearest neighbor search. seed, 4(2):1. Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, Song Han, and Maosong Sun. Infll… — InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory
  49. Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, 2024b. find
  50. Lee Xiong, Chenyan Xiong, Ye Li, Kwok-Fung Tang, Jialin Liu, Paul N. Bennett, Junaid Ahmed, and Arnold Overwijk. Approximate nearest neighbor negative contrastive learning for dense text retrieval. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-…
  51. OpenReview.net, 2021. URL https://openreview.net/forum? id=zeFrfgyZln. ↗
  52. Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. In Hugo Larochelle, Marc’Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsua…
  53. Hailin Zhang, Xiaodong Ji, Yilin Chen, Fangcheng Fu, Xupeng Miao, Xiaonan Nie, Weipeng Chen, and Bin Cui. Pqcache: Product quantization-based kvcache for long context llm inference. ArXiv preprint, abs/2407.12820, 2024a. URL https://arxiv.org/abs/2407.12820. ↗
  54. Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Hao, Xu Han, Zhen Thai, Shuo Wang, Zhiyuan Liu, and Maosong Sun. ∞Bench: Extending long context evaluation beyond 100K tokens. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the… find
  55. Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark W. Barrett, Zhangyang Wang, and Beidi Chen. H2O: heavy-hitter oracle for efficient generative inference of large language models. In Advances in Neural Information Process… find
  56. Yinmin Zhong, Shengyu Liu, Junda Chen, Jianbo Hu, Yibo Zhu, Xuanzhe Liu, Xin Jin, and Hao Zhang. DistServe: Disaggregating prefill and decoding for goodput-optimized large language model serving. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), pp. 193–210, Santa Cl…
  57. ISBN 978-1-939133-40-3. URL https://www.usenix.org/conference/osdi24/ presentation/zhong-yinmin. ↗

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/liu-2024-retrievalattention-accelerating-long-context-llm/