Introduction
In our last lesson, we implemented prefix caching, a powerful logical optimization for reusing computation across requests with shared prompts. However, our implementation relied on torch.cat and copy.deepcopy, which are inefficient from a low-level memory management perspective. In a high-throughput server, constantly allocating, copying, and deallocating variably sized tensors for the KV cache creates a significant problem: memory fragmentation. This wasted memory directly limits the number of concurrent requests you can serve, thus capping your throughput.
This lesson tackles that fundamental systems problem head-on. Your learning outcome is to explain the memory fragmentation problem in naive KV cache allocation and the core concept of paged memory management. We will dissect why traditional memory allocation fails for LLM inference and explore the elegant solution inspired by a classic computer science concept you are very familiar with: virtual memory in operating systems.
The Problem: Memory Fragmentation
When serving LLMs, the exact output length for any given request is unpredictable. This dynamic nature is the root cause of memory management challenges. Early serving systems tried to handle this by pre-allocating a single, contiguous block of GPU memory for each request, large enough to hold the KV cache for its maximum possible sequence length.
This seemingly safe approach leads to massive memory inefficiency. Let's watch a segment from a talk by the creators of vLLM that introduces this problem.
Fast LLM Serving with vLLM and PagedAttention
This clip from the Anyscale channel introduces the KV cache's role in inference and clearly visualizes three distinct types of memory waste that arise from naive allocation strategies.
Watch from 02:58 to 05:08. Pay close attention to the breakdown of memory waste into three categories: internal fragmentation, reservation, and external fragmentation. The video quantifies the impact, showing that 60-80% of memory can be wasted.
As the video explained, this waste can be categorized into two main types of fragmentation:
-
Internal Fragmentation: This is memory that is inside an allocated block but is unused. If you allocate space for 2048 tokens but the model only generates 50, the memory for the remaining 1998 tokens is wasted until the request is finished. It's allocated to the request but contains no data.
-
External Fragmentation: This occurs when free memory is broken up into small, non-contiguous chunks. Even if there is enough total free memory to satisfy a new request, if no single contiguous block is large enough, the allocation fails. This is a classic problem in systems with dynamic memory allocation.
To solidify your understanding of external fragmentation, let's explore a detailed walkthrough. Given your background, the analogy to a classic buddy memory allocator will be very direct.
Paged Attention from First Principles: A View Inside vLLM
The blog 'Paged Attention from First Principles' provides an excellent, detailed explanation of how external fragmentation occurs in the context of serving LLMs.
Read the section titled 'Fragmentation in continuous batching'. The text walks through a step-by-step example with a buddy allocator, clearly demonstrating how a series of allocations and deallocations can lead to a state where a large request cannot be fulfilled, despite ample total free memory.
The consequence of this fragmentation is severe: a significant portion of your expensive GPU VRAM sits idle, unable to be used. This drastically reduces the batch size, which in turn kills throughput and drives up the cost per token. The core limitation is the requirement that a request's KV cache must be stored in a contiguous block of memory. The solution is to remove that requirement.
The Analogy: Operating System Virtual Memory
The problem of managing a limited physical memory space to accommodate multiple processes with large, dynamic memory needs was solved decades ago in operating systems with virtual memory and paging. The core idea is to decouple a program's logical view of its memory from the physical layout.
Since you have a deep background in computer science, this concept will be a straightforward and powerful analogy for understanding the solution to KV cache fragmentation.
Paged Attention from First Principles: A View Inside vLLM
Let's quickly refresh the key concepts of OS-level paging. The same blog post we just looked at has a concise and accurate summary.
Read the sections 'Paging as the Analogy' and 'Operating Systems analogy'. Focus on the mapping: virtual address space -> pages, physical memory -> page frames, and the role of the page table and MMU. This establishes the conceptual framework we'll apply to the KV cache.
The Solution: Paged Memory Management for KV Cache
The vLLM paper introduced a groundbreaking approach by applying this exact OS concept to GPU memory management for the KV cache. This technique is enabled by an algorithm called PagedAttention.
Here is how the concepts map directly:
| Operating System Concept | LLM KV Cache Concept | Description |
|---|---|---|
| Virtual Address Space | Logical Token Sequence | A request's view of its KV cache as a single, continuous sequence of tokens. |
| Physical Memory | GPU VRAM | The physical memory on the accelerator where the cache is actually stored. |
| Page | Logical Block | A fixed-size group of tokens in the logical sequence. |
| Page Frame | Physical Block | A fixed-size, contiguous chunk of GPU VRAM that can store one logical block. |
| Page Table | Block Table | A per-request data structure that maps a request's logical blocks to physical blocks in VRAM. |
By partitioning the KV cache into fixed-size blocks, the memory manager no longer deals with variable-size allocations. It becomes a simple block manager.
- External fragmentation is eliminated: Because all physical blocks are the same size, any free block can be used for any request. There are no "holes" that are too small.
- Internal fragmentation is minimized: Waste only occurs in the very last block of a sequence. If your block size is 16 tokens, you waste at most 15 token slots' worth of memory per request, a negligible amount compared to pre-allocating for thousands of tokens.
Let's see this in action.
Fast LLM Serving with vLLM and PagedAttention
The same Anyscale talk now presents the solution, PagedAttention. The animations here are invaluable for visualizing how logical blocks are mapped to non-contiguous physical blocks.
Watch from 06:02 to 09:54. Observe how PagedAttention operates on non-contiguous memory, and how the block table is used to translate between the logical view (a request's sequence) and the physical view (scattered blocks in VRAM). Note the 'on-demand' allocation, which is key to avoiding waste.
This paged memory layout allows multiple requests to efficiently and flexibly share the GPU memory, with their physical blocks interleaved.
To get the most rigorous understanding, let's go to the source: the vLLM paper itself. It formalizes these concepts and shows how the attention mechanism is adapted to work with this new memory layout.
[PDF] Efficient Memory Management for Large Language Model Serving ...
The vLLM paper, 'Efficient Memory Management for Large Language Model Serving with PagedAttention', formally details the memory manager and the attention algorithm.
Please review the following sections: Section 4.1, 'PagedAttention': Understand how the attention formula is adapted for block-wise computation. Section 4.2, 'KV Cache Manager': This section formally describes the analogy to virtual memory, explaining the roles of logical blocks, physical blocks, and block tables. Section 4.3 and Figure 6: Walk through the example in 'Decoding with PagedAttention and vLLM'. Figure 6 is a crucial diagram that shows the state of the block table as tokens are generated and new blocks are allocated on demand.
In essence, PagedAttention is the custom CUDA kernel that "knows" how to read the block table. For each token, it fetches the scattered physical blocks corresponding to the context sequence and performs the attention calculation as if they were contiguous. This combination of a smart memory manager and a compatible attention kernel is what allows vLLM to achieve near-zero memory waste and, consequently, state-of-the-art throughput.
Conclusion
In this lesson, we diagnosed a critical bottleneck in naive LLM serving: memory fragmentation. By drawing a powerful analogy to virtual memory in operating systems, we uncovered the principles of paged memory management as the solution.
Key Takeaways:
- Naive KV Cache Allocation, which uses contiguous memory blocks pre-sized for the maximum sequence length, leads to severe internal and external fragmentation, wasting the majority of available VRAM.
- Paged Memory Management, inspired by OS paging, resolves this. It divides the KV cache into fixed-size blocks (like pages).
- A per-request block table (like a page table) maps logical blocks in a sequence to non-contiguous physical blocks in GPU VRAM.
- This approach eliminates external fragmentation and minimizes internal fragmentation, allowing for up to 96% memory utilization.
- PagedAttention is the attention algorithm designed to work with this non-contiguous memory layout by using the block table to gather the necessary keys and values for computation.
Preview of the Next Lesson:
We have now established the "why" and "what" of paged memory management. The next logical step is the "how". In the next lesson, you will get hands-on and implement a paged KV cache allocator that manages non-contiguous physical memory blocks for storing logical token sequences. You will build the core data structures, including the block manager and block tables, setting the foundation for integrating this advanced memory system into an inference loop.