TL;DR — You will build a paged attention KV cache from scratch in pure Python, with real, runnable code that demonstrates memory management, algorithm design, and transformer internals. This project signals systems engineering depth to hiring managers and can be extended into a production-grade inference engine.

Most transformer inference engines hide the KV cache behind abstractions, but the leap from “it works” to “it scales” lives in how you manage that cache. By implementing paged attention—a technique that eliminates fragmentation and wasted memory—you show you understand both the algorithm and the systems underneath. This guide walks you through a complete, buildable implementation, then maps out a roadmap to turn it into a senior-level portfolio piece.

Why This Project Stands Out on a CV

A paged attention KV cache is a rare project that exercises multiple skills hiring managers value:

  • Memory management and virtualization: You are implementing a custom allocator that maps logical sequence positions to physical memory blocks, mirroring operating system paging. This directly translates to backend roles at companies like Meta, Google, or Cloudflare.
  • Algorithm design for attention: You will write the core attention computation that underpins every modern LLM, showing you understand the math and the computational bottlenecks.
  • End-to-end ML systems: By integrating the cache with a transformer layer, you demonstrate you can ship a working inference pipeline, not just isolated algorithms.
  • Production orientation: The extension roadmap (persistence, scaling, observability) proves you think beyond a toy project, which is exactly what senior engineers do.

If you are targeting roles in ML infrastructure, AI platform engineering, or high-performance backend systems, this project gives you concrete talking points and code to back them up.

Architecture Overview

The system has four core components:

  1. Physical Block Pool — A pre-allocated array of fixed-size memory blocks. Each block stores key and value vectors for a single transformer layer and a single sequence. Blocks are the unit of allocation and deallocation.
  2. Block Allocator — Manages a free list of unused blocks. When a sequence needs more space for new tokens, it either appends to the last block if space remains or grabs a fresh block from the pool. On sequence completion, all its blocks are returned to the free list.
  3. Logical-to-Physical Mapping — A dictionary keyed by (layer_id, sequence_id) that maintains the ordered list of physical block indices for that sequence. This indirection is what makes paging possible: logical contiguous sequences are stored in non-contiguous physical blocks.
  4. Paged Attention Engine — Computes attention by first gathering the key/value vectors from the mapped blocks into a contiguous temporary array, then applying the standard dot-product attention with scaling and softmax. The “paged” aspect means this gather is cheap and cache-friendly.

In production, this is exactly what vLLM’s PagedAttention does, but we build the core logic ourselves.

+------------------+     +-------------------+
|   Sequence 1     |     |   Sequence 2      |
|  (layer 0)       |     |  (layer 0)        |
|  Block 3 -> 7 -> |     |  Block 1 -> 4 ->  |
|  Block 2         |     |  Block 9          |
+------------------+     +-------------------+
         |                       |
         v                       v
+-------------------------------+
|      Physical Block Pool      |
|  [0] [1] [2] [3] [4] [5] ...  |
|   free  used  free used used free|
+-------------------------------+

Building It Step by Step

Step 1: Define the PagedKVCache class

We start with the data structures. For clarity, we use Python lists of lists; in a real system you would replace these with NumPy arrays or torch tensors.

# kv_cache.py
import math
from typing import List, Tuple, Optional

class PagedKVCache:
    def __init__(self, block_size: int = 16, num_blocks: int = 1024):
        self.block_size = block_size
        self.num_blocks = num_blocks
        # Physical storage: each block holds a list of key vectors and value vectors.
        # Initially all blocks are free (None).
        self.k_blocks: List[Optional[List[List[float]]]] = [None] * num_blocks
        self.v_blocks: List[Optional[List[List[float]]]] = [None] * num_blocks
        # Free list for block allocation
        self.free_blocks: List[int] = list(range(num_blocks))
        # Mapping: (layer_id, seq_id) -> list of physical block indices
        self.mapping: dict = {}

    def _allocate_block(self) -> int:
        """Grab a free block index and initialise its storage."""
        if not self.free_blocks:
            raise RuntimeError("KV cache out of memory – increase num_blocks or evict.")
        block_id = self.free_blocks.pop()
        self.k_blocks[block_id] = []
        self.v_blocks[block_id] = []
        return block_id

    def _free_block(self, block_id: int) -> None:
        """Return a block to the free list."""
        self.k_blocks[block_id] = None
        self.v_blocks[block_id] = None
        self.free_blocks.append(block_id)

    def append(self, layer_id: int, seq_id: int, k: List[float], v: List[float]) -> None:
        """
        Append a single token's key and value to the specified sequence.
        If the last block is full, allocate a new one.
        """
        key = (layer_id, seq_id)
        if key not in self.mapping:
            self.mapping[key] = []

        blocks = self.mapping[key]
        if blocks and len(self.k_blocks[blocks[-1]]) < self.block_size:
            # There is room in the last block – append directly.
            self.k_blocks[blocks[-1]].append(k)
            self.v_blocks[blocks[-1]].append(v)
        else:
            # Need a fresh block.
            block_id = self._allocate_block()
            self.k_blocks[block_id].append(k)
            self.v_blocks[block_id].append(v)
            blocks.append(block_id)

    def get_kv(self, layer_id: int, seq_id: int) -> Tuple[List[List[float]], List[List[float]]]:
        """
        Return the full key and value sequences for a given layer and sequence,
        gathered from all physical blocks in order.
        """
        key = (layer_id, seq_id)
        if key not in self.mapping:
            return [], []
        k_all: List[List[float]] = []
        v_all: List[List[float]] = []
        for block_id in self.mapping[key]:
            k_all.extend(self.k_blocks[block_id])
            v_all.extend(self.v_blocks[block_id])
        return k_all, v_all

    def free_sequence(self, layer_id: int, seq_id: int) -> None:
        """
        Release all blocks belonging to a sequence. Call this when generation
        finishes or the sequence is evicted.
        """
        key = (layer_id, seq_id)
        if key in self.mapping:
            for block_id in self.mapping[key]:
                self._free_block(block_id)
            del self.mapping[key]

Step 2: Implement paged attention

Attention is the heart of the transformer. We compute scaled dot-product attention, but we explicitly gather the KV from the paged blocks first.

# attention.py
def paged_attention(
    query: List[float],
    k_cache: List[List[List[float]]],  # list of blocks, each block is list of key vectors
    v_cache: List[List[List[float]]],  # same shape as k_cache
) -> List[float]:
    """
    Compute scaled dot-product attention over a paged KV cache.
    `query` is a single vector (the current token's query).
    Returns the weighted sum of value vectors.
    """
    d_k = len(query)

    # Flatten the paged blocks into contiguous sequences.
    k_flat: List[List[float]] = []
    v_flat: List[List[float]] = []
    for block in k_cache:
        k_flat.extend(block)
    for block in v_cache:
        v_flat.extend(block)

    if not k_flat:
        return [0.0] * d_k

    # 1. Compute raw scores: Q·K^T / sqrt(d_k)
    scale = 1.0 / math.sqrt(d_k)
    scores = []
    for k_vec in k_flat:
        dot = sum(q * k for q, k in zip(query, k_vec))
        scores.append(dot * scale)

    # 2. Softmax
    max_score = max(scores)
    exp_scores = [math.exp(s - max_score) for s in scores]
    sum_exp = sum(exp_scores)
    weights = [e / sum_exp for e in exp_scores]

    # 3. Weighted sum of values
    output = [0.0] * d_k
    for i, w in enumerate(weights):
        for j in range(d_k):
            output[j] += w * v_flat[i][j]

    return output

Step 3: Wire it into a minimal transformer layer

To prove the cache works end-to-end, we simulate a single transformer layer. In a real model you would load pretrained weights; here we use random projections to focus on the cache logic.

# transformer_layer.py
import random
from kv_cache import PagedKVCache
from attention import paged_attention

class MiniTransformerLayer:
    def __init__(self, d_model: int = 64, num_heads: int = 4):
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        # Random projection matrices – in production these are learned.
        self.W_q = [[random.random() for _ in range(d_model)] for _ in range(d_model)]
        self.W_k = [[random.random() for _ in range(d_model)] for _ in range(d_model)]
        self.W_v = [[random.random() for _ in range(d_model)] for _ in range(d_model)]
        self.W_o = [[random.random() for _ in range(d_model)] for _ in range(d_model)]

    def _linear(self, x: List[float], W: List[List[float]]) -> List[float]:
        """Simple matrix-vector multiplication."""
        return [sum(x_i * w_i for x_i, w_i in zip(x, row)) for row in W]

    def forward(self, x: List[float], layer_id: int, seq_id: int, cache: PagedKVCache) -> List[float]:
        """
        x: input token embedding (size d_model).
        Returns the output after attention + residual (we skip FFN for brevity).
        """
        # Project to Q, K, V
        q = self._linear(x, self.W_q)
        k = self._linear(x, self.W_k)
        v = self._linear(x, self.W_v)

        # Store KV in the cache
        cache.append(layer_id, seq_id, k, v)

        # Retrieve full KV for this sequence
        k_blocks, v_blocks = cache.get_kv(layer_id, seq_id)

        # Compute attention (simplified: single-head, so we use full vectors)
        attn_out = paged_attention(q, k_blocks, v_blocks)

        # Output projection
        output = self._linear(attn_out, self.W_o)

        # Residual connection
        return [x_i + o_i for x_i, o_i in zip(x, output)]

Step 4: Run a generation loop

Now we simulate autoregressive generation: feed one token at a time, let the cache grow, and observe the output.

# demo.py
from kv_cache import PagedKVCache
from transformer_layer import MiniTransformerLayer
import random

def main():
    d_model = 32
    cache = PagedKVCache(block_size=8, num_blocks=64)
    layer = MiniTransformerLayer(d_model=d_model, num_heads=1)

    # Start with a random initial token
    token = [random.random() for _ in range(d_model)]
    seq_id = 1
    layer_id = 0

    print("Generating 20 tokens with paged KV cache...")
    for step in range(20):
        output = layer.forward(token, layer_id, seq_id, cache)
        # Use output as next input (simplified: no token embedding lookup)
        token = output
        # Print the first 4 dimensions as a proxy for the "token"
        print(f"Step {step:2d}: {[round(v, 4) for v in token[:4]]} ...")

    # Verify cache state
    k, v = cache.get_kv(layer_id, seq_id)
    print(f"\nCache now holds {len(k)} key/value vectors for seq {seq_id}.")
    print(f"Physical blocks used: {len(cache.mapping[(layer_id, seq_id)])}")

if __name__ == "__main__":
    main()

Running and Testing It

Save the four files (kv_cache.py, attention.py, transformer_layer.py, demo.py) in a directory and run:

python demo.py

You should see 20 lines of generated vectors and a final summary showing the cache holds 20 key/value vectors distributed across a few physical blocks (since block_size=8, you will use 3 blocks). To verify correctness, add a unit test that compares the output of paged_attention with a naive full-matrix attention implementation for the same KV pairs.

# test_paged_attention.py
import math
from attention import paged_attention

def naive_attention(query, k_flat, v_flat):
    d_k = len(query)
    scale = 1.0 / math.sqrt(d_k)
    scores = [sum(q * k for q, k in zip(query, k)) * scale for k in k_flat]
    max_s = max(scores)
    exp_s = [math.exp(s - max_s) for s in scores]
    weights = [e / sum(exp_s) for e in exp_s]
    out = [0.0] * d_k
    for i, w in enumerate(weights):
        for j in range(d_k):
            out[j] += w * v_flat[i][j]
    return out

# Test with random data
import random
d = 8
q = [random.random() for _ in range(d)]
k = [[random.random() for _ in range(d)] for _ in range(10)]
v = [[random.random() for _ in range(d)] for _ in range(10)]

# Simulate paged blocks: split into two blocks of 5
k_blocks = [k[:5], k[5:]]
v_blocks = [v[:5], v[5:]]

out_paged = paged_attention(q, k_blocks, v_blocks)
out_naive = naive_attention(q, k, v)

assert len(out_paged) == len(out_naive)
for a, b in zip(out_paged, out_naive):
    assert abs(a - b) < 1e-9, f"Mismatch: {a} vs {b}"
print("Paged attention matches naive implementation.")

Run it with python test_paged_attention.py. If the assertion passes, your paged gather and attention math are correct.

Extending It: Your Roadmap to Senior-Level

A toy implementation is a starting point. The following upgrades turn it into a portfolio piece that speaks to production readiness:

  1. Persistent KV cache with SQLite — Store KV pairs on disk and reload for multi-session serving. Reason: Real applications serve thousands of concurrent sessions; in-memory only is not viable.
  2. Distributed KV cache via gRPC — Shard sequences across multiple nodes using a consistent hash ring. Reason: Horizontal scaling is required when a single machine cannot hold the working set.
  3. Observability with Prometheus metrics — Expose cache hit rate, allocation latency, block utilisation, and attention throughput. Reason: Engineers cannot debug or optimise what they cannot measure.
  4. Fault tolerance with replication — Replicate each block to a secondary node and use a lightweight consensus protocol (e.g., Raft) for failover. Reason: Production inference services must survive node failures without losing active sessions.
  5. Benchmarking suite using Locust — Simulate concurrent users with varying sequence lengths and measure p99 latency. Reason: Hiring managers expect evidence of performance under load, not just correctness.
  6. Integration with Hugging Face Transformers — Hook your cache into the forward method of a real model (e.g., GPT-2) by replacing the default KV cache. Reason: Demonstrates you can retrofit your solution into an established ecosystem, which is how adoption happens.

Each of these can be tackled incrementally; even implementing one or two will elevate the project from “class assignment” to “production prototype.”

Key Takeaways

  • A paged attention KV cache is a systems problem disguised as an algorithm: you are building a custom virtual memory manager for transformer states.
  • The core logic—block allocation, logical-to-physical mapping, and paged gather—is only a few dozen lines of Python but exercises memory management, algorithm design, and ML systems integration.
  • This project signals to hiring managers that you understand transformer internals and the engineering challenges of scaling them.
  • The extension roadmap (persistence, distribution, observability, fault tolerance, benchmarking, real-model integration) shows you can think beyond a demo to a production service.
  • Always test with a naive implementation to verify correctness; numerical equivalence is your first and most important validation.

Further Reading