PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Memory Management for Large Language Model Serving with PagedAttention

Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, Ion Stoica

arXiv:2309.06180v1cs.LGcs.DC

TL;DR

LLM serving is limited by memory-bound generation and inefficient KV-cache allocation that wastes memory through fragmentation. The paper introduces PagedAttention and vLLM to manage KV caches in paged, non-contiguous blocks, achieving 2-4× higher throughput than state-of-the-art systems.

  • Problem

    Existing LLM serving systems waste KV-cache memory through internal and external fragmentation, limiting efficient batching and throughput.

  • Method

    PagedAttention stores KV caches in non-contiguous blocks, while vLLM uses block-level memory management and scheduling for high-throughput serving.

  • Results

    2-4× higher throughput than state-of-the-art systems is achieved by vLLM across evaluated models and workloads.

  • Takeaways & Limitations

    PagedAttention enables more requests to fit in GPU memory for batching, improving LLM serving throughput.

  • Takeaways & Limitations

    The paging techniques are not generally applicable to workloads with static tensor shapes or compute-bound performance, where they may degrade performance.

Abstract

from arXiv · show

High throughput serving of large language models (LLMs) requires batching sufficiently many requests at a time. However, existing systems struggle because the key-value cache (KV cache) memory for each request is huge and grows and shrinks dynamically. When managed inefficiently, this memory can be significantly wasted by fragmentation and redundant duplication, limiting the batch size. To address this problem, we propose PagedAttention, an attention algorithm inspired by the classical virtual memory and paging techniques in operating systems. On top of it, we build vLLM, an LLM serving system that achieves (1) near-zero waste in KV cache memory and (2) flexible sharing of KV cache within and across requests to further reduce memory usage. Our evaluations show that vLLM improves the throughput of popular LLMs by 2-4$\times$ with the same level of latency compared to the state-of-the-art systems, such as FasterTransformer and Orca. The improvement is more pronounced with longer sequences, larger models, and more complex decoding algorithms. vLLM's source code is publicly available at https://github.com/vllm-project/vllm

1 Introduction

LLM serving is expensive and memory-bound because autoregressive generation underutilizes GPUs, while inefficient KV-cache management limits batch size and throughput. PagedAttention addresses fragmentation and sharing limitations through paged KV-cache blocks, forming the basis of vLLM’s near-zero-waste serving system.

  • Motivation: Autoregressive token generation makes LLM serving memory-bound, underutilizing GPU computation and limiting serving throughput.Each request repeatedly generates tokens until termination, conditioning on the prompt and previously generated output tokens.
  • KV-cache management: Inefficient KV-cache management limits batch size and consequently LLM throughput because request states grow dynamically while model weights remain constant.For a 13B-parameter LLM on an NVIDIA A100 with 40GB RAM, approximately 65% is allocated to static model weights and close to 30% to dynamic request states.
  • KV-cache management: Existing systems suffer internal and external fragmentation by storing each request’s dynamically changing KV cache in contiguous memory sized for its maximum length.Pre-allocating for a maximum such as 2048 tokens can waste substantial space when the actual request is shorter.
  • KV-cache management: Existing systems cannot share KV-cache memory across sequences generated by parallel sampling or beam search because those sequences occupy separate contiguous spaces.These decoding algorithms generate multiple outputs per request whose sequences can partially share their KV cache.
  • PagedAttention: PagedAttention divides each request’s KV cache into fixed-token blocks that need not be stored contiguously, enabling more flexible memory management.The approach is inspired by operating-system virtual memory and paging, which address memory fragmentation and sharing.
  • vLLM: vLLM builds a high-throughput distributed serving engine on PagedAttention and achieves near-zero KV-cache memory waste through block-level management and preemptive scheduling.It supports GPT, OPT, and LLaMA models of varying sizes, including models exceeding a single GPU’s memory capacity.

2 Background

LLMs model token sequences autoregressively with Transformer self-attention and serve requests through parallel prompt processing followed by sequential generation that reuses cached key and value vectors. Iteration-level batching improves serving flexibility by adding and removing requests after each generation iteration.

  • Language Modeling: Language models factorize the joint probability of tokens into conditional probabilities, so each new token depends on the preceding sequence.This autoregressive decomposition is expressed as P(x) = P(x1) · P(x2 | x1) · · · P(xn | x1, . . ., xn−1).
  • Transformers: Transformers are the standard large-scale architecture, with self-attention computing outputs from queries, preceding keys, and value vectors.Other listed Transformer components are applied independently at each position, whereas self-attention connects positions through attention computation.
  • LLM Serving: LLM services generate output tokens conditioned on input prompts, treating the concatenated prompt and output as a sequence.A request supplies prompt tokens (x1, . . ., xn), and the service generates (xn+1, . . ., xn+T).
  • Generation Phases: Generation has a parallel prompt phase and a sequential autoregressive phase that caches prior key and value vectors for future tokens.The prompt phase uses matrix-matrix multiplication over known prompt tokens, while each later iteration computes only the new key and value vectors.
  • Batching and Scheduling: Iteration-level scheduling updates batches after each iteration by removing completed requests and adding new ones, reducing admission delay compared with request-level batching.A new request can be processed after waiting for a single iteration rather than for the entire batch to complete.

3 Memory Challenges in LLM Serving

LLM serving throughput is constrained by GPU memory, especially the dynamically growing KV cache. Existing contiguous, statically pre-allocated management wastes memory and cannot flexibly support complex decoding or variable request lengths.

  • Memory-bound throughput: LLM serving throughput is memory-bound because GPU capacity, particularly KV-cache storage, limits the number of requests that can be batched.Fine-grained batching reduces computation waste but does not remove the KV-cache memory constraint.
  • Large KV cache: 1.6 GB is the maximum KV-cache memory required for one 2048-token OPT-13B request.Each token requires 800 KB, calculated from two vectors, hidden size 5120, 40 layers, and FP16 storage.
  • Complex decoding algorithms: Complex decoding algorithms require flexible KV-cache sharing because sharing the prompt cache can reduce memory usage across multiple samples.The prompt portion accounts for 12% of total KV-cache memory in the cited experiment, while sharing patterns evolve during decoding.
  • Scheduling for unknown input & output lengths: Variable input and output lengths require memory management to accommodate diverse prompts and expanding KV caches during generation.Growing output caches may exhaust memory needed by incoming requests or ongoing generation, requiring scheduling decisions.
  • Fragmentation and static allocation: Existing systems statically pre-allocate contiguous KV-cache chunks, causing waste from reserved future-token slots, internal fragmentation, and external fragmentation.Compaction is impractical for massive KV caches and still cannot enable decoding-specific memory sharing.

4 Method

vLLM introduces PagedAttention, which stores each request’s KV cache in non-contiguous fixed-size blocks mapped through logical-to-physical block tables. This design enables dynamic allocation, sharing, copy-on-write, mixed decoding, and CPU swapping for efficient memory management.

  • 4.1 PagedAttention: PagedAttention partitions each sequence’s KV cache into fixed-size blocks that may reside in non-contiguous physical memory.The kernel identifies and fetches blocks separately during attention computation.
  • 4.2 KV Cache Manager: Logical-to-physical block tables let vLLM grow KV caches dynamically without reserving all maximum-length positions in advance.Each entry records the physical block corresponding to a logical block and its filled positions.
  • 4.2 KV Cache Manager: Allocating a new physical block only after earlier blocks fill limits each request’s waste to one block and allows more requests to fit in memory.The resulting higher memory utilization supports larger batches and improves throughput.
  • 4.4 Memory Management for Decoding: vLLM shares prompt KV-cache blocks across multiple output samples and uses copy-on-write for the final logical block, greatly reducing memory usage for long prompts.The same physical block-sharing mechanism reduces KV-cache copying across beam candidates.
  • 4.4 Memory Management for Decoding: A common logical-to-physical mapping layer lets vLLM simultaneously process requests using different decoding methods despite their diverse memory-sharing patterns.The model and execution kernel operate on physical block lists while the mapping layer hides the sharing complexity.
  • 4.5 Variable-Length Sequences: When GPU physical blocks are exhausted, vLLM evicts selected sequences and transfers their KV-cache blocks to CPU memory using a CPU block allocator.This implements swapping for managing KV-cache capacity across GPU and CPU memory.

5 Implementation

vLLM combines an OpenAI-compatible FastAPI frontend with a GPU inference engine, custom CUDA kernels, and Python control components. It supports decoding algorithms through sequence fork, append, and free operations.

  • System architecture: vLLM provides a FastAPI frontend extending the OpenAI API, with per-request sampling controls including maximum sequence length and beam width k.The engine comprises 8.5K lines of Python and 2K lines of C++/CUDA code.
  • System architecture: The implementation develops the scheduler and block manager in Python while using custom CUDA kernels for key operations.PyTorch and Transformers support the system, while NCCL provides tensor communication across distributed GPU workers.
  • GPU kernels: PagedAttention is optimized with GPU kernels because its memory access patterns are inefficiently supported by existing systems.The kernels fuse reshape with block write and block read with attention to reduce overhead and optimize memory access.
  • Decoding algorithms: vLLM implements decoding algorithms with fork, append, and free methods that create sequences, add tokens, and delete sequences.Parallel sampling forks multiple output sequences from one input, appends tokens each iteration, and frees sequences meeting a stopping condition.

6 Evaluation

The evaluation tests vLLM across multiple model sizes, workloads, and decoding methods, showing higher sustainable request rates and substantial KV-cache savings than Orca baselines. PagedAttention is especially beneficial for long prompts, shared prefixes, parallel sampling, and beam search.

  • Evaluation setup: The evaluation uses OPT models with 13B, 66B, and 175B parameters plus LLaMA-13B on A2 instances with NVIDIA A100 GPUs.Workloads are synthesized from ShareGPT and Alpaca datasets, whose average input and output lengths differ by 8.4× and 5.8×, respectively.
  • Basic sampling: 1.7×–2.7× higher request rates are sustained by vLLM than Orca (Oracle), and 2.7×–8× higher than Orca (Max), at similar latencies on ShareGPT.For OPT-13B, vLLM processes 2.2× more simultaneous requests than Orca (Oracle) and 4.3× more than Orca (Max).
  • Memory sharing: vLLM’s memory sharing yields 6.1%–9.8% savings for parallel sampling and 37.6%–55.2% for beam search on Alpaca.On ShareGPT, savings reach 16.2%–30.5% for parallel sampling and 44.3%–66.3% for beam search.
  • Cross-request prefix sharing: With shared translation prefixes, vLLM achieves 1.67× higher throughput than Orca (Oracle) for the one-shot workload.The experiment uses LLaMA-13B and WMT16 English-to-German translation requests with one-shot or few-shot prefixes.
  • Chatbot workload: 2× higher request rates are sustained by vLLM than the three Orca baselines on the chatbot workload.PagedAttention handles long prompts by resolving memory fragmentation and reservation, while Orca reserves output space using buddy allocation.

7 Ablation Studies

The ablations identify performance costs and tuning tradeoffs in PagedAttention and vLLM. Dynamic block mapping adds attention-kernel overhead, while block size and recovery mechanism choices substantially affect performance.

  • PagedAttention overhead: 20–26% higher attention-kernel latency results from PagedAttention’s block-table accesses, extra branches, and variable-length handling versus FasterTransformer.The overhead arises in GPU operations involving stored KV-cache blocks, including reads, writes, and attention.
  • Block-size ablation: Block size balances GPU parallelism, internal fragmentation, and KV-cache sharing probability.Too-small blocks can underutilize GPU parallelism, whereas too-large blocks increase fragmentation and reduce sharing opportunities.
  • Block-size ablation: Block sizes 16–128 perform best on ShareGPT, while Alpaca favors sizes 16 and 32 and degrades substantially with larger blocks.The Alpaca degradation occurs because sequences become shorter than the block sizes; evaluations use basic sampling under fixed request rates.
  • Recovery mechanisms: Small block sizes make swapping excessively costly because numerous small CPU–GPU transfers limit effective PCIe bandwidth.The study compares recomputation and swapping through end-to-end evaluations and microbenchmarks.

8 Discussion

The discussion argues that paging benefits LLM serving because KV-cache allocation is dynamic and constrained by GPU capacity, but may not generalize to workloads with static tensor shapes. It also highlights vLLM’s LLM-specific techniques for reducing paging overhead and recovering evicted states.

  • Applying virtual memory and paging to other GPU workloads: Paging suits LLM serving because output lengths are unknown in advance, requiring dynamic memory allocation under GPU-memory capacity constraints.The passage contrasts this with DNN training, where tensor shapes are typically static and memory allocation can be optimized.
  • LLM-specific optimizations: vLLM augments virtual memory and paging with application-specific semantics, including an all-or-nothing swap-out policy for request token states.Processing a request requires all corresponding token states to remain stored in GPU memory.
  • LLM-specific optimizations: vLLM recovers evicted blocks through recomputation, an approach described as infeasible in operating systems.
  • LLM-specific optimizations: vLLM reduces paging’s memory-indirection overhead by fusing GPU memory-access kernels with attention and other operations.

9 Related Work

Prior work spans general model serving, transformer-specific optimizations, and memory-reduction techniques for training and inference. vLLM complements Orca’s scheduling approach and introduces block-level memory management for online serving.

  • General model serving systems: General model serving systems study batching, caching, placement, and scheduling for serving single or multiple models.Examples include Clipper, TensorFlow Serving, Nexus, InferLine, Clockwork, DVABatch, and REEF.
  • Specialized serving systems for transformers: Transformer-serving systems use GPU kernel optimizations, advanced batching, model parallelism, and parameter sharing; Orca is most relevant to vLLM.
  • Comparison to Orca: Orca increases throughput through iteration-level scheduling, whereas vLLM increases memory utilization so more requests fit into GPU memory.The techniques are complementary: Orca schedules and interleaves requests, while PagedAttention reduces fragmentation and enables sharing.
  • Memory optimizations: Existing memory optimizations include swapping and recomputation for training, weight and token-state swapping for inference, tensor-lifetime optimization, and tiled attention kernels.FlexGen does not target online serving, OLLA lacks fine-grained block-level management and online serving, and FlashAttention reduces attention-memory peaks and I/O costs.

10 Conclusion

The paper proposes PagedAttention and presents vLLM, a high-throughput LLM serving system with efficient memory management. It adapts virtual memory and copy-on-write techniques to manage KV caches and support various decoding algorithms.

  • 10 Conclusion: PagedAttention stores attention keys and values in non-contiguous paged memory.The algorithm is inspired by operating-system memory management.
  • 10 Conclusion: vLLM is a high-throughput LLM serving system enabled by PagedAttention’s efficient memory management.
  • 10 Conclusion: Virtual memory and copy-on-write techniques are adapted to efficiently manage KV caches and handle various decoding algorithms in LLM serving.
Loading 2309.06180v1…