Paged KV Cache · AI infra interview lesson

Why It Matters

At four tokens per block, five tokens need two blocks. Logical blocks 0 and 1 might occupy physical blocks 7 and 2. Locate the fifth token’s slot before studying allocation and sharing.

A concrete number: LLaMA-2-7B (32 layers, 32 KV heads, head_dim=128, FP16) stores 2 * 32 * 32 * 128 * 2 = 524,288 bytes = 512 KiB of KV per token. Reserving 4096 tokens costs 2 GiB per request; if requests average 600 tokens, about 85% of it is wasted. With paging, each sequence wastes at most block_size - 1 tokens: with vLLM's default of 16 that is at most 7.5 MiB. nano-vllm's default model, Qwen3-0.6B (28 layers, 8 KV heads, head_dim=128, BF16), needs 2 * 28 * 8 * 128 * 2 = 114,688 bytes = 112 KiB per token, so one 256-token block is 28 MiB.

Walk the Shapes First

block_size=4, H_kv=2, d=4, 8 physical blocks in total. Two sequences: A has 6 tokens with block_table=[7, 2]; B has 3 tokens with block_table=[5]. One prefill, then two decode steps:

Rules read off the table: the slot of position p is always block_table[p // bs] * bs + p % bs; slot_mapping only covers the tokens computed in this step, while attention reads keys over the whole context; a new block is needed exactly at the step where len % bs == 1.

Running the examples: execute this page in order in one Python environment. Later snippets may reuse earlier definitions. CPU examples validate calculations and state changes, not GPU timing. Download the notebook for the complete exercise workspace.

Core Implementations

Step 1: slot_mapping and store_kvcache

What is new: "where does token p's KV go" becomes one integer slot, and freshly computed K/V are written into the paged cache by slot; a slot of -1 is skipped.

Formula: slot = block_table[p // bs] * bs + p % bs. nano-vllm's prepare_prefill does not divide per token; it emits range(slot_start, slot_end) per block: the first block may start at start % bs (prefix hit, second half of a chunked prefill), and the last block ends at end - i * bs. For the write, k_cache is viewed as [num_blocks * bs, H_kv * d], matching the slot * D addressing in the Triton kernel.

import torch
torch.manual_seed(0)

def compute_slot_mapping(block_table, start, end, block_size):
    slots = []
    if start >= end:
        return slots
    start_block, end_block = start // block_size, (end + block_size - 1) // block_size
    for i in range(start_block, end_block):
        base = block_table[i] * block_size
        lo = start % block_size if i == start_block else 0                 # first block may start mid-block
        hi = end - i * block_size if i == end_block - 1 else block_size    # last block may be partial
        slots.extend(range(base + lo, base + hi))
    return slots

def store_kvcache(key, value, k_cache, v_cache, slot_mapping):
    N, H, d = key.shape
    kc, vc = k_cache.view(-1, H * d), v_cache.view(-1, H * d)
    mask = slot_mapping != -1                                              # -1 = padding, do not write
    idx = slot_mapping[mask].long()
    kc[idx] = key.reshape(N, -1)[mask]
    vc[idx] = value.reshape(N, -1)[mask]

assert compute_slot_mapping([7, 2], 0, 6, 4) == [28, 29, 30, 31, 8, 9]   # sequence A from the table
for bs in [4, 16, 256]:
    bt = [7, 2, 5]
    for start, end in [(0, 2 * bs + 2), (bs // 2 + 1, 2 * bs + 1), (bs, 2 * bs), (3, 3)]:
        assert compute_slot_mapping(bt, start, end, bs) == [bt[p // bs] * bs + p % bs for p in range(start, end)]
    H, d, n = 2, 4, 2 * bs + 2
    k_cache, v_cache = torch.zeros(8, bs, H, d), torch.zeros(8, bs, H, d)
    key, value = torch.randn(n, H, d), torch.randn(n, H, d)
    slots = torch.tensor(compute_slot_mapping(bt, 0, n, bs), dtype=torch.int32)
    slots[1] = -1
    store_kvcache(key, value, k_cache, v_cache, slots)
    for p in range(n):
        want = torch.zeros(H, d) if p == 1 else key[p]
        assert torch.equal(k_cache[bt[p // bs], p % bs], want)
    assert (k_cache.abs().sum((-1, -2)) != 0).sum() == n - 1

Step 2: paged_attention_decode

What is new: during decode each sequence has exactly one query, and its K/V must be gathered through block_table from physical blocks in arbitrary order; only the front of the last block is valid; with GQA, query head h uses KV head h // (Hq // H_kv).

Formula: n_blk = ceil(ctx_len / bs), K = k_cache[block_table[:n_blk]].reshape(-1, H_kv, d)[:ctx_len], then ordinary softmax(q K^T * scale) V. The -1 padding only appears after n_blk, so slice before indexing. The reference is F.scaled_dot_product_attention on the same K/V stored contiguously.

import torch
import torch.nn.functional as F
torch.manual_seed(0)

def paged_attention_decode(q, k_cache, v_cache, block_tables, context_lens, scale):
    B, Hq, d = q.shape
    bs, Hkv = k_cache.shape[1], k_cache.shape[2]
    out = torch.empty_like(q)
    for b in range(B):
        n = int(context_lens[b])
        blocks = block_tables[b, :(n + bs - 1) // bs].long()      # drop the -1 padding first
        K = k_cache[blocks].reshape(-1, Hkv, d)[:n]               # cut the unused tail of the last block
        V = v_cache[blocks].reshape(-1, Hkv, d)[:n]
        K, V = K.repeat_interleave(Hq // Hkv, 1), V.repeat_interleave(Hq // Hkv, 1)
        p = (torch.einsum("hd,nhd->hn", q[b], K) * scale).softmax(-1)
        out[b] = torch.einsum("hn,nhd->hd", p, V)
    return out

for bs in [4, 16, 256]:
    Hq, Hkv, d, nb = 4, 2, 8, 12
    lens = [1, bs, 2 * bs + 3]
    k_cache, v_cache = torch.randn(nb, bs, Hkv, d), torch.randn(nb, bs, Hkv, d)   # pre-filled with garbage
    free = torch.randperm(nb).tolist()
    tables, Ks, Vs = [], [], []
    for L in lens:
        bt = [free.pop() for _ in range((L + bs - 1) // bs)]
        K, V = torch.randn(L, Hkv, d), torch.randn(L, Hkv, d)
        for p in range(L):
            k_cache[bt[p // bs], p % bs], v_cache[bt[p // bs], p % bs] = K[p], V[p]
        tables.append(bt); Ks.append(K); Vs.append(V)
    M = max(map(len, tables))
    block_tables = torch.tensor([t + [-1] * (M - len(t)) for t in tables], dtype=torch.int32)
    q = torch.randn(len(lens), Hq, d)
    out = paged_attention_decode(q, k_cache, v_cache, block_tables, torch.tensor(lens), d ** -0.5)
    for b in range(len(lens)):
        K = Ks[b].repeat_interleave(Hq // Hkv, 1).transpose(0, 1)       # [Hq, L, d]
        V = Vs[b].repeat_interleave(Hq // Hkv, 1).transpose(0, 1)
        ref = F.scaled_dot_product_attention(q[b][:, None], K, V, scale=d ** -0.5)[:, 0]
        assert torch.allclose(out[b], ref, atol=1e-5)

Step 3: varlen attention and cu_seqlens

What is new: in prefill the sequences in a batch have different lengths; (Pro lesson)

Step 4: prepare_prefill / prepare_decode

What is new: turning the scheduler's list of Sequence s into the metadata attention needs. (Pro lesson)

Append slots during decode

Prefill processes a span; decode usually processes one new token per request. (Pro lesson)

Step 5: BlockManager: ref counts and chained-hash prefix caching

What is new: who decides which physical blocks go into block_table . (Pro lesson)

Separate allocation, sharing and release

Track a physical block from free to allocated to shared. Return it to the free pool only after its final reference is released. (Pro lesson)

In Real Systems

nano-vllm nanovllm/engine/block_manager.py : Block ( ref_count , hash , token_ids ) and BlockManager ( compute_hash uses xxhash.xxh64 , upda… (Pro lesson)

Common Pitfalls

Writing the decode slot as block_table[-1] * bs + last_block_num_tokens → the new token lands one slot too far and the next step's attention… (Pro lesson)

Say It in the Interview

PagedAttention treats the KV cache like virtual memory: it is split into fixed-size physical blocks, and each sequence keeps a block table m… (Pro lesson)

Hands-on Notebook

The notebook uses nano-vllm's real interfaces, parametrizes block_size (4 / 16 / 256), and checks every exercise against SDPA or a per-token… (Pro lesson)