Efficient Memory Management for LLM Serving
Introduced PagedAttention and the vLLM engine, applying operating system paging concepts to the KV cache to eliminate memory fragmentation and drastically increase serving throughput.
Paper: Efficient Memory Management for Large Language Model Serving with PagedAttention
Authors: Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Hao Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, Ion Stoica · 2023
Read the paperThe Problem
While FlashAttention optimized training, deploying LLMs for serving (running inference for many concurrent users) faced a completely different bottleneck: the KV cache memory.
When an LLM generates text, it must store the Key and Value vectors for all previous tokens in memory. The problem is that the length of the output is unpredictable. Before PagedAttention, serving engines pre-allocated a large, contiguous chunk of memory for every incoming request, assuming it might generate up to its maximum limit.
This resulted in massive internal and external memory fragmentation. If a request only generated 10 tokens but was allocated space for 1000, 990 slots were wasted. The researchers found that existing systems wasted up to 60-80% of their KV cache memory. Because GPU memory is limited, this fragmentation severely restricted the number of users that could be served simultaneously (batch size), making LLM hosting incredibly expensive.
The Idea
The UC Berkeley researchers looked to classical operating systems for the solution. In an OS, physical memory is divided into fixed-size "pages." A process's virtual memory appears contiguous, but the OS maps it to non-contiguous physical pages dynamically as needed, eliminating fragmentation.
The authors applied this exact concept to the KV cache, inventing PagedAttention.
How It Works
- Blocks: Instead of storing a sequence's KV cache in one massive contiguous tensor, PagedAttention divides the KV cache into fixed-size "blocks" (e.g., each block holds KV vectors for 16 tokens).
- Block Table: The engine maintains a block table mapping logical token sequences to physical blocks in GPU memory, just like an OS page table.
- Dynamic Allocation: When a request starts, it is only allocated one block. As the LLM generates tokens and fills that block, the engine dynamically allocates a new, non-contiguous physical block and updates the table.
- Memory Sharing: Because the memory is block-based, multiple requests can easily share blocks. If two users send identical prompts (or if generating multiple completions for the same prompt), the engine maps them to the exact same physical blocks, saving massive amounts of memory.
Why It Mattered
PagedAttention virtually eliminated memory fragmentation. By freeing up 60-80% of the KV cache memory, the system could batch significantly more concurrent requests together.
The researchers open-sourced their engine as vLLM, which immediately delivered 2x to 4x throughput improvements over existing state-of-the-art serving engines (like HuggingFace TGI). It became the industry standard engine for open-weight LLM inference almost overnight.
What Came After
vLLM and PagedAttention are now foundational to modern AI infrastructure. The concept has been widely adopted and integrated into competing serving engines (like TensorRT-LLM and LMDeploy). Continuous improvements have focused on supporting complex routing (MoE) and handling prefix-caching across massive clusters.