TL;DR — Build a pluggable KV cache eviction engine that lets you extend a model’s effective context length while bounding GPU memory. The implementation is pure Python, runs on a single machine, and demonstrates cache coherence, policy swapping, and benchmarking—skills that map directly to senior systems roles.

In production LLM serving, the KV cache is the primary memory consumer; without a smart eviction policy you either OOM or truncate context, hurting quality. This post walks you through a from‑scratch KV cache manager that supports interchangeable eviction strategies (LRU, FIFO, and a simple learned score), integrates with a minimal transformer forward pass, and ships with tests that prove correctness. You’ll end up with a portfolio piece that showcases systems design, algorithmic thinking, and hands‑on Python engineering.

Why This Project Stands Out on a CV

  • Systems‑oriented thinking – You’re not just calling an API; you’re designing a memory subsystem with clear contracts, policy abstraction, and observability hooks.
  • Algorithmic depth – Implementing eviction policies forces you to reason about cache hits/misses, temporal locality, and trade‑offs between latency and accuracy.
  • Production relevance – Real‑world LLM deployments (e.g., OpenAI, Anthropic, internal corporate models) rely on KV cache management; showing you can build it signals readiness for ML infrastructure roles.
  • Scalability narrative – The architecture is written so you can later plug in a distributed cache (Redis, Memcached) or a GPU‑aware allocator, giving you a story for senior‑level talks.
  • Testing rigor – Unit tests for cache correctness, benchmark scripts for throughput, and a CI‑friendly workflow demonstrate engineering discipline.

Architecture Overview

The engine is composed of four loosely coupled components:

  1. KV Store – A dictionary‑like structure that holds (key, value) pairs for each token position. It supports get, set, and evict.
  2. Eviction Policy – An interface with a single method should_evict(store, key) -> bool. Concrete implementations include LRU, FIFO, and ScoreBased.
  3. Cache Manager – Orchestrates reads/writes, invokes the policy, and emits metrics (hits, misses, evictions).
  4. Transformer Adapter – A thin wrapper that calls the model’s forward pass, intercepts the KV tensors, and routes them through the cache manager.

A simplified diagram (textual):

+----------------+      +-------------------+
|  Transformer   | ---> |  Cache Manager    |
|  (forward)     |      |  - KV Store       |
+----------------+      |  - Eviction Policy|
                        |  - Metrics        |
                        +-------------------+

The manager exposes a context_window parameter that limits the maximum number of cached tokens; when the limit is exceeded, the policy decides which entries to drop.

Building It Step by Step

1. Project Skeleton

mkdir kv_cache_engine
cd kv_cache_engine
python -m venv venv
source venv/bin/activate
pip install torch pytest

Create kv_cache/__init__.py, kv_cache/store.py, kv_cache/policy.py, kv_cache/manager.py, and tests/test_cache.py.

2. KV Store

# kv_cache/store.py
from collections import OrderedDict
from typing import Any, Tuple

class KVStore:
    """Thread‑safe OrderedDict for KV pairs."""
    def __init__(self, max_size: int):
        self.max_size = max_size
        self._store = OrderedDict()
        self._lock = __import__('threading').Lock()

    def get(self, key: int) -> Tuple[Any, Any] | None:
        with self._lock:
            if key in self._store:
                # Move to end to support LRU
                self._store.move_to_end(key)
                return self._store[key]
            return None

    def set(self, key: int, value: Tuple[Any, Any]) -> None:
        with self._lock:
            if key in self._store:
                self._store.move_to_end(key)
            self._store[key] = value
            if len(self._store) > self.max_size:
                # Eviction handled by policy; here we just pop the first item
                self._store.popitem(last=False)

    def evict(self, key: int) -> None:
        with self._lock:
            self._store.pop(key, None)

    def __len__(self) -> int:
        return len(self._store)

3. Eviction Policy Interface

# kv_cache/policy.py
from abc import ABC, abstractmethod
from .store import KVStore

class EvictionPolicy(ABC):
    """Strategy pattern for deciding which KV entries to drop."""
    @abstractmethod
    def should_evict(self, store: KVStore) -> list[int]:
        """Return a list of keys that should be evicted now."""

LRU Implementation

class LRUPolicy(EvictionPolicy):
    def should_evict(self, store: KVStore) -> list[int]:
        # The KVStore already moves accessed keys to end;
        # the first items are least recently used.
        if len(store) > store.max_size:
            # Return the oldest entry
            return [next(iter(store._store))]
        return []

FIFO Implementation

class FIFOPolicy(EvictionPolicy):
    def should_evict(self, store: KVStore) -> list[int]:
        if len(store) > store.max_size:
            # First inserted is at the beginning
            return [next(iter(store._store))]
        return []

Score‑Based (Learned) Policy

class ScorePolicy(EvictionPolicy):
    def __init__(self, scores: dict[int, float]):
        self.scores = scores  # key -> importance score

    def should_evict(self, store: KVStore) -> list[int]:
        if len(store) <= store.max_size:
            return []
        # Sort keys by ascending score; evict the lowest
        sorted_keys = sorted(store._store.keys(), key=lambda k: self.scores.get(k, 0.0))
        # Evict enough to bring size back under limit
        to_evict = sorted_keys[: len(store) - store.max_size]
        return to_evict

4. Cache Manager

# kv_cache/manager.py
from .store import KVStore
from .policy import EvictionPolicy
from collections import defaultdict
import time

class CacheManager:
    def __init__(self, max_tokens: int, policy: EvictionPolicy):
        self.store = KVStore(max_size=max_tokens)
        self.policy = policy
        self.metrics = defaultdict(int)

    def retrieve(self, key: int):
        start = time.perf_counter()
        hit = self.store.get(key)
        self.metrics['latency_retrieve'] += time.perf_counter() - start
        if hit:
            self.metrics['hits'] += 1
        else:
            self.metrics['misses'] += 1
        return hit

    def insert(self, key: int, kv_pair):
        self.store.set(key, kv_pair)
        # After insertion, ask policy if eviction is needed
        evict_keys = self.policy.should_evict(self.store)
        for k in evict_keys:
            self.store.evict(k)
            self.metrics['evictions'] += 1

    def snapshot(self):
        return dict(self.metrics)

5. Transformer Adapter (Minimal)

# kv_cache/adapter.py
import torch
from .manager import CacheManager

class KVCacheAdapter:
    """Wraps a transformer forward pass and routes KV through CacheManager."""
    def __init__(self, model, max_tokens: int, policy):
        self.model = model
        self.cache = CacheManager(max_tokens, policy)

    def forward(self, input_ids: torch.LongTensor):
        # In a real model you would hook into the attention layers.
        # Here we simulate KV extraction for illustration.
        batch_size, seq_len = input_ids.shape
        # Dummy KV generation (replace with actual model hooks)
        keys = torch.randn(batch_size, seq_len, 64)
        values = torch.randn(batch_size, seq_len, 64)

        # For each token position, cache the KV pair
        for pos in range(seq_len):
            key = pos  # using position as key for simplicity
            kv = (keys[:, pos, :], values[:, pos, :])
            self.cache.insert(key, kv)

        # Continue with model forward
        outputs = self.model(input_ids)
        return outputs

6. Putting It Together

# example_run.py
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from kv_cache.adapter import KVCacheAdapter
from kv_cache.policy import LRUPolicy

model = GPT2LMHeadModel.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

adapter = KVCacheAdapter(model, max_tokens=512, policy=LRUPolicy())

prompt = "Once upon a time,"
input_ids = tokenizer(prompt, return_tensors='pt')['input_ids']

with torch.no_grad():
    adapter.forward(input_ids)

print(adapter.cache.snapshot())

Running and Testing It

  1. Install dependencies – pip install torch transformers pytest
  2. Run the example – python example_run.py should print a metrics snapshot showing hits, misses, and evictions (initially zeros, but will populate after repeated calls).
  3. Unit tests – pytest tests/test_cache.py verifies that:
    • The store respects max_size.
    • LRU policy evicts the least recently used entry.
    • FIFO policy evicts the oldest entry.
    • Score policy evicts the lowest‑scoring keys.

Example test snippet:

# tests/test_cache.py
from kv_cache.store import KVStore
from kv_cache.policy import LRUPolicy, FIFOPolicy, ScorePolicy
from kv_cache.manager import CacheManager

def test_lru_eviction():
    store = KVStore(max_size=2)
    policy = LRUPolicy()
    manager = CacheManager(2, policy)
    manager.insert(1, ('k1', 'v1'))
    manager.insert(2, ('k2', 'v2'))
    # Access key 1 to make it recently used
    manager.retrieve(1)
    # Insert new key, should evict key 2
    manager.insert(3, ('k3', 'v3'))
    assert manager.store.get(2) is None
    assert manager.store.get(1) is not None
    assert manager.store.get(3) is not None

Running the test suite should produce all green, proving correctness.

Extending It: Your Roadmap to Senior-Level

  1. Persisted Cache – Serialize the KV store to disk (e.g., using pickle or torch.save) and load on startup. Why it matters: Enables warm‑start across model restarts, reducing latency for repeated prompts.
  2. Distributed Cache – Replace the in‑memory KVStore with a Redis cluster or a custom gRPC service. Why it matters: Allows horizontal scaling across multiple GPU workers, each sharing a global cache.
  3. Observability – Export metrics (hits, misses, eviction count) to Prometheus via a /metrics endpoint and visualize in Grafana. Why it matters: Production systems need visibility to tune cache size and policy.
  4. Fault Tolerance – Implement a write‑ahead log (WAL) for the cache so that after a crash you can recover recent entries. Why it matters: Guarantees durability without sacrificing too much throughput.
  5. Adaptive Policy – Train a small neural network to predict eviction scores based on token position, frequency, and attention weights. Why it matters: Moves from static heuristics to learned, context‑aware eviction, improving effective context length.
  6. Benchmark Harness – Create a script that measures end‑to‑end latency, memory usage, and perplexity impact across different policies and cache sizes. Why it matters: Provides data‑driven evidence for choosing a policy in a real deployment.

Key Takeaways

  • A KV cache manager is the cornerstone of scalable LLM serving; mastering it signals deep systems expertise.
  • Pluggable eviction policies (LRU, FIFO, learned) let you trade off between simplicity and context quality.
  • Real‑world impact comes from adding persistence, distribution, observability, and fault tolerance.
  • The project is a springboard: you can evolve it into a production‑grade cache service or integrate it with existing model servers.
  • Demonstrating this on a CV shows you can design, implement, and benchmark memory‑critical components.

Further Reading