Visualize PagedAttention block layout

When I first read the vLLM paper, the comparison to operating system virtual memory clicked immediately. Traditional KV cache management pre-allocates a contiguous chunk of GPU memory for each request at its maximum possible sequence length. This wastes memory the same way pre-allocated fixed-size arrays waste heap space. PagedAttention fixes this the same way virtual memory fixed physical memory: indirection through a page table.

Today I want to make this concrete by building a block allocator and visualizing exactly what happens in GPU memory as requests come and go.

The block abstraction

In PagedAttention, GPU memory for the KV cache is divided into fixed-size blocks. Each block holds the key and value tensors for a fixed number of tokens (the block size, typically 16). A request's KV cache is a linked list of these blocks, managed through a block table that maps logical block indices to physical block locations.

# Block structure for a single layer
# block_size = 16 tokens
# num_heads = 32
# head_dim = 128
# dtype = float16

# Each block stores:
#   key:   [block_size, num_heads, head_dim] = [16, 32, 128] float16
#   value: [block_size, num_heads, head_dim] = [16, 32, 128] float16
# Size per block per layer: 2 * 16 * 32 * 128 * 2 bytes = 262,144 bytes = 256 KB

# For 32 layers: 256 KB * 32 = 8 MB per block
# With 80 GB HBM: roughly 10,000 blocks available

The block allocator

The block allocator is essentially a free list, the same data structure you would find in a simple memory allocator. It maintains a pool of physical block IDs and hands them out on demand.

class BlockAllocator:
    def __init__(self, num_blocks: int, block_size: int = 16):
        self.block_size = block_size
        self.free_blocks = list(range(num_blocks))
        self.ref_counts = [0] * num_blocks  # for copy-on-write sharing

    def allocate(self) -> int:
        if not self.free_blocks:
            raise RuntimeError("Out of KV cache blocks")
        block_id = self.free_blocks.pop()
        self.ref_counts[block_id] = 1
        return block_id

    def free(self, block_id: int):
        self.ref_counts[block_id] -= 1
        if self.ref_counts[block_id] == 0:
            self.free_blocks.append(block_id)

    def share(self, block_id: int):
        """Increment ref count for copy-on-write sharing."""
        self.ref_counts[block_id] += 1

    @property
    def num_free(self) -> int:
        return len(self.free_blocks)

Block tables: the page table equivalent

Each active request has a block table, a list of physical block IDs that map to its logical token positions. When the attention kernel needs to read the KV cache for token position 42, it looks up block_table[42 // block_size] to find the physical block, then indexes into position 42 % block_size within that block.

class RequestState:
    def __init__(self, request_id: int, prompt_len: int, allocator: BlockAllocator):
        self.request_id = request_id
        self.num_tokens = prompt_len
        self.block_table = []
        self.allocator = allocator

        # Allocate blocks for the prompt
        num_blocks_needed = (prompt_len + allocator.block_size - 1) // allocator.block_size
        for _ in range(num_blocks_needed):
            self.block_table.append(allocator.allocate())

    def append_token(self):
        """Called after each decode step."""
        self.num_tokens += 1
        # Check if we need a new block
        if self.num_tokens % self.allocator.block_size == 1 and self.num_tokens > 1:
            self.block_table.append(self.allocator.allocate())

    def release(self):
        """Free all blocks when request completes."""
        for block_id in self.block_table:
            self.allocator.free(block_id)
        self.block_table = []

Visualizing the layout

Let me trace through a scenario with 20 physical blocks and three requests arriving at different times:

allocator = BlockAllocator(num_blocks=20, block_size=16)

# Request A arrives: 40 tokens prompt (needs 3 blocks: ceil(40/16))
req_a = RequestState(request_id="A", prompt_len=40, allocator=allocator)
# Block table: [0, 1, 2]       Free: 17 blocks

# Request B arrives: 24 tokens prompt (needs 2 blocks)
req_b = RequestState(request_id="B", prompt_len=24, allocator=allocator)
# Block table: [3, 4]          Free: 15 blocks

# Request C arrives: 50 tokens prompt (needs 4 blocks)
req_c = RequestState(request_id="C", prompt_len=50, allocator=allocator)
# Block table: [5, 6, 7, 8]    Free: 11 blocks

# GPU memory layout after initial allocation:
# [A][A][A][B][B][C][C][C][C][ ][ ][ ][ ][ ][ ][ ][ ][ ][ ][ ]
#  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19

# A generates 10 more tokens (now 50 total, needs 4th block)
for _ in range(10):
    req_a.append_token()
# A's block table: [0, 1, 2, 9]   (block 9 allocated for overflow)

# B finishes. Release its blocks.
req_b.release()
# Blocks 3 and 4 return to free list

# GPU memory layout after B finishes:
# [A][A][A][ ][ ][C][C][C][C][A][ ][ ][ ][ ][ ][ ][ ][ ][ ][ ]
#  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19
#
# Note: A's blocks are NOT contiguous! Blocks 0,1,2 and 9.
# This is fine because the block table handles the indirection.

This is the key insight. The KV cache for a single request does not need to be contiguous in physical memory. The block table provides the indirection layer that lets blocks be scattered anywhere in the GPU memory pool. This eliminates fragmentation waste entirely.

Copy-on-write for beam search and prefix sharing

The ref_count mechanism enables copy-on-write (CoW) sharing. When two requests share a common prefix (like the same system prompt), their block tables can point to the same physical blocks. The blocks are only copied when one request needs to modify them.

# Two requests share the same 32-token system prompt
# Instead of allocating 4 blocks total (2 per request),
# allocate 2 blocks and share them:

sys_prompt_blocks = [allocator.allocate(), allocator.allocate()]

# Request D points to shared blocks + its own
req_d_table = sys_prompt_blocks + [allocator.allocate()]
for b in sys_prompt_blocks:
    allocator.share(b)  # ref_count goes to 2

# Request E points to the same shared blocks + its own
req_e_table = sys_prompt_blocks + [allocator.allocate()]

# Memory saved: 2 blocks (the shared prefix)
# This is prefix caching at the memory level

How the attention kernel uses the block table

The attention kernel receives the block table as an input alongside the query vector. For each attention head, it iterates over the blocks in the table, loading 16 key-value pairs at a time. The FlashAttention-style tiling maps naturally onto blocks: each block is one tile of the KV cache.

In vLLM's implementation, the PagedAttention CUDA kernel takes a block_tables tensor of shape [batch_size, max_num_blocks] and a context_lens tensor that tells it how many tokens are valid. The kernel handles the gather from non-contiguous physical blocks transparently.

PagedAttention is the most impactful single idea in LLM serving since continuous batching. It solves the memory fragmentation problem the same way virtual memory solved it for CPUs, and like virtual memory, once you have it, you never want to go back.

Tomorrow I will look at the tradeoff between TTFT and throughput as batch size changes, the operational question that all of this memory management is ultimately in service of.