How to Implement Paged Attention in PyTorch and vLLM

When you size a KV cache for a 32-layer GQA model (8 KV heads, head_dim 128, FP16) at max sequence length 4096, each sequence reserves 2 × 32 × 4096 × 8 × 128 × 2 bytes - 512 MB. Assume 256 concurrent sequences pre-allocating at max length, and you've committed 128 GB of theoretical reservation before loading a single weight. Most of those slots stay empty. Learning how to implement paged attention means replacing that contiguous slab with fixed-size blocks allocated on demand and reclaimed when sequences finish.

The data flow hides behind CUDA kernels in production systems like vLLM. This tutorial builds a minimal block-based KV-cache allocator in PyTorch, traces a decode step through the block table, then maps every component to vLLM's production implementation.

Quick answer

Paged attention splits each sequence's KV cache into fixed-size blocks managed by a block table mapping logical token positions to physical GPU memory slots. Implement it by building a block allocator with a free-block stack, maintaining per-sequence block tables for index translation, and gathering K/V slices during attention. Block size controls the fragmentation-versus-lookup-overhead trade-off.

How to implement paged attention: the minimal data flow

PagedAttention manages KV-cache memory layout. The dot-product math stays identical. Every decode step follows the same path:

1. Reserve a logical position - the sequence length increments. If the current block is full, the allocator pops one from the free pool.
2. Translate through the block table - logical_block = token_pos // block_size maps to a physical block ID.
3. Write K and V - the vectors land at offset token_pos % block_size inside the physical block.
4. Read all blocks for attention - the kernel walks the block table, gathers every physical block for the sequence, and runs scaled dot-product attention across the full KV cache.

The tutorial builds five components: a pool of physical KV blocks, a free-block list, per-sequence block tables, a sequence-length tracker, and a reference attention function. No custom CUDA code is required. Production systems fuse the gather-and-attend step into a single CUDA kernel, avoiding the contiguous copy entirely.

Size the paged KV cache and choose a block size

Contiguous KV-cache memory per sequence follows:

2 × layers × max_tokens × kv_heads × head_dim × dtype_bytes

Using the 32-layer GQA model from the introduction (8 KV heads, head_dim 128, FP16), that's 512 MB per sequence at max length 4096.

Paged allocation replaces max_tokens with actual tokens generated. With block_size=16, a full-length sequence needs ceil(4096/16) = 256 blocks. A sequence at 200 tokens uses only ceil(200/16) = 13 blocks.

Block-table metadata adds negligible overhead - one int32 per entry. Smaller blocks eliminate internal fragmentation but multiply table entries and scatter memory reads. Larger blocks improve locality but waste the unfilled tail. vLLM defaults to 16 tokens per block [1], keeping fragmentation under one block per sequence while limiting gather overhead during flash-attention-compatible decoding.

Build the block allocator and block tables in PyTorch

The allocator below uses deterministic seeds, explicit shapes, and a device flag you flip between "cpu" and "cuda". Key and value tensors are each shaped [num_physical_blocks, num_layers, block_size, kv_heads, head_dim]. Four invariants hold throughout: every physical block has exactly one owner, no allocated block appears in the free list, every block-table entry references a valid physical block ID, and freeing a sequence returns all its blocks.

Represent blocks, sequences, and free capacity

The block table maps each sequence ID and logical block index to a physical block ID. Sequence lengths track token counts separately so logical_block = pos // block_size and offset = pos % block_size give the write location. Free blocks sit in a Python list used as a stack - pop() to acquire, append() to release - giving O(1) allocation.

Allocate, append, and free KV-cache blocks

import torch

class PagedKVCache:
    def __init__(self, num_blocks, num_layers, block_size, kv_heads, head_dim, device="cpu"):
        self.block_size = block_size
        self.k = torch.zeros(num_blocks, num_layers, block_size, kv_heads, head_dim, device=device)
        self.v = torch.zeros_like(self.k)
        self.free = list(range(num_blocks - 1, -1, -1))
        self.tables = {}
        self.seq_lens = {}

    def add_sequence(self, seq_id):
        self.tables[seq_id] = []
        self.seq_lens[seq_id] = 0

    def append_token(self, seq_id, k_vec, v_vec):
        pos = self.seq_lens[seq_id]
        if pos % self.block_size == 0:
            if not self.free:
                raise RuntimeError("OOM: no free blocks")
            self.tables[seq_id].append(self.free.pop())
        phys = self.tables[seq_id][pos // self.block_size]
        offset = pos % self.block_size
        self.k[phys, :, offset] = k_vec  # shape: [num_layers, kv_heads, head_dim]
        self.v[phys, :, offset] = v_vec
        self.seq_lens[seq_id] = pos + 1

    def free_sequence(self, seq_id):
        self.free.extend(self.tables.pop(seq_id))
        del self.seq_lens[seq_id]

Production allocators extend this with copy-on-write for prefix-shared blocks, speculative preallocation, and LRU eviction.

Decode tokens through paged KV-cache lookups

A fused CUDA paged-attention kernel walks block-table indirections while computing attention - no contiguous copy materializes. The reference implementation below reconstructs full KV tensors explicitly as a correctness oracle.

Trace a block-table lookup during one decode step

With block_size=4, sequence 0 has 7 tokens (indices 0-6) and block table [5, 2]. Token index 6 maps to logical block 6 // 4 = 1, physical block 2, offset 6 % 4 = 2. The next append is token index 7: logical block 1, offset 3 - still room in physical block 2. Token index 8 triggers 8 % 4 == 0, so the allocator pops a new block (say physical block 9), extending the table to [5, 2, 9]. After the write, sequence length is 9.

Compute reference attention without hiding the layout

The query tensor q has shape [num_layers, 1, q_heads, head_dim] - one query position per decode step, with q_heads potentially larger than num_kv_heads for GQA models.

def reference_paged_attention(cache, seq_id, q, num_kv_heads):
    # q shape: [num_layers, 1, q_heads, head_dim]
    seq_len = cache.seq_lens[seq_id]
    blocks = cache.tables[seq_id]
    head_dim = cache.k.shape[-1]
    num_layers = cache.k.shape[1]

    # Gather K/V into contiguous tensors [num_layers, seq_len, kv_heads, head_dim]
    k_full = torch.zeros(num_layers, seq_len, num_kv_heads, head_dim, device=cache.k.device)
    v_full = torch.zeros_like(k_full)
    for pos in range(seq_len):
        phys = blocks[pos // cache.block_size]
        off = pos % cache.block_size
        k_full[:, pos] = cache.k[phys, :, off]
        v_full[:, pos] = cache.v[phys, :, off]

    # GQA head expansion
    q_heads = q.shape[2]
    repeats = q_heads // num_kv_heads
    k_exp = k_full.unsqueeze(3).expand(-1, -1, -1, repeats, -1).reshape(num_layers, seq_len, q_heads, head_dim)
    v_exp = v_full.unsqueeze(3).expand(-1, -1, -1, repeats, -1).reshape(num_layers, seq_len, q_heads, head_dim)

    scale = head_dim ** -0.5
    scores = torch.einsum("lqhd,lkhd->lhqk", q, k_exp) * scale
    attn = torch.softmax(scores, dim=-1)
    return torch.einsum("lhqk,lkhd->lqhd", attn, v_exp)

For correctness tests use torch.allclose with atol=1e-5 for FP32, atol=1e-3 for FP16, and atol=5e-3 for BF16.

Test and benchmark paged attention against a contiguous KV cache

Three test categories cover correctness:

• Write-read roundtrips: store known KV pairs within a block, at a block boundary, and across noncontiguous blocks; assert retrieved values match.
• Output equivalence: feed identical random Q/K/V (same torch.manual_seed) to reference_paged_attention and contiguous F.scaled_dot_product_attention; check torch.allclose at dtype-appropriate tolerance.
• Lifecycle invariants: free sequences and confirm blocks return to the free list; exhaust capacity and catch RuntimeError; interleave sequences and verify zero leaked blocks.

For benchmarks, separate prompt-fill from single-token decode, call torch.cuda.synchronize() around timed regions, and record device, dtype, sequence lengths, and block size.

Criterion	Contiguous KV cache	Paged KV cache
Allocation shape	One slab per sequence at max length	Fixed-size blocks on demand
Internal fragmentation	Up to max_len − actual_len per sequence	At most block_size − 1 per sequence
Growth behavior	Pre-allocated; cannot grow	Grows one block at a time
Batching flexibility	Bounded by worst-case reservation	Admits more concurrent sequences
Lookup cost	Direct indexing	Block-table translation per token
Implementation complexity	Minimal	Allocator + block tables + gather logic

Latency improves only when a fused CUDA paged-attention kernel (like vLLM's paged_attention_v1) eliminates the explicit gather [2]. The pure-PyTorch reference copies scattered blocks before computing attention, adding overhead that can make it slower per step despite better memory utilization.

Map the PyTorch design to vLLM PagedAttention

Every tutorial component has a vLLM counterpart spread across several subsystems:

• Physical tensor pool → Cache engines pre-allocate KV block tensors at startup, sized by gpu_memory_utilization and block_size.
• Free-block list → The block manager tracks free and reference-counted physical blocks with copy-on-write: shared prefix blocks get a refcount increment, copied only on the first divergent write.
• Block tables → Slot mappings pass as integer tensors to the model runner each step.
• Reference attention → Fused paged_attention_v1 / paged_attention_v2 CUDA kernels read block tables directly during attention, skipping the gather loop.
• Allocator lifecycle → The scheduler decides which requests to admit, preempt, or evict, then instructs the block manager before each forward pass. This enables continuous batching - requests enter and leave the running batch every iteration.

Early vLLM releases used a single BlockSpaceManager; later versions split block management into separate implementations with different eviction and prefix-caching strategies. Verify symbols against your installed vLLM release (v0.8+) before subclassing cache-management code.

PagedAttention vs sliding-window attention, RadixAttention, and SSMs

PagedAttention changes where KV vectors live in memory. Sliding-window attention changes which KV vectors exist - each query attends only to the most recent w tokens. The two compose: a sliding window caps retained history while paged allocation manages the physical blocks that hold it.

RadixAttention, used in SGLang, organizes cached blocks in a radix tree keyed by token content. High prefix overlap and rapid request churn favor the tree; workloads with unique long contexts work fine with flat block tables.

State space models (Mamba, RWKV) carry fixed-size hidden states instead of a growing KV cache, so paged block allocation has nothing to manage.

Workload	Recommended approach
Standard transformer serving	PagedAttention with block tables
Bounded-context models (Mistral sliding window)	Sliding window + paged allocation
Prefix-heavy (shared prompts, multi-turn)	RadixAttention or PagedAttention with prefix caching
Non-attention architectures (Mamba, RWKV)	Fixed recurrent state; no KV allocator needed

FAQ

How does paged attention work?

Paged attention splits each sequence's KV cache into fixed-size blocks stored at arbitrary physical GPU locations. A per-sequence block table maps logical token positions to physical block IDs. During decoding, the attention kernel walks the block table to gather K/V vectors and computes standard scaled dot-product attention. Blocks are allocated on demand and reclaimed when a sequence finishes.

How do you implement paged attention?

Build a pre-allocated tensor pool of physical KV blocks, a free-block stack for O(1) allocation, a per-sequence block table translating logical positions to physical IDs, and a decode function that gathers K/V slices through the table before computing attention. Allocate a new block when token_pos % block_size == 0; free all blocks on sequence completion.

What is paged attention in vLLM?

vLLM's PagedAttention implements block-based KV-cache management with a block manager that tracks free and reference-counted physical blocks, supporting copy-on-write for shared prefixes. Fused CUDA kernels read block tables directly during attention, eliminating the explicit gather step. The scheduler coordinates block allocation with continuous batching each iteration.

What is paged attention in LLM inference?

Paged attention is a memory management technique for transformer KV caches. Instead of reserving a contiguous buffer at maximum context length per sequence, it allocates small fixed-size blocks on demand as tokens generate. This eliminates most wasted pre-allocated memory, letting serving systems fit more concurrent sequences in the same GPU.

References

1. Group divisible designs with block size five: direct constructions - Anthony D. Forbes (2022)
2. Self-Attention through Kernel-Eigen Pair Sparse Variational Gaussian Processes - Yingyi Chen, Qinghua Tao, Francesco Tonin et al. (2024)