Our first walk of mlc-ai/mlc-llm told the compile-and-serve story at the repository scale: passes that fuse weight dequantization into matmuls, quantization schemes, the async Python engine, and a note that a compact C++ core does the real work. That C++ core is where the interesting scheduling happens, and it deserves the kind of attention TurboMind got from us in the lmdeploy follow-up. Every inference loop on the planet has the same shape — pick a batch, run prefill or decode, sample, stream deltas back — and the differences between frameworks live in how each decides what a “step” is, how requests move between queues, and what happens when you want to draft five tokens and throw four away. mlc-llm answers all three with an unusually clean vocabulary: one background thread, a step function that selects exactly one EngineAction per iteration, and a request-state machine whose draft token slots make speculative decoding a bookkeeping exercise rather than a special case. As always, this is an educational read of published Apache-2.0 source on the main branch; the first post at /mlc-llm/ covers the compiler side we deliberately skip here.

The architectural headline is that the whole serving engine is a two-thread conversation. Requests arrive from Python as instructions into a queue owned by the background thread; the background thread runs step after step, and every delta — every new token — flows back through a callback loop into Python without either thread ever blocking the other. Everything else in cpp/serve exists to keep that conversation moving: the actions are the verbs, the state is the ledger, and the sampler is the arithmetic.

Overview of the serve core: Python engines hand instructions to the ThreadedEngine, its background loop drives EngineImpl's step dispatcher, the EngineAction family covers prefill, decode, and speculative paths, EngineState and RequestState track every request, and the prefix cache plus GPU sampler close the loop.

Reading the overview from left to right:

• engine_base.py is the Python side of the callback bridge: the async stream callback turns C++ deltas into asyncio wakeups, with a sync twin at line 548.
• sync_engine.py exposes add_request and an explicit step for callers who want to own the loop.
• threaded_engine.cc owns RunBackgroundLoop and RunBackgroundStreamBackLoop, fed by an instruction queue at line 356.
• engine.cc is the dispatcher: EngineImpl::Step, beside AddRequest at line 668 and HandleSpecialRequests at 535.
• engine_state.cc keeps the queues and hands out per-request states plus the running-state list at 33.
• action.h declares the entire action vocabulary in one line: Step takes the state and returns finished requests.
• new_request_prefill.cc starts running requests; batch_decode.cc keeps them running one token at a time.
• auto_spec_decode.cc routes between draft, verify, and jump-forward when a draft model is configured.
• request_state.cc is the ledger: CommitToken, RollbackTokens at 70, AddDraftToken at 85.
• prefix_cache.cc reuses work through ExtendSequence, RollBackSequence at 177, and RecycleSequence at 191.
• gpu_sampler.cc does the arithmetic: top-p renormalization and batched sampling at 190.
• logit_processor.cc shapes logits before sampling with bias, penalties, and grammar masks.

A step is an action

The design’s one big idea is easiest to see in the dispatcher. EngineImpl::Step does no model work itself: it checks whether there are new requests to absorb or dead requests to clean, picks exactly one action object, calls its Step, and hands the returned list of finished request ids to the state for cleanup. Which action gets picked is a matter of one if-chain: if any waiting request can join the batch, prefill; if a speculative configuration is active, the speculative action; otherwise decode. Each action is a self-contained class. NewRequestPrefillActionObj::Step walks the waiting queue, consults the prefix cache for reusable blocks, and packs as many prompts as fit into a chord-shaped prefill batch. BatchDecodeActionObj::Step is the treadmill’s heartbeat: it assembles the running batch, runs one forward step on the model, then samples the whole batch in one call at line 170. Underneath both, function_table.cc is the FFI registry that made the model callable: Init and _InitFunctions at 207 load the TVM runtime functions — prefill, decode, kv cache ops — that every action calls by name, and CopyToWorker0 shuttles host data to the device. Where actions need each other’s plumbing, action_commons.cc shares the batch-assembly and output-merge helpers rather than duplicating them.

The two-thread story is what makes this elegant rather than merely neat. Python calls into ThreadedEngine only to push an instruction — add request, abort request, reload — and the instruction queue at line 356 is the single crossing point. The background loop drains it, runs steps, and parks deltas in an output queue; the stream-back loop at 191 waits on that queue and fires the callback into Python. Two condition variables, no locks held during model execution, and the vtable at line 386 is the whole surface Python ever sees. It is a lighter pattern than the process-per-worker lattice we mapped in TensorRT-LLM, and it means a reload never tears down the Python event loop.

The lifecycle bookkeeping around that loop is worth its own paragraph, because it shows how much production reality the dispatcher absorbs. Engine construction is Engine::Create at line 1016, which funnels through AutoDecideEngineConfig at line 893 — the stage that fills in defaults the user left blank and picks a mode single-GPU versus speculative versus disaggregated. Requests that are not ordinary completions are intercepted before they reach the queues: HandleSpecialRequests at line 535 answers debug and control instructions inline, and HandleDisaggRequest at line 553 recognizes requests that arrived from a remote prefill node in disaggregated serving. Failures stream back through the same callback as successes — StreamBackError at line 529 wraps an error and a finish reason into a normal-looking delta so the Python side never needs a second error channel. On the far side of the bridge, process_chat_completion_request at line 687 is where OpenAI-shaped JSON becomes a C++ Request, and its stream-output twin at line 781 converts deltas back into chat-completions chunks. The wire dialect changes nothing in the core — which is precisely the point of the layering.

Speculative drafts and the state machine

Speculative decoding is where this architecture pays off. Without a draft model, a step is a decode. With one, AutoSpecDecodeActionObj::Step becomes a tiny state machine of its own: it alternates BatchDraftActionObj::Step, where the small model commits a block of k draft tokens per request, with BatchVerifyActionObj::Step, where the large model runs all drafts in parallel and the sampler accepts the longest correct prefix. A jump-forward action at line 36 handles the structured-output case — when a grammar or a JSON schema says the next five characters are already decided, the engine jumps them without a forward pass at all. The EAGLE path swaps in different actions — eagle_batch_draft.cc and eagle_batch_verify.cc — with feature vectors and tree-shaped verification, and the dispatcher at line 78 wraps whichever family you configured.

What makes draft-then-verify cheap to reason about is the request state. RequestModelStateNode::AddDraftToken parks each draft in a numbered slot; CommitToken records accepted tokens and RollbackTokens rewinds the ones the verifier rejected, with RemoveAllDraftTokens reporting which KV cache slots to recycle. The same ledger serves non-speculative paths: grammar-constrained decoding asks for a token bitmask at line 48, and every streamed delta is composed by GetDeltaRequestReturn at line 175. The sampler underneath is quietly sophisticated: BatchRenormalizeProbsByTopP renormalizes on-device before BatchSampleTokensWithProbAfterTopP at line 190 draws from the surviving mass, and GenerateUniformSamples at line 480 pre-generates random numbers on GPU so the per-request RNG state never round-trips through the host. Bias and penalties arrive earlier through InplaceUpdateLogits, UpdateWithPenalty at 270, and UpdateWithMask at 372.

The prefix cache underneath

Serving quality often depends on what you do not recompute, and this repo’s answer is a paged radix tree in C++. PrefixCacheImpl exposes four operations: ExtendSequence claims blocks for growing tokens, CommitSequenceExtention at 154 makes a speculative extension durable, RollBackSequence at 177 releases the tail after verification rewinds, and RecycleSequence at 191 marks a finished conversation for lazy reuse. The tree itself, radix_tree.cc, indexes sequences at block granularity so two conversations sharing a system prompt share those KV blocks for free — the same idea we saw in SGLang’s radix cache and in the Python-side tree the first tour showed, now living close to the C++ actions that use it. TryFreeMemory at line 225 is the pressure valve that evicts leaves when a new prefill needs room.

From Install to Your First Model

None of this needs configuration: python -m mlc_llm serve builds it and runs it, and create_sync or AsyncThreadedEngine expose the same core in-process. The visible knobs are speculative — a draft model plus spec decode mode flips the action dispatcher — and the debug surface: the event trace recorder emits per-request spans you can read in Chrome’s trace viewer, which is the fastest way to see the prefill-chord and draft-verify alternation in a real log.

Honest limits: we read the single-GPU path; disaggregated prefill and decode across nodes exists in the same actions folder but we did not trace it, and the embedding engine shares none of this. The sampler reads correct but we benchmarked nothing — top-p behavior on GPU is exactly where silent numerical drift hides. And the action set is still growing; jump-forward and EAGLE both landed recently, which is why this layer rewards a second reading. For context: vLLM solves speculative decoding at the Python scheduler layer, SGLang shares the radix-tree idea from the same lineage, LM Studio-style local stacks wrap engines rather than reimplement them, and LocalAI routes to backends it does not own. mlc-llm’s one-action-per-step C++ core is the most legible of these — you can hold the whole token treadmill in your head, which is exactly why it was worth a second visit.

Enjoyed this post? Never miss out on future posts by following us