Memory Optimization in LLM Serving with PagedAttention

Introduction

Welcome to our next lesson on Inference at Scale. In the previous lesson, we established that while the KV cache is a brilliant solution to the quadratic complexity of autoregressive generation, it introduces a significant memory bottleneck. We saw how naive memory allocation—pre-allocating a large, contiguous block of VRAM for each request's maximum possible sequence length—is extremely wasteful and limits system throughput.

This lesson directly addresses that problem, focusing on the learning outcome: Explain how specialized serving libraries (e.g., vLLM, TGI) optimize memory management for LLM inference using techniques like PagedAttention.

We will explore how modern serving systems overcome the memory challenges of the KV cache. The central idea we'll discuss, PagedAttention, is a powerful example of applying a classic computer science concept—virtual memory and paging from operating systems—to solve a critical problem in a new domain. Given your background, the parallels to OS memory management will be immediately apparent and provide a solid theoretical foundation for understanding these advanced systems.

The Problem Revisited: Memory Waste in LLM Serving

Before diving into the solution, let's precisely define the sources of memory inefficiency in traditional serving systems that use contiguous allocation for the KV cache.

As the video explains, the total memory waste is substantial, often leaving only 20-40% of the allocated KV cache space actually in use. Let's formalize those concepts:

• Internal Fragmentation: This occurs when a pre-allocated memory block is larger than the actual data it needs to hold. In LLM inference, a system might allocate space for 2048 tokens, but the generation stops after 300. The remaining 1748 token slots are wasted within that allocation.
• Reservation: This is a form of waste where memory is allocated for future tokens in a sequence. While this memory will eventually be used, it remains idle for the duration of the generation, preventing other requests from using it.
• External Fragmentation: When requests with different maximum lengths are processed, the memory allocator ends up with free memory chunks of various sizes scattered throughout the address space. A new request might not fit into any single one of these chunks, even if the total free memory is sufficient.

These issues severely limit the number of requests that can be batched together, directly capping the system's throughput.

The Solution: PagedAttention and vLLM

To overcome these limitations, researchers from UC Berkeley developed vLLM, a high-throughput serving engine built around an algorithm called PagedAttention. The core idea is to abandon contiguous memory allocation in favor of a paged memory model, directly analogous to how modern operating systems manage memory.

Let's explore the theory behind this approach by reading from the original vLLM paper.

How PagedAttention Works

The vLLM system virtualizes the KV cache by introducing several key concepts borrowed from operating systems:

1. KV Blocks: The entire GPU memory space for the KV cache is partitioned into fixed-size blocks. These are the equivalent of physical pages. A typical block might hold the KV pairs for 16 or 32 tokens.
2. Logical vs. Physical Blocks: A sequence's KV cache is viewed as a contiguous sequence of logical blocks. However, these logical blocks are mapped to non-contiguous physical blocks in GPU memory.
3. Block Table: For each request, vLLM maintains a block table—the equivalent of a page table—that stores the mapping from its logical blocks to the physical blocks in memory.

This architecture fundamentally changes how memory is managed:

• On-Demand Allocation: Physical blocks are allocated only when they are needed, as the sequence generation progresses. There is no need to pre-allocate memory for the maximum sequence length.
• Elimination of Fragmentation: Since all physical blocks are the same size, there is no external fragmentation. Internal fragmentation is limited to only the very last block of a sequence, making the waste negligible.

The video you started with explains this mechanism very clearly.

Of course, this non-contiguous memory layout requires an attention kernel that knows how to operate on it. PagedAttention is not just the memory manager; it's also a custom CUDA kernel that can gather the scattered KV blocks to perform the attention calculation correctly.

The result of this design is near-zero memory waste, which dramatically increases memory utilization. As the video highlights, this allows for a 3-5x increase in the effective batch size, leading to a major boost in throughput.

Advanced Capabilities: Memory Sharing with Copy-on-Write

The benefits of PagedAttention go beyond just eliminating fragmentation. The abstraction of logical-to-physical block mapping enables highly efficient memory sharing for complex decoding strategies.

Consider parallel sampling, where multiple output sequences are generated from the same prompt. With contiguous allocation, the KV cache for the prompt would have to be duplicated for each output. With PagedAttention, all output sequences can simply map their logical prompt blocks to the same set of physical blocks.

A copy-on-write mechanism is used to handle modifications. If a sequence needs to append a new token to a shared block, vLLM first checks the block's reference count. If the count is greater than one, it allocates a new block, copies the content, and updates the sequence's block table to point to the new copy before writing. This is precisely how fork() is implemented efficiently in operating systems.

This same principle applies to more complex strategies like beam search and even allows for sharing common prefixes (e.g., system prompts) across different users.

The following resources detail how this sharing is implemented.

Broader Context: vLLM, TGI, and TensorRT-LLM

While vLLM's PagedAttention is a revolutionary approach, it's important to understand it within the broader ecosystem of serving libraries.

• vLLM: Primarily focuses on maximizing throughput by optimizing memory management via PagedAttention and using a sophisticated continuous batching scheduler.
• HuggingFace Text Generation Inference (TGI): A production-focused server that prioritizes stability, ease of use, and broad model support. It uses a more traditional memory allocation scheme but implements dynamic batching and other optimizations. It often provides lower time-to-first-token at the expense of maximum throughput.
• NVIDIA TensorRT-LLM: An optimization library from NVIDIA that compiles models into highly efficient engines. It also includes an implementation of paged KV cache (sometimes called in-flight batching) and supports features like INT8/FP8 quantization to further reduce memory and accelerate computation on NVIDIA hardware.

The choice between these systems involves design trade-offs:

In essence, vLLM's PagedAttention represents a deep, architectural optimization for memory efficiency, while TGI offers a more feature-rich, general-purpose solution. The trend, however, is clear: efficient, paged KV cache management is becoming a standard feature in high-performance serving frameworks, as evidenced by its inclusion in NVIDIA's own TensorRT-LLM.

Conclusion

In this lesson, we delved into the sophisticated memory management techniques used by modern LLM serving libraries. We saw that the memory waste inherent in naive, contiguous KV cache allocation can be almost entirely eliminated by applying principles from operating system design.

Key Takeaways:

• Problem: Traditional KV cache management suffers from severe internal and external fragmentation, wasting a significant portion of valuable GPU memory and limiting throughput.
• Solution (PagedAttention): By partitioning the KV cache into non-contiguous, fixed-size blocks and managing them with a logical-to-physical mapping system (block tables), libraries like vLLM eliminate fragmentation and allocate memory on demand.
• OS Analogy: PagedAttention is a direct application of virtual memory, paging, and copy-on-write concepts to the domain of LLM inference.
• Benefits: This approach leads to near-optimal memory utilization, enabling larger batch sizes and 2-4x higher throughput. It also facilitates efficient memory sharing for complex decoding algorithms like beam search.
• System Trade-offs: Different serving libraries make different design choices. vLLM is highly optimized for throughput via PagedAttention, while TGI prioritizes production features and ease of deployment, resulting in different performance profiles.

Preview of the Next Lesson:

Now that we've seen how libraries like vLLM optimize the execution of a single model, our next lesson will zoom out to the system level. We will explore the role of a dedicated inference server like NVIDIA Triton, which orchestrates multiple models, manages request queues, and handles the complexities of deploying and scaling these powerful libraries in a real-world production environment.