PagedAttention: Optimizing KV Cache in vLLM

In our last few lessons, we've focused on optimizing inference performance by shrinking the model itself through quantization, comparing methods like AWQ and GPTQ. You've benchmarked these and seen firsthand the significant gains in throughput and reductions in memory. However, model size is only one part of the memory story. Today, we turn our attention to the other major memory consumer in LLM inference: the KV cache.

This lesson explores PagedAttention, a groundbreaking memory management technique implemented in vLLM. You will learn how it solves the critical problem of memory fragmentation in the KV cache, which is often the primary bottleneck preventing higher throughput.

Your extensive background in site reliability and platform engineering, particularly with operating systems, will be a significant asset here. The core idea of PagedAttention is directly inspired by the virtual memory and paging systems you're already familiar with, providing a powerful analogy for understanding this optimization.

1. The Problem: A Leaky Bucket of GPU Memory

As we've discussed, the KV cache is essential for efficient decoding. It stores the intermediate key and value vectors for every token in a sequence, preventing costly re-computation at each generation step.

However, traditional LLM serving systems manage this cache inefficiently. For each incoming request, they allocate a single, large, contiguous slab of GPU memory to hold the entire potential KV cache for that request. This "pre-allocation" strategy is necessary because the final output length is unknown, but it leads to massive memory waste.

This video from IBM Technology provides a quick, high-level overview of the challenges in LLM serving, including this memory problem.

The memory waste described in the video manifests in a few distinct ways. The following video from Anyscale, the company behind vLLM, provides a more detailed breakdown.

This extreme inefficiency is the core problem PagedAttention was designed to solve. When memory is wasted, you can't fit as many requests into a batch, which directly limits your system's overall throughput.

2. The Solution: Virtual Memory for the KV Cache

The creators of vLLM recognized that this memory fragmentation problem is not new. Operating systems solved a very similar issue decades ago with virtual memory and paging. PagedAttention applies this same fundamental concept to the KV cache.

Instead of allocating one large, contiguous slab per request, PagedAttention partitions the entire KV cache space on the GPU into many small, fixed-size blocks.

This simple but powerful idea is explained in detail in the following article.

Let's summarize how this "paged" approach solves the fragmentation problems:

• No External Fragmentation: Because all blocks are the same size, any freed block can immediately be reused by any new request. The memory space never becomes a "Swiss cheese" of unusable holes.
• Minimal Internal Fragmentation: Memory is allocated "just-in-time" as a sequence generates new tokens. Waste is now confined to the very last block of a sequence. With a typical block size of 16 tokens, the maximum waste per request is less than 16 tokens' worth of memory, compared to thousands in the old system.
• No Reservation Waste: Memory isn't pre-allocated for the maximum possible length. Blocks are only assigned from the free pool as they are needed.

The result is a dramatic increase in memory efficiency. The vLLM paper reports that memory waste drops from 60-80% down to less than 4%. This efficiency gain allows for much larger batch sizes, directly leading to higher throughput.

3. Making it Work: The PagedAttention Kernel

Now, you might be wondering: "How can the attention mechanism work if the keys and values for a sequence are scattered all over GPU memory?"

This is the "compute" part of the solution. The standard attention mechanism in libraries like PyTorch expects key and value tensors to be contiguous. PagedAttention introduces a custom attention kernel that is "block-aware."

The PagedAttention kernel simply reads the request's block table to find the physical memory addresses of the KV blocks. It then iterates through these blocks in their logical order, performing the attention calculation piece by piece. A running softmax normalization ensures that the final mathematical result is exactly the same as it would be with a contiguous memory layout.

The key takeaway is that PagedAttention is a memory layout optimization; it does not change the fundamental mathematics of the attention mechanism itself.

4. Advanced Benefit: Zero-Cost Sharing with Copy-on-Write

The block-based architecture unlocks another powerful feature that is particularly relevant to your work as a systems engineer: nearly free memory sharing.

Consider use cases like parallel sampling (generating multiple candidate outputs from one prompt) or beam search. In these scenarios, multiple output sequences share a common prefix (the prompt). In a traditional system, you would have to duplicate the entire KV cache for the prompt for each new sequence, which is incredibly wasteful.

PagedAttention solves this using a copy-on-write mechanism, a concept that will be very familiar from your experience with Linux memory management and the fork() system call.

Here's how it works:

1. Sharing: When multiple sequences share a prompt, their block tables simply point to the same physical memory blocks. A reference count for each physical block is incremented. No data is copied.
2. Divergence: When one of the sequences generates a new, unique token, it needs to write to the last shared block.
3. Copy-on-Write: Before writing, the system checks the block's reference count. If it's greater than one, it allocates a new block, copies the contents of the shared block to it, and updates the diverging sequence's block table to point to this new block. The write then proceeds on the new, private block.

This block-level copy-on-write is vastly more efficient than copying the entire multi-gigabyte KV cache history, dramatically reducing memory overhead for complex sampling strategies.

Conclusion

In this lesson, you've seen how PagedAttention fundamentally redesigns KV cache memory management, moving from a wasteful contiguous allocation model to a highly efficient, block-based system inspired by OS virtual memory.

Here are the key takeaways:

• The Problem: Traditional KV cache management suffers from severe internal and external fragmentation, wasting 60-80% of allocated GPU memory and limiting throughput.
• The Analogy: PagedAttention solves this by applying the principles of virtual memory and paging to the KV cache.
• The Mechanism: It partitions the cache into fixed-size blocks, uses a per-request block table to map logical to physical blocks, and allows for non-contiguous storage. This nearly eliminates fragmentation.
• The Benefits: This leads to drastically improved memory utilization, allowing for larger batch sizes and 2-4x higher throughput. It also enables highly efficient memory sharing for complex decoding strategies via a copy-on-write mechanism.

PagedAttention isn't just an incremental improvement; it's a foundational change that enables much of the high performance we see in modern serving engines like vLLM. It is the infrastructure that allows other advanced optimizations to be built on top.

In our next lesson, we will explore one such optimization: speculative decoding. You will see how a small, fast "draft" model can be used to accelerate generation from a large, powerful model, and you'll appreciate how PagedAttention's efficient sharing capabilities make such a technique practical.