Building an Inference Engine

By now we have all the pieces: tokenization, embeddings, positional encoding, attention, RMSNorm, feed-forward networks, the language model head, and sampling. This chapter assembles those pieces into an actual inference engine. The engine is not a new mathematical idea. It is the operational system that places tensors in memory, runs the forward pass on hardware, maintains caches, and serves tokens to callers with acceptable latency and cost.

From Chapter 1 pseudocode to a real engine

In Chapter 1 the entire system fit into a short sketch: load weights, tokenize prompt, run a forward pass, sample a token, append, repeat. That sketch was correct, but it hid all the engineering. The real engine expands each function into a subsystem with concrete data structures, memory lifetimes, and hardware choices. The running example remains “What is the capital of France?”, but now we care less about the linguistic answer and more about what the machine must do to produce it quickly.

weights    = load_model_tensors_from_disk()
tokenizer  = load_tokenizer()
state      = allocate_runtime_state(weights, max_batch, max_seq)
prompt_ids = tokenizer.encode("What is the capital of France?")

prefill(state, prompt_ids)            // processes all prompt tokens
while !stop:
    logits   = decode_one_token(state) // uses KV cache
    next_id  = sample(logits)
    emit(next_id)
    state.append(next_id)

This looks small because the control loop really is small. The complexity is in how the functions are implemented, how much memory they consume, and how efficiently they map to the target hardware.

Plain English architecture

The most useful split is prefill versus decode. Prefill ingests the entire prompt and populates the KV cache. Decode then runs one position at a time, reusing cached keys and values so each step only computes work for the newest token. Most engines succeed or fail based on how well they manage that split.

A minimal but complete C-like reference implementation

typedef struct {
    Tensor tok_embeddings, rms_att, wq, wk, wv, wo;
    Tensor rms_ffn, w_gate, w_up, w_down;
    Tensor final_rms, lm_head;
    int n_layers, n_heads, n_kv_heads, d_model, d_head, d_ff, vocab;
} Model;

typedef struct {
    Tensor* k_cache;   // [layer][max_seq][n_kv_heads][d_head]
    Tensor* v_cache;   // [layer][max_seq][n_kv_heads][d_head]
    float* x;          // [d_model]
    float* xb;         // [d_model]
    float* q;          // [n_heads * d_head]
    float* k;          // [n_kv_heads * d_head]
    float* v;          // [n_kv_heads * d_head]
    float* att;        // [n_heads * max_seq]
    float* logits;     // [vocab]
    int pos, max_seq;
} RunState;

void forward_token(Model* m, RunState* s, int token_id) {
    lookup_embedding(s->x, m->tok_embeddings, token_id);
    for (int l = 0; l < m->n_layers; ++l) {
        rmsnorm(s->xb, s->x, m->rms_att[l]);
        matmul_qkv(s->q, s->k, s->v, s->xb, m->wq[l], m->wk[l], m->wv[l]);
        apply_rope(s->q, s->k, s->pos, m->n_heads, m->d_head, m->n_kv_heads);
        store_kv(s->k_cache[l], s->v_cache[l], s->pos, s->k, s->v);
        attention_decode(
            s->xb, s->q,
            s->k_cache[l], s->v_cache[l],
            s->att, s->pos + 1,
            m->n_heads, m->n_kv_heads, m->d_head);
        matmul_add_residual(s->x, s->xb, m->wo[l]);

        rmsnorm(s->xb, s->x, m->rms_ffn[l]);
        ffn_gated_residual(
            s->x, s->xb,
            m->w_gate[l], m->w_up[l], m->w_down[l],
            m->d_model, m->d_ff);
    }
    rmsnorm(s->xb, s->x, m->final_rms);
    matmul_vocab(s->logits, s->xb, m->lm_head, m->vocab, m->d_model);
}

void generate(Model* m, Tokenizer* tok, Sampler* sampler, const char* prompt) {
    int ids[MAX_PROMPT];
    int n = tokenize(tok, prompt, ids);
    RunState s = alloc_state(m, MAX_SEQ);

    for (int i = 0; i < n; ++i) {     // prefill
        s.pos = i;
        forward_token(m, &s, ids[i]);
    }

    int next = sample(sampler, s.logits);
    while (!should_stop(sampler, next, &s)) {
        emit_token(tok, next);
        s.pos += 1;
        forward_token(m, &s, next);   // decode one position
        next = sample(sampler, s.logits);
    }
    free_state(&s);
}

This omits many production concerns — batching, allocator reuse, quantization kernels, multi-GPU partitioning, grammar masks, speculative decoding — but it is complete enough to reveal the control structure. The engine is a loop around forward_token plus careful management of persistent state.

Step 1: load model weights from disk

The first subsystem is file I/O plus metadata interpretation. In modern LLM stacks, two common formats are safetensors and GGUF. Safetensors is a simple, safe tensor container widely used in Python-centric ecosystems. GGUF is popular in local runtimes and often packages quantized tensors plus tokenizer metadata in one format. Whatever the container, the engine must answer the same questions: what tensors exist, what are their shapes, what dtype are they stored in, and where should they live at runtime?

There are two broad strategies. Load-into-memory reads all tensors eagerly into RAM or VRAM. Memory-map keeps the file mapped and faults pages in on demand, which can reduce startup time and support very large models on systems with enough virtual address space. For GPU inference, many engines still end up copying tensors into device memory, but a memory-mapped host representation can remain useful for staging or tiered offload.

The important systems principle is that a tensor’s stored layout is not always its best execution layout. Some engines transpose, pack, or tile weights during load so later matmuls are faster. That increases startup cost but lowers per-token latency. This is one of the first places where serving is clearly not just “run the training code in inference mode.”

Step 2: initialize the tokenizer

The engine’s front door is the tokenizer. It loads merge rules, vocabulary tables, special tokens, and any chat-template behavior. Mechanically, tokenization is separate from neural inference, but operationally it is part of the engine because request latency starts at the moment text arrives. A slow tokenizer can dominate small prompts on CPU-only deployments.

For the prompt “What is the capital of France?”, the tokenizer might produce a sequence conceptually similar to [“What”, “ is”, “ the”, “ capital”, “ of”, “ France”, “?”], but the actual token IDs depend on the vocabulary. The engine must never assume word boundaries. All downstream logic — KV cache lengths, stop sequences, grammar constraints, penalties — works on token IDs and token byte strings, not on idealized words.

Step 3: the prefill pass

Prefill runs the prompt tokens through the model before any new tokens are generated. If the prompt length is T, prefill processes positions 0..T−1, computes attention across the full prompt, and stores keys and values for every layer and position. The final prompt token’s hidden state produces the logits for the first generated token.

Prefill is where GPUs usually shine. Sequence dimension exists, so matrix multiplications can be large and parallel. Attention cost is still heavier here because the prompt positions attend across the whole prompt, but the arithmetic intensity is higher than decode. In practice, this means time-to-first-token is often dominated by prompt length and prefill efficiency.

Prefill attention score shape per layer:
[n_heads, T, d_head] × [n_heads, d_head, T] → [n_heads, T, T]

That square [T, T] attention structure is why long prompts are expensive. The exact memory strategy varies with attention kernel choice, but the conceptual scaling is still there: longer prompts imply more pairwise token interactions.

Step 4: the decode loop

After prefill, the engine enters decode. This is the hot user-facing loop: one token at a time, one forward pass per token, until a stop condition is reached. The control flow is simple:

while request.is_active:
    logits = forward_one_position(last_token, kv_cache)
    next   = sample(logits, sampler_state)
    emit(next)
    append_to_sequence(next)
    if stop(next): break

The simplicity is deceptive. Every iteration reads a large fraction of the model weights, touches the KV cache for every layer, runs the output head over the vocabulary, and then calls the sampler. The serial dependency is absolute: token t+1 cannot be produced until token t is chosen. That is why interactive latency matters so much more than raw throughput for chat.

Step 5: KV cache management

Without a KV cache, decode would be catastrophically wasteful. You would recompute keys and values for the entire prefix on every new token. With the cache, each layer stores prior keys and values once, then the new query attends against that stored history. Compute drops dramatically, and the engine becomes viable.

KV bytes ≈ batch × n_layers × seq_len × 2 × n_kv_heads × d_head × bytes_per_element

The factor of 2 is for keys plus values. If you serve multiple requests, sequence length grows for each live request, and the cache can become a major memory consumer even when weights dominate the baseline footprint. This is why grouped-query attention and multi-query attention matter operationally: fewer KV heads mean a smaller cache.

Paged KV caches are increasingly popular because they decouple logical sequence growth from physical contiguity. Instead of reserving one giant contiguous buffer per request, the engine allocates fixed-size blocks and maps logical positions to blocks. That helps continuous batching and reduces fragmentation at scale.

Step 6: memory planning

The inference engine is a memory planner as much as a math engine. Before serving anything, you want a budget for weights, KV cache, activations, temporary workspaces, and safety margin. A practical first-order calculator is:

Total bytes
  ≈ weight_bytes
  + kv_cache_bytes
  + activation_bytes
  + workspace_bytes
  + fragmentation_margin

weight_bytes    = params × bytes_per_param
kv_cache_bytes  = batch × n_layers × seq_len × 2 × n_kv_heads × d_head × bytes_per_kv
activation_bytes≈ batch × max_live_tokens × d_model × bytes_per_act × small_constant

For a 7B-parameter model at 16-bit weights, weight_bytes ≈ 14 GB before overhead. A quantized 4-bit representation cuts that dramatically, but activations and KV cache still need real memory in some runtime dtype. The point is not to memorize one number. The point is to think in classes of memory: static weights, dynamic per-request state, and transient workspace.

Good engines separate persistent allocations from per-step scratch. Persistent allocations include weights, tokenizer tables, request metadata, and the KV cache. Scratch includes softmax buffers, temporary projections, and reduction workspaces. Reusing scratch buffers matters because allocator churn in a per-token loop is self-inflicted latency.

Where do the tensors actually live?

“Memory planning” becomes much clearer when you think in placement tiers rather than one giant pool. Cold storage is the model file on disk. Warm host memory may hold mapped tensors, tokenizer tables, and staging buffers. Hot accelerator memory holds the tensors that must be touched every token: active weights, KV cache blocks, live activations, and scratch space. Some engines also use a middle tier, such as pinned host memory, for fast transfers or partial offload.

Weights are mostly read-only and reused across every request, so they want stable placement and cache-friendly layout. The KV cache is per-request but extremely hot during decode, so moving it off device can destroy latency. Tokenizer metadata and request logs can stay on CPU. Once you frame the engine this way, many design decisions become ordinary locality questions rather than AI-specific mysteries.

Step 7: batching multiple requests

Single-request generation is the simplest mental model, but production engines almost always batch. The two important forms are prefill batching and continuous decode batching. Prefill batching groups many incoming prompts and processes them together. Continuous batching keeps a live set of decode requests on device and admits new work between steps as slots become available.

Continuous batching is operationally powerful because different requests are at different sequence positions. One request may be at token 8, another at token 300, another may have just finished prefill. The scheduler’s job is to keep the device busy while preserving each request’s causal order and sampler state. This often means a slot table, a request queue, and a mapping from logical request positions to physical KV pages.

loop:
    admit_finished_prefills_into_decode_batch()
    retire_completed_requests()
    compact_or_remap_kv_pages_if_needed()
    run_one_decode_step_for_all_live_requests()
    sample_one_token_per_live_request()
    stream_outputs_to_clients()

Batching is where “the engine” becomes unmistakably different from “the model.” The model is the same set of weights. The engine is the scheduler deciding how many requests share the hardware, how fairness works, and whether you optimize for best latency, best throughput, or some negotiated middle ground.

Request lifecycle, admission control, and eviction

Once multiple requests exist, the engine needs a lifecycle model. A request arrives, tokenizes, waits for admission, runs prefill, enters decode, streams tokens, then finishes, cancels, or times out. At each phase the engine must know what resources are reserved and what can be reclaimed. Good serving stacks make these transitions explicit because hidden state is where memory leaks and starvation bugs breed.

Admission control usually balances queue depth, free KV pages, and latency goals. If the device is nearly full, adding one more long prompt can hurt everyone already decoding. Eviction is simpler when a request completes normally: free its KV cache pages and sampler state immediately. Harder cases include client disconnects, streaming errors, and partial cancellations mid-batch. Robust engines treat cleanup as a first-class path, not an afterthought.

Step 8: quantization integration

A serving engine that ignores quantization ignores reality. Weight quantization changes the economics of deployment by reducing memory footprint and memory traffic. The engine must therefore know not only how to read quantized tensors, but also how to execute them efficiently.

The central design choice is whether to dequantize on the fly inside a fused matmul kernel or to pre-dequantize weights into a wider runtime format. Pre-dequantization simplifies compute but increases memory traffic and storage. On-the-fly dequantization keeps weights compact in memory and is often faster overall because decode is so bandwidth-sensitive, but it requires specialized kernels for each quantization scheme.

Quantization also affects memory planning. A 4-bit 7B model is dramatically smaller than a 16-bit 7B model, but the KV cache may still sit in 16-bit or 8-bit formats depending on kernel support and quality targets. Reducing weight precision does not automatically shrink every other tensor in the engine.

Execution target tabs: the same engine, different machines

Performance profiling: what to measure

Three metrics matter immediately: time to first token (TTFT), tokens per second during decode, and throughput across many concurrent requests. TTFT is mostly a prefill story plus scheduling delay. Tokens per second is mainly a decode story. Throughput reflects batching efficiency, admission control, and whether the hardware stays busy.

Profile prefill and decode separately. If prefill is slow, look at prompt batching, attention kernels, and tokenization overhead. If decode is slow, look at memory bandwidth, cache locality, KV layout, kernel launch overhead, and vocabulary projection cost. Different phases fail for different reasons.

Also keep a correctness profile, not just a speed profile. Compare a handful of known prompts against a trusted reference engine and store logits or sampled outputs for regression checks. Fast wrong answers are still wrong. Especially after adding quantized kernels, fused paths, or paged KV logic, you want a way to tell whether an optimization preserved semantics or quietly changed them.

In practice, a tiny golden set goes a long way: one short factual prompt, one long prompt, one structured-output prompt, and one repetition-stress prompt. If all four stay stable across engine changes, you have at least a minimal guardrail while iterating on performance.

Why decode becomes memory-bound: arithmetic intensity

Arithmetic intensity is roughly FLOPs / bytes moved. High intensity means you do a lot of computation per byte fetched; low intensity means data movement dominates. Prefill has decent intensity because batches and sequence lengths are larger. Decode is the opposite: one new token, lots of layers, huge weight matrices, and small activation shapes. The engine keeps streaming weights from memory for relatively little arithmetic each step.

This is why decode optimization often feels like storage engineering rather than pure linear algebra. Quantization helps because smaller weights reduce bytes moved. Fused kernels help because they reduce redundant reads and writes of activations. Better batching helps because multiple requests can reuse the same weight reads across more live work.

Kernel fusion opportunities

Engines look for opportunities to fuse adjacent operations so intermediate tensors never leave fast memory. Common fusions include RMSNorm + QKV projections, dequantize + matmul, RoPE + attention score preparation, bias + activation + gate multiply, and softmax + masking in attention or sampling. Each fusion reduces memory traffic and kernel launch overhead.

The trade-off is code complexity. Fused kernels are harder to write, harder to validate, and more specific to a hardware target. The right approach is usually layered: keep a clear reference path for correctness and add specialized fused paths only where profiling shows they matter.

Tensor shapes across the live engine

Tracking shapes honestly helps answer operational questions. For example: why does the output head remain expensive even for one token? Because [d_model] × [d_model, vocab] still touches a vocabulary-wide matrix. Why does the KV cache explode with long contexts? Because its sequence dimension grows with every generated token across every live request.

Operation and bound classification

Caching beyond KV: what else should persist?

The obvious cache is KV, but engines also benefit from caching tokenizer artifacts, packed/transposed weight layouts, compiled graph variants, grammar parser states, sampler history, and reusable scratch buffers. If request prompts repeat, some systems even cache prompt-prefix KV states so shared system prompts do not require full prefill every time. Prefix caching can be a large win for chat workloads with stable system instructions.

The principle is simple: if a value is expensive to compute, reused frequently, and valid across requests or across decode steps, ask whether it should persist. Just be honest about invalidation rules. Bad cache invalidation in an inference engine means wrong outputs, not merely missed performance.

Analogies that keep the systems story honest

There are two especially good analogies. First, the engine as a database execution engine: model files are tables, tokenization is parsing, prefill is the first large query stage, KV cache is a materialized working set, continuous batching is the scheduler, and sampling is the final row selection. Second, the engine as a deep pipeline machine: weights are large read-only memories, activations are live registers, and the decode loop is a serial dependency chain with extreme sensitivity to cache and bandwidth behavior.

The analogy to avoid is “it’s just one function call.” Architecturally, yes. Operationally, no. A mature engine is a loader, planner, scheduler, kernel dispatcher, allocator, stream manager, cache manager, and sampler wrapped around one conceptual forward pass.

Common mistakes when building your own engine

If you were implementing this yourself...

Start with the simplest correct path: CPU, float32 or float16, single request, greedy sampling, no quantization, contiguous KV cache, and a straightforward loader. Once outputs match a trusted reference, add one acceleration feature at a time: better kernels, quantized weights, paged KV cache, batching, prefix caching, speculative decoding, grammar constraints. At every step, keep a correctness harness that compares logits or sampled outputs on a tiny prompt like “What is the capital of France?”.

Implementation order:
1. Load weights + tokenizer metadata correctly
2. Run single-token forward pass against a known reference
3. Add prompt prefill and verify logits after the last prompt token
4. Add KV cache and verify decode matches no-cache reference
5. Add sampler and stop conditions
6. Add memory reuse and scratch buffers
7. Add batching / scheduling
8. Add quantized kernels and profile again

The durable lesson of this whole book is that transformer inference is not one giant mystery. It is a composition of well-defined numerical steps and equally well-defined systems problems. Loading, placement, batching, caching, and data movement determine whether the elegant math becomes a usable engine. By this point, the box is open.