Elamur Research
☰
home / papers / 2024 / Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference

Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference

Written for the Research library (claude-opus-5-5) from the paper text, in our own words: the license "arXiv — non-exclusive license to distribute" does not allow redistributing the text itself. Original — arxiv.org/pdf/2406.10774

Summary

Generating text with very long contexts is slow because every new token requires reading the full KV cache from memory. For Llama 7B at 32K tokens, that cache reaches about 16GB. Earlier eviction methods such as H2O, TOVA and StreamingLLM permanently drop tokens judged unimportant from past attention. The authors show that which tokens matter shifts with each new query, so a dropped token may be needed later.

Quest keeps the whole cache but reads only the relevant parts. The cache is split into pages, and each page stores the per-dimension minimum and maximum of its Key vectors. At every decoding step, the current Query is combined with these bounds to estimate an upper limit on attention for each page. Exact attention is then computed only over the Top-K highest-scoring pages.

Tests on PG19, passkey retrieval and LongBench show better accuracy than baselines at equal token budgets, along with large latency gains. This makes long-context serving cheaper without retraining.

Key points

  • At 32K context with a 2048-token budget, self-attention runs up to 7.03× faster than FlashInfer. End-to-end decoding is 1.74× faster with FP16 weights and 2.23× faster with 4-bit AWQ weights.
  • Evaluations used LongChat-7b-v1.5-32k and Llama-2-7B, with token budgets from 256 to 4K. Outside the first two layers, more than 90% of KV tokens could be skipped with almost no change in PG19 perplexity.
  • In passkey retrieval, Quest's recall of the top-attended tokens stays close to full attention, while H2O's recall drops noticeably.
  • Under a lossless-accuracy target on LongBench, NarrativeQA needed about 5K tokens with Quest versus 14K with TOVA. Attention ran 3.82× faster on GovReport and 4.54× faster on TriviaQA.
  • The paper appeared at ICML 2024, and the code is published at github.com/mit-han-lab/Quest.

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

Fields: Other

Cited by

196 papers of our library cite it (of the 140673 whose bibliographies we read); 2 of them have a page here, 194 not yet.

Cite

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

@misc{tang2024quest,
  title = {{Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference}},
  author = {Tang, Jiaming and Zhao, Yilong and Zhu, Kan and Xiao, Guangxuan and Kasikci, Baris and Han, Song},
  year = {2024},
  doi = {10.48550/arxiv.2406.10774},
  eprint = {2406.10774},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2406.10774}
}

References

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

  1. Anthropic. Introducing the next generation of Claude. https://www.anthropic.com/news/claude-3-family, 2024. [Accessed 28-05-2024]. ↗
  2. Bai, Y., Lv, X., Zhang, J., Lyu, H., Tang, J., Huang, Z., Du, Z., Liu, X., Zeng, A., Hou, L., Dong, Y., Tang, J., and Li, J. Longbench: A bilingual, multitask benchmark for long context understanding, 2023. find our card
  3. Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022. find
  4. Dasigi, P., Lo, K., Beltagy, I., Cohan, A., Smith, N. A., and Gardner, M. A dataset of information-seeking questions and answers anchored in research papers, 2021. find our card
  5. Ge, S., Zhang, Y., Liu, L., Zhang, M., Han, J., and Gao, J. Model tells you what to discard: Adaptive kv cache compression for llms, 2024. find our card
  6. Huang, L., Cao, S., Parulian, N., Ji, H., and Wang, L. Efficient attentions for long document summarization. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 1419– 1436, Online, June 2021. Associati… find our card
  7. Joshi, M., Choi, E., Weld, D., and Zettlemoyer, L. TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. In Barzilay, R. and Kan, M.-Y. (eds.), Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 160… find our card
  8. Kočiský, T., Schwarz, J., Blunsom, P., Dyer, C., Hermann, K. M., Melis, G., and Grefenstette, E. The NarrativeQA reading comprehension challenge. Transactions of the Association for Computational Linguistics, 6:317–328, 2018. doi: 10.1162/tacl a 00023. URL https://aclanthology.org/Q18-1023. ↗ our card
  9. Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., and Stoica, I. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023. find
  10. Li, D., Shao, R., Xie, A., Sheng, Y., Zheng, L., Gonzalez, J. E., Stoica, I., Ma, X., and Zhang, H. How long can open-source llms truly promise on context length?, June 2023. URL https://lmsys.org/blog/2023-06-29-longchat. ↗
  11. Liu, H., Yan, W., Zaharia, M., and Abbeel, P. World model on million-length video and language with blockwise ringattention, 2024a. find
  12. Liu, X., Yan, H., Zhang, S., An, C., Qiu, X., and Lin, D. Scaling laws of rope-based extrapolation, 2024b. find
  13. NVIDIA. Nvidia ada lovelace professional gpu architecture. https://images.nvidia.com/aem-dam/en-zz/Solutions/technologies/NVIDIA-ADA-GPU-PROVIZ-Architecture-Whitepaper\_1.1.pdf, 2023. [Accessed 28-05-2024]. ↗
  14. NVIDIA. Nvbench: Nvidia’s benchmarking tool for gpus, 2024. Available online: https://github.com/NVIDIA/nvbench. ↗
  15. OpenAI. New models and developer products announced at devday. https://openai.com/blog/new-models-and-developer-products-announced-at-de OpenAI, November 2023. Accessed: 2024-01-31. ↗
  16. OpenAI. Introducing gpt-4o: our fastest and most affordable flagship model. https://platform.openai. com/docs/models, 2024. [Accessed 28-05-2024]. ↗
  17. Oren, M., Hassid, M., Adi, Y., and Schwartz, R. Transformers are multi-state RNNs, 2024. URL https://arxiv.org/abs/2401.06104. arXiv:2401.06104. ↗
  18. Peng, B., Quesnelle, J., Fan, H., and Shippole, E. Yarn: Efficient context window extension of large language models, 2023. find
  19. Rae, J. W., Potapenko, A., Jayakumar, S. M., Hillier, C., and Lillicrap, T. P. Compressive transformers for long-range sequence modelling. arXiv preprint, 2019. URL https://arxiv.org/abs/1911.05507. ↗ our card
  20. Ribar, L., Chelombiev, I., Hudlass-Galley, L., Blake, C., Luschi, C., and Orr, D. Sparq attention: Bandwidth-efficient llm inference, 2023. find
  21. Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding, 2023. find
  22. Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lample, G. Llama: Open and efficient foundation language models, 2023. find
  23. Tworkowski, S., Staniszewski, K., Pacek, M., Wu, Y., Michalewski, H., and Miłoś, P. Focused transformer: Contrastive training for context scaling, 2023. find
  24. Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Efficient streaming language models with attention sinks. arXiv, 2023. find
  25. Yang, Z., Qi, P., Zhang, S., Bengio, Y., Cohen, W. W., Salakhutdinov, R., and Manning, C. D. Hotpotqa: A dataset for diverse, explainable multi-hop question answering, 2018. find our card
  26. Ye, Z., Lai, R., Lu, R., Lin, C.-Y., Zheng, S., Chen, L., Chen, T., and Ceze, L. Cascade inference: Memory bandwidth efficient shared prefix batch decoding. https://flashinfer.ai/2024/01/08/cascade-inference.html, Jan 2024. URL https://flashinfer.ai/2024/01/08/cascade-inference.html. Accessed on 202… ↗
  27. Zhang, J., Naruse, A., Li, X., and Wang, Y. Parallel topk algorithms on gpu: A comprehensive study and new methods. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC ’23, New York, NY, USA, 2023a. Association for Computing Machinery.… — [doi:10.1145/3581784.3607062](https://doi.org/10.1145/3581784.3607062) find
  28. Zhang, Z., Sheng, Y., Zhou, T., Chen, T., Zheng, L., Cai, R., Song, Z., Tian, Y., Ré, C., Barrett, C., Wang, Z., and Chen, B. H 2 o: Heavy-hitter oracle for efficient generative inference of large language models, 2023b. find
  29. Zhao, Y., Lin, C.-Y., Zhu, K., Ye, Z., Chen, L., Zheng, S., Ceze, L., Krishnamurthy, A., Chen, T., and Kasikci, B. Atom: Low-bit quantization for efficient and accurate llm serving, 2024. find