18. Engine v1: load Qwen3 and chat

From model to engine

You have a model (Chapter 17), a sampler (8) and a cache (16). An engine joins them under a contract:

• Input: token IDs plus generation settings (temperature, top-p, max new tokens, stop tokens, seed).
• Output: generated token IDs, a finish reason (stop or length), and measurements.
• Owned state: the KV cache for the request, created when the request starts and released when it ends, even if it ends with an error.

Text appears only at the boundary. Turning a chat into token IDs and IDs back into text is the tokenizer’s job, and keeping it outside the engine core has practical benefits. The engine is testable with integer inputs alone, the same engine serves different tokenizers, and token-level outputs make parity checks exact.

Loading Qwen3-0.6B

Download the checkpoint (if you didn’t in Chapter 1):

uv pip install -r optional-requirements.txt
hf download Qwen/Qwen3-0.6B --local-dir models/Qwen3-0.6B

Because your parameter names match the checkpoint’s exactly (Chapter 17), loading is a checked copy:

@torch.no_grad()
def load_qwen3(directory, device="cpu", dtype=torch.bfloat16):
    """Dense Qwen3: our parameter names equal the checkpoint's, so this is a checked copy.  (Your engine: Chapter 18)"""
    cfg = Qwen3Config.from_hf(read_config(directory))
    with torch.device("meta"):
        model = Qwen3(cfg)                      # no memory yet: parameters are shapes only
    model = model.to_empty(device=device).to(dtype)
    if cfg.tie_word_embeddings:                 # to_empty breaks the tie; restore it
        model.lm_head.weight = model.model.embed_tokens.weight
    destinations = dict(model.named_parameters())
    loaded, head_alias = set(), None
    for name, value in snapshot_tensors(directory):
        if name == "lm_head.weight" and cfg.tie_word_embeddings:
            head_alias = value
            continue
        if name not in destinations:
            raise ValueError(f"Unexpected tensor {name}")
        assign(destinations[name], value, name)
        loaded.add(name)
    missing = set(destinations) - loaded
    if missing:
        raise ValueError(f"Missing tensors: {sorted(missing)[:5]}")
    if head_alias is not None and not torch.equal(head_alias.to(model.lm_head.weight), model.lm_head.weight):
        raise ValueError("Tied lm_head and embed_tokens disagree")
    return model.eval()

Three engineering points:

• Build on the meta device first. Qwen3(cfg) on the CPU would allocate and randomly initialize 596M parameters, only to overwrite them. Built under torch.device("meta"), parameters are shapes without storage. to_empty(device=...) then allocates uninitialized memory directly on the target device, and the loader fills it shard by shard. Peak memory is about one model plus one shard, not two models.
• Restore weight tying. to_empty creates fresh storage for every parameter, which silently unties the head from the embedding. Re-tie them, and if the checkpoint also contains lm_head.weight, verify it equals the embedding rather than letting file order decide which copy wins.
• Refuse what you don’t implement. Qwen3Config.from_hf rejects other model types, scaled RoPE and sliding windows. A loader that “mostly works” on an unsupported variant is worse than one that refuses.

The text boundary

The tokenizer comes from the checkpoint’s own files. Your engine uses Hugging Face’s implementation through AutoTokenizer, loaded with local_files_only=True so nothing is fetched or executed remotely. (Chapter 4’s BPE explains what it does; reading Qwen’s tokenizer.json with your own BPE is a stretch exercise.)

Three details decide whether chat works at all:

1. The chat template wraps messages in <|im_start|>role ... <|im_end|> markers and ends with the generation prompt <|im_start|>assistant\n (Chapter 4).
2. Thinking mode. Qwen3 can “think” in a <think>...</think> block before answering. enable_thinking=False in the template turns it off for short, direct answers; True gives better reasoning at the cost of many more tokens.
3. Stop tokens. The assistant ends its turn with <|im_end|> (ID 151645); <|endoftext|> (151643) ends a document. Stop on both, or the model will continue with an imaginary next user turn.

The engine

class LLM:
    def __init__(self, model, tokenizer=None, manifest=None):
        self.model = model.eval()
        self.tokenizer = tokenizer
        self.device = next(model.parameters()).device
        self.manifest = manifest or {}

    @classmethod
    def from_pretrained(cls, directory, device=None, dtype="bf16", **load_options):
        device = device or default_device()
        model = load_model(directory, device, parse_dtype(dtype), **load_options)
        tokenizer = None
        try:   # The tokenizer is the only third-party piece, used at the text boundary.
            from transformers import AutoTokenizer
            tokenizer = AutoTokenizer.from_pretrained(directory, local_files_only=True)
        except Exception:
            pass
        config = read_config(directory)
        manifest = {"model_dir": str(directory), "model_type": config.get("model_type"),
                    "dtype": str(parse_dtype(dtype)), "device": str(device),
                    "torch": torch.__version__, "python": platform.python_version(),
                    "gpu": torch.cuda.get_device_name() if str(device).startswith("cuda") else None}
        return cls(model, tokenizer, manifest)

    def encode(self, text, chat=False, thinking=False):
        if self.tokenizer is None:
            raise RuntimeError("No tokenizer: pass token IDs directly")
        if chat:
            text = self.tokenizer.apply_chat_template([{"role": "user", "content": text}], tokenize=False,
                                                      add_generation_prompt=True, enable_thinking=thinking)
        return self.tokenizer(text, add_special_tokens=False).input_ids

    def default_stop_ids(self):
        if self.tokenizer is None or self.tokenizer.eos_token_id is None:
            return ()
        ids = {self.tokenizer.eos_token_id}
        for token in ("<|im_end|>", "<|endoftext|>"):
            index = self.tokenizer.convert_tokens_to_ids(token)
            if isinstance(index, int) and index >= 0 and index != self.tokenizer.unk_token_id:
                ids.add(index)
        return tuple(ids)

    @torch.inference_mode()
    def stream(self, prompt_ids, config=GenerationConfig()):
        """Yield token IDs as they are produced; afterwards self.last_result holds the summary.  (Your engine: Chapter 18)

        Measures time to first token (prefill + first sample) and the mean time per later token.
        """
        ids = torch.tensor([list(prompt_ids)], device=self.device)
        total = ids.shape[1] + config.max_new_tokens
        if total > self.model.context_limit:
            raise ValueError("Prompt plus max_new_tokens exceed the model's context")
        generator = torch.Generator(device=self.device).manual_seed(config.seed)
        policy = dict(temperature=config.temperature, top_k=config.top_k, top_p=config.top_p,
                      min_p=config.min_p, generator=generator)
        cache = self.model.new_cache(1, total)
        synchronize(self.device)
        start = time.perf_counter()
        logits = self.model(ids, cache)
        produced, finish, first_token_time = [], "length", None
        try:
            for step in range(config.max_new_tokens):
                token = sample(logits[:, -1], **policy)
                value = int(token)                    # host sync: we must know it to stream it
                if first_token_time is None:
                    first_token_time = time.perf_counter()
                produced.append(value)
                yield value
                if value in config.stop_ids:
                    finish = "stop"
                    break
                if step + 1 < config.max_new_tokens:
                    logits = self.model(token, cache)
        finally:
            end = time.perf_counter()
            decode_steps = max(len(produced) - 1, 0)
            self.last_result = GenerationResult(produced, finish, metrics={
                "prompt_tokens": ids.shape[1], "new_tokens": len(produced),
                "ttft_ms": round(1000 * ((first_token_time or end) - start), 3),
                "ms_per_output_token": round(1000 * (end - (first_token_time or end)) / decode_steps, 3) if decode_steps else None,
                "tokens_per_s": round(len(produced) / (end - start), 2) if produced else 0.0})

    def generate(self, prompt, config=None, chat=False):
        """prompt is text (needs a tokenizer) or a list of token IDs."""
        config = config or GenerationConfig(stop_ids=self.default_stop_ids())
        prompt_ids = self.encode(prompt, chat=chat) if isinstance(prompt, str) else list(prompt)
        for _ in self.stream(prompt_ids, config):
            pass
        result = self.last_result
        if self.tokenizer is not None:
            result.text = self.tokenizer.decode(result.token_ids, skip_special_tokens=True)
        return result

    def run_manifest(self, prompt_ids, config):
        """Everything needed to reproduce (and compare) a run."""
        return {**self.manifest, "prompt_token_ids": list(prompt_ids), "generation": asdict(config)}

stream is the heart: prefill the prompt, sample, then decode one token at a time, yielding each token as soon as it exists so a UI can display it. The try/finally records the result even if the consumer stops reading early (a user pressing “stop” is a normal event, not an error). It measures two numbers users feel:

• TTFT (time to first token): from the start of the request to the first sampled token. That’s prefill plus one sample. It grows with prompt length.
• Time per output token (TPOT): the average time between later tokens. That’s decode. It’s nearly constant, rising slowly as the KV cache grows.

run_manifest records everything needed to reproduce a run: model directory and type, dtype, device, library versions, prompt token IDs and generation settings. Save one next to every result you want to keep or compare. Later, “it got faster” or “it got worse” questions become answerable.

Chat with it

python run.py chat --model-dir models/Qwen3-0.6B --prompt "Why is the sky blue? Answer in two sentences."

The text streams token by token, followed by the measurements, like this (exact text and timings will vary):

The sky appears blue because molecules in the atmosphere scatter shorter (blue) wavelengths of sunlight
much more strongly than longer (red) ones. ...
{"prompt_tokens": 23, "new_tokens": 41, "ttft_ms": ..., "ms_per_output_token": ..., "tokens_per_s": ..., "finish": "stop"}

Tip

If the model rambles, repeats the question or starts a fake “user:” turn, check the boundary before suspecting the model: was the generation prompt added, and are both stop IDs in stop_ids? llm.encode(text, chat=True) and tokenizer.convert_ids_to_tokens(ids) show exactly what the model sees.

A ladder of evidence

“It produces English” is weak evidence. A model with a subtly wrong RoPE still writes fluent sentences. Climb this ladder instead, and stop at the first rung that fails:

1. Unit operations against small hand-checked cases: Chapter 17’s tests.
2. Internal consistency: full, chunked and token-by-token forwards agree (Chapter 16’s tests). This catches position and cache bugs.
3. Independent parity on the real checkpoint: load the same snapshot into Hugging Face Transformers and into your engine, in FP32, and compare logits on fixed token IDs. Expect a maximum absolute difference around 1e-4 to 1e-3.
4. Layer-by-layer localization when rung 3 fails: register forward hooks on both models’ decoder layers, and compare the residual stream after each layer. The first layer that disagrees contains the bug.
5. Behavior: identical greedy text over several prompts, in FP32.

import torch
from transformers import AutoModelForCausalLM
from engine.loaders import load_qwen3

path = "models/Qwen3-0.6B"
ref = AutoModelForCausalLM.from_pretrained(path, dtype=torch.float32, attn_implementation="eager").eval()
mine = load_qwen3(path, device="cpu", dtype=torch.float32)
ids = torch.randint(0, 151936, (1, 33))
with torch.no_grad():
    diff = (mine(ids) - ref(ids).logits).abs()
print("max abs logit difference:", diff.max().item())

For BF16, compare like with like: your BF16 against Transformers’ BF16, both judged against FP32 (the three-comparison rule from Chapter 13). Greedy tokens can legitimately diverge after a few dozen steps in BF16, because a near-tie between two tokens flips on rounding. Compare logits, not long generations.

Memory budget

For Qwen3-0.6B in BF16 on a single request:

The logits row is a surprise to most people. Prefill only needs the last position’s logits, so computing all 4,096 × 151,936 of them wastes memory and a large matmul. Production engines slice the hidden states to the last position before the head. That’s a one-line stretch exercise below.

Measure against the ceiling

Use your Chapter 10 tools. Predict, then measure:

• Decode ceiling: decode_ceiling(1.19e9, bandwidth), about 230 tokens/s on DGX Spark and about 840 on an RTX 4090.
• Achieved: ms_per_output_token from the engine.

Engine v1 typically reaches a small fraction of the ceiling on a GPU. A 0.6B model’s decode step is hundreds of small kernels, each waiting on Python, plus one int(token) synchronization per step, so the GPU idles between kernels. A profile (torch.profiler) shows the gaps clearly. Chapter 19 removes them.

The Rust engine does the same, on the CPU

The Rust track loads the same snapshot and generates on the CPU. Its weights stay in BF16 and are widened inside a multi-threaded mat-vec, which is a direct application of the memory-bound decode model:

#![allow(unused)]
fn main() {
    /// Process one token at position cache.len and return the next-token logits.
    pub fn forward_token(&self, token: u32, cache: &mut KvCache) -> Vec<f32> {
        let c = &self.cfg;
        let pos = cache.len;
        let d = c.head_dim;
        let mut x = self.embed.row(token as usize);
        for (li, layer) in self.layers.iter().enumerate() {
            let h = rmsnorm(&x, &layer.input_norm, c.eps);
            let (mut q, mut k, v) = (layer.q.matvec(&h), layer.k.matvec(&h), layer.v.matvec(&h));
            for head in q.chunks_mut(d) {                    // per-head RMSNorm, then rotate
                head.copy_from_slice(&rmsnorm(head, &layer.q_norm, c.eps));
                rope(head, pos, c.theta);
            }
            for head in k.chunks_mut(d) {
                head.copy_from_slice(&rmsnorm(head, &layer.k_norm, c.eps));
                rope(head, pos, c.theta);
            }
            let (keys, values) = cache.append(li, &k, &v);
            let attended = causal_attention(&q, keys, values, 1, pos + 1, c.heads, c.kv_heads, d);
            for (xi, oi) in x.iter_mut().zip(layer.o.matvec(&attended)) {
                *xi += oi;
            }
            let h = rmsnorm(&x, &layer.post_norm, c.eps);
            let gated: Vec<f32> = layer.gate.matvec(&h).iter().zip(layer.up.matvec(&h)).map(|(g, u)| silu(*g) * u).collect();
            for (xi, mi) in x.iter_mut().zip(layer.down.matvec(&gated)) {
                *xi += mi;
            }
        }
        cache.len += 1;
        let h = rmsnorm(&x, &self.norm, c.eps);
        self.head.as_ref().unwrap_or(&self.embed).matvec(&h)
    }
}

#![allow(unused)]
fn main() {
/// Prefill the prompt token by token (simple, not fast), then decode `new_tokens` tokens.
pub fn generate(model: &Qwen3, prompt: &[u32], new_tokens: usize, policy: &crate::sampling::Policy,
                rng: &mut crate::rng::Rng, stop: &[u32]) -> Vec<u32> {
    let mut cache = model.new_cache(prompt.len() + new_tokens);
    let mut logits = vec![];
    for &t in prompt {
        logits = model.forward_token(t, &mut cache);
    }
    let mut out = vec![];
    for _ in 0..new_tokens {
        let next = crate::sampling::sample(&logits, policy, rng) as u32;
        out.push(next);
        if stop.contains(&next) || out.len() == new_tokens {
            break;
        }
        logits = model.forward_token(next, &mut cache);
    }
    out
}
}

cd rust
IDS=$(python tokenize_ids.py encode --model-dir ../models/Qwen3-0.6B "Why is the sky blue?")
OUT=$(cargo run --release -- generate --model-dir ../models/Qwen3-0.6B --ids "$IDS" --new-tokens 48)
python tokenize_ids.py decode --model-dir ../models/Qwen3-0.6B "$OUT"

Its cargo test reproduces the Python reference’s logits at every position to 1e-4, for both F32 and BF16 checkpoints. Compare its CPU tokens/s with your machine’s memory bandwidth divided by 1.19 GB.

Build it

pytest tests/test_ch18_engine.py
python run.py chat --impl engine --model-dir models/Qwen3-0.6B

Stretch exercises

1. ★ Compute logits only for the last position during prefill (return_hidden=True, slice, then lm_head). Measure the memory saved on a 4,096-token prompt. Where: the prefill model call in LLM.stream in engine/engine.py.
2. ★★ Implement rung 4: a function that registers hooks on every decoder layer of both models and prints the first layer whose output differs by more than a tolerance. Where: experiments/ch18.py (create it), using engine.qwen3 and the corresponding Transformers model.
3. ★★ Add a --thinking flag to run.py chat and compare the answer quality and token count of thinking and non-thinking modes on five reasoning questions. Where: the argument parser and cmd_chat in run.py; pass it as thinking=... to LLM.encode (which forwards it to the chat template).
4. ★★★ Load Qwen’s tokenizer.json with your own BPE implementation: parse the vocabulary and merges, implement the byte-to-unicode mapping and the pre-tokenizer regex (use the third-party regex package), and check that you reproduce AutoTokenizer’s IDs on 1,000 lines of text. Where: add a checkpoint-tokenizer loader in engine/tokenizer.py; compare it in experiments/ch18.py (create it).

Check your understanding

1. Why should you compare intermediate layer outputs before generated text?
2. Why can a model that fits in memory still fail to allocate for a long prompt?
3. Why are “tokens emitted” and “tokens processed into the cache” different counts?
4. Why does to_empty require you to re-tie the embedding and head?
5. What two stop tokens does a Qwen3 chat need, and what happens with only one?

Going deeper

• BALLM Chapter 5 §§5.4-5.5 for the GPT-2 version of loading pretrained weights; Raschka’s ch05/11_qwen3 notebook loads real Qwen3 weights into a from-scratch model.
• PMPP §20.4 and §20.6 for cache sizing and decode intensity.
• The Qwen3 model card (chat template, thinking mode, recommended sampling settings: temperature 0.7, top-p 0.8 for non-thinking mode).
• vLLM’s LLM class and SGLang’s Engine: the production versions of this chapter’s interface.