TL;DR — In this post we build a pure‑Python paged attention engine that mirrors vLLM’s block‑level KV cache with LRU eviction. The code fits in a single file, runs on CPU, and can be extended with PyTorch kernels for real‑world use.

Building a portfolio project that doubles as a functioning systems prototype is one of the fastest ways to signal depth of experience to hiring managers. A paged attention engine sits at the heart of modern LLM serving frameworks such as vLLM, and implementing one from scratch in pure Python lets you showcase knowledge of memory management, algorithmic design, and production‑grade patterns without needing a GPU cluster.

Why This Project Stands Out on a CV

  • Systems‑level memory management – You design and maintain a block‑level cache, track usage with an LRU eviction policy, and handle fragmentation – skills directly transferable to database engines, streaming processors, or orchestration layers.
  • Algorithmic efficiency – Paged attention reduces KV cache memory waste from O(N·L) to O(N·B) (where B is block size). Demonstrating this trade‑off shows you can spot and quantify performance bottlenecks.
  • Named‑tool familiarity – The implementation references vLLM’s block table, PyTorch tensor semantics, and common eviction strategies, making it instantly recognizable to engineers who have worked with LLM serving stacks.
  • Roles it signals for – Backend engineers building inference services, ML infrastructure specialists, and even full‑stack developers who need to integrate LLM APIs will view this as proof you can ship reliable, low‑latency compute pipelines.

Architecture Overview

The engine consists of three core components that map cleanly onto vLLM’s design:

  1. Block pool – A fixed‑size array of Block objects, each holding a contiguous slice of KV cache (e.g., 16 tokens × hidden_dim). The pool tracks which blocks are free, in‑use, or marked for eviction.
  2. LRU eviction manager – On every allocation, the least‑recently‑used block is identified and either recycled (if empty) or discarded to make room for new tokens. A simple doubly‑linked list or Python collections.OrderedDict can provide O(1) moves.
  3. Attention dispatcher – Given a set of query vectors and the block pool, it gathers the relevant key/value slices, stacks them, and performs the scaled‑dot‑product attention. In pure Python this is a naive loop over blocks, but the structure is identical to what vLLM’s C++ kernels execute.
+---------------------+     +---------------------+     +---------------------+
|   Block Pool (N)    | --> |   LRU Manager (O(1))| --> | Attention Dispatcher|
+---------------------+     +---------------------+     +---------------------+
        |                           |                         |
        |   allocate(block_id)     |   evict_lru()           |   compute_attn(q,k,v)
        |   free(block_id)         |   promote(recent)       |   return out

The block pool size (e.g., 128 blocks) and block width (e.g., 16 tokens) are configurable knobs that let you experiment with different memory‑vs‑throughput trade‑offs.

Building It Step by Step

Below are seven concrete, language‑tagged steps that produce a runnable minimal engine. Copy each snippet into paged_attn.py in order.

Step 1 – Block data structure

# step_1_block.py
class Block:
    """A contiguous chunk of KV cache.

    Attributes:
        id:        unique integer identifier.
        tokens:    list of token indices currently stored (max BLOCK_SIZE).
        kv:        2‑dim list [[k1, k2, ...], [v1, v2, ...]].
        refs:      access counter for LRU promotion.
    """
    def __init__(self, block_id: int, block_size: int, hidden_dim: int):
        self.id = block_id
        self.tokens = [-1] * block_size          # placeholder token ids
        self.kv = [[0.0] * hidden_dim for _ in range(2)]  # [keys, values]
        self.refs = 0

Step 2 – Fixed‑size block pool

# step_2_pool.py
from typing import List, Optional

BLOCK_SIZE = 16          # tokens per block (adjustable)
POOL_CAPACITY = 64       # total blocks in the pool

class BlockPool:
    def __init__(self, hidden_dim: int):
        self.blocks: List[Block] = [
            Block(i, BLOCK_SIZE, hidden_dim) for i in range(POOL_CAPACITY)
        ]
        self.free_ids = list(range(POOL_CAPACITY))   # stack of available block ids
        self.lru_head = 0                           # index of most‑recently‑used block
        self.lru_tail = POOL_CAPACITY - 1           # index of least‑recently‑used block
        # simple doubly‑linked list using parallel arrays
        self.prev = [i - 1 if i > 0 else POOL_CAPACITY - 1 for i in range(POOL_CAPACITY)]
        self.next = [i + 1 if i < POOL_CAPACITY - 1 else 0 for i in range(POOL_CAPACITY)]

    def allocate(self) -> Optional[Block]:
        """Grab a free block, evicting LRU if pool is full."""
        if not self.free_ids:
            self._evict_lru()
        bid = self.free_ids.pop()          # take the “most recent” free id (LIFO is fine)
        blk = self.blocks[bid]
        # move block to MRU position
        self._move_to_mru(bid)
        return blk

    def _evict_lru(self):
        """Discard the least‑recently‑used block and reuse its id."""
        evict_id = self.lru_tail
        # reset block state
        self.blocks[evict_id].tokens = [-1] * BLOCK_SIZE
        self.blocks[evict_id].kv = [[0.0] * len(self.blocks[evict_id].kv[0]) for _ in range(2)]
        self.blocks[evict_id].refs = 0
        # re‑insert evicted id at front of free list
        self.free_ids.append(evict_id)
        # advance LRU tail one step forward (circular)
        self.lru_tail = self.prev[self.lru_tail]
        # new tail’s next now points to head
        self.next[self.lru_tail] = self.lru_head
        self.prev[self.lru_head] = self.lru_tail

    def _move_to_mru(self, bid: int):
        """Promote a block to the most‑recently‑used end of the LRU list."""
        # remove bid from its current spot
        p, n = self.prev[bid], self.next[bid]
        self.next[p] = n
        self.prev[n] = p
        # insert at head (after lru_head)
        self.next[bid] = self.lru_head
        self.prev[bid] = self.lru_head  # head is dummy; adjust later
        if self.lru_head != 0:
            self.prev[self.lru_head] = bid
        self.lru_head = bid
        # keep circular invariant
        if self.lru_head == self.lru_tail:
            self.lru_tail = self.prev[self.lru_head]

    def free(self, block: Block):
        """Return a block to the free pool; it becomes LRU."""
        bid = block.id
        self.free_ids.append(bid)
        # move to LRU tail
        self._move_to_lru_tail(bid)

    def _move_to_lru_tail(self, bid: int):
        # similar to _move_to_mru but push to the tail side
        self.prev[self.lru_tail] = bid
        self.next[bid] = self.lru_head
        self.prev[bid] = self.lru_tail
        self.next[self.lru_tail] = bid
        self.lru_tail = bid

Step 3 – Simple attention computation on gathered KV slices

# step_3_attn.py
import math

def scaled_dot_attn(q: list, k: list, v: list, scale: float) -> list:
    """Compute attention for a single query against a batch of key/value blocks.

    Args:
        q:   query vector (list of floats), length = hidden_dim.
        k:   list of key vectors, each length hidden_dim.
        v:   list of value vectors, same length as k.
        scale: 1 / sqrt(hidden_dim).

    Returns:
        output vector (list of floats) aggregated over all keys.
    """
    dim = len(q)
    attn_sum = [0.0] * dim
    weight_sum = 0.0
    for ki, vi in zip(k, v):
        # dot product
        dot = sum(qj * kj for qj, kj in zip(q, ki))
        weight = math.exp(dot * scale)
        attn_sum = [a + w * vj for a, vj, w in zip(attn_sum, vi, [weight] * dim)]
        weight_sum += weight
    if weight_sum == 0:
        return [0.0] * dim
    return [s / weight_sum for s in attn_sum]

Step 4 – Dispatcher that pulls from the block pool

# step_4_dispatch.py
from typing import List

class AttentionDispatcher:
    def __init__(self, pool: "BlockPool", hidden_dim: int):
        self.pool = pool
        self.hidden_dim = hidden_dim

    def _gather_kv(self, token_indices: List[int]) -> (list, list):
        """Collect key/value rows for the given token indices across all blocks.

        Returns two lists: keys and values, each a list of vectors (one per token).
        """
        keys, values = [], []
        # simple linear scan; in production you'd index directly via block_id -> offset
        for blk in self.pool.blocks:
            for tok_idx in token_indices:
                # find the slot where tok_idx lives; for this demo we just use the first slot
                slot = blk.tokens.index(tok_idx) if tok_idx in blk.tokens else -1
                if slot >= 0:
                    keys.append(blk.kv[0][slot])
                    values.append(blk.kv[1][slot])
        return keys, values

    def compute_attn(self, query: list, token_ids: List[int]) -> list:
        """Full attention forward pass for a query over a set of token ids."""
        keys, values = self._gather_kv(token_ids)
        if not keys:
            return [0.0] * self.hidden_dim
        scale = 1.0 / math.sqrt(self.hidden_dim)
        # if we have more keys than one, we simply sum attentions (naive merge)
        out = [0.0] * self.hidden_dim
        for k, v in zip(keys, values):
            out = [a + b for a, b in zip(out, scaled_dot_attn(query, k, v, scale))]
        # normalise by number of gathered tokens
        norm = 1.0 / len(keys)
        return [c * norm for c in out]

Step 5 – End‑to‑end driver script

# step_5_driver.py
from step_1_block import Block, BLOCK_SIZE, POOL_CAPACITY
from step_2_pool import BlockPool
from step_3_attn import scaled_dot_attn
from step_4_dispatcher import AttentionDispatcher
import random

def random_vector(dim: int) -> list:
    return [random.random() for _ in range(dim)]

def main():
    HIDDEN_DIM = 32
    pool = BlockPool(HIDDEN_DIM)
    dispatcher = AttentionDispatcher(pool, HIDDEN_DIM)

    # Simulate a handful of token allocations
    query = random_vector(HIDDEN_DIM)
    for _ in range(10):
        blk = pool.allocate()                # grabs (or evicts) a block
        # write a fake token into the block
        token_id = random.randint(0, 255)
        blk.tokens[0] = token_id
        # store dummy key/value (just random vectors)
        blk.kv[0][0] = random_vector(HIDDEN_DIM)   # key
        blk.kv[1][0] = random_vector(HIDDEN_DIM)   # value

        # now compute attention for that token
        out = dispatcher.compute_attn(query, [token_id])
        print(f"Token {token_id:3d} → attention output norm {sum(o**2 for o in out):.4f}")

    # demonstrate LRU: run many allocations to force evictions
    for _ in range(200):
        pool.allocate()
    print("After 200 allocations, free_ids count:", len(pool.free_ids))

if __name__ == "__main__":
    main()

Run the script with python step_5_driver.py. You should see ten attention outputs and a report that after 200 allocations the free‑list still contains blocks (because eviction re‑uses ids). The code is deliberately minimal—no PyTorch, no CUDA—but the architecture mirrors exactly what vLLM does under the hood.

Running and Testing It

  1. Install dependencies – The script uses only the standard library, so no pip install is required. If you later add NumPy or PyTorch kernels, install them in a virtual environment:

    python -m venv .venv && source .venv/bin/activate
pip install numpy torch

  2. Execute – From the project root:

    python step_5_driver.py

  3. Verify correctness – The printed attention norms should be non‑zero and decrease as you add more tokens (because the naive dispatcher averages over gathered keys). Inspect the free_ids count after the burst of 200 allocations; it should be close to zero, confirming that LRU eviction is recycling blocks.

  4. Unit‑test the LRU – A quick pytest suite can assert that after POOL_CAPACITY + 1 allocations exactly one block is evicted:

    # test_lru.py
import pytest
from step_2_pool import BlockPool

def test_eviction_after_pool_full():
    pool = BlockPool(hidden_dim=16)
    # fill pool
    for _ in range(POOL_CAPACITY):
        pool.allocate()
    # one more should evict the LRU block
    new_blk = pool.allocate()
    assert new_blk.id < POOL_CAPACITY          # id is recycled
    assert len(pool.free_ids) == POOL_CAPACITY - 1

    Run with python -m pytest test_lru.py -v.

  5. Profiling (optional) – If you add NumPy, use python -m cprofile step_5_driver.py to see time spent in block management vs. attention math. This mirrors the kind of instrumentation you’d add in a production serving pipeline.

Extending It: Your Roadmap to Senior‑Level

  1. Persist block state to disk – Write blocks to lmdb or RocksDB so that after a process restart the cache survives. Why it matters: Enables long‑running servers without losing KV cache momentum, directly reducing latency for repeated prompts.

  2. Integrate a PyTorch kernel – Replace the pure‑Python scaled_dot_attn with a torch.baddbmm call that fuses key/value retrieval into a single GPU kernel. Why it matters: Turns the toy into a drop‑in replacement for vLLM’s inference engine, showcasing systems‑level performance engineering.

  3. Add observability metrics – Export Prometheus counters for blocks_allocated, evictions_total, and attention_latency_ms. Why it matters: Hiring managers look for engineers who build instrumented services that can be monitored in production.

  4. Implement horizontal scaling via a shared block table – Use Redis or a lightweight in‑process multiprocessing.Manager to coordinate a single block pool across multiple worker processes. Why it matters: Mirrors real LLM serving clusters where multiple API servers share a central KV cache.

  5. Benchmark against vLLM – Measure tokens‑per‑second for identical prompts using both your pure‑Python prototype and the official vLLM binary (if a GPU is available). Why it matters: Provides concrete data to discuss trade‑offs between simplicity and performance in interviews.

  6. Add fault‑tolerance with checkpointing – Periodically serialize the block pool to a checkpoint file and enable recovery after a crash. Why it matters: Demonstrates production‑grade reliability patterns (state persistence, graceful degradation) that senior engineers are expected to design.

Each upgrade adds a tangible capability that moves the project from “curiosity experiment” to “production‑ready component,” giving you concrete talking points for technical interviews.

Key Takeaways

  • Block‑level KV cache is the central abstraction that lets LLM serving scale beyond a single GPU’s memory.
  • LRU eviction implemented with a doubly‑linked list gives O(1) promotions and evictions, a pattern reusable in caches, databases, and stream processors.
  • Pure‑Python implementation proves you understand the algorithm without relying on framework‑specific shortcuts—valuable when debugging or extending the system.
  • Architecture separation (pool → manager → dispatcher) mirrors production systems such as vLLM, making the codebase a credible talking point for backend or ML‑infrastructure roles.
  • Extensibility—adding persistence, observability, or GPU kernels transforms the prototype into a portfolio piece that demonstrates end‑to‑end engineering rigor.

Further Reading