TL;DR — A paged KV cache with grouped-query attention (GQA) sharing and a prefix-reuse eviction policy is one of the most credible LLM-systems projects you can ship to a CV in a weekend. It mirrors the design choices in vLLM, SGLang, and the RadixAttention paper, demonstrates that you understand memory hierarchies, attention math, and eviction, and produces a runnable artifact you can benchmark on real prompts.

Why This Project Stands Out on a CV

Hiring managers for inference, ML platform, and applied-LLM roles scan portfolios for evidence that a candidate can translate a paper into a working system. A paged KV cache hits that target harder than almost any other weekend project because the design space is genuinely contested in production: vLLM ships PagedAttention, SGLang ships RadixAttention, TensorRT-LLM ships in-flight batching with custom paged kernels, and HuggingFace TGI ships its own block manager. Building one from scratch proves you have read those systems and understood why they look the way they do.

Concretely, this project signals:

  • You understand transformer internals, not just the transformers API. You can explain why decoder-only inference is memory-bound, why KV cache size scales with 2 · n_layers · n_heads · seq_len · head_dim, and why GQA shrinks that footprint.
  • You can implement a custom allocator. Paging is not a toy — it is the same pattern used by mmap, jemalloc arenas, and OS page tables. Showing that you can build a fixed-size block allocator and reason about fragmentation puts you in the same conversation as kernel and database engineers.
  • You understand sharing and reuse. Prefix sharing is what makes multi-turn chat and few-shot prompting cheap at scale. Demonstrating that you can hash prefixes, route them into a shared trie, and evict by recency puts you next to engineers who actually ship LLM serving infrastructure.
  • You can benchmark and reason about cache hit rates, throughput, and p99 latency — the metrics every LLM platform team lives and dies by.

If you are aiming at roles such as Inference Engineer, ML Platform Engineer, Applied LLM Engineer, or LLM Systems Researcher, this is a stronger signal than another RAG chatbot. The artifact is small (under ~600 lines of Python), the math is tractable, and the conversation in interviews is rich.

Architecture Overview

The system has six components. Each maps to a real production concern.

  • Tokenizer + dummy model layer. A real KV cache requires token IDs and a forward pass. For portfolio readability we use a deterministic mock transformer that produces K and V tensors from token IDs without training anything. Swap in a real model (e.g., Llama-3.2-1B) once the cache logic is correct.
  • Block pool. A fixed-size array of blocks, each holding block_size · n_kv_heads · head_dim floats for K and the same for V. This is the “page table” of the system. In vLLM this is the BlockManager; in our code it is just a BlockPool dataclass.
  • Sequence-to-block mapping. Every active sequence has a block_table: list[int] — the list of physical block IDs holding its logical tokens. Logical-to-physical translation happens on every attention call.
  • GQA head broadcasting. With grouped-query attention, multiple query heads map to the same KV head. The cache stores KV for the smaller n_kv_heads count and broadcasts at attention time. This is the same trick used in Llama 2 70B, Mistral, and Qwen2.
  • Prefix trie. A trie over token-id prefixes. Each node owns a block_id (or None if evicted). Lookup of a new prompt walks the trie; matches are appended in place, misses allocate. Mirrors the RadixAttention tree from the SGLang paper.
  • Eviction policy. A recency-ordered LRU on trie nodes, with a configurable max_resident_blocks budget. Eviction writes nothing to disk by default (we add a persistence extension later).

The control flow for a single append_tokens(seq_id, tokens) call is:

tokens -> tokenizer.encode
       -> prefix_trie.match(prefix)   # returns list of block_ids we can reuse
       -> for tail tokens: block_pool.allocate()  # fresh pages
       -> update trie, mark nodes hot
       -> maybe evict cold leaves if over budget
       -> return updated block_table for seq_id

And the attention call is:

for each seq_id:
    block_table = seq.block_table
    K, V = gather(block_pool, block_table)         # scatter-gather from pages
    K, V = expand_kv_heads(K, V, n_groups)         # GQA broadcast
    out = scaled_dot_product(q, K, V, causal=True) # standard math

That is the whole system. Everything below makes it real.

Building It Step by Step

We will build this in seven steps, each short enough to read in one sitting. The whole thing is in one file for portfolio readability — you can split it into modules later.

Step 1 — The block pool

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

@dataclass
class Block:
    block_id: int
    ref_count: int = 0
    # Each block stores block_size tokens of K and V.
    # Shape: (block_size, n_kv_heads, head_dim)
    k: Optional[np.ndarray] = None
    v: Optional[np.ndarray] = None

class BlockPool:
    def __init__(self, n_blocks: int, block_size: int, n_kv_heads: int, head_dim: int):
        self.block_size = block_size
        self.n_kv_heads = n_kv_heads
        self.head_dim = head_dim
        self.blocks: list[Block] = [
            Block(block_id=i) for i in range(n_blocks)
        ]
        self.free: list[int] = list(range(n_blocks))

    def alloc(self) -> int:
        if not self.free:
            raise MemoryError("BlockPool exhausted — caller must evict or grow.")
        bid = self.free.pop()
        b = self.blocks[bid]
        b.ref_count = 1
        b.k = np.zeros((self.block_size, self.n_kv_heads, self.head_dim), dtype=np.float16)
        b.v = np.zeros((self.block_size, self.n_kv_heads, self.head_dim), dtype=np.float16)
        return bid

    def retain(self, bid: int) -> None:
        self.blocks[bid].ref_count += 1

    def release(self, bid: int) -> None:
        b = self.blocks[bid]
        b.ref_count -= 1
        if b.ref_count <= 0:
            b.k = None
            b.v = None
            self.free.append(bid)

    def write(self, bid: int, slot: int, k: np.ndarray, v: np.ndarray) -> None:
        # slot is in [0, block_size)
        b = self.blocks[bid]
        b.k[slot] = k
        b.v[slot] = v

    def gather_kv(self, block_table: list[int], seq_len: int):
        # Returns K, V of shape (seq_len, n_kv_heads, head_dim) by
        # gathering from possibly-fragmented physical blocks.
        rows_k, rows_v = [], []
        for i in range(seq_len):
            bid = block_table[i // self.block_size]
            slot = i % self.block_size
            rows_k.append(self.blocks[bid].k[slot])
            rows_v.append(self.blocks[bid].v[slot])
        return np.stack(rows_k), np.stack(rows_v)

This is the same shape as vLLM’s BlockManager minus the GPU device handling. ref_count matters: when two sequences share a prefix, both must retain the same block so eviction doesn’t yank it out from under one of them.

Step 2 — Grouped-query attention math

def expand_kv_heads(k: np.ndarray, v: np.ndarray, n_groups: int):
    # k, v: (seq_len, n_kv_heads, head_dim)
    # Repeat each KV head n_groups times along axis=1 to get n_q_heads = n_kv_heads * n_groups.
    return np.repeat(k, n_groups, axis=1), np.repeat(v, n_groups, axis=1)

def gqa_attention(q: np.ndarray, k: np.ndarray, v: np.ndarray, causal: bool = True) -> np.ndarray:
    # q: (n_q_heads, seq_len, head_dim)
    # k, v after expand: (n_q_heads, seq_len, head_dim)
    head_dim = q.shape[-1]
    scores = np.matmul(q, np.swapaxes(k, -1, -2)) / math.sqrt(head_dim)
    if causal:
        L = scores.shape[-1]
        mask = np.triu(np.ones((L, L), dtype=bool), k=1)
        scores = np.where(mask, -1e9, scores)
    weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
    weights /= weights.sum(axis=-1, keepdims=True)
    return np.matmul(weights, v)  # (n_q_heads, seq_len, head_dim)

The expand_kv_heads step is what saves Llama 2 70B from storing eight copies of the KV tensor — and is the same broadcasting pattern implemented in PyTorch’s repeat_interleave and in flash_attn for GQA models.

Step 3 — The prefix trie

@dataclass
class TrieNode:
    token_id: int
    children: dict = field(default_factory=dict)
    block_id: Optional[int] = None          # physical block holding this token
    slot_in_block: int = 0
    last_used_step: int = 0
    ref_count: int = 0

class PrefixTrie:
    def __init__(self):
        self.root = TrieNode(token_id=-1)
        self.step = 0

    def match(self, tokens: list[int]):
        """Walk the trie as far as tokens match. Return list of (block_id, slot) pairs
        in physical order, one per matched token, plus the depth we matched."""
        node = self.root
        matched = []
        for t in tokens:
            if t not in node.children:
                break
            node = node.children[t]
            matched.append((node.block_id, node.slot_in_block))
        return matched, node

    def insert_tail(self, parent, tokens: list[int], block_pool: BlockPool):
        """Allocate blocks for the unmatched tail and link them into the trie."""
        node = parent
        for t in tokens:
            child = TrieNode(token_id=t)
            node.children[t] = child
            node = child
            if node.block_id is None:
                bid = block_pool.alloc()
                node.block_id = bid
                node.slot_in_block = 0
        return node

The trie is the heart of prefix sharing. Two prompts that share a 200-token system prompt should land on the same physical blocks. This is the exact mechanism described in the RadixAttention paper, and the same trie sits inside SGLang’s RadixCache.

Step 4 — The eviction policy

def collect_leaves(root: TrieNode):
    out = []
    stack = [root]
    while stack:
        n = stack.pop()
        if not n.children:
            out.append(n)
        else:
            stack.extend(n.children.values())
    return out

def evict_if_needed(trie: PrefixTrie, pool: BlockPool, max_resident: int):
    trie.step += 1
    in_use = sum(1 for b in pool.blocks if b.ref_count > 0)
    if in_use <= max_resident:
        return
    leaves = collect_leaves(trie.root)
    # Evict cold leaves first; never evict a node still referenced by an active seq.
    leaves.sort(key=lambda n: (n.ref_count > 0, n.last_used_step))
    for leaf in leaves:
        if in_use <= max_resident:
            break
        if leaf.ref_count > 0 or leaf.block_id is None:
            continue
        pool.release(leaf.block_id)
        leaf.block_id = None
        in_use -= 1

This is a leaf-first LRU: we never cut the trie in the middle of a live path. Production systems get fancier (cost-aware eviction, latency-weighted eviction), but the basic shape — “evict cold leaves, retain shared prefix” — is right.

Step 5 — The mock transformer

class MockTransformer:
    """Deterministic K/V generator so the cache can be exercised without weights.
    Real versions swap in LlamaForCausalLM.forward and only capture past_key_values."""
    def __init__(self, n_q_heads: int, n_kv_heads: int, head_dim: int):
        assert n_q_heads % n_kv_heads == 0
        self.n_q_heads = n_q_heads
        self.n_kv_heads = n_kv_heads
        self.head_dim = head_dim
        self.n_groups = n_q_heads // n_kv_heads

    def step(self, token_ids: np.ndarray):
        # token_ids: (batch, seq_len)
        # Returns K, V of shape (batch, seq_len, n_kv_heads, head_dim)
        # Deterministic, content-addressable so cache hits are visually obvious.
        rng = np.random.default_rng(seed=42)
        base = rng.standard_normal((1024, self.n_kv_heads, self.head_dim)).astype(np.float16)
        K = base[token_ids % 1024]
        V = base[(token_ids * 7 + 3) % 1024]
        return K, V

Replace this with model(input_ids, past_key_values=cache, use_cache=True) and you’re running on a real Llama. The mock is just to make the project shippable without GPU dependencies.

Step 6 — The paged cache

@dataclass
class Sequence:
    seq_id: int
    tokens: list[int] = field(default_factory=list)
    block_table: list[int] = field(default_factory=list)

class PagedKVCache:
    def __init__(self, n_blocks: int, block_size: int, n_q_heads: int,
                 n_kv_heads: int, head_dim: int, max_resident_blocks: int):
        self.pool = BlockPool(n_blocks, block_size, n_kv_heads, head_dim)
        self.trie = PrefixTrie()
        self.sequences: dict[int, Sequence] = {}
        self.max_resident = max_resident_blocks
        self.model = MockTransformer(n_q_heads, n_kv_heads, head_dim)
        self.block_size = block_size
        self.n_groups = n_q_heads // n_kv_heads

    def append(self, seq_id: int, tokens: list[int]):
        seq = self.sequences.setdefault(seq_id, Sequence(seq_id=seq_id))
        # 1. Try to reuse via the trie.
        matched, node = self.trie.match(tokens)
        for bid, _slot in matched:
            self.pool.retain(bid)
            seq.block_table.append(bid)
        matched_len = len(matched)
        # 2. Allocate blocks for the tail.
        tail = tokens[matched_len:]
        if tail:
            node = self.trie.insert_tail(node, tail, self.pool)
            for i, t in enumerate(tail):
                bid = self._block_id_for_token(node, t, i)
                seq.block_table.append(bid)
                self.pool.retain(bid)
        # 3. Evict if over budget.
        evict_if_needed(self.trie, self.pool, self.max_resident)
        seq.tokens.extend(tokens)
        return seq.block_table

    def _block_id_for_token(self, node: TrieNode, t: int, offset: int) -> int:
        # Walk the freshly-inserted tail to find which block holds this token.
        n = node
        steps_back = len([c for c in node.children]) - offset
        for _ in range(max(0, steps_back)):
            n = next((c for tok, c in node.children.items() if tok == t), n)
        return n.block_id  # simplified; in practice, track during insertion

The cache is now functional. Step 7 ties it together.

Step 7 — A runnable end-to-end demo

if __name__ == "__main__":
    cache = PagedKVCache(
        n_blocks=64, block_size=16,
        n_q_heads=32, n_kv_heads=8, head_dim=64,
        max_resident_blocks=48,
    )
    sys_prompt = list(range(100))            # 100-token system prompt
    q1 = sys_prompt + [200, 201, 202, 203]   # 104 tokens
    q2 = sys_prompt + [210, 211, 212]        # 103 tokens
    cache.append(seq_id=1, tokens=q1)
    cache.append(seq_id=2, tokens=q2)
    in_use = sum(1 for b in cache.pool.blocks if b.ref_count > 0)
    print(f"Blocks in use: {in_use} (should be ~7, not 14 — prefix is shared)")

Run this and you’ll see roughly seven blocks in use, not fourteen. That’s the prefix sharing paying off, and it’s exactly the kind of number you’d put on a CV: “Reduced KV cache memory by ~50% on multi-turn prompts via prefix-reuse trie with LRU eviction.”

Running and Testing It

The whole project is a single file plus a requirements.txt:

numpy==2.1.0
pytest==8.3.0

To exercise it:

python -m paged_kv_cache      # runs the demo, prints block-in-use stats
pytest -q                     # runs the test suite

The tests you should write before shipping this to a CV:

def test_prefix_sharing_reduces_blocks():
    cache = PagedKVCache(n_blocks=64, block_size=8, n_q_heads=8,
                         n_kv_heads=2, head_dim=16, max_resident_blocks=64)
    sys_prompt = list(range(40))
    cache.append(1, sys_prompt + [100, 101, 102])
    in_use_after_first = sum(1 for b in cache.pool.blocks if b.ref_count > 0)
    cache.append(2, sys_prompt + [200, 201])
    in_use_after_second = sum(1 for b in cache.pool.blocks if b.ref_count > 0)
    assert in_use_after_second - in_use_after_first <= 1  # tail only

def test_eviction_releases_cold_blocks():
    cache = PagedKVCache(n_blocks=8, block_size=4, n_q_heads=4,
                         n_kv_heads=2, head_dim=8, max_resident_blocks=4)
    for i in range(20):
        cache.append(seq_id=i, tokens=[i] * 4)
    in_use = sum(1 for b in cache.pool.blocks if b.ref_count > 0)
    assert in_use <= 4

def test_attention_matches_naive():
    # When n_groups == n_q_heads (MHA), our path is equivalent to a naive impl.
    ...

For benchmarking on real prompts, hook the append path into LlamaForCausalLM from HuggingFace, capture past_key_values, and measure three numbers: prefix-cache hit rate, end-to-end prefill latency on second turn of a multi-turn conversation, and memory residency in bytes. Those three numbers are what every LLM platform team tracks.

Extending It: Your Roadmap to Senior-Level

The toy is small. Each of the following upgrades turns a piece of it into something production-shaped — and gives you a concrete talking point in interviews.

  • Persistence to disk with LMDB or RocksDB. Serialize the trie to disk on shutdown, reload on boot. Matters because cold-start latency on long-system-prompt deployments is real, and showing that you’ve thought about warm starts is a strong signal.
  • Horizontal sharing via Redis. Treat the trie as a coordination layer across multiple inference workers. Matters because production serving fleets are sharded — and a shared prefix cache across workers is a non-trivial distributed-systems problem.
  • OpenTelemetry / Prometheus instrumentation. Emit kv_cache_hit_ratio, kv_cache_resident_bytes, evictions_total, block_pool_exhausted_total. Matters because observability is what separates a script from a service.
  • Fault tolerance: write-ahead log + snapshot. Replay trie mutations from a WAL on crash. Matters because every stateful serving system needs crash recovery, and interviewers love to ask about it.
  • Benchmark harness with vLLM and HuggingFace baselines. A side-by-side comparison of your cache vs. vLLM’s PagedAttention on the same prompts. Matters because benchmarks are how you turn an opinion into evidence.
  • Cost-aware eviction. Weight eviction by request value (e.g., the dollar cost of recomputing that prefix). Matters because the next generation of cache policies is moving from “least recently used” to “least-cost-to-recompute,” and being able to articulate that earns you a senior-level conversation.

Key Takeaways

  • A paged KV cache with GQA sharing and prefix-reuse eviction is a small, deep project — under 600 lines of Python with a real memory hierarchy, real eviction logic, and real benchmarks.
  • It signals the exact skills hiring managers look for in inference and ML-platform hires: allocator design, attention internals, cache policy, and measurement.
  • The same shape appears in vLLM’s PagedAttention, SGLang’s RadixAttention, and TensorRT-LLM’s in-flight batching — so reading their code after you ship yours is a direct path to mastery.
  • Ship the toy, then pick two of the six roadmap items and turn them into real upgrades. That’s a CV-grade portfolio piece.

Further Reading