If you have been following this journey, you know that the KV cache is what makes decode memory-bound. Every request accumulates key and value tensors at every layer, and they stay in GPU memory for the entire lifetime of the generation. With a 7B model serving hundreds of concurrent requests, the KV cache can easily consume more memory than the model weights themselves.
The naive approach is to pre-allocate a contiguous tensor for each request's maximum possible sequence length. This wastes enormous amounts of memory because most requests never reach their maximum length. vLLM's PagedAttention solved this with a paged memory approach, and today I am building the block allocator that makes it possible.
The block abstraction
Instead of allocating one big contiguous buffer per request, we divide KV cache memory into fixed-size blocks, each holding a fixed number of tokens' worth of keys and values. A typical block size is 16 tokens. For a model with 32 layers, 32 KV heads, and head dimension 128 in FP16, one block per layer stores:
# Per-layer block size in bytes:
# 16 tokens * 128 head_dim * 2 bytes (FP16) * 2 (K and V) = 8,192 bytes = 8 KB
# Across 32 layers: 8 KB * 32 = 256 KB per logical block
BLOCK_SIZE = 16 # tokens per block
NUM_LAYERS = 32
NUM_KV_HEADS = 32
HEAD_DIM = 128
DTYPE_SIZE = 2 # FP16
bytes_per_block_per_layer = BLOCK_SIZE * HEAD_DIM * DTYPE_SIZE * 2 # K and V
total_bytes_per_block = bytes_per_block_per_layer * NUM_LAYERSWith 24 GB of GPU memory reserved for KV cache, we can fit roughly 24 GB / 256 KB = 98,304 blocks, which at 16 tokens per block supports about 1.5 million tokens across all concurrent requests.
The block allocator
The allocator is conceptually simple: maintain a free list of block IDs and hand them out on demand. Each request gets a block table, which is a list of physical block IDs that maps logical token positions to physical memory locations.
from dataclasses import dataclass, field
from typing import Dict, List, Optional
from collections import OrderedDict
@dataclass
class BlockAllocator:
num_blocks: int
block_size: int
free_blocks: List[int] = field(default_factory=list)
# Maps request_id to its list of allocated block IDs
block_tables: Dict[str, List[int]] = field(default_factory=dict)
# Track how many tokens are used in the last block of each request
last_block_usage: Dict[str, int] = field(default_factory=dict)
def __post_init__(self):
# Initialize free list with all block IDs
self.free_blocks = list(range(self.num_blocks - 1, -1, -1))
def allocate(self, request_id: str) -> int:
"""Allocate a single block for a request. Returns block ID."""
if not self.free_blocks:
raise MemoryError("No free blocks available")
block_id = self.free_blocks.pop()
if request_id not in self.block_tables:
self.block_tables[request_id] = []
self.last_block_usage[request_id] = 0
self.block_tables[request_id].append(block_id)
return block_id
def free(self, request_id: str):
"""Free all blocks belonging to a request."""
if request_id in self.block_tables:
self.free_blocks.extend(self.block_tables[request_id])
del self.block_tables[request_id]
del self.last_block_usage[request_id]
def append_token(self, request_id: str) -> int:
"""Record a new token for a request, allocating a new block if needed.
Returns the physical block ID where the token should be written."""
usage = self.last_block_usage.get(request_id, 0)
if usage == 0 or usage >= self.block_size:
# Need a new block
block_id = self.allocate(request_id)
self.last_block_usage[request_id] = 1
return block_id
else:
self.last_block_usage[request_id] = usage + 1
return self.block_tables[request_id][-1]
@property
def num_free(self) -> int:
return len(self.free_blocks)
@property
def utilization(self) -> float:
return 1.0 - (len(self.free_blocks) / self.num_blocks)The block table: virtual to physical mapping
The block table is the heart of paged KV cache. When the attention kernel needs to read the KV cache for a given request at token position t, it looks up which physical block holds that token:
def get_physical_location(self, request_id: str, token_pos: int):
"""Map a logical token position to a physical (block_id, offset) pair."""
block_index = token_pos // self.block_size
offset = token_pos % self.block_size
block_id = self.block_tables[request_id][block_index]
return block_id, offsetThis indirection is exactly what operating systems do with virtual memory pages. The attention kernel receives a list of physical block IDs and offsets, and gathers the KV vectors from non-contiguous memory. This is what makes PagedAttention different from standard attention: it reads from scattered blocks instead of a contiguous buffer.
Eviction: what happens when memory runs out
When the free list is empty and a new request arrives, we need to evict. The question is which request to evict. I implemented three policies to compare:
class EvictionPolicy:
"""Manages eviction when KV cache memory is exhausted."""
def __init__(self, policy: str = "lru"):
self.policy = policy
# Track access order for LRU
self.access_order: OrderedDict[str, float] = OrderedDict()
def record_access(self, request_id: str, timestamp: float):
"""Record that a request was accessed (generated a token)."""
if request_id in self.access_order:
self.access_order.move_to_end(request_id)
self.access_order[request_id] = timestamp
def select_victim(self, block_tables: Dict[str, List[int]]) -> Optional[str]:
"""Select a request to evict based on the policy."""
if not block_tables:
return None
if self.policy == "lru":
# Evict the least recently accessed request
for request_id in self.access_order:
if request_id in block_tables:
return request_id
return None
elif self.policy == "largest":
# Evict the request using the most blocks
return max(block_tables, key=lambda r: len(block_tables[r]))
elif self.policy == "shortest_remaining":
# Evict the request closest to its max_tokens limit
# (requires external info, simplified here)
return min(block_tables, key=lambda r: len(block_tables[r]))LRU eviction is intuitive, but in LLM serving, a request that has not generated tokens recently might be waiting for a long prompt to prefill. Evicting it would waste all the prefill work. Production systems like vLLM use a priority-based approach that considers whether a request is in prefill or decode, how much work has been done, and whether the request can be swapped to CPU memory instead of being fully evicted.
Swapping: eviction without loss
True eviction means killing a request and losing all its progress. A better approach is swapping: copy the request's KV cache blocks from GPU memory to CPU memory, freeing the GPU blocks for other requests. When the swapped request gets scheduled again, copy its blocks back.
class SwapManager:
"""Manages GPU-to-CPU and CPU-to-GPU KV cache block swaps."""
def __init__(self, allocator: BlockAllocator, cpu_blocks: int):
self.gpu_allocator = allocator
self.cpu_allocator = BlockAllocator(cpu_blocks, allocator.block_size)
# Maps request_id to its CPU block table (when swapped out)
self.swapped_requests: Dict[str, List[int]] = {}
def swap_out(self, request_id: str) -> List[tuple]:
"""Swap a request's KV cache from GPU to CPU.
Returns list of (gpu_block, cpu_block) pairs for the copy kernel."""
gpu_blocks = self.gpu_allocator.block_tables[request_id]
swap_pairs = []
for gpu_block in gpu_blocks:
cpu_block = self.cpu_allocator.allocate(request_id)
swap_pairs.append((gpu_block, cpu_block))
self.swapped_requests[request_id] = list(
self.cpu_allocator.block_tables[request_id]
)
self.gpu_allocator.free(request_id)
return swap_pairs # caller issues async memcpy for each pair
def swap_in(self, request_id: str) -> List[tuple]:
"""Swap a request's KV cache from CPU back to GPU."""
cpu_blocks = self.swapped_requests[request_id]
swap_pairs = []
for cpu_block in cpu_blocks:
gpu_block = self.gpu_allocator.allocate(request_id)
swap_pairs.append((cpu_block, gpu_block))
self.cpu_allocator.free(request_id)
del self.swapped_requests[request_id]
return swap_pairsThe swap pairs are handed to an async CUDA memcpy, so the actual data transfer overlaps with computation on other requests. The PCIe bus between CPU and GPU becomes the bottleneck here, which is why systems try to minimize swapping and use it only as a last resort.
Fragmentation and compaction
Unlike contiguous allocation, block-based allocation eliminates external fragmentation entirely. Every block is the same size, so any free block can serve any request. Internal fragmentation (wasted space in the last partially-filled block) is bounded by one block per request, which is at most 16 tokens' worth of memory.
This is a massive improvement over the naive approach. With contiguous allocation, a request that might generate up to 2048 tokens must reserve 2048 tokens of memory upfront, even if it only generates 50. With blocks, it allocates 4 blocks (64 tokens) and wastes at most 14 tokens in the last block.
The block allocator turned KV cache management from a hard memory planning problem into a simple bookkeeping problem. Pre-allocate the blocks, manage a free list, and let requests grow incrementally. The rest is just accounting.
What I learned
Building this allocator made the vLLM architecture click for me in a way that reading the paper did not. The key insights:
- Blocks eliminate fragmentation. Any block can serve any request, so memory utilization is near-optimal.
- Block tables are tiny. A request with 4096 tokens has a block table of 256 integers. The metadata overhead is negligible.
- Swapping is better than eviction. CPU memory is cheap and the PCIe copy is fast enough that swapping beats recomputing from scratch.
- The eviction policy matters less than you think. With good block utilization, you hit memory limits much later, and swapping handles the overflow gracefully.
Tomorrow I am building on top of this allocator to add prefix caching with hash deduplication, where multiple requests that share a common prompt prefix can share the same physical KV cache blocks.