KV Cache Optimization and PagedAttention

When an LLM generates text, it produces one token at a time, and each new token attends to every previous token. Naively, that would mean recomputing the attention 'keys' and 'values' (the K and V vectors) for the entire sequence on every single step — quadratic, wasteful work. The KV cache is the fix: after computing K and V for a token once, you store them and reuse them for all future tokens. Think of it as mise en place in a kitchen — you chop the onions once at the start and keep them in a bowl, rather than re-chopping for every dish.

The catch is that this cache is large and grows with every token. Its size is roughly: 2 (K and V) x num_layers x num_kv_heads x head_dim x sequence_length x batch_size x bytes_per_value. For a big model with long sequences and a large batch, the KV cache can consume more GPU memory than the model weights themselves. So how you MANAGE that memory becomes the throughput bottleneck.

The classic problem is fragmentation. Traditional serving reserved one big contiguous block of memory per request, sized for the maximum possible sequence length. But most requests are shorter than the max, so huge chunks sit reserved-but-unused — and because the free space is scattered in awkward-sized gaps, you can't fit new requests into it. Studies found naive serving wasted 60-80% of KV cache memory this way.

PagedAttention (introduced by vLLM) borrows the idea of virtual memory and paging from operating systems. Instead of one contiguous reservation per request, the KV cache is split into small fixed-size 'blocks' (pages). A request's tokens are stored in whatever blocks are free, and a per-request 'block table' maps logical positions to physical blocks — exactly like an OS page table. Because blocks are small and allocated on demand, fragmentation nearly disappears (waste drops to under ~4%), so you can fit far more concurrent requests in the same GPU memory. That directly raises throughput.

A bonus: blocks can be SHARED. If many requests start with the same system prompt, they can point to the same physical blocks for that shared prefix (prefix caching / copy-on-write) instead of each storing its own copy — saving even more memory. This is why PagedAttention is the headline feature that made vLLM's throughput dramatically higher than naive Hugging Face serving.

# A tiny model of WHY the KV cache dominates memory, and how
# paged (block-based) allocation reduces waste vs contiguous.

def kv_cache_bytes(num_layers, num_kv_heads, head_dim,
                   seq_len, batch, bytes_per_val=2):
    """Approx KV cache size in bytes (2 = K and V)."""
    return 2 * num_layers * num_kv_heads * head_dim * seq_len * batch * bytes_per_val

gb = kv_cache_bytes(32, 8, 128, seq_len=4096, batch=16) / 1e9
print(f"KV cache for 16 concurrent 4k-token requests: ~{gb:.1f} GB")

# Contiguous reservation wastes memory when requests are short;
# paged allocation only uses the blocks it needs.
BLOCK = 16  # tokens per page (vLLM default is 16)

def contiguous_waste(actual_len, max_len):
    return max_len - actual_len  # reserved-but-unused slots

def paged_waste(actual_len, block=BLOCK):
    import math
    reserved = math.ceil(actual_len / block) * block
    return reserved - actual_len  # at most block-1 slots wasted

print("contiguous waste (200-token req, 4096 max):", contiguous_waste(200, 4096))
print("paged waste       (200-token req):", paged_waste(200))
# ~3896 wasted slots vs at most 15 — the fragmentation PagedAttention removes