AI/ML systems deep dive

PagedAttention: how an LLM KV cache became a virtual-memory problem

During autoregressive generation, every new token needs the attention keys and values produced by all earlier tokens. A naïve server reserves one contiguous GPU region per request, wastes space on sequences that finish early, and cannot pack a busy batch well. PagedAttention fixes the shape of that problem by storing the KV cache in fixed-size GPU blocks and giving each sequence a page table.

AI/ML systems deep dive · vLLM’s block manager, GPU page tables, prefix sharing, allocator failure modes, isolation, and a reproducible simulator

2026-09-05

The concrete problem

A decoder-only transformer generates one token at a time. At step t, each attention head needs the key and value vectors for positions 0…t−1. Recomputing those vectors would make generation needlessly expensive, so the server caches them. This cache is the KV cache: for every layer, request, token position, and KV head, it stores a key vector and a value vector.

The cache grows with the output. The server usually does not know the final sequence length, so it faces a memory-allocation problem inside GPU memory:

Reserving each request’s maximum possible length avoids moving a live cache, but wastes memory when the request stops at an ordinary length. Reserving only the current length requires growing or copying the cache. Either way, a fragmented batch has less usable capacity than its total free bytes suggest.

The engineering target is not merely “make attention faster.” It is: keep many unpredictable sequences resident, append one token without relocating old tokens, and let the GPU read a logically contiguous history from physically scattered memory.

Virtual memory is the useful analogy

PagedAttention divides a sequence into logical blocks of a fixed number of tokens, commonly called the block_size. A request’s logical block 0, block 1, and so on are mapped to arbitrary physical KV-cache blocks. The mapping is the request’s block table.

logical sequence positions:  0  1  2  3 | 4  5  6  7 | 8  9
logical blocks:                  block 0     block 1     block 2
                                   │           │           │
block table:                    [ 17,         4,          29 ]
                                   │           │           │
GPU cache blocks:               physical 17  physical 4  physical 29

The analogy is precise enough to guide design, but it has limits. An operating system maps virtual pages through a hardware page-table walk and can fault pages independently. PagedAttention keeps the mapping in scheduler metadata and passes it to a custom attention kernel; there is no GPU page fault that fetches a missing KV block from disk. If a block is not resident, the scheduler must stop, evict, swap, or reject work before launching the kernel.

Why fixed blocks improve utilization

With a block size of B, a sequence of length L consumes ceil(L/B) blocks. Only its final block can have unused token slots, so internal fragmentation is bounded by roughly one block per sequence. A contiguous allocator can leave holes much larger than that after requests with different lifetimes finish.

contiguous allocator, capacity = 16 token slots
A uses [0..5]    B uses [6..9]    C uses [10..15]
B finishes → free [6..9], but a new 6-token request cannot fit

paged allocator, block_size = 4
A → [2, 8]        B → [1]        C → [7, 5]
B finishes → block 1 is immediately reusable
new request → [1, ...] without moving A or C

The remaining waste is explicit and measurable: the unused tail of the last block. A smaller block reduces tail waste but increases block-table entries and indirection overhead. A larger block reduces metadata and can improve locality, but strands more unused slots when a sequence ends.

The major win is packing: the scheduler can admit a request whenever enough free physical blocks exist, regardless of whether those blocks are adjacent. The cache’s logical history remains ordered through the table.

Data and control flow in a vLLM-style server

The names vary across versions and backends, but the path is stable:

There are two related mappings. The block table answers “which physical block contains logical block i?” The slot mapping answers “which physical token slot should receive the newly computed K/V for this token?” Confusing them is a classic implementation bug: the read path can look correct for old tokens while newly appended tokens are written to the wrong location.

Conceptually, the kernel computes the physical location for a logical token position p as:

logical_block = p // block_size
offset         = p % block_size
physical_block = block_table[logical_block]
    cache_index    = physical_block * block_size + offset

Real kernels also index layer, K/V choice, KV head, head dimension, sequence stride, data type, and sometimes a sliding-window boundary. They use specialized CUDA loads and reductions rather than performing this calculation in a slow per-element high-level loop.

Important implementation landmarks

vLLM is the reference open-source implementation behind the original PagedAttention paper. The exact directory layout evolves, so treat these as the contracts to trace rather than as a promise that every release has identical helper names.

A useful reading order is allocator first, runner metadata second, CUDA kernel third. Reading only the kernel hides the ownership rules; reading only the allocator hides the physical layout contract. The interface between them is the block table plus the cache tensor shape.

Prefix caching and copy-on-write

Once the cache is block-addressed, a shared prompt prefix can be represented by the same physical blocks for multiple requests. A prefix cache typically hashes a full logical block plus model- and request-relevant context, then maps a matching hash to a physical block. The hit avoids recomputing prompt attention and avoids allocating duplicate K/V bytes.

Sharing is safe only while the block is immutable. When two requests diverge after a shared prefix, their new tokens go into different blocks. If an operation needs to mutate an already shared block—beam search and some speculative-decoding paths are examples—the block manager must copy it first, decrement the old reference count, and give the writer an exclusive block.

Prefix caching changes the performance model: cache hits save prefill compute and memory bandwidth, while a poor hash key can produce incorrect results or leak one tenant’s prompt-derived state to another.

Tradeoffs and failure modes

• Indirection cost. Every attention read follows a block-table mapping. The cost is usually worth the utilization gain, but it competes with memory bandwidth and can matter for short contexts or kernels with poor locality.
• Block-size tuning. Small blocks reduce tail waste and make scheduling flexible, but enlarge metadata and may reduce coalesced access. Large blocks do the reverse. The best value depends on model head size, context distribution, and GPU architecture.
• Block exhaustion. Total free bytes are not the right metric; the scheduler needs enough free blocks, and usually enough blocks for all layers’ cache layout. A request can be paused even when a dashboard shows nonzero GPU memory.
• Stale mappings. Freeing a block before the last kernel using it completes can make a later request read another request’s K/V. Lifetime must cover the asynchronous CUDA work, not just the CPU scheduling event.
• Prefix-cache poisoning. If the cache key omits tokenizer settings, model revision, LoRA identity, or tenant policy, a byte-identical-looking prefix can refer to semantically different computation.
• Sharing races. A shared block must be read-only until copy-on-write. Mutating it in place silently changes another sequence’s attention history, a correctness failure that may appear only under beam search or concurrent scheduling.
• Kernel-contract mismatch. A block table built for one block size, number of KV heads, or layout cannot safely be passed to a kernel compiled for another. The result may be wrong answers rather than a clean error.
• Preemption and swap. Moving blocks to CPU or another GPU turns a simple allocator into a transfer protocol. Stalls, stale device pointers, and incomplete restoration become part of token latency.

Security and isolation implications

The KV cache is derived from user prompts and model state. Treating it as anonymous scratch memory is unsafe in a multi-tenant service.

Debugging a bad answer or a sudden throughput collapse

1. Validate the logical-to-physical mapping on CPU. For a tiny sequence, print every position’s block, offset, and physical slot. Compare it with the GPU metadata immediately before the attention launch.
2. Separate allocator pressure from kernel performance. Record free-block count, allocated blocks per request, tail waste, prefix hit rate, preemption count, and kernel time. A lower tokens-per-second rate can be caused by waiting for blocks, not slower attention math.
3. Force block reuse. Run short requests that finish in a pattern designed to reuse the same physical blocks. Use deterministic token inputs and verify that a new request’s logits do not change when old requests are added or removed.
4. Disable prefix caching. If the wrong answer disappears, compare cache keys, reference counts, copy-on-write transitions, and model/adaptor identity before blaming the attention kernel.
5. Check asynchronous lifetimes. Add CUDA synchronization only as a diagnostic. If synchronization “fixes” corruption, inspect whether the CPU freed or reassigned a block before the prior stream finished reading it.
6. Compare against a contiguous reference. For a very small batch, materialize each sequence’s K/V history contiguously and compare logits with the paged path within the expected floating-point tolerance.

Compact lab: simulate blocks and prove the mapping

Goal: compare a contiguous allocator with a fixed-block allocator, then verify that a scattered block table reconstructs the same logical KV history.

Prerequisites: Python 3.8 or newer. The lab uses only the standard library and runs on a CPU; it models addresses and data, not CUDA performance.

python3 - <<'PY'
from math import ceil

BLOCK = 4
capacity = 12

def paged_slots(block_table, length):
    return [block_table[p // BLOCK] * BLOCK + (p % BLOCK)
            for p in range(length)]

# A physical cache has 3 blocks. Requests receive non-adjacent blocks.
cache = [None] * capacity
tables = {"A": [2, 0], "B": [1]}

for token, slot in enumerate(paged_slots(tables["A"], 6)):
    cache[slot] = f"A{token}"
for token, slot in enumerate(paged_slots(tables["B"], 3)):
    cache[slot] = f"B{token}"

reconstructed_a = [cache[s] for s in paged_slots(tables["A"], 6)]
reconstructed_b = [cache[s] for s in paged_slots(tables["B"], 3)]
print("A physical slots:", paged_slots(tables["A"], 6))
print("B physical slots:", paged_slots(tables["B"], 3))
print("A logical history:", reconstructed_a)
print("B logical history:", reconstructed_b)

# Compare block consumption with contiguous max-length reservations.
lengths = [6, 3, 5]
print("paged blocks:", sum(ceil(n / BLOCK) for n in lengths))
print("contiguous slots at max length 8:", len(lengths) * 8)
PY

Expected result:

A physical slots: [8, 9, 10, 11, 0, 1]
B physical slots: [4, 5, 6]
A logical history: ['A0', 'A1', 'A2', 'A3', 'A4', 'A5']
B logical history: ['B0', 'B1', 'B2']
paged blocks: 5
contiguous slots at max length 8: 24

The logical histories are contiguous even though A occupies physical slots 8–11 and then 0–1. The allocator needs five blocks for the three actual lengths, while reserving eight contiguous slots for each request would require twenty-four slots.

Revealing variations

Change tables["A"] to [2, 1]. B now shares physical block 1, so the printed histories expose overwritten data: this is why the block manager must never hand a live block to two writers. Then change it to [2, 0, 1] and increase A’s length to 9; the third logical block shows that a block table can grow without moving the first two blocks.

For a GPU-backed experiment, install a matching vLLM release and use its documented metrics and profiling tools, but do not infer production behavior from a CPU toy model: the real costs include CUDA launch shape, memory bandwidth, kernel specialization, scheduler delay, and GPU-to-CPU transfer time.