TL;DR — This project implements a continuous batching LLM serving engine in pure Python, featuring a custom KV cache and token streaming. It demonstrates real systems skills such as memory management, concurrency, and latency optimization, making it a compelling portfolio piece for engineering roles.

In the race to ship generative AI features, the inference engine is often the bottleneck that determines whether a product feels responsive or sluggish. Building a performant serving stack from scratch forces you to confront low‑level concerns—memory layout, scheduling, and network I/O—that many engineers only read about. This guide walks you through constructing a minimal yet functional continuous batching LLM server with a custom KV cache and token streaming, all in pure Python. The result is a portfolio project that signals deep systems understanding to hiring managers looking for evidence of production‑ready engineering.

Why This Project Stands Out on a CV

  • End‑to‑end ownership – You design the request scheduler, KV cache, and streaming output, showcasing ability to ship a complete subsystem rather than a single function.
  • Memory management expertise – Implementing a custom KV cache forces you to reason about allocation, reuse, and eviction policies, a skill directly transferable to high‑performance backend services.
  • Concurrency & scheduling – Continuous batching requires non‑blocking I/O and work‑stealing patterns, demonstrating familiarity with async Python and real‑time scheduling.
  • Latency optimization – Token streaming and incremental decoding highlight an obsession with time‑to‑first‑token and throughput, metrics that matter in production LLM APIs.
  • Signal for roles – This project speaks to ML infrastructure, platform engineering, and backend positions where performance and reliability are paramount.

Architecture Overview

The engine is composed of four loosely coupled components:

  1. Request Scheduler – Accepts HTTP/gRPC requests, validates them, and enqueues them into a priority queue.
  2. Continuous Batcher – Periodically pulls batches from the queue, groups sequences by length, and forms a single forward‑pass batch.
  3. KV Cache Manager – Maintains a per‑sequence key‑value store, allocating memory blocks and recycling them when sequences finish.
  4. Token Streamer – Emits decoded tokens back to the client as they become available, using Server‑Sent Events or WebSocket frames.

A simplified flow:

Client → Scheduler → Batcher → KV Cache → Model → Token Streamer → Client

The batcher runs in an event loop, while the model forward pass executes on the main thread (or a thread pool for CPU). The KV cache is a list of pre‑allocated tensors; each sequence gets a slice that grows as new tokens are generated.

Building It Step by Step

1. Project Skeleton

mkdir llm-serving
cd llm-serving
python -m venv venv
source venv/bin/activate
pip install torch fastapi uvicorn

Create the following files:

llm-serving/
├── main.py
├── scheduler.py
├── batcher.py
├── kv_cache.py
└── streamer.py

2. KV Cache Implementation

The KV cache stores keys and values for each sequence. We use a simple list of tensors, each representing a block of tokens.

# kv_cache.py
import torch
from typing import List, Tuple

class KVCache:
    def __init__(self, max_seq_len: int, num_layers: int, hidden_dim: int, block_size: int = 16):
        self.max_seq_len = max_seq_len
        self.num_layers = num_layers
        self.hidden_dim = hidden_dim
        self.block_size = block_size
        # Pre-allocate blocks: shape (num_blocks, block_size, hidden_dim)
        num_blocks = (max_seq_len + block_size - 1) // block_size
        self.k_blocks = torch.zeros(num_blocks, block_size, hidden_dim)
        self.v_blocks = torch.zeros(num_blocks, block_size, hidden_dim)
        self.next_block = 0  # simple linear allocator

    def allocate(self, seq_len: int) -> int:
        """Reserve enough blocks for seq_len tokens; returns starting block index."""
        needed = (seq_len + self.block_size - 1) // self.block_size
        if self.next_block + needed > self.k_blocks.shape[0]:
            raise RuntimeError("KV cache exhausted")
        start = self.next_block
        self.next_block += needed
        return start

    def write(self, start_block: int, offset: int, k: torch.Tensor, v: torch.Tensor):
        """Write new KV entries into the cache."""
        block_idx = start_block + offset // self.block_size
        intra_block_offset = offset % self.block_size
        self.k_blocks[block_idx, intra_block_offset:intra_block_offset + k.shape[0]] = k
        self.v_blocks[block_idx, intra_block_offset:intra_block_offset + v.shape[0]] = v

    def read(self, start_block: int, length: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """Read `length` consecutive KV entries starting at start_block."""
        end_block = start_block + (length + self.block_size - 1) // self.block_size
        k = self.k_blocks[start_block:end_block].reshape(-1, self.hidden_dim)[:length]
        v = self.v_blocks[start_block:end_block].reshape(-1, self.hidden_dim)[:length]
        return k, v

3. Continuous Batcher

The batcher collects pending sequences, groups them by length, and forms a single forward‑pass batch.

# batcher.py
import asyncio
from typing import List, Dict
from kv_cache import KVCache

class ContinuousBatcher:
    def __init__(self, kv_cache: KVCache, max_batch_tokens: int = 2048):
        self.kv_cache = kv_cache
        self.max_batch_tokens = max_batch_tokens
        self.queue: List[Dict] = []

    def enqueue(self, seq_id: int, prompt: List[int]):
        self.queue.append({"seq_id": seq_id, "prompt": prompt, "generated": []})

    def _group_by_length(self) -> Dict[int, List[Dict]]:
        groups: Dict[int, List[Dict]] = {}
        for item in self.queue:
            length = len(item["prompt"]) + len(item["generated"])
            groups.setdefault(length, []).append(item)
        return groups

    def _form_batch(self) -> List[Dict]:
        groups = self._group_by_length()
        batch = []
        total_tokens = 0
        # Simple greedy packing: take longest sequences first
        for length in sorted(groups.keys(), reverse=True):
            for item in groups[length]:
                if total_tokens + length <= self.max_batch_tokens:
                    batch.append(item)
                    total_tokens += length
                else:
                    break
        # Remove selected items from queue
        for item in batch:
            self.queue.remove(item)
        return batch

    async def get_batch(self) -> List[Dict]:
        await asyncio.sleep(0)  # yield to event loop
        return self._form_batch()

4. Token Streamer

Streaming is achieved by yielding tokens as they are decoded, using an async generator.

# streamer.py
import asyncio
from typing import AsyncGenerator, List

class TokenStreamer:
    def __init__(self):
        self.queue = asyncio.Queue()

    async def push(self, token: int):
        await self.queue.put(token)

    async def stream(self) -> AsyncGenerator[int, None]:
        while True:
            token = await self.queue.get()
            if token is None:  # sentinel
                break
            yield token

5. Wiring Everything Together

# main.py
import asyncio
import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from batcher import ContinuousBatcher
from kv_cache import KVCache
from streamer import TokenStreamer

app = FastAPI()

# Assume a dummy model for demonstration
class DummyModel:
    def __init__(self):
        self.device = "cpu"

    def forward(self, input_ids: torch.Tensor, kv_cache: KVCache, start_block: int):
        # Simulate a forward pass; replace with real model
        logits = torch.randn(input_ids.shape[0], 1)  # dummy logits
        next_token = torch.argmax(logits, dim=-1)
        # Simulate KV cache update
        batch_size, seq_len = input_ids.shape
        k = torch.randn(batch_size, seq_len, 64)
        v = torch.randn(batch_size, seq_len, 64)
        kv_cache.write(start_block, 0, k[0], v[0])
        return next_token

model = DummyModel()
kv_cache = KVCache(max_seq_len=512, num_layers=2, hidden_dim=64)
batcher = ContinuousBatcher(kv_cache)

@app.post("/generate")
async def generate(prompt: str):
    # Tokenize prompt (simplified)
    input_ids = torch.tensor([ord(c) for c in prompt])
    seq_id = id(prompt)
    batcher.enqueue(seq_id, input_ids.tolist())
    streamer = TokenStreamer()

    async def event_generator():
        while True:
            batch = await batcher.get_batch()
            if not batch:
                await asyncio.sleep(0.01)
                continue
            # For simplicity, process first item
            item = batch[0]
            start_block = kv_cache.allocate(len(item["prompt"]))
            # Prefill
            logits = model.forward(torch.tensor(item["prompt"]).unsqueeze(0), kv_cache, start_block)
            token = logits.item()
            item["generated"].append(token)
            await streamer.push(token)
            # Decode loop
            for _ in range(20):  # max new tokens
                next_input = torch.tensor([token])
                logits = model.forward(next_input.unsqueeze(0), kv_cache, start_block)
                token = logits.item()
                item["generated"].append(token)
                await streamer.push(token)
                if token == 0:  # EOS
                    break
            await streamer.push(None)  # signal end
            break

    return StreamingResponse(event_generator(), media_type="text/event-stream")

Note: The above code is intentionally minimal. Replace DummyModel with a real transformer (e.g., a tiny GPT‑2) and integrate proper tokenization via transformers.

Running and Testing It

  1. Start the server
uvicorn main:app --reload
  1. Send a request
curl -N -X POST http://localhost:8000/generate -d '{"prompt":"Hello world"}'

You should see a stream of token IDs (or decoded text if you add a tokenizer) printed as Server‑Sent Events.

  1. Validate performance

Add a simple benchmark script:

# bench.py
import asyncio
import time
import httpx

async def send_request():
    async with httpx.AsyncClient() as client:
        start = time.time()
        async with client.stream("POST", "http://localhost:8000/generate", json={"prompt": "The quick brown fox"}) as resp:
            async for line in resp.aiter_lines():
                pass
        return time.time() - start

async def main():
    latencies = await asyncio.gather(*[send_request() for _ in range(10)])
    print(f"Average latency: {sum(latencies)/len(latencies):.2f}s")

asyncio.run(main())

Run with python bench.py. Aim for sub‑second time‑to‑first‑token after scaling up the model.

Extending It: Your Roadmap to Senior‑Level

  1. Persistent KV cache with Redis – Store KV blocks in Redis to enable request affinity and reduce recomputation across sessions, crucial for multi‑turn conversations.
  2. Horizontal scaling via Ray – Distribute batches across multiple workers using Ray’s distributed scheduler, allowing the engine to handle thousands of concurrent requests.
  3. Observability with OpenTelemetry – Emit latency, throughput, and cache‑hit metrics to a Prometheus backend; add tracing to pinpoint bottlenecks in the batcher or model forward.
  4. Fault‑tolerant request queue – Replace the in‑memory list with a durable Kafka topic, ensuring no requests are lost during worker restarts.
  5. Adaptive batching heuristics – Implement a dynamic batch‑size policy based on queue depth and token length variance, improving tail latency under load.
  6. Model quantization & CUDA integration – Load a quantized model (e.g., GPT‑Q) and offload KV cache to GPU memory, dramatically increasing throughput for production traffic.

Key Takeaways

  • A custom KV cache forces you to master memory allocation, reuse, and eviction—core skills for any systems engineer.
  • Continuous batching demonstrates ability to balance throughput and latency through intelligent scheduling.
  • Token streaming showcases an understanding of client‑side expectations and network‑efficient data transfer.
  • The project maps directly to ML infrastructure, platform, and backend roles where performance is measured in milliseconds.
  • Extending it with persistence, scaling, and observability turns a prototype into a production‑ready service.

Further Reading