S20 · The Complete Engine · Build an LLM Inference Engine

Motivation

The problem

You have nineteen working parts and no engine. Assembling them takes more than importing everything into one file: the parts have ordering constraints, shared state, and a few real conflicts.

The prefix cache and the block manager both own blocks. The scheduler and the speculative decoder both decide how many tokens a step produces. CUDA graphs want fixed shapes and continuous batching refuses to provide them. Getting these to coexist is the work of building an inference engine.

Core idea

The solution

One step() function, called in a loop, with a strict ordering. Everything else is a component it calls.

the complete engine looppython

def step(self) -> list[Output]:
    # 1. SCHEDULE — who runs, and with what token budget       (S10, S11)
    batch = self.scheduler.schedule()
    if not batch:
        return []

    # 2. MEMORY — resolve prefix hits, then allocate blocks    (S06, S07)
    for req in batch.new_requests:
        reused, start = self.prefix_cache.match(req.token_ids)
        req.block_table = reused + self.blocks.allocate(req, from_token=start)

    # 3. PROPOSE — optional draft tokens for verification      (S14)
    if self.spec_enabled and batch.is_decode_only:
        batch = self.proposer.attach_drafts(batch, k=self.adaptive_k())

    # 4. EXECUTE — one forward pass over a ragged batch        (S03, S12, S13)
    runner = self.graph_runner if batch.can_use_graph else self.eager_runner
    logits = runner.forward(batch)

    # 5. SAMPLE — per-request params, grammar masks, verify    (S04, S15, S14)
    tokens = self.sampler.sample(logits, batch.sampling_metadata)
    tokens = self.verifier.accept(tokens, batch)     # no-op without speculation

    # 6. RETIRE — stop checks, free blocks, publish to streams (S19)
    return self.postprocess(batch, tokens)

The ordering is not stylistic

Scheduling must precede allocation, because the scheduler decides who gets memory. Prefix matching must precede allocation, because a hit changes how many blocks are needed. Speculation follows scheduling, since it consumes the budget the scheduler set. Retirement comes last, so that freed blocks are available to the next step's scheduler. Reorder any of these and you get subtle, load-dependent bugs.

Mechanics

How it works

Diagramevery chapter, in one picture

Layers top to bottom: serving, scheduling, memory and model, output. The dashed line back to the API is the streaming path.

The conflicts, and how real engines resolve them

CUDA graphs vs continuous batching. Graphs need fixed shapes; batches change every step. Resolution: capture buckets and pad, and fall back to eager mode for prefill.
Speculation vs batching. Both spend spare compute. Resolution: make the speculation length a function of current batch size — aggressive when idle, off when saturated.
Prefix cache vs KV capacity. Blocks retained for a possible future request are blocks a running request cannot have. Resolution: reference counting plus leaf-only LRU, with cached-but-unreferenced blocks first in line for eviction.
Chunked prefill vs prefix caching. They compose well, but only if a chunk whose blocks are already cached is skipped entirely rather than recomputed.

What to build in what order

If you are writing your own engine, the payoff ordering is roughly:

KV cache — nothing else matters until this exists.
Continuous batching + paged memory — the two together, since neither works well alone.
Prefix caching — largest win per line of code, if your traffic has shared prefixes.
Chunked prefill — cheap, and it fixes the worst tail-latency behaviour.
Quantization — when memory or bandwidth is the binding constraint.
Everything else — measure first. CUDA graphs and speculation win big at low concurrency and almost nothing at high.

Explore

Try it

Every technique in the course, as a toggle. The multipliers come from published measurements; what matters is their relative magnitude. Turn features off one at a time and see which ones you would miss.

Simulatorthe complete engineturn everything off, then add features back in order

target concurrency32GPU memory80 GB

throughput: 173.2k tok/s
TTFT: 148 ms
inter-token latency: 0.6 ms
concurrent requests: 32

the engine, feature by feature

cumulative throughput

naive loop (S01)21

+ KV cache882

+ PagedAttention2.3k

+ prefix caching3.4k

+ INT4 weights + INT8 KV7.2k

+ continuous batching23.1k

+ chunked prefill25.9k

+ FlashAttention35.0k

+ CUDA graphs + fusion41.2k

+ speculative decoding43.3k

throughput before this feature
added by this feature

Overall 2062× against the naive loop from S01. The multipliers are illustrative and drawn from published measurements; read the ordering and the relative magnitudes rather than the third significant figure. Notice how much of the total comes from the first two features.

memory is still the constraint

weights

32 × KV

KV per request: 0.36 GB
requests that fit: 175
memory efficiency: 18%

Turn off PagedAttention and quantization and watch how many requests stop fitting. Nineteen chapters in, the conclusion from S05 still holds: the KV cache is the scarce resource, and throughput is whatever fits.

The instructive experiment: set concurrency to 1 and note how much speculative decoding and CUDA graphs contribute. Then set it to 128. Their contribution nearly vanishes, and continuous batching and paging carry everything. Most disagreements about which optimisation matters are really disagreements about the concurrency people have in mind.

Build it

Implementation

The reference implementation for this chapter is the whole engine: about 420 lines, importing the components from the previous nineteen files. It exposes a benchmark harness so you can reproduce the simulator's ablation on your own machine.

code/s20_complete_engine.py (excerpt)python

def ablate(workload, features):
    """Turn one feature off at a time and measure what it was worth."""
    baseline = benchmark(Engine(**features), workload)
    rows = []

    for name in features:
        reduced = {**features, name: False}
        result = benchmark(Engine(**reduced), workload)
        rows.append((name, baseline.throughput / result.throughput))

    return sorted(rows, key=lambda r: -r[1])   # most valuable first

A feature can score below 1.00× and that is the answer, not a bug

The runnable ablation reports speculative decoding as a net loss at max_num_seqs=64; the engine is faster without it. That is correct. Speculation spends spare compute, and at that concurrency there is none spare. Any ablation you run is a statement about your workload and your concurrency, not a universal ranking.

Run it locally

Assembles every component from S01–S19 into one engine, runs a synthetic chat workload through it, and performs a leave-one-out ablation reporting what each feature is worth on your machine, plus a memory sweep showing where requests stop fitting.

$ python code/s20_complete_engine.py

Expect: An end-to-end throughput and latency report, an ablation table ranking the features (with speculative decoding scoring below 1× at high concurrency), and a KV-pool sweep where 256 blocks completes 16/120 requests and 2048 completes all of them.

Only NumPy is required — setup instructions.

Production notes

Where to go next

You can now read the real engines. Concretely, here is where each chapter lives in vLLM's source:

vllm/v1/core/sched/scheduler.py — S10, S11
vllm/v1/core/kv_cache_manager.py — S06, S07
vllm/v1/worker/gpu_model_runner.py — S09, S13
vllm/v1/sample/ — S04, S15
vllm/v1/spec_decode/ — S14
vllm/distributed/ — S17, S18

At the small end of the spectrum, picoLM is ~2,944 lines of C implementing S01–S05, S08 and S12 with no dependencies. You can read it in an afternoon, and it is a good check on whether you understood the forward pass.

Final exercises

1
Run the ablation on your own hardware and workload. The ordering will differ from the simulator's, and understanding why it differs is the real exercise.
2
Port the engine to a real model — load Llama-3.2-1B weights, replace the toy forward pass with the S03 implementation, and verify output against HuggingFace generate() at temperature 0.
3
Pick the one thing this course omitted that matters for your use case — multimodal inputs, LoRA adapter serving, embeddings, or encoder-decoder models — and add it. The scheduler and memory manager you built will mostly accommodate it. How far they stretch is a good test of the abstractions.

Continue

What you built

Twenty chapters, and the arc is short enough to state in one paragraph. An inference engine is a loop that calls a model. The loop is slow because it recomputes the past, so you cache it. The cache becomes the scarce resource, so you page it, share it and compress it. The GPU is idle during decode, so you batch, and batching needs a scheduler, and the scheduler needs chunking to stay fair. Then you spend the remaining idle capacity on speculation, split the model across GPUs when it stops fitting, split the workload across machines when the two phases conflict, and put an API in front of it.

Everything else in this field is a refinement of one of those moves.

Check yourself

Questions

Answer before you look. A wrong answer tells you more than a right one, because the explanation points at the part of the chapter to reread.

Check yourselfQuestion 1 of 6

In the engine's step(), why must prefix matching happen before block allocation?

Score 0/6