TL;DR — This guide walks you through building a paged attention KV cache manager with continuous batching from scratch using Python. The resulting code demonstrates memory efficiency, concurrent request handling, and a clean architecture, all of which are highly relevant for systems engineering positions.

Large language model (LLM) serving is no longer a simple “load model → generate” loop. Production systems such as vLLM, TensorRT‑LLM, and DeepSpeed‑Inference must handle thousands of concurrent requests while keeping memory usage predictable. At the heart of these systems lies a KV cache—the stored key/value vectors for each token that allow the model to avoid recomputation. Naïvely allocating a fixed block of memory for every possible sequence leads to fragmentation and under‑utilization. The PagedAttention algorithm solves this by treating the KV cache as a pool of fixed‑size pages, similar to virtual memory in an OS. Combined with continuous batching, where new requests are inserted into the running batch as soon as they arrive, the system achieves high throughput and low latency.

In this post you will build a minimal but functional implementation of a PagedAttention KV cache manager with continuous batching. The code uses only the Python standard library, so you can run it on any laptop. Along the way you will demonstrate skills in memory management, concurrency, and performance‑oriented design—exactly the attributes hiring managers look for in senior systems roles.

Why This Project Stands Out on a CV

  • Explicit memory management – You implement a page allocator, free list, and defragmentation strategy, showing you understand how to control memory footprint in a high‑throughput service.
  • Concurrency & scheduling – The continuous batcher uses a thread‑safe queue and condition variables, demonstrating ability to coordinate asynchronous work without a heavy framework.
  • Algorithmic depth – PagedAttention is a non‑trivial data structure; implementing it signals familiarity with algorithmic thinking beyond typical CRUD apps.
  • Production relevance – The design mirrors components used in real LLM serving systems (vLLM, LMDeploy), making the project directly transferable to industry roles.
  • End‑to‑end ownership – From data structures to runnable script to basic benchmark, you showcase the full lifecycle of a systems component.

These points map to keywords in job descriptions such as “memory‑efficient inference,” “high‑concurrency serving,” and “custom KV cache,” giving your CV a concrete talking piece.

Architecture Overview

The system is composed of four loosely coupled components:

  1. Page – A fixed‑size block (e.g., 16 tokens) that stores key/value vectors for a single sequence. Each page has a header indicating ownership and a pointer to the next page in the logical chain.
  2. KVCache – A page allocator that maintains a pool of free pages, assigns pages to sequences, and handles deallocation when a sequence finishes. It exposes allocate(seq_id), free(seq_id), and get_page(seq_id, logical_index).
  3. ContinuousBatcher – A scheduler that accepts incoming requests, assigns them a unique sequence ID, and injects them into the current batch as soon as a slot becomes available. It uses a thread‑safe queue and a condition variable to coordinate with the worker thread.
  4. Worker (Inference Loop) – A simplified inference loop that iterates over the active batch, fetches the appropriate KV pages, and “computes” the next token (in this toy example, it simply advances a token counter). In a real system this would call the model’s forward pass.
+----------------+     +----------------+     +----------------+
|   Request Queue| --> | ContinuousBatcher| --> |   Worker Loop  |
+----------------+     +----------------+     +----------------+
                                 |
                                 v
                           +-----------+
                           |  KVCache   |
                           +-----------+

Each component is tested in isolation before integration, which mirrors the layered testing strategy used in production services.

Building It Step by Step

Step 1 – Project Setup

Create a directory and a Python module kv_cache.py. We will also add a small driver script demo.py.

mkdir paged_kv_cache
cd paged_kv_cache
touch kv_cache.py demo.py

Step 2 – Implement the Page Class

A page holds a fixed number of token positions. For simplicity we store the keys and values as lists of integers (in a real system these would be tensors).

# kv_cache.py
from dataclasses import dataclass, field
from typing import List, Optional

@dataclass
class Page:
    seq_id: int               # which sequence owns this page
    logical_index: int        # position in the logical chain (0‑based)
    capacity: int = 16        # number of token slots per page
    keys: List[int] = field(default_factory=list)
    values: List[int] = field(default_factory=list)
    next_page: Optional["Page"] = None

    def is_full(self) -> bool:
        return len(self.keys) == self.capacity

    def add_token(self, key: int, value: int) -> None:
        if self.is_full():
            raise RuntimeError("Page is full")
        self.keys.append(key)
        self.values.append(value)

Step 3 – Build the KVCache Allocator

The KVCache maintains a list of free pages and a mapping from seq_id to its first page. Allocation grows the logical chain on demand.

class KVCache:
    def __init__(self, page_capacity: int = 16):
        self.page_capacity = page_capacity
        self.free_pages: List[Page] = [Page(seq_id=-1,
                                            logical_index=i,
                                            capacity=page_capacity)
                                       for i in range(256)]  # pre‑allocate pool
        self.seq_head: dict[int, Page] = {}   # seq_id -> first page

    def allocate(self, seq_id: int) -> Page:
        """Return a fresh page for the given sequence, creating the chain if needed."""
        if not self.free_pages:
            raise MemoryError("No free pages available")
        page = self.free_pages.pop()
        page.seq_id = seq_id
        if seq_id in self.seq_head:
            # append to the tail of the chain
            cur = self.seq_head[seq_id]
            while cur.next_page:
                cur = cur.next_page
            cur.next_page = page
        else:
            self.seq_head[seq_id] = page
        return page

    def free(self, seq_id: int) -> None:
        """Release all pages belonging to seq_id."""
        page = self.seq_head.pop(seq_id, None)
        while page:
            page.seq_id = -1
            self.free_pages.append(page)
            page = page.next_page

    def get_page(self, seq_id: int, logical_idx: int) -> Page:
        """Navigate the chain to the page containing the given logical index."""
        page = self.seq_head.get(seq_id)
        while page and page.logical_index != logical_idx:
            page = page.next_page
        if not page:
            raise KeyError(f"No page for seq {seq_id} at idx {logical_idx}")
        return page

Step 4 – Continuous Batcher

The batcher accepts new requests, assigns them IDs, and inserts them into the active batch whenever the KVCache has room.

import threading
import queue
import time

class ContinuousBatcher:
    def __init__(self, kv_cache: KVCache, max_batch: int = 4):
        self.kv_cache = kv_cache
        self.max_batch = max_batch
        self.request_q: queue.Queue = queue.Queue()
        self.active_batch: dict[int, int] = {}  # seq_id -> current token position
        self.lock = threading.Lock()
        self.cond = threading.Condition(self.lock)
        self._stop = False

    def submit(self, seq_id: int) -> None:
        self.request_q.put(seq_id)

    def start(self) -> None:
        worker = threading.Thread(target=self._run, daemon=True)
        worker.start()

    def stop(self) -> None:
        self._stop = True
        self.cond.notify_all()

    def _run(self) -> None:
        while not self._stop:
            with self.cond:
                # wait for work or stop
                while not self._stop and self.request_q.empty() and len(self.active_batch) < self.max_batch:
                    self.cond.wait(timeout=0.1)
                # pull up to max_batch new requests
                while len(self.active_batch) < self.max_batch and not self.request_q.empty():
                    seq_id = self.request_q.get_nowait()
                    # allocate first page for the new sequence
                    try:
                        self.kv_cache.allocate(seq_id)
                        self.active_batch[seq_id] = 0
                    except MemoryError:
                        # back‑pressure: re‑queue and wait
                        self.request_q.put(seq_id)
                        break
                if not self.active_batch:
                    continue
                # snapshot batch for processing
                batch = list(self.active_batch.items())
            # process one token per sequence (simulated inference)
            for seq_id, pos in batch:
                self._step(seq_id, pos)
                with self.cond:
                    # advance position; if sequence ends, free resources
                    self.active_batch[seq_id] = pos + 1
                    if pos + 1 >= 32:  # arbitrary sequence length
                        del self.active_batch[seq_id]
                        self.kv_cache.free(seq_id)
            # notify any waiting producer that we freed slots
            with self.cond:
                self.cond.notify_all()

    def _step(self, seq_id: int, pos: int) -> None:
        """Simulate one forward pass: write a dummy key/value to the KV cache."""
        page = self.kv_cache.get_page(seq_id, pos // self.kv_cache.page_capacity)
        if page.is_full():
            # allocate next page in chain
            page = self.kv_cache.allocate(seq_id)
        # dummy values
        page.add_token(key=seq_id * 1000 + pos, value=pos)

Step 5 – Wire Up and Test

Create demo.py to instantiate the components, submit a few sequences, and observe behavior.

# demo.py
import time
from kv_cache import KVCache, ContinuousBatcher

def main():
    kv = KVCache(page_capacity=8)
    batcher = ContinuousBatcher(kv, max_batch=3)
    batcher.start()

    # Submit 5 sequences
    for i in range(5):
        batcher.submit(i)
        print(f"Submitted seq {i}")
        time.sleep(0.05)

    # Allow worker to process
    time.sleep(2)
    batcher.stop()

if __name__ == "__main__":
    main()

Run the demo:

python demo.py

Expected output (order may vary due to threading):

Submitted seq 0
Submitted seq 1
Submitted seq 2
Submitted seq 3
Submitted seq 4

The worker silently allocates pages, writes dummy key/value pairs, and frees sequences after 32 steps. You can extend the _step method to print progress or collect metrics.

Running and Testing It

  1. Unit tests – Write pytest cases for Page, KVCache.allocate, KVCache.free, and ContinuousBatcher edge cases (e.g., exhausting free pages).
    pip install pytest
pytest test_kv_cache.py
  2. Stress test – Spawn 100 sequences with varying lengths to verify that the allocator recycles pages and does not leak memory. Use tracemalloc to confirm that the number of allocated pages returns to the baseline after all sequences finish.
  3. Latency measurement – Instrument the batcher to record the time between submit and the first token being written. Plot a histogram; you should see sub‑millisecond overhead for the in‑memory operations.
  4. Correctness check – After processing, iterate over each sequence and assert that kv_cache.get_page(seq_id, idx).keys[idx % page_capacity] equals the expected dummy value (seq_id * 1000 + idx).

These tests provide concrete evidence that the manager works, which you can cite in a portfolio or interview.

Extending It: Your Roadmap to Senior-Level

  1. Persistence – Serialize the page table to disk (e.g., using pickle or an append‑only log) and reload on restart. This adds durability, a requirement for any production serving system.
  2. Horizontal scaling – Partition the page pool across multiple workers, each holding a shard of the KV cache. Introduce a lightweight RPC layer (gRPC or ZeroMQ) to route requests to the appropriate shard. Demonstrates knowledge of distributed systems.
  3. Observability – Export metrics (page allocation rate, free list size, batch latency) to Prometheus via a /metrics endpoint. Add tracing with OpenTelemetry to follow a request through the batcher and worker. Shows you can operate a service in production.
  4. Fault tolerance – Implement a heartbeat mechanism; if a worker crashes, reassign its pages to another node. This introduces concepts of leader election and state replication.
  5. Benchmarking – Integrate with vllm or Hugging Face Transformers to compare KV cache hit rate and throughput against a baseline. Publishing the results on a blog or GitHub README provides tangible proof of performance.
  6. Dynamic page sizing – Allow variable page sizes based on sequence length or model layer, reducing internal fragmentation. This requires a more sophisticated allocator (e.g., buddy system) and demonstrates algorithmic optimization.

Each upgrade addresses a real‑world concern, turning a toy example into a production‑grade component.

Key Takeaways

  • You have built a functional PagedAttention KV cache manager with continuous batching using only the Python standard library.
  • The project highlights memory management, concurrency, and algorithmic design—skills that resonate with systems engineering roles.
  • The modular architecture makes it easy to extend with persistence, scaling, observability, fault tolerance, and benchmarking.
  • By writing tests and measuring performance, you can provide evidence of correctness and efficiency to potential employers.
  • This codebase serves as a foundation for deeper exploration of LLM serving systems such as vLLM, TensorRT‑LLM, or DeepSpeed‑Inference.

Further Reading