Continuous Batching Engine

The project idea came out of a conversation with Claude about what actually makes engines like vLLM and TGI fast, which is what pointed me at scheduling over model quality as the interesting problem. Claude was also involved with the design too, for things like the KV-cache block manager and the admission control policy, it laid out a few viable approaches with their tradeoffs each time, but the choice of which to build, and every line of the implementation, was mine.

I wanted this project to be about AI itself. When I looked at how vLLM and TGI actually get their throughput, the thing that stuck with me was that almost none of it comes from the model. It's the scheduling. How do you keep a GPU busy when a hundred people are asking for different amounts of text at different times?

That question is what got me building this. It's a continuous-batching engine, with a request queue, admission control, and a scheduler that packs many in-flight generations into one batched step, running on top of a transformer I wrote myself in plain NumPy so I'd actually understand what the KV-cache underneath it is doing, instead of treating the model as a black box.

The engine has a request queue with admission control, a scheduler that packs many in-flight generations into one batched step, and a KV-cache block manager, all running on top of a transformer I wrote myself in plain NumPy rather than a pretrained model. Weights are randomly initialized by default, since the point of the project is serving mechanics and throughput, not text quality, and ModelWeights.load() is there if I ever want to drop in trained weights later.

A KV-cache is only worth anything if decoding with it gives you the exact same result as recomputing the whole sequence from scratch, so the first real test I wrote was to test that. I ran a prompt through prefill(), separately decoded the last token using a cache from the first N-1, and asserted the logits matched. Extending that to a ragged batch, two sequences of different lengths decoded together in one call, is what actually caught bugs in my padding and masking logic.

The hard part I found for continuous batching is that sequences sharing a decode step almost never have the same length, someone's three tokens in while someone else just joined. I pad each sequence's cached K and V up to the batch's current max length and mask out the padding in the attention scores, so it's still one batched matmul instead of a loop per sequence, the same technique real batched-decode implementations use before something like PagedAttention avoids the padding waste entirely.

My first version of step() admitted waiting requests and then immediately ran a decode step on the whole batch, including the ones that had just joined. That meant a brand new request got two tokens in its first tick while everyone else got one, which quietly broke the guarantee the whole design rests on, that every sequence advances by exactly one token per tick. My scheduler tests caught it immediately by asserting exact token counts after a single step(), and the fix was to decode only the sequences that were already running.

step() is a plain synchronous method on purpose, since it's just NumPy matmuls with nothing to await, but that meant I couldn't use asyncio.Queue for streaming tokens back to a client the way I first tried. On Python 3.9 it binds to whatever event loop is running the moment it's constructed, and submit() has no guarantee one exists yet. I ran into that as a very confusing test failure before I understood what was happening. The fix ended up simpler than the original design, RequestHandle just polls the request's growing token list instead.

None of this matters if batching doesn't measurably beat serving requests one at a time, so bench.py runs the exact same model and weights both ways. On my machine, 24 concurrent requests generating 40 tokens each finish about 2.3x faster batched than sequential, and 64 requests at 50 tokens each lands around 2.2x. The peak batch size the benchmark reports is capped below the request count when the cache pool is smaller than demand, which is the admission control actually doing its job.

The benchmark runs two passes over the same model with the same randomly initialized weights: one that serves each request in isolation and waits for it to finish before starting the next, and one that feeds them all through the continuous-batching engine. The only variable is the serving strategy. Below are the outputs for a small and a larger run, showing how the throughput gap holds as the request count scales up.

The peak concurrent batch size stays at 21 regardless of how many total requests are submitted because it's set by the KV cache block pool, not the request count. With the default 64 blocks at size 16, the pool saturates at around 21 in-flight sequences given the prompt and generation length in use. Once it's full, the scheduler queues new requests and admits them one by one as running sequences finish and free their blocks. To raise the ceiling you'd pass a larger --num-blocks value; to see every request admitted at once you'd keep the request count below the pool's capacity.