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.
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.
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
| Concurrency | Throughput (tok/s) | TTFT mean (s) |
|---|---|---|
| 1 | 44.7 | 0.07 |
| 4 | 83.8 | 0.27 |
| 8 | 113.5 | 0.54 |
| 16 | 108.5 | 1.27 |
Interpretation
At equal settings the engine matches Hugging Face one-at-a-time (44.7 vs 40.0 tok/s) and reaches about 2.8x throughput once batching 8 requests (113.5 tok/s). The gain is pure scheduling and memory efficiency, not faster math.
Speculative decoding is correct (identical greedy output, 54% of drafts accepted) but about 0.7x slower on CPU, where the draft forward is not cheap relative to the target. It stays in the codebase as a correctness-verified feature. A comparison against vLLM is deliberately absent because this machine has no GPU, these are honest CPU numbers.
Technical Details
Tech stack:
What Was Learned
- Paged attention is bookkeeping, not math: The attention kernel barely changes; the allocator, refcounting, and hashing are where the design lives.
- Continuous batching works, even on CPU: going from 1 to 8 concurrent requests lifts throughput from 44.7 to 113.5 tok/s, about 2.5x its own single-request rate and about 2.8x the 40.0 tok/s Hugging Face baseline. The gain comes entirely from not wasting compute on padding.
- Measure honestly on your hardware: Speculative decoding's speedup is conditional on draft ≪ target cost. On CPU it goes backwards, more useful to know than to assume GPU numbers transfer.
- Correctness first, performance second: Prefix-cache correctness with in-place rewrites required deregistering content hashes before mutation. Cumulative capacity checks had to span the whole batch, not per-sequence.
- Tokenizer boundaries matter: UTF-8 can split across tokens in byte-level BPE, so detokenization must hold back partial tails until they decode cleanly.