TL;DR — Building a paged KV cache from scratch teaches you the exact memory-management and block-allocation patterns that power vLLM and other high-throughput inference engines. This project demonstrates virtual-memory paging, cache coherence, and systems-level Python engineering — a rare and impressive signal on any CV.

Why This Project Stands Out on a CV

Most portfolio projects for ML engineers involve fine-tuning a Hugging Face model or building a chatbot wrapper around an API. Those are fine for juniors, but they signal consumption, not construction. A paged KV cache is different — it sits at the intersection of deep learning systems, operating systems concepts, and high-performance engineering.

Here is what this project demonstrates to a hiring manager:

  • Systems-level thinking: You understand that inference is not just about models — it is about memory layout, allocation strategies, and waste minimization. The paged KV cache directly borrows from virtual memory paging in OS kernels, a concept that separates engineers who understand where computation happens from those who only understand what computation runs.
  • Deep learning infrastructure fluency: You can name-drop and implement the core idea behind PagedAttention (the paper behind vLLM), which is one of the most cited systems innovations in LLM serving.
  • Production-quality Python: You write clean abstractions — classes for blocks, allocators, and cache managers — rather than notebook scripts. This signals you can work in a codebase, not just prototype in isolation.
  • Performance awareness: By benchmarking allocation strategies and measuring fragmentation, you show you care about the gap between “it works” and “it works at scale.”

The roles this signals: ML Systems Engineer, Inference Optimization Engineer, Backend Engineer for AI Platforms, and Technical Lead positions where you need to bridge model research and serving infrastructure.

Architecture Overview

The project decomposes into four tightly-coupled components, each with a single responsibility:

┌─────────────────────────────────────────────────────┐
│                  PagedKVCache                        │
│  ┌──────────┐  ┌──────────────┐  ┌───────────────┐  │
│  │ Block    │  │ Block        │  │ Sequence      │  │
│  │ Pool     │  │ Allocator    │  │ Manager       │  │
│  │          │  │              │  │               │  │
│  │ [Block]  │  │ free_list    │  │ seq_id →      │  │
│  │ [Block]  │  │ used_map     │  │   [block_ids] │  │
│  │ [Block]  │  │              │  │               │  │
│  │ [Block]  │  │ allocate()   │  │ append()      │  │
│  │ [Block]  │  │ deallocate() │  │ get_kv()      │  │
│  └──────────┘  └──────────────┘  └───────────────┘  │
│         │                │              │             │
│         └────────────────┴──────────────┘             │
│                       │                               │
│              ┌────────▼────────┐                      │
│              │ Attention      │                      │
│              │ Forward Pass   │                      │
│              │ (consumer)     │                      │
│              └─────────────────┘                      │
└─────────────────────────────────────────────────────┘
  • Block Pool: A pre-allocated array of fixed-size blocks. Each block holds a slice of keys and values for all heads in a single transformer layer. The block size is configurable (e.g., 16 or 32 tokens).
  • Block Allocator: Manages free and used blocks using a first-fit strategy. It mirrors a buddy allocator or a simple free-list allocator from OS kernels. When a sequence needs more space, the allocator hands it blocks; when a sequence finishes, those blocks return to the free pool.
  • Sequence Manager: Tracks which blocks belong to which sequence ID and maintains the logical position within each sequence. This is the bookkeeping layer that maps token positions to physical block locations.
  • Attention Consumer: A minimal multi-head attention function that reads from the paged cache to compute attention scores, demonstrating that the cache is actually usable for inference.

The key design decision: blocks are fixed-size pages. Unlike a contiguous buffer where you must pre-allocate the maximum sequence length, paged allocation means a short sequence wastes almost no memory and a long sequence simply uses more pages. This is the core insight from PagedAttention — treat KV cache memory like virtual memory, and you eliminate internal fragmentation.

Building It Step by Step

Step 1: Define the Block Structure

Each block stores keys and values for one layer, one head dimension slice, and a fixed number of token positions. We use NumPy arrays for the actual tensor storage but keep all management logic in pure Python.

import numpy as np
from dataclasses import dataclass, field
from typing import Optional

@dataclass
class KVCacheBlock:
    """A single page in the paged KV cache."""
    block_size: int = 16
    num_heads: int = 8
    head_dim: int = 64
    dtype: np.dtype = np.float32

    keys: np.ndarray = field(default_factory=lambda: None)
    values: np.ndarray = field(default_factory=lambda: None)
    in_use: bool = False
    seq_id: Optional[int] = None
    start_pos: int = 0  # first token position this block covers

    def __post_init__(self):
        if self.keys is None:
            self.keys = np.zeros((self.block_size, self.num_heads, self.head_dim), dtype=self.dtype)
            self.values = np.zeros((self.block_size, self.num_heads, self.head_dim), dtype=self.dtype)

    def is_full(self) -> bool:
        return self.start_pos + self.block_size > self._max_pos_seen

    def available_slots(self) -> int:
        return self.block_size - self.start_pos

Step 2: Build the Block Allocator

The allocator maintains a free list (stack of available block indices) and a mapping from sequence IDs to their allocated blocks. This is the heart of the system — every allocation and deallocation runs in O(1).

class BlockAllocator:
    """Manages block allocation and deallocation, like a kernel's page allocator."""

    def __init__(self, total_blocks: int, block_size: int, num_heads: int, head_dim: int):
        self.total_blocks = total_blocks
        self.block_size = block_size
        self.free_list: list[int] = list(range(total_blocks))  # stack of free block indices
        self.used_map: dict[int, list[int]] = {}  # seq_id -> [block_indices]
        self.block_meta: dict[int, KVCacheBlock] = {}  # block_idx -> block object

        # Pre-allocate all blocks
        for idx in range(total_blocks):
            self.block_meta[idx] = KVCacheBlock(block_size, num_heads, head_dim)

    def allocate(self, seq_id: int, num_blocks: int = 1) -> list[int]:
        """Allocate `num_blocks` for `seq_id`. Returns list of block indices."""
        if len(self.free_list) < num_blocks:
            raise MemoryError(
                f"KV cache out of memory: need {num_blocks} blocks, "
                f"only {len(self.free_list)} free of {self.total_blocks} total."
            )

        allocated = []
        for _ in range(num_blocks):
            block_idx = self.free_list.pop()
            block = self.block_meta[block_idx]
            block.in_use = True
            block.seq_id = seq_id
            block.start_pos = len(self.used_map.get(seq_id, [])) * self.block_size
            allocated.append(block_idx)

        self.used_map.setdefault(seq_id, []).extend(allocated)
        return allocated

    def deallocate(self, seq_id: int) -> None:
        """Return all blocks belonging to `seq_id` to the free pool."""
        if seq_id not in self.used_map:
            return
        for block_idx in self.used_map[seq_id]:
            block = self.block_meta[block_idx]
            block.in_use = False
            block.seq_id = None
            block.start_pos = 0
            self.free_list.append(block_idx)
        del self.used_map[seq_id]

    def fragment_ratio(self) -> float:
        """Fraction of total blocks that are free but scattered (simplified metric)."""
        free_count = len(self.free_list)
        return 1.0 - (free_count / self.total_blocks) if self.total_blocks > 0 else 0.0

    def stats(self) -> dict:
        return {
            "total_blocks": self.total_blocks,
            "free_blocks": len(self.free_list),
            "used_blocks": sum(len(v) for v in self.used_map.values()),
            "active_sequences": len(self.used_map),
            "fragment_ratio": self.fragment_ratio(),
        }

Step 3: Build the Paged KV Cache

The cache ties everything together. It exposes append and retrieve operations that the attention engine calls.

class PagedKVCache:
    """Main interface: a paged KV cache for transformer inference."""

    def __init__(self, num_layers: int, total_blocks: int = 256,
                 block_size: int = 16, num_heads: int = 8, head_dim: int = 64):
        self.num_layers = num_layers
        self.block_size = block_size
        self.allocator = BlockAllocator(total_blocks, block_size, num_heads, head_dim)
        self.seq_positions: dict[int, int] = {}  # seq_id -> current token position

    def ensure_blocks(self, seq_id: int, needed_blocks: int) -> None:
        """Allocate more blocks if the current allocation is insufficient."""
        current_blocks = len(self.allocator.used_map.get(seq_id, []))
        if current_blocks < needed_blocks:
            self.allocator.allocate(seq_id, needed_blocks - current_blocks)

    def append(self, seq_id: int, layer: int, key_vec: np.ndarray, val_vec: np.ndarray) -> None:
        """Append a single token's key and value to the cache for a given layer."""
        pos = self.seq_positions.get(seq_id, 0)
        block_idx = pos // self.block_size
        offset = pos % self.block_size

        # Ensure enough blocks exist
        needed_blocks = block_idx + 1
        self.ensure_blocks(seq_id, needed_blocks)

        # Write into the correct block
        block_indices = self.allocator.used_map[seq_id]
        physical_block = self.allocator.block_meta[block_indices[block_idx]]
        physical_block.keys[offset] = key_vec
        physical_block.values[offset] = val_vec

        self.seq_positions[seq_id] = pos + 1

    def retrieve(self, seq_id: int, layer: int) -> tuple[np.ndarray, np.ndarray]:
        """Retrieve all keys and values for a sequence up to its current length."""
        if seq_id not in self.allocator.used_map:
            return np.array([]), np.array([])

        block_indices = self.allocator.used_map[seq_id]
        pos = self.seq_positions[seq_id]
        num_tokens = pos

        all_keys = []
        all_values = []
        for block_idx in block_indices:
            block = self.allocator.block_meta[block_idx]
            tokens_in_block = min(self.block_size, num_tokens - block.start_pos)
            all_keys.append(block.keys[:tokens_in_block])
            all_values.append(block.values[:tokens_in_block])

        return np.concatenate(all_keys, axis=0), np.concatenate(all_values, axis=0)

    def free_sequence(self, seq_id: int) -> None:
        """Release all memory for a completed sequence."""
        self.allocator.deallocate(seq_id)
        self.seq_positions.pop(seq_id, None)

    def cache_stats(self) -> dict:
        return {
            "allocator": self.allocator.stats(),
            "active_sequences": len(self.seq_positions),
        }

Step 4: Wire Up a Minimal Attention Consumer

This is the proof that the cache actually works for inference. We implement scaled dot-product attention that reads from the paged cache.

def scaled_dot_product_attention(
    query: np.ndarray,          # (num_heads, head_dim)
    cache: PagedKVCache,
    seq_id: int,
    layer: int,
    scale: float = 0.125        # 1/sqrt(head_dim) approximation
) -> np.ndarray:
    """Compute attention output for a single query token using paged KV cache."""
    keys, values = cache.retrieve(seq_id, layer)

    if len(keys) == 0:
        return np.zeros((cache.allocator.block_meta[0].num_heads, cache.allocator.block_meta[0].head_dim))

    # keys: (seq_len, num_heads, head_dim), query: (num_heads, head_dim)
    # Compute attention scores: (num_heads, seq_len)
    scores = np.einsum('hd,shd->hs', query, keys) * scale

    # Softmax over sequence dimension
    exp_scores = np.exp(scores - np.max(scores, axis=1, keepdims=True))
    attn_weights = exp_scores / np.sum(exp_scores, axis=1, keepdims=True)

    # Weighted sum of values: (num_heads, head_dim)
    output = np.einsum('hs,shd->hd', attn_weights, values)
    return output

Step 5: Create a Minimal Transformer Block That Uses It

A single transformer block with self-attention that actually feeds into and reads from the paged cache.

class MiniTransformerBlock:
    """A single transformer layer with paged KV cache integration."""

    def __init__(self, cache: PagedKVCache, layer_idx: int,
                 hidden_dim: int = 512, num_heads: int = 8, head_dim: int = 64):
        self.cache = cache
        self.layer_idx = layer_idx
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.head_dim = head_dim

        # Simple random weight matrices (in production, these would be trained)
        self.w_q = np.random.randn(hidden_dim, num_heads * head_dim) * 0.02
        self.w_k = np.random.randn(hidden_dim, num_heads * head_dim) * 0.02
        self.w_v = np.random.randn(hidden_dim, num_heads * head_dim) * 0.02
        self.w_o = np.random.randn(num_heads * head_dim, hidden_dim) * 0.02

    def forward(self, seq_id: int, hidden_state: np.ndarray) -> np.ndarray:
        """Forward pass: compute attention using paged KV cache."""
        # Project to Q, K, V
        q = hidden_state @ self.w_q  # (hidden_dim,) -> (num_heads * head_dim,)
        k = hidden_state @ self.w_k
        v = hidden_state @ self.w_v

        # Reshape into heads
        q_heads = q.reshape(self.num_heads, self.head_dim)
        k_vec = k.reshape(self.num_heads, self.head_dim)
        v_vec = v.reshape(self.num_heads, self.head_dim)

        # Append to paged cache
        self.cache.append(seq_id, self.layer_idx, k_vec, v_vec)

        # Compute attention from cache
        attn_output = scaled_dot_product_attention(q_heads, self.cache, seq_id, self.layer_idx)

        # Output projection
        return attn_output.reshape(-1) @ self.w_o

Running and Testing It

Install the dependency and run a smoke test that proves the cache works end-to-end:

pip install numpy

Create a test script test_paged_cache.py:

import numpy as np
from paged_kv_cache import PagedKVCache, MiniTransformerBlock, BlockAllocator

def test_basic_allocation():
    """Test that blocks are allocated and deallocated correctly."""
    allocator = BlockAllocator(total_blocks=64, block_size=16, num_heads=8, head_dim=64)

    # Allocate for sequence 1
    blocks = allocator.allocate(seq_id=1, num_blocks=3)
    assert len(blocks) == 3
    assert all(allocator.block_meta[b].in_use for b in blocks)
    assert allocator.stats()["free_blocks"] == 61

    # Deallocate
    allocator.deallocate(seq_id=1)
    assert allocator.stats()["free_blocks"] == 64
    assert len(allocator.used_map) == 0
    print("✓ Basic allocation test passed")

def test_paged_cache_append_and_retrieve():
    """Test that tokens appended to the cache can be retrieved."""
    cache = PagedKVCache(num_layers=2, total_blocks=128, block_size=16)

    seq_id = 42
    for i in range(35):  # More than one block worth of tokens
        key_vec = np.random.randn(8, 64).astype(np.float32)
        val_vec = np.random.randn(8, 64).astype(np.float32)
        cache.append(seq_id, layer=0, key_vec=key_vec, val_vec=val_vec)

    keys, values = cache.retrieve(seq_id, layer=0)
    assert keys.shape[0] == 35  # 35 tokens retrieved
    assert keys.shape[1:] == (8, 64)
    print(f"✓ Append/retrieve test passed: retrieved {keys.shape[0]} tokens")

def test_transformer_block():
    """End-to-end test: a mini transformer block produces valid output."""
    cache = PagedKVCache(num_layers=1, total_blocks=64)
    block = MiniTransformerBlock(cache, layer_idx=0)

    hidden = np.random.randn(512).astype(np.float32)
    output = block.forward(seq_id=7, hidden_state=hidden)

    assert output.shape == (512,)
    assert not np.isnan(output).any()
    print(f"✓ Transformer block test passed: output shape {output.shape}")

def test_memory_exhaustion():
    """Verify that out-of-memory raises a clear error."""
    cache = PagedKVCache(num_layers=1, total_blocks=4)

    seq_id = 1
    try:
        for i in range(100):  # Far exceeds 4 blocks
            cache.append(seq_id, layer=0,
                         key_vec=np.random.randn(8, 64).astype(np.float32),
                         val_vec=np.random.randn(8, 64).astype(np.float32))
        assert False, "Should have raised MemoryError"
    except MemoryError as e:
        print(f"✓ Memory exhaustion test passed: {e}")

if __name__ == "__main__":
    test_basic_allocation()
    test_paged_cache_append_and_retrieve()
    test_transformer_block()
    test_memory_exhaustion()
    print("\nAll tests passed. Your paged KV cache is working.")

Run it:

python test_paged_cache.py

Expected output:

✓ Basic allocation test passed
✓ Append/retrieve test passed: retrieved 35 tokens
✓ Transformer block test passed: output shape (512,)
✓ Memory exhaustion test passed: KV cache out of memory: need 100 blocks, only 0 free of 4 total.

All tests passed. Your paged KV cache is working.

To profile memory behavior, add a simple benchmark:

import time

def benchmark_allocation(total_blocks=1024, num_seqs=50, tokens_per_seq=128):
    cache = PagedKVCache(num_layers=2, total_blocks=total_blocks, block_size=16)
    start = time.time()

    for seq_id in range(num_seqs):
        for token in range(tokens_per_seq):
            cache.append(seq_id, layer=0,
                         key_vec=np.random.randn(8, 64).astype(np.float32),
                         val_vec=np.random.randn(8, 64).astype(np.float32))

    elapsed = time.time() - start
    stats = cache.cache_stats()
    print(f"Allocated {num_seqs} sequences × {tokens_per_seq} tokens in {elapsed:.3f}s")
    print(f"Cache stats: {stats}")

benchmark_allocation()

Extending It: Your Roadmap to Senior-Level

The toy above proves the concept. Here are six concrete upgrades that transform it into something that would impress in a production environment:

  1. Persistence via memory-mapped files. Use Python’s mmap module to persist KV blocks to disk so that sequences can be swapped out when GPU memory is full. This is the exact mechanism vLLM uses for CPU offloading, and implementing it shows you understand memory hierarchy and I/O bottlenecks.

  2. Thread-safe concurrent allocation. Add a threading.Lock or use multiprocessing.Manager to make the block allocator safe for concurrent request servicing. Production inference servers handle dozens of requests simultaneously; a race condition in block allocation corrupts the entire cache.

  3. Observability with Prometheus metrics. Instrument allocate(), deallocate(), and append() with counters, histograms, and gauges, then expose them via a /metrics endpoint using the prometheus_client library. This signals you understand that systems without observability are systems that fail silently.

  4. Block compaction and defragmentation. Implement a background routine that moves blocks between sequences to reduce fragmentation — analogous to page compaction in Linux kernels. This directly addresses the fragment_ratio metric you already have and shows you care about long-running system health.

  5. Fault tolerance with write-ahead logging. Before writing a key/value block, append the operation to a WAL (a simple append-only file). On restart, replay the log to reconstruct the cache state. This is a fundamental pattern from distributed systems and databases that very few ML engineers can demonstrate.

  6. Integration with Hugging Face transformers. Replace the random weight matrices in MiniTransformerBlock with actual model layers from transformers, and wire the paged cache as the attention KV store. This bridges the gap between your toy and a real inference engine, making the project immediately demonstrable in an interview.

Key Takeaways

  • A paged KV cache applies operating-system virtual memory paging concepts to transformer inference, eliminating the internal fragmentation that plagues contiguous KV cache allocation.
  • The core implementation requires only three components: a block pool, a block allocator, and a sequence manager — each with a clear single responsibility.
  • Building this from scratch demonstrates systems-level thinking, memory management expertise, and production-grade Python skills that distinguish you from candidates who only use high-level APIs.
  • The project is immediately extensible into production-flavored territory through persistence, observability, fault tolerance, and integration with real transformer models.
  • Benchmarking and fragmentation metrics give you concrete numbers to discuss in interviews, showing you think about performance, not just correctness.
  • The architecture mirrors real systems like vLLM’s PagedAttention, making it a credible signal that you understand the infrastructure behind modern LLM serving.

Further Reading