Building Our First Inference Engine

The sampler in this chapter’s original engine had a bug that no test caught and no output revealed. Step \(k\) of a request with seed \(s\) used the random seed \(s + k\). That looks harmless, and it means that request 2’s step 0 draws the same random number as request 1’s step 1, request 3’s step 0 the same as request 2’s step 1, and so on. For 100 requests of 32 steps each, 3,200 draws needed 3,200 independent random streams, and the scheme provided 131. When I fed two requests the same scores at every step, the output of seed 2 was an exact one-step-shifted copy of seed 1 in all 31 compared positions.

Nothing crashed. The results looked random, were reproducible, and passed the determinism test, because deterministic is exactly what they were. The defect showed only when two requests were compared with each other, which is the one comparison a single-request test never makes. It is also exactly the situation a serving system creates when it hands consecutive seeds to consecutive users.

This is a chapter about building an engine, and the defect is the argument for how to do it. An engine is a set of replaceable parts behind interfaces. Every part that gets replaced, whether for speed or for a different model, needs a reference to be checked against, a record of exactly what ran, and tests that look at more than one request.

Where this sits

Chapters 1 to 5 each added a piece: the request record, the artefact manifest, the sampling configuration, the hand-written reference forward pass, and the generation loop with stopping and streaming. This chapter puts them together as the Inference Lab engine, version 0, and closes Part I.

Version 0 is deliberately naive. It has no KV cache, no batching, no concurrency and no GPU. Its job is not to be fast. Its job is to be the thing Parts II and III make faster without changing what it computes. A claim of the form “this optimisation is faster” is only meaningful against a baseline that is correct, reproducible and measured, and this chapter produces all three.

For the Pi series, this is the point at which the book stops being about one model and becomes about a runtime. The agent, retrieval and memory volumes call a model through some interface. The engine here is the local, inspectable version of that interface: a request goes in, a result with usage, timing and provenance comes out, and the model behind it can be swapped.

Parts and contracts

The engine has five parts, each with a contract narrow enough to test and replace.

The shape matters more than the contents. A faster attention kernel changes a backend and nothing else. A different decoding strategy changes a sampler. A new stop condition changes the loop. If swapping one part forces edits to another, the boundary was in the wrong place.

The request is validated when it is constructed

Request is a frozen dataclass that checks its own invariants: a positive token limit, a positive temperature unless the rule is greedy, a top-k of at least one, a top-p in \((0, 1]\), non-empty stop strings. A request that violates them cannot exist, so no later stage needs to defend against it. The alternative is validation at the point of use, which means every sampler, loop and backend discovers a bad request in its own way, and a few discover it late, in the middle of a generation.

The rule is carried as Chapter 3’s SamplingConfig: greedy flag, temperature, top-k, top-p, seed. Temperature zero is not an error and not a division; it is spelled greedy=True.

The backend hides everything model-specific

The backend interface is “ids in, scores out”, plus enough identity to say what ran. There are three implementations:

• TableBackend, the Chapter 1 seeded transition table. It is fast and model-free, and every test that does not need a real model uses it.
• LibraryBackend, Hugging Face GPT-2 evaluated by the library.
• OracleBackend, the same weights evaluated by Chapter 4’s hand-written forward pass.

Each backend returns a manifest: implementation and version, model, snapshot, dtype, device, file hashes and the parameter digest from Chapter 2. The engine hashes the manifest into a manifest_sha256 and puts both in every result.

The oracle backend exists to answer a question that deserves to be asked of every pair of runtimes: are they the same? The answer needs two parts. Are the outputs the same? And are the runtimes the same? These are different questions. I ran both backends through the same engine on four prompts for 32 greedy tokens each. The outputs were identical in all four, with no divergence index. The parameter digests were identical too, since they load the same weights. The manifest hashes were different, because one manifest says the library computed the result and the other says a hand-written function did.

That last difference is deliberate, and it is the book’s rule that two runtimes must not be collapsed merely because their outputs agree. Agreement on four prompts is evidence of equivalence, but it is not equivalence. A future prompt might sit on a near-tie where float32 rounding order decides the token, and the two would diverge. The record preserves the fact that they were two runtimes, so a divergence, when it comes, can be traced.

Two samplers, one of them trusted

Chapter 3’s sampler works on Python lists. That is exactly right for clarity and quite wrong for 50,257 logits, as the measurement below shows. The engine therefore has two samplers behind one interface: ListSampler, the reference, and TorchSampler, a vectorised version. The second is acceptable only because it is checked against the first.

The check is a differential test. Both samplers are run on the same logits with the same configuration and seed, and the chosen tokens compared:

The five configurations are greedy, plain temperature, temperature with top-k, temperature with top-p, and all three together. Making the two agree needed care in three places: sorts must be stable so ties break towards the lower index in both; the vectorised version works in float64 to match Python’s arithmetic; and both draw the same uniform random number from the same derived seed and invert the same cumulative distribution.

Zero mismatches in 3,200 comparisons is evidence that the vectorised sampler implements the reference, not a proof. A cumulative sum accumulated in a different order could in principle disagree at a boundary by one unit in the last place and flip a token. The test suite runs a smaller version of the comparison on every change, which will catch a regression in the tie-breaking, though not every rounding case.

The seed defect, in detail

The fix for the opening defect is a function, derive_seed(seed, step), that hashes the pair into a 64-bit integer. The experiment compared the old and the new scheme on identical logits [2.0, 1.0, 0.5, 0.1] at every step, which isolates the random stream from everything else:

With the identical logits, two independent streams would agree at a given step with probability \(\sum_i p_i^2 = 0.40\). Sixteen of 31 is 52 per cent, which is 1.3 standard errors above chance for 31 trials: consistent with independence, not proof of it. The old scheme’s 31 of 31 is not subtle.

Why did every earlier test pass? Because each of them asked “is this request reproducible?”, and a defective seed scheme is perfectly reproducible. The regression test now asserts the property that was actually violated: distinct requests get distinct random streams. It also records the old scheme’s 131-stream figure so that the defect is documented in the test and not only in this chapter.

The wider lesson is about what tests are for. A determinism test checks that a system is a function. It cannot check that the function is a good one.

What the engine returns

Result carries what the earlier chapters promised it would:

Result(
    text, generated_ids, stop_reason,
    usage   = {prompt_tokens, completion_tokens, positions_processed},
    timing  = {total_seconds, forward_seconds, select_seconds, other_seconds},
    steps   = [...per-step records...],
    provenance = {backend, manifest_sha256, manifest, sampler, sampling, loop},
)

Three details are worth defending.

Usage includes positions_processed. Completion tokens say what the user asked for. Positions processed says what the machine did. In the naive loop they differ by a factor that grows with length (Chapter 5), and a result that reports only the first hides the cost that Chapter 10 removes.

Timing is partitioned and the parts must sum to the total. Forward time, selection time and everything else add up to wall time by construction, and a test asserts it. The remainder, “other”, is what the engine itself costs. For a 32-token request it was about 2 ms, which is small enough to be reassuring and large enough to be worth watching once the forward pass gets faster.

Provenance is not decoration. It records what a reader needs to reproduce the result: the full manifest, its hash, the sampler implementation, the exact selection configuration and the loop. It does not claim more than it knows. It says the loop has no KV cache, because that is part of what ran.

Where a step’s time goes

I measured the engine on GPT-2 for 32 tokens with each sampler, under greedy selection and under a realistic rule (temperature 0.8, top-k 40, top-p 0.9):

The reference sampler is correct and nearly a third of the step under a normal sampling configuration. The vectorised one costs 5 per cent. Chapter 1 reported that sampling was negligible, and it was, for greedy argmax. The conclusion does not carry over to the full pipeline of temperature, sort, cumulative sum and draw, run over a 50,000-entry vocabulary in Python. Both samplers produced the same text prefix in the same configuration, as the differential test says they should.

The lesson is not that the list sampler is bad. It is the reference, and it was never meant to be fast. The lesson is that an estimate of where time goes made for one configuration was applied to another, and the measurement is what corrected it. It also bears on what to expect later: when the forward pass gets ten times faster, a 5 ms selection step stops being negligible.

The baseline

The numbers Part II will try to beat, for the library and oracle backends through the engine with the vectorised sampler, 32 greedy tokens from the 17-token prompt, one warm-up request discarded, three runs each:

Both processed 1,040 positions to produce 32 tokens, against 48 distinct tokens in total, a waste factor of 22. Within each backend the output was identical across runs.

The spread between runs is the finding that needs the most care. The oracle’s three runs span 48 per cent, and in the equivalence experiment earlier in the same session it took about as long as the library. I do not trust these timings enough to say which backend is faster, and I am not claiming one is. This is what an uncontrolled benchmark looks like: one machine, other work happening, three samples. Chapter 40 is about designing benchmarks that can support a ranking. This one can support only a rough scale: about 100 to 200 ms per token for this model on one CPU thread without a cache, which is the number to beat.

What the engine does not do

A list worth reading as a plan for the rest of the book.

• No KV cache. Every step reprocesses the whole prefix. Chapter 10.
• No batching or concurrency. One request at a time. Chapters 15, 26 and 27.
• No GPU. Every measurement is CPU. Chapter 7 onwards.
• No streaming transport. The loop can emit text deltas, but nothing sends them over a connection. Chapter 27.
• No cancellation. A request cannot be stopped between steps from outside. Chapter 30.
• No memory accounting. The result does not say how much memory the request used. Chapter 10 onwards.
• No model compatibility checks. The engine trusts that tokeniser and weights belong together, which Chapter 2’s manifest makes checkable but does not yet enforce.

What the numbers do not show

• One model, one prompt, one machine. GPT-2 small on one CPU thread. None of the timings transfers.
• Noisy timing. The baseline spread is large. It supports an order of magnitude and no ranking.
• Sampler equivalence is evidence, not proof. 3,200 comparisons with no mismatch does not rule out a rare boundary case.
• Backend equivalence is on four prompts. Identical output over 32 tokens is not identical output everywhere.
• The seed fix shows independence only weakly. 31 comparisons cannot distinguish a good scheme from a mediocre one. What it rules out is the specific aliasing defect.
• Nothing here measures quality. The text GPT-2 produces is of the standard Chapter 3 described.

Inference Lab increment

The Lab now has an engine. Its public surface is small:

engine = Engine(backend, sampler)
result = engine.generate(Request(prompt, max_new_tokens=32, sampling=SamplingConfig(...)))

Everything the earlier chapters added is in it: the request record (Chapter 1), the manifest (Chapter 2), the sampling configuration (Chapter 3), the oracle as a backend (Chapter 4) and the loop with stop strings and streaming (Chapter 5). The rule for Parts II and III is simple: a change is accepted when the engine’s output is unchanged against the reference and its measured cost is lower, and the result record says so. When a change alters the output, it is a different model and the record must say that too.

What you should now be able to do

• Design an engine as replaceable parts with narrow, testable contracts, and say what each one owns.
• Validate a request at construction time so that no later stage defends against bad input.
• Check a faster component against a slower reference with a differential test, and say what the test does and does not prove.
• Explain why two runtimes with identical outputs should still have distinct records.
• Find a defect that single-request tests cannot see by comparing requests with each other.
• Decompose a request’s time into forward, selection and overhead, and report positions processed alongside tokens.
• State a baseline honestly, including how noisy it is.

Key terms

• Engine: the assembled runtime that turns a request into a result.
• Backend: the model-specific part, behind “ids in, scores out”.
• Manifest: the hashable description of what a backend executes.
• Provenance: the part of a result that says what produced it.
• Differential test: running a reference and a candidate on the same input and comparing outputs.
• Oracle: a slow, trusted implementation used as the reference.
• Seed aliasing: distinct requests sharing parts of their random streams.
• Baseline: the measurement later changes are compared against.
• Positions processed: the number of token positions a request pushed through the model.

Reader exercises

1. Reintroduce seed + step and find the smallest change to the test suite that detects it from a single request. Is there one?
2. Add a Backend that wraps a second model, such as a small model you have cached. What has to change in the engine? What has to change in the manifest?
3. Make the oracle and library backends diverge by switching one to a different dtype. At what token does the first difference appear? What does the record say about it?
4. Extend the differential test to compare the two samplers on logits that contain exact ties. Do they agree?
5. Add peak memory to timing or usage. Which measurement would you use, and what would it miss?
6. Run the baseline ten times on a quiet machine, then with another heavy process running. How large is the effect on the median? What does that suggest for benchmark design?
7. Write the smallest change that would let two requests share one forward pass, and list what in the record would have to change.

References

• Efficient Memory Management for Large Language Model Serving with PagedAttention — a systems paper showing that an engine’s performance depends on how the parts around the model are designed.
• SGLang: Efficient Execution of Structured Language Model Programs — a systems paper in which scheduling and reusable state sit outside the model.
• deepseek-ai/DeepInfer at 774ce244 — a pinned first-party repository; its documentation states that its public release has model and hardware limits.

Evidence in this chapter

Implementation status

• Manuscript: content/books/inference/06-chapter.md
• Runnable examples: examples/ch06/engine.py (original toy engine), examples/ch06/lab_engine.py (the Lab engine)
• Tests: tests/test_ch06_engine.py, tests/test_ch06_lab_engine.py (request validation, seed regression, sampler differential test, engine behaviour on the table backend; the model commands are not exercised by tests)
• Evidence: research/ch06/sources.md, research/ch06/results.json, research/evidence-register.md
• Verification command: python -m unittest discover -s tests -v

The chapter is a draft with draft = false. It makes no claim about cached, batched, GPU or production behaviour. Part II begins with the operation that dominates the cost of everything above: matrix multiplication, and what it takes to run it on a GPU.