# How to Implement Paged Attention in PyTorch and vLLM

> Learn how to implement paged attention in PyTorch with a block KV cache, decoding tests, memory math, benchmarks, and a clear mapping to vLLM.

- Author: Swarnava Dutta (https://swarnava.dev)
- Published: 2026-09-03
- Tags: Vllm, Pytorch, Cuda
- Reading time: 8 min (1765 words)
- Canonical: https://swarnava.dev/blogs/implement-paged-attention-pytorch

---

![Illustration of how to implement paged attention: A long wall of uniform lockers spans the frame; on the left, mismatched](/images/blogs/implement-paged-attention-pytorch-hero.webp)

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](/blogs/kv-cache-optimization-llm-inference).

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.

![Flow diagram showing how to implement paged attention decode path from logical token position through block table to physical KV block](/images/blogs/implement-paged-attention-pytorch-diagram-1.jpg "The decode path from token position to attention read")

## 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]](#ref-1), keeping fragmentation under one block per sequence while limiting gather overhead during [flash-attention-compatible](/blogs/how-flash-attention-works) 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.

![Diagram of how to implement paged attention allocator with free-block pool assigning physical blocks into per-sequence block tables](/images/blogs/implement-paged-attention-pytorch-diagram-2.jpg "Free-block pool feeding per-sequence block tables")

### 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

```python
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](/blogs/prefix-caching-vs-kv-cache) 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](/blogs/grouped-query-attention-vs-mqa) models.

```python
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]](#ref-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](/blogs/sglang-vs-vllm-deployment), 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](https://arxiv.org/abs/2211.14124v3) - Anthony D. Forbes (2022)
2. [Self-Attention through Kernel-Eigen Pair Sparse Variational Gaussian Processes](https://arxiv.org/abs/2402.01476v2) - Yingyi Chen, Qinghua Tao, Francesco Tonin et al. (2024)