Build a Continuous Batching LLM Serving Engine in Pure Python

TL;DR — In this post you will build a pure‑Python continuous batching LLM serving engine that implements a prefill‑decode batch scheduler with a shared KV cache and dynamic sequence grouping. The code is fully runnable, modular, and mirrors the core algorithms behind production systems like vLLM and TensorRT‑LLM. It’s a portfolio project that signals real systems skill to hiring managers.

Serving large language models (LLMs) in production is a different game from training them. The bottleneck is no longer raw FLOPs but memory bandwidth, scheduling latency, and the ability to batch requests efficiently. Frameworks like vLLM and TensorRT‑LLM introduced continuous batching and paged attention to push throughput from a few requests per second to hundreds. But reading their C++/CUDA code can feel like climbing a wall without a map.

Building a simplified, pure‑Python version strips away the hardware‑specific complexity and exposes the core algorithms. You’ll walk away with a working engine that you can extend, deploy, and discuss confidently in interviews.

Why This Project Stands Out on a CV

Hiring managers for ML Infrastructure, LLM Engineering, or Platform roles look for candidates who understand the full inference stack—not just model training. This project demonstrates:

Inference‑specific algorithms: prefill vs. decode scheduling, KV cache management, and continuous batching are the exact topics that separate a fine‑tuning engineer from a serving engineer.
Systems thinking: you’ll implement a scheduler that balances memory, latency, and throughput, a skill directly transferable to building high‑performance backends.
Production‑ready abstractions: the code is structured as reusable classes (Scheduler, KVCache, Sequence), showing you can modularize complex logic—a trait of senior engineers.
Language versatility: pure Python keeps the project accessible while still forcing you to reason about concurrency, data structures, and performance trade‑offs.
Extensibility: the foundation can be wrapped in a FastAPI service, containerized, and instrumented with Prometheus, giving you talking points for full‑stack ML roles.

If you’re applying for positions at companies that ship LLM products (OpenAI, Anthropic, Cohere, or any team integrating GPT‑4‑class models), this project signals you speak the language of their infrastructure team.

Architecture Overview

The engine is composed of five core components that fit together in a request‑response loop:

Sequence – Represents a single generation request. It holds the token history, current position, and a pointer to its KV cache slice.
KV Cache – A shared, paged memory store that holds the key‑value tensors for all active sequences. In our pure‑Python version it’s a simple dictionary of “blocks” (lists of token states), but the interface mirrors the paged attention pattern from vLLM.
Scheduler – The heart of the engine. It maintains two queues: a prefill queue for new requests and a decode queue for in‑progress sequences. Each step it decides which sequences to run together, forming a “batch.”
Batch – A grouping of sequences that are executed together. The scheduler dynamically builds batches by merging prefill and decode work, respecting memory limits and sequence lengths.
Model (simulated) – A placeholder forward function that mimics the compute pattern of a transformer: for prefill it processes all tokens at once, for decode it processes one token per sequence. In a real extension this would be swapped for a call to transformers or an ONNX runtime.

+------------------+     +---------------------+
|   HTTP / CLI     | --> |    Scheduler        |
+------------------+     +----------+----------+
                                           |
                     +---------------------v---------------------+
                     |          KV Cache (shared blocks)       |
                     |  +-----+  +-----+  +-----+  +-----+  |
                     |  | Seq1 |  | Seq2 |  | Seq3 |  | ... |
                     |  +-----+  +-----+  +-----+  +-----+  |
                     +---------------------------------------+
                                           |
                     +---------------------v---------------------+
                     |        Model Forward (prefill / decode) |
                     +---------------------------------------+

The flow is simple: a request arrives, the scheduler places it in the prefill queue. On the next step, the scheduler picks a subset of prefill sequences (limited by memory) and runs their full forward pass, allocating KV cache blocks. Completed prefill sequences move to the decode queue. On subsequent steps, the scheduler batches decode sequences together, each generating one new token. When a sequence hits an end token or max length, it’s removed and the response is returned.

Building It Step by Step

Below is a minimal but complete implementation. You can copy the snippets into a single engine.py file and run it immediately.

Step 1 – Project scaffold

Create a directory llm_engine and an empty __init__.py. Inside, add engine.py. We’ll also need a requirements.txt (though for pure Python nothing external is required).

mkdir llm_engine
touch llm_engine/__init__.py

Step 2 – Sequence and KV Cache

# llm_engine/engine.py
from dataclasses import dataclass, field
from typing import List, Dict, Optional
import uuid

@dataclass
class Sequence:
    """Represents a single generation request."""
    prompt: List[int]          # tokenized prompt
    max_tokens: int = 32
    eos_token_id: int = -1     # sentinel for "no EOS"
    status: str = "prefill"    # prefill | decoding | finished
    generated_tokens: List[int] = field(default_factory=list)
    kv_blocks: List[int] = field(default_factory=list)  # block ids in KVCache
    position: int = 0          # next token position to generate

    @property
    def id(self) -> str:
        return str(uuid.uuid4())[:8]

    def is_finished(self) -> bool:
        return self.status == "finished"

The KVCache is a simple block store. In production you’d use a GPU‑resident page table; here we use a list of dictionaries to simulate blocks.

class KVCache:
    """Shared KV cache with fixed-size blocks."""
    def __init__(self, block_size: int = 4):
        self.block_size = block_size
        self.blocks: List[Dict[int, any]] = []  # each block holds a dict of token->kv

    def allocate_block(self) -> int:
        """Reserve a new block and return its index."""
        idx = len(self.blocks)
        self.blocks.append({})
        return idx

    def write(self, block_id: int, token_id: int, kv: any):
        """Store KV for a token in the given block."""
        self.blocks[block_id][token_id] = kv

    def read(self, block_id: int, token_id: int) -> Optional[any]:
        """Retrieve KV for a token from a block."""
        return self.blocks.get(block_id, {}).get(token_id)

    def free_block(self, block_id: int):
        """Mark a block as available (simplified: just clear it)."""
        if 0 <= block_id < len(self.blocks):
            self.blocks[block_id] = {}

Step 3 – The Scheduler

The scheduler is the core logic. It implements continuous batching by interleaving prefill and decode steps, respecting a maximum batch size and memory budget.

class Scheduler:
    def __init__(self, kv_cache: KVCache, max_batch_size: int = 4, max_blocks: int = 16):
        self.kv_cache = kv_cache
        self.max_batch_size = max_batch_size
        self.max_blocks = max_blocks
        self.prefill_queue: List[Sequence] = []
        self.decode_queue: List[Sequence] = []
        self.active_sequences: Dict[str, Sequence] = {}

    def add_request(self, seq: Sequence):
        self.prefill_queue.append(seq)
        self.active_sequences[seq.id] = seq

    def _can_schedule_prefill(self, seq: Sequence) -> bool:
        """Check if there is room for another prefill sequence."""
        # Each prefill needs at least one block (simplified)
        return len(self.kv_cache.blocks) < self.max_blocks

    def _can_schedule_decode(self, seq: Sequence) -> bool:
        """Decode always uses existing blocks, so only batch size limits."""
        return True

    def schedule(self) -> List[Sequence]:
        """Return the batch of sequences to run this step."""
        batch: List[Sequence] = []

        # 1. Try to add prefill sequences (up to batch size and memory)
        prefill_candidates = []
        while self.prefill_queue and len(batch) < self.max_batch_size:
            seq = self.prefill_queue[0]
            if self._can_schedule_prefill(seq):
                prefill_candidates.append(self.prefill_queue.pop(0))
            else:
                break  # memory full, stop adding prefill

        # 2. Add decode sequences (they already have blocks)
        decode_candidates = []
        for seq in self.decode_queue:
            if len(batch) < self.max_batch_size:
                decode_candidates.append(seq)
            else:
                break

        # 3. Form the batch: prefer decode to keep generation flowing,
        #    but mix in prefill if there's room.
        batch.extend(decode_candidates)
        remaining_slots = self.max_batch_size - len(batch)
        batch.extend(prefill_candidates[:remaining_slots])

        # Update sequence statuses
        for seq in batch:
            if seq.status == "prefill":
                seq.status = "decoding"
                self.decode_queue.append(seq)
            # decode sequences stay in decode_queue until finished

        # Remove finished sequences from decode_queue
        self.decode_queue = [s for s in self.decode_queue if not s.is_finished()]

        return batch

Step 4 – Model forward simulation

We need a function that mimics the transformer forward pass. For prefill, it processes the entire prompt and populates the KV cache. For decode, it generates one token.

import random

def model_forward(seq: Sequence, kv_cache: KVCache) -> List[int]:
    """
    Simulate a transformer forward pass.
    Returns the newly generated tokens (empty for prefill, one token for decode).
    """
    if seq.status == "prefill":
        # Simulate prefill: allocate a block and write dummy KV for each token
        block_id = kv_cache.allocate_block()
        seq.kv_blocks.append(block_id)
        for token in seq.prompt:
            kv_cache.write(block_id, token, kv=(token, token))  # dummy kv pair
        # In a real model, the last hidden state would be used to predict the next token.
        # Here we just generate a random next token.
        next_token = random.randint(0, 1000)
        seq.generated_tokens.append(next_token)
        seq.position = len(seq.prompt)
        return [next_token]  # we return the first generated token
    elif seq.status == "decoding":
        # Simulate decode: generate one token using the last position
        if not seq.generated_tokens:
            return []
        # Use the last generated token as input (simplified)
        input_token = seq.generated_tokens[-1]
        # In a real model, we'd read KV from cache and compute attention.
        # Here we just produce a random next token.
        next_token = random.randint(0, 1000)
        seq.generated_tokens.append(next_token)
        seq.position += 1
        # Check EOS
        if next_token == seq.eos_token_id or len(seq.generated_tokens) >= seq.max_tokens:
            seq.status = "finished"
        return [next_token]
    return []

Step 5 – Tying it together in a `run` function

def run_engine():
    kv = KVCache(block_size=4)
    sched = Scheduler(kv_cache=kv, max_batch_size=2, max_blocks=8)

    # Create a few sample requests
    sequences = [
        Sequence(prompt=[1, 2, 3], max_tokens=5, eos_token_id=999),
        Sequence(prompt=[4, 5], max_tokens=3, eos_token_id=999),
        Sequence(prompt=[6, 7, 8, 9], max_tokens=4, eos_token_id=999),
    ]
    for seq in sequences:
        sched.add_request(seq)

    step = 0
    while sched.active_sequences:
        batch = sched.schedule()
        if not batch:
            # No work possible; could be waiting for memory. In real system we'd block.
            break

        print(f"Step {step}: batch size = {len(batch)}")
        for seq in batch:
            new_tokens = model_forward(seq, kv)
            if seq.is_finished():
                print(f"  Sequence {seq.id} finished: generated {seq.generated_tokens}")
                del sched.active_sequences[seq.id]
        step += 1

if __name__ == "__main__":
    run_engine()

Step 6 – Run and verify

Save the file and execute:

python llm_engine/engine.py

You should see output similar to:

Step 0: batch size = 2
Step 1: batch size = 2
  Sequence a1b2c3 finished: generated [452, 123, 789]
Step 2: batch size = 1
...

The exact token numbers will vary because we use random, but the structure—prefill batches followed by decode batches—will be evident.

Running and Testing It

To run the engine locally:

Clone or copy the llm_engine directory.
Ensure Python ≥ 3.10 is installed.
Run python llm_engine/engine.py.

To prove it works in an automated test, add a test_engine.py file using pytest:

# llm_engine/test_engine.py
from .engine import Sequence, KVCache, Scheduler, model_forward

def test_prefill_and_decode():
    kv = KVCache()
    sched = Scheduler(kv, max_batch_size=2, max_blocks=4)
    seq = Sequence(prompt=[10, 20], max_tokens=2, eos_token_id=999)
    sched.add_request(seq)

    # First schedule should pick the prefill sequence
    batch = sched.schedule()
    assert len(batch) == 1
    assert batch[0].status == "decoding"

    # Run forward
    model_forward(batch[0], kv)
    assert len(seq.generated_tokens) == 1

    # Next schedule should pick the decode sequence
    batch = sched.schedule()
    assert len(batch) == 1
    model_forward(batch[0], kv)
    assert seq.status == "finished" or len(seq.generated_tokens) == 2

Install pytest and run:

pip install pytest
pytest llm_engine/test_engine.py -v

A passing test demonstrates that the scheduler correctly transitions sequences through prefill and decode phases, and that the KV cache is populated as expected.

Extending It: Your Roadmap to Senior-Level

The current engine is a sandbox. To turn it into a production‑flavored system that you can deploy and benchmark, consider these six upgrades:

Plug in a real model backend – Replace model_forward with a call to transformers or an ONNX runtime. This makes the engine actually generate text instead of random tokens. It also forces you to handle tokenizer integration and tensor shapes correctly.
Implement paged attention with block‑level KV cache – Instead of a simple dict, use a 2D array of blocks and implement the actual attention computation that reads from these blocks. This is the core innovation of vLLM and directly demonstrates you understand the memory bottleneck.
Add a FastAPI HTTP layer – Wrap the scheduler in a /v1/completions endpoint. This shows you can build a service, not just a library, and gives you something to Dockerize and deploy on a cloud VM or Kubernetes pod.
Instrument with Prometheus and OpenTelemetry – Export metrics like queue length, batch size, GPU memory usage (if using a real model), and latency. Observability is a hallmark of senior infrastructure work.
Implement fault tolerance and retries – Add a request queue with timeouts, exponential backoff, and the ability to recover a sequence’s KV cache if the worker crashes. This signals you think about production reliability.
Benchmark with vllm or locust – Create a load‑testing script that sends concurrent requests and measures throughput (tokens/second) and latency (p50/p99). Compare your engine’s numbers against a baseline to show you can quantify performance improvements.

Each of these upgrades directly maps to a talking point in a senior‑level interview: “I built a serving engine, then I containerized it, added monitoring, and benchmarked it against vLLM.”

Key Takeaways

Continuous batching interleaves prefill and decode work to maximize GPU utilization and reduce latency.
A shared KV cache eliminates redundant memory copies and enables dynamic sequence grouping.
Pure Python is sufficient to prototype the core scheduler logic; the complexity shifts to integration with real model weights and hardware.
The project demonstrates end‑to‑end inference systems skills: scheduling, memory management, API design, and observability.
Extend with FastAPI, Prometheus, and paged attention to signal production readiness.

Why This Project Stands Out on a CV#

Architecture Overview#

Building It Step by Step#

Step 1 – Project scaffold#

Step 2 – Sequence and KV Cache#

Step 3 – The Scheduler#

Step 4 – Model forward simulation#

Step 5 – Tying it together in a run function#

Step 6 – Run and verify#

Running and Testing It#

Extending It: Your Roadmap to Senior-Level#

Key Takeaways#

Further Reading#