Make the Engine Talk

An inference engine has one job. Given a prompt, it produces the next token, sticks it on the end, and does the whole thing again, until the model decides it is finished. Attention, the KV cache, the kernels, the quantization — all the machinery in this book exists to make that one loop correct and fast.

So before we build any of it, let us run the finished engine and watch it work. You will end this chapter with a real model generating real text on your own machine, and a clear picture of the loop that everything else plugs into.

forge ships a small model, Qwen3-0.6B, and all the code needed to run it. The model, the weight loader, the tokenizer, and the generation loop are forge’s own code, written over the course of this book. No Hugging Face model class does the real work behind the curtain, and no serving framework. It is all here.

1.1 Run it

Clone forge at the commit in SOURCE_PIN, download the Qwen3-0.6B weights, and point FORGE_WEIGHTS at them. The whole program is short.

import os
import torch
from forge_engine.config import QWEN3_0_6B
from forge_engine.model import Qwen3
from forge_engine.weights import load_checkpoint
from forge_engine.tokenizer import Qwen3Tokenizer
from forge_engine.sampling import generate

weights = os.environ.get("FORGE_WEIGHTS", "weights")
model = Qwen3(QWEN3_0_6B)
load_checkpoint(model, f"{weights}/qwen3-0.6B-base.pth")
model.to("cpu").eval()
tokenizer = Qwen3Tokenizer(f"{weights}/tokenizer-base.json")

prompt = "The capital of France is"
text, _ = generate(model, tokenizer, prompt, torch.device("cpu"), max_new_tokens=20)
print(prompt + text)

Listing 1.1 · The driver. It wires forge’s own model, weight loader, tokenizer, and generate together, then prints the result.

It loads the weights into forge’s model, wraps the tokenizer, and asks for twenty tokens after the prompt. Greedy decoding is the default, so there are no sampling knobs to set. On a laptop CPU it finishes in a few seconds.

Listing 1.2 · A from-scratch engine continuing the prompt. Each sentence names a new country and its capital, until the twenty-token budget stops it mid-sentence.

The model was never handed a table of capitals. It produced that text one token at a time, and the twenty-token budget cut it off partway through the next sentence. Every one of those tokens came out of the loop we read next.

1.2 The loop behind it

generate in the driver is a thin wrapper. The real work is stream, the decode loop in forge’s sampling.py.

@torch.inference_mode()
def stream(
    model: LanguageModel,
    tokenizer: Tokenizer,
    prompt: str,
    device: torch.device,
    max_new_tokens: int = 256,
    temperature: float = 0.0,
    top_p: float | None = None,
    generator: torch.Generator | None = None,
    stats: Stats | None = None,
) -> Iterator[str]:
    if top_p is not None and not temperature:
        raise ValueError("top_p requires temperature > 0; greedy decoding would ignore it")

    st = stats if stats is not None else Stats()
    ids = tokenizer.encode(prompt)
    st.prompt_tokens = len(ids)
    x = torch.tensor(ids, device=device).unsqueeze(0)

    cache = KVCache(model.config.n_layers)
    model.reset_kv_cache()
    model.eval()

    prefill_start = time.perf_counter()
    logits = model(x, cache=cache)[:, -1]
    st.prefill_s = time.perf_counter() - prefill_start

    decode_start = time.perf_counter()
    for _ in range(max_new_tokens):
        if temperature:
            # softmax in fp32: over a 151,936-wide vocabulary bf16 lifts the tail
            probas = torch.softmax(logits.float() / temperature, dim=-1)
            probas = top_p_filter(probas, top_p)
            # multinomial has diverged and crashed on accelerator backends for
            # large draws; the round trip through cpu is cheap
            nxt = torch.multinomial(probas.cpu(), num_samples=1, generator=generator).to(device)
        else:
            nxt = torch.argmax(logits, dim=-1, keepdim=True)

        token = int(nxt.item())
        if token == tokenizer.eos_token_id:
            st.stopped = "eos"
            break

        st.generated += 1
        yield tokenizer.decode([token])
        logits = model(nxt, cache=cache)[:, -1]

    st.decode_s = time.perf_counter() - decode_start

Read it from the top. The st. assignments and time.perf_counter() calls are timing counters the benchmark chapters read later; ignore them here. The work is this: it turns the prompt into token ids and runs one forward pass over all of them at once — the prefill — which fills the KV cache and hands back the logits for the last position. Then it enters the loop. Each pass reads those logits, picks a token, and stops if the token is the end-of-sequence marker. Otherwise it decodes the token to text, yields it, and runs the model on that single new token to get the next set of logits.

What matters is the shape: one forward pass to prime the cache, then one cheap pass per token, each feeding the next. Everything expensive in an inference engine happens inside this loop, which is why most of this book is about the three calls it makes.

1.3 What the rest of Part I builds

The loop is complete, but three of its calls are still closed boxes. Part I opens them in the order the loop reaches for them.

• The forward pass, model(x, cache=cache) — the transformer that turns tokens into scores. The next chapter assembles it and proves it correct against an independent implementation; the chapters after open its four mechanisms — RMSNorm, rotary position embeddings, grouped-query attention, and the feed-forward block — one at a time.
• The cache, KVCache — the memory that lets each new token attend to the ones before it without recomputing them, and the first place the engine’s footprint grows with the conversation.
• The sampler — greedy, temperature, and top-p decoding, and why that softmax runs in fp32 over a vocabulary this wide.

By the end of Part I, the loop you just read runs on code you built and can follow line by line.