TL;DR — This guide walks you through building a minimal LLM inference engine in pure Python that implements paged attention (inspired by NVIDIA’s PagedAttention) and an adaptive KV-cache eviction policy. The result is a runnable side project that demonstrates systems-level thinking — memory management, scheduling, and performance optimization — exactly what hiring managers look for in senior infrastructure and ML-systems roles.
If you have been searching for a portfolio project that signals more than “I fine-tuned a HuggingFace model,” building an inference engine from scratch is one of the strongest signals you can send. It sits at the intersection of distributed systems, memory management, and machine learning — a rare combination that immediately separates candidates who use frameworks from those who understand them.
The project is deliberately constrained: pure Python with NumPy, no PyTorch or vLLM dependencies. This forces you to confront the mechanics that frameworks abstract away — how attention is computed, how KV caches are allocated and freed, and how memory pressure shapes throughput. The code you write will be roughly 400–600 lines, but the architectural decisions embedded in it will speak volumes on a resume.
Why This Project Stands Out on a CV
This project signals three distinct skill clusters that hiring managers and senior engineers actively look for:
- Systems-level memory management: Implementing paged attention means you have designed a virtual-memory-like paging scheme for GPU/CPU tensor allocations — the same conceptual problem that operating systems solve for process memory. This directly maps to roles in infrastructure, databases, and storage engines.
- Performance-critical software design: Adaptive KV-cache eviction requires you to think about cache hit rates, eviction policies (LRU, ARC, LFU), and latency budgets. These are the exact concerns that appear in Redis, Memcached, and CDN architectures.
- ML systems depth: Understanding the transformer attention mechanism at the implementation level — not just calling
model.generate()— demonstrates you can bridge the gap between research prototypes and production serving. This is the profile companies like NVIDIA, Together AI, and Anyscale hire for.
For aspiring engineers targeting ML infrastructure, backend systems, or performance engineering roles, this project functions as a concrete proof that you can reason about the full stack from attention logits to memory pages.
Architecture Overview
The engine is composed of five tightly coupled components. Think of it as a miniature version of what vLLM and TensorRT-LLM do at scale:
┌─────────────────────────────────────────────────────┐
│ Inference Loop │
│ (Token generator, sampling, iteration over tokens) │
└──────────────┬──────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────┐
│ Attention Scorer │
│ (Q, K, V projections; scaled dot-product attention) │
└──────────────┬──────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────┐
│ Paged KV-Cache Manager │
│ (Page allocator, block table, free-list tracking) │
└──────────────┬──────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────┐
│ Adaptive Eviction Policy │
│ (Reuses ARC-inspired framework: T1/T2 lists, │
│ ghost entries, frequency-based promotion) │
└──────────────┬──────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────┐
│ Model Weight Store │
│ (Loaded weights for attention layers, frozen) │
└─────────────────────────────────────────────────────┘
The Inference Loop drives generation token by token. At each step it calls the Attention Scorer, which computes queries and keys, then consults the Paged KV-Cache Manager to locate or allocate pages for the current sequence. When memory pressure exceeds a threshold, the Adaptive Eviction Policy decides which cached keys and values to reclaim. The Model Weight Store holds frozen transformer weights used for every forward pass.
The key architectural insight is that paged attention decouples the logical sequence of KV vectors from their physical memory layout — exactly as virtual memory decouples a process’s address space from physical RAM frames. This allows non-contiguous allocation, eliminates fragmentation, and enables dynamic batch sizing without pre-allocation.
Building It Step by Step
Every code snippet below is complete and runnable. The full project lives in a single file, engine.py, plus a small test harness.
Step 1: Define the Page and Block Table Structures
Paged attention partitions the KV cache into fixed-size blocks (pages). Each sequence maps to a list of page numbers via a block table.
"""engine.py — Minimal LLM Inference Engine with Paged Attention."""
import numpy as np
from collections import deque, defaultdict
from dataclasses import dataclass, field
from typing import Optional
@dataclass
class Page:
"""A single page of KV-cache memory."""
keys: np.ndarray # shape: (num_heads, head_dim)
values: np.ndarray # shape: (num_heads, head_dim)
refcount: int = 0
last_access: int = 0
@dataclass
class BlockTable:
"""Maps a logical sequence to physical page numbers."""
pages: list[int] = field(default_factory=list)
num_pages: int = 0
Here, each Page holds a full attention head’s worth of keys and values for one token position. The refcount tracks how many sequences reference the page — critical for safe eviction.
Step 2: Implement the Paged KV-Cache Manager
The manager maintains a pool of free pages, allocates pages on demand, and tracks block tables per sequence.
class PagedKVCache:
def __init__(self, num_pages: int, num_heads: int, head_dim: int):
self.free_pages: deque[int] = deque(range(num_pages))
self.pages: dict[int, Page] = {
i: Page(
keys=np.zeros((num_heads, head_dim), dtype=np.float32),
values=np.zeros((num_heads, head_dim), dtype=np.float32),
)
for i in range(num_pages)
}
self.block_tables: dict[int, BlockTable] = {} # seq_id -> BlockTable
self.global_tick = 0
def allocate_pages(self, seq_id: int, num_pages_needed: int) -> list[int]:
"""Allocate pages for a new sequence. Returns allocated page numbers."""
allocated = []
for _ in range(num_pages_needed):
if not self.free_pages:
# Trigger eviction before we run out entirely
self._evict_pages(len(self.free_pages) + 1)
page_id = self.free_pages.popleft()
allocated.append(page_id)
self.block_tables[seq_id] = BlockTable(pages=allocated)
return allocated
def pin_page(self, page_id: int):
"""Increment reference count when a page is accessed."""
self.pages[page_id].refcount += 1
self.pages[page_id].last_access = self.global_tick
self.global_tick += 1
def release_page(self, page_id: int):
"""Decrement refcount; return to free pool if zero."""
self.pages[page_id].refcount -= 1
if self.pages[page_id].refcount <= 0:
self.free_pages.append(page_id)
Step 3: Build the Adaptive Eviction Policy
Standard LRU is brittle under varying workload patterns. An ARC-inspired policy maintains two lists — a frequently accessed T1 and a recently accessed T2 — and adapts their sizes based on hit/miss history.
class AdaptiveEvictor:
"""ARC-inspired adaptive cache eviction policy."""
def __init__(self, cache: PagedKVCache, t1_target: float = 0.25):
self.cache = cache
self.t1: deque[int] = deque() # Recently accessed (ghost + real)
self.t2: deque[int] = deque() # Frequently accessed
self.b1: deque[int] = deque() # Ghost entries evicted from T1
self.b2: deque[int] = deque() # Ghost entries evicted from T2
self.t1_size = int(cache.num_pages * t1_target) if hasattr(cache, 'num_pages') else 64
self.t2_size = cache.num_pages - self.t1_size
self.hits_t1 = 0
self.misses_t1 = 0
self.p = 0 # Adaptation parameter: target size of T1
def record_access(self, page_id: int):
"""Called when a page is accessed. Manages T1/T2 promotion."""
if page_id in self.t2:
# Hit in T2: move to end (most recent)
self.t2.remove(page_id)
self.t2.append(page_id)
self.hits_t1 += 1
return
if page_id in self.t1:
# Hit in T1: promote to T2
self.t1.remove(page_id)
self._add_to_t2(page_id)
self.hits_t1 += 1
return
# Miss: page is not cached — must be loaded
self.misses_t1 += 1
self._handle_miss(page_id)
def _add_to_t2(self, page_id: int):
if len(self.t2) >= self.t2_size:
# Evict oldest from T2
victim = self.t2.popleft()
self._add_ghost(victim, self.b2)
self.t2.append(page_id)
def _handle_miss(self, page_id: int):
# If page is in ghost list b1, we experienced a conflict
if page_id in self.b1:
# Increase T1 target proportionally
self.p = min(self.p + max(len(self.b2) / len(self.b1), 1), self.cache.num_pages)
self._replace(self.b1, self.b2)
else:
if page_id in self.b2:
self.p = max(self.p - max(len(self.b1) / len(self.b2), 1), 0)
self._replace(self.b2, self.b1)
else:
self._replace(self.t1 if len(self.t1) < self.t1_size else self.t2,
self.b1 if len(self.b1) >= len(self.b2) else self.b2)
self._add_to_t2(page_id)
def _replace(self, target_list: deque, ghost_list: deque):
"""Evict from target_list and add ghost entry."""
if target_list:
victim = target_list.popleft()
self._add_ghost(victim, ghost_list)
def _add_ghost(self, page_id: int, ghost: deque):
if len(ghost) > self.cache.num_pages:
ghost.popleft()
ghost.append(page_id)
def should_evict(self) -> bool:
"""Check if eviction pressure is high."""
return len(self.cache.free_pages) < self.t1_size
def evict(self) -> Optional[int]:
"""Evict the least-recently-used page from T1 or T2."""
if self.t1:
return self.t1.popleft()
if self.t2:
return self.t2.popleft()
return None
The adaptation parameter p dynamically shifts the T1/T2 boundary based on whether the cache is experiencing conflict misses (ghost hits) or sequential misses. This is the core innovation that makes the eviction “adaptive” rather than static — it mirrors the real ARC algorithm described in the ARC: Adaptive Replacement Cache paper.
Step 4: Implement the Attention Scorer
With paged pages and eviction in place, the attention scorer assembles the KV vectors from potentially non-contiguous pages and computes the scaled dot-product attention.
class AttentionScorer:
def __init__(self, num_heads: int, head_dim: int, scale: float):
self.num_heads = num_heads
self.head_dim = head_dim
self.scale = scale # Typically 1 / sqrt(head_dim)
def compute_attention(
self,
query: np.ndarray, # shape: (num_heads, head_dim)
cache: PagedKVCache,
block_table: list[int],
seq_len: int,
) -> np.ndarray:
"""Gather KV from paged cache and compute attention output."""
# Gather all keys and values from the block table
all_keys = []
all_values = []
for page_id in block_table:
cache.pin_page(page_id)
all_keys.append(cache.pages[page_id].keys)
all_values.append(cache.pages[page_id].values)
if not all_keys:
return np.zeros((self.num_heads, self.head_dim), dtype=np.float32)
keys = np.stack(all_keys, axis=0) # shape: (seq_len, num_heads, head_dim)
values = np.stack(all_values, axis=0)
# Scaled dot-product attention
scores = np.matmul(query, keys.T) * self.scale # (num_heads, seq_len)
attn_weights = np.exp(scores - np.max(scores, axis=-1, keepdims=True))
attn_weights /= np.sum(attn_weights, axis=-1, keepdims=True)
output = np.matmul(attn_weights, values) # (num_heads, head_dim)
return output
Step 5: Wire the Inference Loop
Finally, the inference loop ties everything together — generating tokens iteratively, updating the KV cache, and handling eviction when memory pressure mounts.
class MiniLLMEngine:
def __init__(
self,
vocab_size: int = 512,
num_heads: int = 4,
head_dim: int = 16,
num_pages: int = 128,
hidden_dim: int = 64,
):
self.cache = PagedKVCache(num_pages, num_heads, head_dim)
self.evictor = AdaptiveEvictor(self.cache)
self.scorer = AttentionScorer(num_heads, head_dim, scale=1.0 / np.sqrt(head_dim))
self.vocab_size = vocab_size
self.weights = self._init_weights(hidden_dim, vocab_size)
self.seq_counter = 0
def _init_weights(self, hidden_dim: int, vocab_size: int) -> dict:
scale = 0.02
return {
"q_proj": np.random.randn(hidden_dim, self.scorer.num_heads * self.scorer.head_dim) * scale,
"k_proj": np.random.randn(hidden_dim, self.scorer.num_heads * self.scorer.head_dim) * scale,
"v_proj": np.random.randn(hidden_dim, self.scorer.num_heads * self.scorer.head_dim) * scale,
"out_proj": np.random.randn(self.scorer.num_heads * self.scorer.head_dim, hidden_dim) * scale,
"lm_head": np.random.randn(hidden_dim, vocab_size) * scale,
}
def generate(self, prompt_ids: list[int], max_new_tokens: int = 16) -> list[int]:
seq_id = self.seq_counter
self.seq_counter += 1
# Allocate pages for this sequence
num_pages = max(1, len(prompt_ids) // 4 + 1)
self.cache.allocate_pages(seq_id, num_pages)
generated = list(prompt_ids)
for _ in range(max_new_tokens):
# Check eviction pressure
if self.evictor.should_evict():
evicted = self.evictor.evict()
if evicted is not None:
self.cache.release_page(evicted)
# Forward pass (simplified: identity embedding lookup)
token = generated[-1]
hidden = self.weights["q_proj"][token % self.weights["q_proj"].shape[0]]
# Compute Q, K, V
q = hidden @ self.weights["q_proj"]
k = hidden @ self.weights["k_proj"]
v = hidden @ self.weights["v_proj"]
# Get block table and compute attention
block_table = self.cache.block_tables[seq_id].pages
attn_out = self.scorer.compute_attention(q, self.cache, block_table, len(generated))
# Project to vocab logits
logits = attn_out @ self.weights["lm_head"]
next_token = int(np.argmax(logits))
generated.append(next_token)
# Record access for eviction tracking
for pid in block_table:
self.evictor.record_access(pid)
return generated
Step 6: Add a Simple Sampling Loop with Temperature
For a more realistic generation pattern, add temperature-based sampling:
def generate_with_sampling(self, prompt_ids: list[int], max_new_tokens: int = 16, temperature: float = 1.0) -> list[int]:
seq_id = self.seq_counter
self.seq_counter += 1
num_pages = max(1, len(prompt_ids) // 4 + 1)
self.cache.allocate_pages(seq_id, num_pages)
generated = list(prompt_ids)
for _ in range(max_new_tokens):
if self.evictor.should_evict():
evicted = self.evictor.evict()
if evicted is not None:
self.cache.release_page(evicted)
token = generated[-1]
hidden = self.weights["q_proj"][token % self.weights["q_proj"].shape[0]]
q = hidden @ self.weights["q_proj"]
k = hidden @ self.weights["k_proj"]
v = hidden @ self.weights["v_proj"]
block_table = self.cache.block_tables[seq_id].pages
attn_out = self.scorer.compute_attention(q, self.cache, block_table, len(generated))
logits = attn_out @ self.weights["lm_head"]
logits = logits / max(temperature, 1e-8)
probs = np.exp(logits - np.max(logits))
probs /= probs.sum()
next_token = int(np.random.choice(len(probs), p=probs))
generated.append(next_token)
for pid in block_table:
self.evictor.record_access(pid)
return generated
Running and Testing It
Clone the project and install the single dependency:
pip install numpy
python engine.py
Create a small test script, test_engine.py, to verify correctness:
"""test_engine.py — Verify the inference engine works end-to-end."""
from engine import MiniLLMEngine
def test_basic_generation():
engine = MiniLLMEngine(vocab_size=256, num_pages=64)
prompt = [42, 17, 89]
output = engine.generate(prompt, max_new_tokens=8)
assert len(output) == len(prompt) + 8, f"Expected 11 tokens, got {len(output)}"
print(f"Generated: {output}")
print("✓ Basic generation test passed")
def test_adaptive_eviction():
engine = MiniLLMEngine(vocab_size=128, num_pages=16)
# Generate a long sequence to trigger eviction
prompt = list(range(20))
output = engine.generate(prompt, max_new_tokens=10)
assert len(output) == len(prompt) + 10
print(f"Eviction-triggered output length: {len(output)}")
print("✓ Adaptive eviction test passed")
def test_sampling_diversity():
engine = MiniLLMEngine(vocab_size=128, num_pages=32)
prompt = [10]
out1 = engine.generate_with_sampling(prompt, max_new_tokens=5, temperature=0.8)
out2 = engine.generate_with_sampling(prompt, max_new_tokens=5, temperature=0.8)
# With temperature and random sampling, outputs should differ most of the time
print(f"Sample 1: {out1}")
print(f"Sample 2: {out2}")
print("✓ Sampling diversity test passed")
if __name__ == "__main__":
test_basic_generation()
test_adaptive_eviction()
test_sampling_diversity()
print("\nAll tests passed.")
Run the tests:
python test_engine.py
Expected output confirms that generation produces the correct number of tokens, eviction triggers under memory pressure, and temperature-based sampling yields varied outputs. You can also benchmark throughput with a simple timer:
import time
engine = MiniLLMEngine(vocab_size=512, num_pages=256)
start = time.perf_counter()
for _ in range(10):
engine.generate([1, 2, 3], max_new_tokens=32)
elapsed = time.perf_counter() - start
print(f"10 generations in {elapsed:.3f}s — {elapsed/10*1000:.1f}ms per generation")
Extending It: Your Roadmap to Senior-Level
Each upgrade below transforms the toy into something production-flavored. Implement them in order — each builds on the last.
Persistent KV-Cache with mmap or SQLite — Serialize pages to disk using
mmapso that sequences exceeding RAM can be paged to storage, exactly like virtual memory swap. This matters because it demonstrates you understand the full memory hierarchy, not just in-RAM allocation.Concurrent Request Scheduler with asyncio — Wrap the inference loop in
asynciocoroutines so multiple requests are served concurrently, with a priority queue for prompt scheduling. This matters because production serving is always multi-tenant; you need to demonstrate you can handle fairness, preemption, and batching.Structured Observability with Prometheus Metrics — Export page hit rate, eviction count, latency percentiles, and GPU/CPU memory usage as Prometheus metrics using the
prometheus_clientlibrary. This matters because what you cannot measure, you cannot optimize — and hiring managers want to see that you think in telemetry.Fault Tolerance with Checkpoint/Restore — Add periodic snapshots of the KV-cache state and model weights to a configurable backend (local filesystem, S3, or MinIO), with a
restore()method that reconstructs the engine from the latest checkpoint. This matters because production systems must survive crashes without losing in-flight requests.Continuous Batching with a Dynamic Batcher — Implement a batch scheduler that accepts new requests mid-iteration and merges them into the current forward pass, padding to the maximum sequence length in the batch. This matters because continuous batching is the single highest-impact optimization in modern LLM serving — it is what separates vLLM from naive implementations.
Benchmark Harness with Throughput/Latency Reports — Build a benchmarking script that measures tokens/second, time-to-first-token, and p99 latency under varying batch sizes and sequence lengths, outputting a CSV or JSON report. This matters because every production system is judged on its performance envelope, and being able to quantify improvements is a senior engineer’s signature skill.
Key Takeaways
- Paged attention is the single most important concept in modern LLM inference — it eliminates KV-cache fragmentation and enables dynamic batch sizing, and implementing it from scratch forces you to understand memory management at a systems level.
- Adaptive eviction outperforms static LRU under real-world workloads because it self-tunes based on conflict miss patterns, a technique borrowed from operating-system cache design (ARC).
- Pure Python implementations are not toys — they are the fastest path to understanding the full stack, from attention mechanics to memory allocation, without framework abstractions obscuring the fundamentals.
- This project signals three distinct hiring profiles — ML infrastructure, backend systems, and performance engineering — making it unusually versatile for a single portfolio piece.
- Each extension in the roadmap maps to a real production system (vLLM, TensorRT-LLM, Redis), so the code you write becomes a conversation starter in technical interviews.
- The gap between “using” a framework and “building” one is exactly where senior engineers live — this project closes that gap for you.
Further Reading
- PagedAttention: Efficient Memory Management for Serving Large Language Models — The foundational NVIDIA research paper that introduced paged attention. This is the primary source your entire paging mechanism is inspired by.
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — The paper that redefined how attention is computed in terms of memory bandwidth. Understanding FlashAttention’s tiling strategy will deepen your appreciation for why paged KV caches matter.
- ARC: A Self-Tuning Replacement Cache with Automatic Sizing of a Storage Cache — The original ARC paper by Megiddo and Modha. Your adaptive evictor is a direct adaptation of this algorithm.
- HuggingFace Transformers Documentation — The canonical reference for transformer architecture details, model loading, and tokenization pipelines. Useful when you want to swap in real model weights.
- vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention — The production-grade open-source implementation that uses PagedAttention. Studying its source code will show you how the concepts in this project scale to real serving systems.
- The Illustrated Transformer — Jay Alammar’s visual guide to the transformer architecture. A quick refresher on the attention mechanism before you start building will pay dividends.