Papers
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 · SOSP 2023, 2023
The problem: the KV cache limits the batch size
Serving a large language model (LLM) with high throughput requires processing many requests together in a batch. The bottleneck, according to the paper, is GPU memory. An autoregressive LLM generates one token at a time, and for each token already seen it stores the attention key and value vectors, the so-called KV cache, so as not to recompute them at every step. In the paper's example, a 40 GB NVIDIA A100 GPU serving a 13-billion-parameter model uses about 65% of its memory for the weights, which stay fixed, and more than 30% for the KV cache; the remainder is small and serves for temporary activations. Since the weights are constant, it is the management of the KV cache that decides how many requests fit in the batch. The size is large: in OPT-13B, a single token takes about 800 KB of KV cache, so a sequence of up to 2048 tokens reaches about 1.6 GB. Moreover, the KV cache grows and shrinks dynamically, and the final length of the output is not known in advance. The authors argue that GPU compute grows faster than memory (from the A100 to the H100, FLOPS rise more than 2x, but maximum memory stays at 80 GB), so memory tends to be an increasingly relevant bottleneck.
Waste from fragmentation in existing systems
The earlier systems cited by the authors, such as FasterTransformer and Orca, store each request's KV cache in a contiguous block of memory, because the deep learning frameworks of the time expect contiguous tensors. Since the final output length is unpredictable, they reserve space up front for the maximum possible request length (for example, 2048 tokens). The paper identifies three sources of waste. The first is reservation: slots set aside for future tokens that have not been generated yet, but that stay occupied and keep other requests from using that space. The second is internal fragmentation: the oversized space left over when the actual output is much shorter than the maximum. The third is external fragmentation, caused by the memory allocator (the paper cites the buddy allocator) when each request reserves a different size. In the authors' measurements with Orca variants (Max, Pow2, and Oracle), only 20.4% to 38.2% of KV cache memory held actual token states; in vLLM, the proposed system, that figure was 96.3%. The authors also note that compacting memory to reduce fragmentation would be impractical, given the volume of KV cache data.
The core idea: PagedAttention and KV cache blocks
The paper proposes PagedAttention, an attention algorithm inspired by the paged virtual memory of operating systems. Each sequence's KV cache is divided into fixed-size blocks, each holding the keys and values of a fixed number of tokens; the analogy is that blocks are pages, tokens are bytes, and requests are processes. Physical blocks do not need to be contiguous in GPU memory: the attention kernel fetches each block separately. vLLM, the serving engine built on this idea, keeps for each request a block table that maps logical blocks (filled from left to right as tokens appear) to physical blocks, also recording how many positions in each block are already filled. A new physical block is allocated only when the previous ones are full, on demand, without reserving the maximum up front. As a result, a request's waste is confined to its last block, and external fragmentation disappears because all blocks have the same size. Block size is a trade-off: blocks that are too small make poor use of GPU parallelism, and large blocks increase internal fragmentation and reduce the chance of sharing. vLLM uses 16 by default, a choice the authors justify with an evaluation of different sizes.
Sharing, continuous batching, and preemption
Because the mapping goes through a table, several logical blocks can point to the same physical block, tracked by a reference counter. In parallel sampling, where one request generates several outputs from the same prompt, the prompt's KV cache is stored only once. When a sequence needs to write to a shared block, vLLM applies copy-on-write at the block level: it allocates a new block, copies the contents and decrements the original's counter. In beam search the sharing is greater and changes throughout decoding; with a shared prefix (for example, an instruction common to several requests), the provider can keep already-computed prefix blocks. This context adds to iteration-level batching, described in the paper as a prior technique (as in Orca): unlike traditional methods, which work at the request level, it operates at the iteration level; after each iteration, completed requests leave the batch and new ones join, and a new request waits only one iteration, not for the entire batch to finish. When the GPU runs out of physical blocks, vLLM uses an FCFS policy and preempts the most recent requests first, with an all-or-nothing policy: it evicts all the blocks of a sequence, and sequences of the same group are preempted together. To recover the blocks, there are two options: swapping to CPU memory or recomputing the KV cache.
What the evaluation measures and the reported results
The authors evaluate vLLM with OPT models of 13B, 66B and 175B parameters and LLaMA 13B, on A100 GPUs on Google Cloud, using workloads synthesized from the ShareGPT and Alpaca datasets, with arrivals modeled by a Poisson distribution. The main metric is normalized latency, the mean across requests of each request's end-to-end latency divided by its output length (as in Orca), measured as a function of request rate: a high-throughput system keeps this latency low even at high rates. The baselines are FasterTransformer and Orca, the latter reimplemented by the authors in three versions (Oracle, Pow2 and Max) according to how much space they reserve. With basic sampling on ShareGPT, vLLM sustained request rates 1.7x to 2.7x higher than Orca (Oracle) and 2.7x to 8x higher than Orca (Max), with similar latencies; against FasterTransformer, up to 22x. In the abstract, the paper cites a 2 to 4x throughput gain at the same latency. The gains grow with more sharing: in beam search, the memory savings on Alpaca were 37.6% to 55.2% (depending on the width tested), and with a shared prefix of 5 examples, throughput was 3.58x that of Orca (Oracle). The cost: the attention kernels were 20% to 26% slower than FasterTransformer's; the authors consider the cost small because it affects only the attention operator, not the others such as the linear layers, and vLLM's end-to-end performance was still superior.
See the ideas at work
Practice the designs that use this paper in the System Design school.