Digitalise Agency

inference-engine

LLM Inference Engine

Some software has to run an AI model and produce its reply token by token, a token being a word or part of one. This is that software, written from scratch. The hard part is serving a crowd at once from a single machine without it sitting idle: memory housekeeping and clever turn-taking. On CPU, serving eight at once lifted it from about 45 tokens a second to about 114.

CPU numbers, read as relative

Runs GPT-2 family models on CPU. Hugging Face is used only to download weights; the forward pass, attention, cache, and scheduler are all implemented here. The benchmarks are single-machine CPU numbers and are honest about what they are.

Paged Attention: the memory trick this engine is built on, animated. Each thread is one person's conversation and the lattice is the model's memory of it, cut into fixed-size pages so many conversations can share one machine; the ghosts racing ahead are guesses at the next words, kept only if they turn out right. Live and interactive: drag it to orbit, scroll or pinch to zoom. Open full screen
Two-panel chart of inference-engine's real serving-sweep benchmark: throughput scaling and time-to-first-token cost of batching.
Measured on an 8-core CPU: peak throughput reaches 113.5 tok/s at 8 concurrent requests, about 2.8x the 40.0 tok/s Hugging Face baseline.
~3e-5max logit diff vs transformers
~2.8xbatched vs sequential throughput
0.54speculative acceptance rate
55tests passing

Problem, solution, result

Problem: LLM inference is latency-bound in production. Serving multiple requests efficiently requires careful memory management, scheduling, and algorithmic tricks. Solution: Built a from-scratch inference engine with paged KV cache, continuous batching, prefix caching, and speculative decoding. Result: ~2.8x throughput improvement on commodity hardware through smart resource allocation, not raw compute.

The Problem

Production LLM serving is constrained by latency and memory. Naive request handling (one request at a time) leaves compute idle. Serving engines solve this through continuous batching, but require careful memory management: cache allocation, eviction, sharing, and rollback under failure. Serving engines are where ML meets systems programming: memory allocators, schedulers, cache coherence, admission control.

The Solution

Built a from-scratch LLM serving engine in Python and PyTorch that implements the core mechanisms needed for efficient inference:

  • Paged KV cache: Fixed-size 16-token blocks with refcounting, per-sequence block tables, and copy-on-write for shared blocks.
  • Prefix caching: Identical prompts reuse cached blocks with content hashing and LRU eviction.
  • Continuous batching: Requests join/leave at token boundaries; preemption lets new requests proceed when memory is full.
  • Speculative decoding: Smaller model drafts tokens, larger model verifies all at once with rejection sampling.
  • Verified correctness: GPT-2 implementation verified against transformers library to 3e-5 max logit diff.
  • OpenAI-compatible API: FastAPI server with /v1/completions, streaming, /models, /healthz, /metrics.

Results: Measured Throughput

Measured on an Intel i7-1185G7 (8 threads), PyTorch CPU float32, greedy decoding, 32 new tokens per request.

Throughput vs concurrency

ConcurrencyThroughput (tok/s)TTFT mean (s)
144.70.07
483.80.27
8113.50.54
16108.51.27

Book a call

Let's talk about what you're building.

Pick a slot below.