Paged KV Cache Allocator

Introduction

In our last lesson, we established the "why" behind paged memory management. We dissected the critical problem of memory fragmentation in naive KV cache allocation and drew a powerful analogy to virtual memory in operating systems. You learned that by partitioning the KV cache into fixed-size blocks and using a per-request block table, we can eliminate fragmentation and dramatically improve memory utilization.

Today, we transition from "why" to "how." Your learning outcome is to implement a paged KV cache allocator that manages non-contiguous physical memory blocks for storing logical token sequences. This is a heavily practical session where you will build the core data structures and logic that act as the "memory operating system" for a high-throughput inference engine. We'll construct the Python classes for blocks, block tables, and the central allocator that brings it all together.

Conceptual Refresher: The Paged KV Cache System

Before we dive into code, let's revisit the high-level architecture of the system we're building. The following video provides an excellent visual walkthrough of the key components: the memory blocks where the KV cache is stored, the page table that maps sequences to blocks, and the free block queue that tracks available memory.

As the video reinforces, our implementation needs three main components:

1. A representation of a physical memory block.
2. A per-sequence block table to map logical tokens to physical blocks.
3. A central block manager (our allocator) to handle the free list and orchestrate allocations.

The Core Data Structures

Let's start by defining the fundamental data structures. The provided blog post on nano-vllm offers clean, minimal dataclass definitions that are perfect for this task.

Based on that reading, we can define our core data structures in Python:

from dataclasses import dataclass, field
from typing import List, Optional, Deque
from collections import deque
import torch




# A physical block of memory in the KV cache.
@dataclass
class Block:
    block_id: int
    ref_count: int = 0 # Useful for sharing, e.g., prefix caching




# A mapping from a logical sequence to physical blocks.
# This is our "page table".
@dataclass
class BlockTable:
    block_ids: List[int] = field(default_factory=list)

    def append_block(self, block_id: int):
        self.block_ids.append(block_id)

The beauty of this design lies in its simplicity. Block is just a handle to a physical block, and BlockTable is a simple list of integers representing the physical block IDs allocated to a sequence. All the complexity is handled by our allocator.

Implementing the Paged KV Cache Allocator

Now we'll build the central allocator class. This class will own the physical cache memory and manage the pool of blocks. We will structure it like the UniversalPagedAttentionCache from the Medium article you've been provided with, as it cleanly separates concerns.

Here is the skeleton of our PagedKVCacheAllocator:

class PagedKVCacheAllocator:
    def __init__(self, num_blocks: int, block_size: int, num_heads: int, head_dim: int, dtype: torch.dtype, device: str = 'cpu'):
        """
        Initializes the allocator.

        Args:
            num_blocks: Total number of physical blocks.
            block_size: The number of tokens each block can hold.
            num_heads: The number of attention heads.
            head_dim: The dimension of each attention head.
            dtype: The data type for the cache tensors.
            device: The device to store the cache on.
        """
        self.num_blocks = num_blocks
        self.block_size = block_size
        self.device = device




        # 1. Allocate the physical cache memory as two large, contiguous tensors
        cache_shape = (num_blocks, num_heads, block_size, head_dim)
        self.k_cache = torch.zeros(cache_shape, dtype=dtype, device=device)
        self.v_cache = torch.zeros(cache_shape, dtype=dtype, device=device)


```grasp
{
  "type": "exercise",
  "id": "77b5c379-b814-4e8d-9fb3-aca16b985877"
}

    # 2. Initialize the block management data structures
    self.blocks = [Block(i) for i in range(num_blocks)]
    self.free_blocks: Deque[int] = deque(range(num_blocks))
    



    # 3. Initialize sequence tracking
    # Maps a sequence ID to its BlockTable
    self.sequence_block_tables: dict[int, BlockTable] = {}

def __repr__(self):
    return f"PagedKVCacheAllocator(free_blocks={len(self.free_blocks)}/{self.num_blocks})"

This `__init__` method sets up the entire system:
1.  It pre-allocates all the GPU memory we will ever need for the KV cache in `self.k_cache` and `self.v_cache`.
2.  It creates a `free_blocks` queue, which acts as our free list, initially containing all block IDs. Using a `deque` gives us efficient O(1) appends and pops.
3.  It prepares a dictionary to hold the `BlockTable` for each active sequence.




#### Allocation and Deallocation

Next, we need methods to allocate blocks to a new sequence and free them when a sequence is complete. The `BlockManager` from the `nano-vllm` resource provides a simple and effective model for this.

```python



# Add these methods to the PagedKVCacheAllocator class

def allocate_for_sequence(self, seq_id: int, num_prompt_tokens: int) -> None:
    """Allocates blocks for a new sequence's prompt."""
    if seq_id in self.sequence_block_tables:
        raise ValueError(f"Sequence {seq_id} is already allocated.")

    num_blocks_needed = (num_prompt_tokens + self.block_size - 1) // self.block_size
    
    if len(self.free_blocks) < num_blocks_needed:
        raise RuntimeError(f"Not enough free blocks. Need {num_blocks_needed}, have {len(self.free_blocks)}.")

    block_table = BlockTable()
    for _ in range(num_blocks_needed):
        block_id = self.free_blocks.popleft()
        self.blocks[block_id].ref_count = 1
        block_table.append_block(block_id)

    self.sequence_block_tables[seq_id] = block_table

def free_sequence(self, seq_id: int) -> None:
    """Frees all blocks associated with a sequence."""
    if seq_id not in self.sequence_block_tables:



        # Can happen if a request is cancelled before being processed
        return

    block_table = self.sequence_block_tables.pop(seq_id)
    for block_id in block_table.block_ids:
        self.blocks[block_id].ref_count -= 1



        # In a more advanced implementation, we would only free if ref_count is 0
        if self.blocks[block_id].ref_count == 0:
            self.free_blocks.append(block_id)

• allocate_for_sequence calculates the number of blocks needed for a prompt, takes them from the free_blocks queue, and creates a BlockTable for the sequence.
• free_sequence does the reverse: it retrieves the block IDs from the sequence's table and returns them to the free_blocks queue.

Mapping Logical Positions to Physical Locations

This is the most crucial part of the allocator. The attention kernel needs to know where in the physical k_cache and v_cache tensors to find the data for a given logical token position. The following method provides this translation service.

The implementation is inspired by the get_block_positions method in the Medium article (LINK).

# Add this method to the PagedKVCacheAllocator class

def get_block_mapping(self, seq_id: int) -> List[int]:
    """
    Returns the list of physical block IDs for a sequence.
    The attention kernel will use this to gather K and V values.
    """
    if seq_id not in self.sequence_block_tables:
        raise ValueError(f"Sequence {seq_id} not found.")
    return self.sequence_block_tables[seq_id].block_ids

In a real system, you would pass this block mapping, along with other metadata like token positions, to a custom CUDA kernel (like PagedAttention). The kernel would then perform the address translation internally: for a token at logical position, it would:

1. Calculate logical_block_idx = position // self.block_size.
2. Calculate offset_in_block = position % self.block_size.
3. Look up physical_block_id = block_mapping[logical_block_idx].
4. Access the data at k_cache[physical_block_id, :, offset_in_block, :].

This indirection is the core of paged memory management. It allows the sequence to be stored in non-contiguous physical blocks.

Dynamic Block Allocation during Decoding

During autoregressive generation, a sequence grows one token at a time. What happens when the last allocated block becomes full? We need a method to append a new block to an existing sequence.

```grasp
{
  "type": "exercise",
  "id": "dfc97315-6bb1-4faa-8df5-a9daea5b7db1"
}

Add this method to the PagedKVCacheAllocator class

def append_block_to_sequence(self, seq_id: int) -> None:
"""Allocates one new block and appends it to a sequence's block table."""
if seq_id not in self.sequence_block_tables:
raise ValueError(f"Sequence {seq_id} not found.")

if not self.free_blocks:
    raise RuntimeError("Out of free blocks.")

block_id = self.free_blocks.popleft()
self.blocks[block_id].ref_count = 1
self.sequence_block_tables[seq_id].append_block(block_id)

In a full inference loop, before generating a new token, the scheduler would check if `(current_length + 1) % block_size == 0`. If so, it would call this method to ensure there is space for the next token. This on-demand allocation is what makes the system so efficient.

Let's watch a final animation that shows this exact process of on-demand allocation.

```grasp
{
  "type": "video",
  "title": "Fast LLM Serving with vLLM and PagedAttention",
  "id": "[LINK](https://www.youtube.com/watch?v=5ZlavKF_98U)",
  "video_id": "5ZlavKF_98U",
  "relevant_section_indices": [
    2,
    3
  ],
  "par_intro": "This clip from the Anyscale talk on vLLM provides a clear animation of the block allocation process during decoding.",
  "par_directions": "Watch from 07:03 to 09:54. The animation shows how tokens are stored in logical blocks, which are mapped to physical blocks. Crucially, it demonstrates what happens when a block becomes full: a new physical block is allocated from the free pool 'on demand' and added to the block table. This directly visualizes the `append_block_to_sequence` logic we just wrote."
}

Conclusion

In this lesson, you have moved from theory to practice by implementing the core of a paged KV cache allocator. You have built the essential data structures and the logic to manage non-contiguous memory blocks, mirroring the advanced memory management found in state-of-the-art inference engines like vLLM.

Key Takeaways:

• The paged KV cache system is built on three pillars: Block data structures, per-sequence BlockTables (page tables), and a central allocator that manages a free list of blocks.
• The allocator's main responsibilities are to allocate blocks for new sequences, free them upon completion, and append new blocks on-demand during generation.
• The mapping from a logical token sequence to scattered physical blocks is handled by the BlockTable, which simply stores an ordered list of physical block IDs.
• The physical cache (k_cache, v_cache) is a large, pre-allocated tensor, and the allocator's job is to manage access to slices (blocks) of this tensor.

Preview of the Next Lesson:

Now that you have built a functional paged KV cache allocator, the next step is to put it to work. In the next lesson, you will integrate the paged KV cache into an inference loop and benchmark the reduction in memory waste under a simulated concurrent workload. This will allow you to empirically verify the benefits of the system you've just constructed and see firsthand how it overcomes the fragmentation that plagues naive allocation strategies.