inference-batching-lab
A decoder-only transformer with a hand-written KV cache, and three request-scheduling strategies benchmarked for real throughput and latency.
Why this project
Once you have a model, serving it well is a scheduling and memory-management problem. Writing the batching logic from scratch means the KV cache, the causal masking and the admission policy are all code I can explain line by line.
The bar: byte-identical output
Batching only pays off if it does not change what gets generated. Every batched strategy has to match one-at-a-time generation exactly, even though requests join and leave the batch at different times. Getting there took six real bugs:
- A causal mask that only worked when query and key lengths matched, so incremental decode could attend to just key zero.
- NaN propagation through left-padded positions whose own query had every key masked out.
- A per-request KV-cache slice using an index meant for the batched tensor.
- A request with
max_new_tokens == 1silently getting one extra token. Only a Hypothesis example found it. - Continuous batching dropping queued requests when every active one finished during prefill.
- New tokens stored at the wrong cache column in batches with different cached lengths.
Measured
On CPU with a 241K-parameter model and 24 requests: sequential 1,917 tok/s, static batching 4,178 tok/s, continuous batching 4,198 tok/s.
Stack
Python · PyTorch · Hypothesis · Pytest
Honest notes
The model is untrained and the machine had no GPU, so this measures scheduling, not model quality. Continuous batching pads rather than pages, and the benchmark submits requests in one burst, which understates its real advantage. The follow-up is paged-speculative-engine.