Elamur Research⌕
☰
home / papers / 2023 / Efficient Memory Management for Large Language Model Serving with PagedAttention

Efficient Memory Management for Large Language Model Serving with PagedAttention

Summary

To make a large language model run fast, requests are processed in batches. The batch size is limited by GPU memory, specifically by the KV cache: for each request it grows and shrinks unpredictably. Existing systems, including FasterTransformer and Orca, reserve in advance a contiguous chunk of memory for each request sized for the maximum possible length. As a result, a lot of memory is lost to internal and external fragmentation, and identical parts of the cache have to be stored multiple times. The authors, from UC Berkeley and other institutions, carried over the idea of paged virtual memory from operating systems to this problem. They proposed the PagedAttention algorithm: the KV cache is split into fixed-size blocks that need not be contiguous in memory and are allocated only when needed. Built on this algorithm is the open-source system vLLM. It has a scheduler capable of preempting requests and a copy-on-write mechanism that lets multiple sequences share blocks. Compared with the best prior systems, vLLM processes 2–4 times more tokens per second at the same latency. The largest gains are observed for long sequences, large models and complex decoding methods. This makes LLM serving noticeably cheaper.

Key points

  • For a 13B-parameter model on an NVIDIA A100 40GB, the weights take up about 65% of memory and the KV cache nearly 30%. Therefore, it is cache management that determines how large a batch can be.
  • According to the authors' measurements, in existing systems only 20.4–38.2% of KV cache memory actually stores token states. In vLLM this share is close to 96%, meaning memory waste is nearly zero.
  • Blocks can be shared both within a single request (parallel sampling, beam search) and across different requests.
  • vLLM supports GPT, OPT and LLaMA models of various sizes, including ones that do not fit in the memory of a single GPU and are therefore distributed across several. The code is open: github.com/vllm-project/vllm. The paper was presented at SOSP '23.
  • Compared with FasterTransformer and Orca, throughput increases 2–4 times, while model accuracy is unchanged.

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: AI & machine learning · Computing & GPUs

Topics Language models & pretraining Efficient models & compression

Cite

Cite as Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang and Ion Stoica. 2023. Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.

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

@misc{kwon2023efficient,
  title = {{Efficient Memory Management for Large Language Model Serving with PagedAttention}},
  author = {Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang and Ion Stoica},
  year = {2023},
  doi = {10.48550/arxiv.2309.06180},
  eprint = {2309.06180},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2309.06180}
}

References

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

  1. Reza Yazdani Aminabadi, Samyam Rajbhandari, Minjia Zhang, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng, Jeff Rasley, Shaden Smith, Olatunji Ruwase, et al. 2022. DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale. arXiv preprint arXiv:2207.00032 (2022). find
  2. Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. 2016. Layer normalization. arXiv preprint arXiv:1607.06450 (2016). our card
  3. Yoshua Bengio, Réjean Ducharme, and Pascal Vincent. 2000. A neural probabilistic language model. Advances in neural information processing systems 13 (2000). find
  4. Ond rej Bojar, Rajen Chatterjee, Christian Federmann, Yvette Graham, Barry Haddow, Matthias Huck, Antonio Jimeno Yepes, Philipp Koehn, Varvara Logacheva, Christof Monz, Matteo Negri, Aurelie Neveol, Mariana Neves, Martin Popel, Matt Post, Raphael Rubino, Carolina Scarton, Lucia Specia, Marco Turchi,… our card
  5. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. Advances in neural information processing systems 33 (2020), 1877–1901. our card
  6. Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Proceedings of Machine Learning and Systems 2 (2020), 497–511. [25] Tom Kilburn, David BG Edwards, Michael J Lanigan, and Frank H Sumner. 1962. One-level storage syst… find

This page in Markdown → https://md.research.elamur.ai/free/en/papers/kwon-2023-efficient-memory-management-large-language/