Skip to content
back to the archive page
#AI

vLLM and PagedAttention: How Ideas from Operating Systems Accelerated LLM Inference (Category AI)

#AI #Research #Software #Architecture #Engineering #RnD

I finally read the entire paper “Efficient Memory Management for Large Language Model Serving with PagedAttention” (Symposium on Operating Systems Principles (SOSP) 2023), the paper that gave rise to vLLM, one of the most popular open-source LLM inference engines. I mentioned it earlier in my Tanenbaum review; now I want to examine the engineering idea itself. It is elegant: the authors made inference several times faster without changing the model or substantially touching the computations—they fixed memory management.

Here is the context. LLM-serving throughput is constrained by batching: a GPU works efficiently when it generates tokens for many requests at once. Batch size, in turn, is limited by memory for the KV cache—the attention keys and values for every token already processed in each request. In the paper’s configuration—OPT-13B with FP16 weights on an A100 with 40 GB—the weights occupy about 65% of memory and the KV cache another roughly 30%. In this setup, the cache determines how many requests fit into the batch.

The Berkeley team, whose authors include Woosuk Kwon, Zhuohan Li, and Ion Stoica, showed that systems of the time—FasterTransformer and Orca—handled this memory wastefully. A request’s KV cache was stored as one contiguous chunk reserved immediately for the maximum length, such as 2,048 tokens, even if the response contained only fifty. The losses came from three sources:

  1. slots reserved “for the future”
  2. internal fragmentation in chunks allocated with spare capacity
  3. external fragmentation in the allocator According to the authors’ measurements, useful data occupied only 20–38% of KV-cache memory; the rest was wasted.

The solution was PagedAttention: virtual memory and paging from operating systems transferred to inference. The KV cache is divided into fixed-size blocks—16 tokens by default—analogous to pages. Logically, the sequence is contiguous; physically, its blocks may live anywhere, while a block table—analogous to a page table—stores the mapping. Memory is allocated as generation proceeds, and only the final block remains partially filled: in the authors’ experiments, KV-cache utilization rose to almost 96%.

A second classic mechanism then comes into play: shared pages and copy-on-write, as with process fork. During parallel sampling, several response variants point to the same physical prompt blocks; in beam search, hypotheses share common prefixes, and when they diverge, one block is copied rather than the entire cache. The authors report memory savings of 6–10% for parallel sampling and 37–55% for beam search.

The vLLM engine was built on top of the algorithm: continuous batching (iterative scheduling from Orca), a block manager, eviction of whole sequences when memory runs short—using swap to CPU memory or recomputation—and custom CUDA kernels for working with scattered blocks. An honest detail is that the attention kernel became 20–26% slower because of access through the block table. Yet the system as a whole delivers 2–4 times the throughput of FasterTransformer and Orca at the same latency, because many more concurrent requests fit into memory. The benefit grows with sequence length and model size.

What happened next is well known: vLLM was released as open source in the summer of 2023, and competitors quickly adopted the idea. Paged KV cache is described in the TensorRT-LLM documentation, while Hugging Face TGI directly uses vLLM’s CUDA kernels. In my view, the approach has effectively become an industry standard.

My main takeaway has two parts. 1️⃣ The bottleneck in LLM inference turned out to be not FLOPs, but memory management, and it can be treated with techniques from an operating-systems textbook that are more than half a century old. 2️⃣ The authors deliberately lost on a local metric (kernel speed) to win on a system metric (throughput at a given latency). Knowing which level to optimize is a sign of mature systems engineering.

#AI #Research #Software #Architecture #Engineering #RnD