How do you build a toy LLM inference engine?

In this article

1. Before you start
2. Why interviewers ask this
3. Toy inference engine from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page is one exercise in the PyTorch and machine learning (ML) coding round guide. If you have not built a text generator yet, the generative pre-trained transformer (GPT) from scratch tutorial shows the model this engine would serve.

Why interviewers ask this

A serving system pays for its graphics processing unit (GPU) time whether or not the GPU is busy, so a candidate who can explain why a batch is re-formed every step, and what limits its size, shows they understand where inference cost comes from. Prompts of this kind are reported by candidates in loops at AI labs and inference companies: write a simplified engine with request batching and a per-batch model call, build a priority queue for a batcher, or design a batching application programming interface (API) with a key-value cache, streaming and overload behavior.

A weak answer waits for N requests and runs them together. A strong answer separates the first pass over the prompt from the one-token steps that follow, schedules per step, and says that memory, not arithmetic, caps the batch. A Staff-level answer adds time to first token, admission control and speculative decoding.

Toy inference engine from first principles

Inference engine

An inference engine is a program that sits between users and a trained model, accepting many generation requests, grouping them into batches, running the model on each batch and returning every result, so that the expensive hardware stays busy while each user gets an answer.

Continuous batching

Continuous batching is a scheduling method that re-forms the batch of running requests after every single model step, so finished requests leave and waiting requests join at once instead of waiting for a whole batch to end. Orca called it iteration-level scheduling.

In plain words

Picture a bus that leaves only when every seat is taken. It returns only when every passenger has reached their stop. Passengers who live nearby sit idle for the whole trip. New riders wait at the stop. Continuous batching is a bus where a seat is refilled as soon as its rider gets off. The analogy breaks in one place. A model step is not a trip along a route. Every running request moves forward exactly one token per step.

In one precise sentence, the engine repeats "choose which requests run, run one step, collect one token for each, drop the finished ones."

Step 1: A request queue and a one-token loop

Generation is a loop of model iterations, and an iteration is one run through all layers of the model; each iteration produces one output token that becomes the next iteration's input (Yu et al., 2022). The first iteration is called prefill: it processes the whole prompt together and yields the first output token. Every later iteration is decode: it feeds one new token and yields the next. Decoding K tokens takes K serial runs of the model (Leviathan et al., 2023).

So the engine needs three parts: a queue of waiting requests, a set of running requests, and a loop. Input to the loop: the running set. Output: one new token per running request. A request leaves when it emits an end token or reaches its token limit. You now know the unit of work, one step, which is what the next step schedules.

Step 2: Re-form the batch after every step

A request-level batcher keeps its batch fixed until every request in it finishes, yet each request may need a different number of iterations, so some finish early and then waste computation (Yu et al., 2022). Iteration-level scheduling fixes this: the scheduler runs only a single iteration of the model on the batch, then decides again. A newly arrived request can then be considered after waiting for only a single iteration.

Take four requests needing 2, 8, 2 and 2 output tokens with two slots (illustrative numbers). Request-level batching takes 10 iterations and keeps slots busy 70% of the time; per-step batching takes 8 iterations at 87.5%.¹ The same trace returns request C at iteration 4 instead of 10.¹ The rule: at the top of every iteration, fill free slots from the queue, run one step, then remove finished requests.

Step 3: The KV cache sets the batch limit

Each request keeps a key-value (KV) cache. The KV cache is a per-request store of the attention keys and values of every token already processed, which lets the model skip recomputing them at each step. That store grows with every generated token. Kwon et al. (2023) state that serving throughput is memory-bound, because this cache limits the batch size.

Their example is a 13-billion-parameter model, whose cache costs 800 KB per token, computed as 2 (keys and values) × 5120 (hidden size) × 40 (layers) × 2 bytes in 16-bit floating point (FP16); one request can need as much as 1.6 GB. The arithmetic gives 819,200 bytes per token.¹ Your engine therefore needs an admission rule: start a request only if its worst-case cache fits the free memory, or the engine stalls mid-generation.

Reserving the maximum length for every request wastes memory, so vLLM stores the cache in fixed-size blocks, an approach called PagedAttention: one can think of blocks as pages and tokens as bytes (Kwon et al., 2023).

Step 4: Speculative decoding generates several tokens per large pass

A decode step for a large model is slow, so speculative decoding runs a cheap draft model several tokens ahead and verifies them with the large model in one pass. The method samples faster without any changes to the outputs and needs no retraining or architecture change (Leviathan et al., 2023).

With draft length γ and an acceptance rate α, the paper's analysis gives an expected (1 − α^(γ+1)) / (1 − α) tokens per large-model pass. For α = 0.8 and γ = 4 (illustrative), that is 3.36 tokens per pass; at α = 0.6 it is 2.31.¹ If each draft step costs 5% of a large step (illustrative), the net speedups are 2.80 and 1.92.¹ The engine must then keep a separate cache for the draft model and roll back rejected tokens. The beam search page covers a different way of exploring several tokens, which is not the same thing.

1 · Queue and loop2 · Per-step batch3 · Memory budget4 · Draft and verifyReplay

Step-by-step description

1. Requests wait; the model runs one step.
2. The batch is re-formed every step.
3. The KV cache caps the batch.
4. Several tokens per large pass.

• The scheduler re-forms the batch after every model step.
• The KV cache budget caps how many requests run together.
• A draft model proposes tokens that the large model verifies in one pass.

Worked example

The trace below uses four requests that need 2, 8, 2 and 2 output tokens (illustrative) and two slots. Each iteration gives every running request one token, and the prompt cost is ignored. A request-level batcher returns results when its whole batch ends. The scheduler is a few lines:

def continuous(lengths, slots):
    queue = list(enumerate(lengths))   # (request id, tokens still needed)
    running, done, step, used = {}, {}, 0, 0
    while queue or running:
        step += 1
        while queue and len(running) < slots:   # refill free slots
            i, n = queue.pop(0)
            running[i] = n
        used += len(running)                    # slot-steps spent
        for i in list(running):
            running[i] -= 1                     # one token each
            if running[i] == 0:
                done[i] = step                  # leaves immediately
                del running[i]
    return step, used, done

The loop refills slots first, then spends one step, then removes finished requests; its complexity is one pass over the running set per step. The static baseline groups requests in arrival order and runs each group for the length of its longest member.¹

The memory side uses an illustrative budget of 12 GB for the cache. At 2,048 tokens per request, the cache is 1,677,721,600 bytes per request, so 7 requests fit; at 256 tokens each, 57 fit.¹ With blocks of 16 tokens (illustrative), a 250-token request uses 16 blocks and wastes 6 token slots; reserving 2,048 slots for a request that uses 256 wastes 87.5% of its reservation.¹

Text description of the diagram

• Request C returns at iteration 4 instead of 10.
• Mean return iteration falls from 9 to 5 in the illustrative trace.

Trade-offs and edge cases

Per-step scheduling is not free. The scheduler runs between every pair of steps, so it must be cheap, and requests of different lengths cannot be batched for every operation: Orca applies batching only to a selected set of operations, a design it calls selective batching (Yu et al., 2022). The attention mask page shows why each request must attend only to its own tokens.

Text description of the diagram

• Iteration-level batching re-forms the batch after every step.
• A finished slot is refilled on the next step.

Other cases an interviewer probes. Time to first token is the wait from submitting a request to receiving its first output token; a long prompt prefill in the same batch delays everyone's next token, so engines cap the tokens processed per step (a token budget). Overload needs admission control: reject or delay requests when the queue or the cache is full, rather than letting latency grow without bound. A priority queue needs a lock and a tiebreak counter so equal priorities keep arrival order.

Sampling, if you write it, must use a numerically stable softmax. If the model does not fit on one GPU, inference engines usually split the model with tensor parallelism; the fully sharded data parallel page shows the related training-side sharding. Note that the idea of per-step scheduling comes from Orca; vLLM's contribution is paged memory management. The design axes above are this page's synthesis of what interviewers discuss, not a published rubric.

How to answer it in an interview

The spoken structure: the loop, the per-step batch, the memory limit, then the optional speedup.

Common mistakes

• "I wait until I have N requests, then run them together." (wrong) That is request-level batching, which holds finished requests and delays new ones. (right) "I re-form the batch every iteration."
• "Throughput is limited by compute." (wrong) Serving throughput is memory-bound because the cache limits the batch. (right) "The cache budget limits how many requests run together."
• "Speculative decoding approximates the large model." (wrong) It changes speed, not the output distribution. (right) "The large model verifies every drafted token."
• "Continuous batching is vLLM's invention." (wrong) Orca introduced iteration-level scheduling; vLLM added paged cache memory. (right) "Orca scheduled per iteration; vLLM paged the cache."

Follow-up questions

How do you implement a thread-safe, priority-based request queue for an inference batcher?

Protect a heap with a lock, or use Python's queue.PriorityQueue, and store tuples of priority, an arrival counter and the request so ties keep arrival order. The batcher thread pops until its slots or token budget are full.

How do you batch embedding requests with batch-size and token constraints?

Add requests to a batch until either the maximum request count or the total token budget would be exceeded. Sorting by length reduces padding, and worker threads can pull finished batches from a shared queue. There is no growing cache, so one pass per batch is enough.

What happens when the engine is overloaded?

Admission control rejects or delays new requests once the queue is long or the cache budget is full, and returns a clear error so clients can retry. Without it, queue wait grows and every request slows down.

How does speculative decoding change the output?

It does not change the distribution: the large model accepts or rejects each drafted token, so the result matches sampling from the large model alone, only faster (Leviathan et al., 2023).

What would you stream to the user?

Send each token as its iteration finishes, since every request gets one token per step. Time to first token then depends on queue wait plus prefill, and later latency on the step time.

Sources

1. Orca: A Distributed Serving System for Transformer-Based Generative Models — USENIX OSDI '22 (Yu et al.), 11 Jul 2022
2. Efficient Memory Management for Large Language Model Serving with PagedAttention (arXiv 2309.06180) — arXiv (Kwon et al.), SOSP 2023, 12 Sep 2023
3. Fast Inference from Transformers via Speculative Decoding (arXiv 2211.17192) — arXiv (Leviathan, Kalman, Matias), ICML 2023, 30 Nov 2022