The forward pass end to end: a Llama-style engine in 100 lines

An inference engine is a program that turns weights and token ids into more token ids, quickly. Before speed, correctness: this lesson implements the complete forward pass of a modern decoder (Llama-style) in plain Python, so you can see every step an engine like llama.cpp performs.

#What a forward step does

For one new token at position pos:

text

x = embedding[token]
for each layer:
    h = RMSNorm(x)
    q, k, v = h·Wq, h·Wk, h·Wv ;  rotate q and k by RoPE(pos)
    append k, v to this layer's KV cache
    for each query head: attend over all cached keys (its shared KV head), mix values
    x = x + (attended · Wo)                       # residual
    h = RMSNorm(x)
    x = x + Down( SiLU(Gate·h) ⊙ (Up·h) )          # SwiGLU feed-forward, residual
logits = Output · RMSNorm(x)

It is exactly the architecture from Modern upgrades, written for inference only: no gradients, one token at a time, weights read-only.

#The code

python

"""A tiny Llama-style inference engine in pure Python: RMSNorm, RoPE, grouped-query attention, SwiGLU, KV cache."""
import math, random


class Config:
    vocab, d, n_layer, n_head, n_kv, hidden, max_seq = 40, 16, 2, 4, 2, 24, 64

    @property
    def head_dim(self):
        return self.d // self.n_head


def mat(rng, r, c, std=0.15):
    return [[rng.gauss(0, std) for _ in range(c)] for _ in range(r)]


def init_weights(cfg, seed=0):
    rng = random.Random(seed)
    hd = cfg.head_dim
    W = {"emb": mat(rng, cfg.vocab, cfg.d), "out": mat(rng, cfg.vocab, cfg.d), "norm": [1.0] * cfg.d, "layers": []}
    for _ in range(cfg.n_layer):
        W["layers"].append({
            "wq": mat(rng, cfg.n_head * hd, cfg.d), "wk": mat(rng, cfg.n_kv * hd, cfg.d), "wv": mat(rng, cfg.n_kv * hd, cfg.d),
            "wo": mat(rng, cfg.d, cfg.n_head * hd), "gate": mat(rng, cfg.hidden, cfg.d), "up": mat(rng, cfg.hidden, cfg.d),
            "down": mat(rng, cfg.d, cfg.hidden), "n1": [1.0] * cfg.d, "n2": [1.0] * cfg.d})
    return W


def matvec(M, x):
    return [sum(a * b for a, b in zip(row, x)) for row in M]


def rmsnorm(x, g, eps=1e-5):
    s = (sum(v * v for v in x) / len(x) + eps) ** -0.5
    return [v * s * gi for v, gi in zip(x, g)]


def softmax(xs):
    m = max(xs)
    e = [math.exp(v - m) for v in xs]
    s = sum(e)
    return [v / s for v in e]


def silu(z):
    return z / (1 + math.exp(-z))


def rope(vec, pos, hd, base=10000.0):
    out = vec[:]
    for h in range(0, len(vec), hd):
        for i in range(0, hd, 2):
            th = pos * base ** (-i / hd)
            c, s = math.cos(th), math.sin(th)
            a, b = vec[h + i], vec[h + i + 1]
            out[h + i], out[h + i + 1] = a * c - b * s, a * s + b * c
    return out


class KVCache:
    def __init__(self, cfg):
        self.k = [[] for _ in range(cfg.n_layer)]
        self.v = [[] for _ in range(cfg.n_layer)]

    def __len__(self):
        return len(self.k[0])


def forward_token(cfg, W, token, pos, cache):
    hd, group = cfg.head_dim, cfg.n_head // cfg.n_kv
    x = W["emb"][token][:]
    for l, L in enumerate(W["layers"]):
        h = rmsnorm(x, L["n1"])
        q = rope(matvec(L["wq"], h), pos, hd)
        k = rope(matvec(L["wk"], h), pos, hd)
        v = matvec(L["wv"], h)
        cache.k[l].append(k)                       # append this position's K and V to the cache
        cache.v[l].append(v)
        att = []
        for head in range(cfg.n_head):
            kvh = head // group                    # grouped-query: several query heads share one KV head
            qs = q[head * hd:(head + 1) * hd]
            scores = [sum(a * b for a, b in zip(qs, kk[kvh * hd:(kvh + 1) * hd])) / math.sqrt(hd) for kk in cache.k[l]]
            w = softmax(scores)
            att += [sum(w[t] * cache.v[l][t][kvh * hd + i] for t in range(len(w))) for i in range(hd)]
        x = [a + b for a, b in zip(x, matvec(L["wo"], att))]
        h = rmsnorm(x, L["n2"])
        ff = [silu(g) * u for g, u in zip(matvec(L["gate"], h), matvec(L["up"], h))]
        x = [a + b for a, b in zip(x, matvec(L["down"], ff))]
    return matvec(W["out"], rmsnorm(x, W["norm"]))


def argmax(xs):
    return max(range(len(xs)), key=lambda j: xs[j])


def generate(cfg, W, prompt, n_new, temperature=0.0, seed=0):
    rng = random.Random(seed)
    cache = KVCache(cfg)
    logits = None
    for pos, t in enumerate(prompt):               # prefill
        logits = forward_token(cfg, W, t, pos, cache)
    out = []
    for i in range(n_new):                         # decode
        if temperature == 0:
            tok = argmax(logits)
        else:
            tok = rng.choices(range(cfg.vocab), weights=softmax([l / temperature for l in logits]))[0]
        out.append(tok)
        logits = forward_token(cfg, W, tok, len(prompt) + i, cache)
    return out

#The correctness test

Generation has two phases (roofline): prefill feeds the prompt tokens and fills the cache, then decode produces one token per step. To prove the cache is correct, compare it with a reference that recomputes everything from scratch at every step:

python

cfg = Config(); W = init_weights(cfg)
prompt = [3, 14, 15, 9, 2, 6]
fast = generate(cfg, W, prompt, 10)
seq = prompt[:]
for _ in range(10):                      # reference: no cache reuse
    c = KVCache(cfg)
    for pos, t in enumerate(seq): lg = forward_token(cfg, W, t, pos, c)
    seq.append(argmax(lg))
print(fast == seq[len(prompt):])         # True
print(2 * cfg.n_layer * cfg.n_kv * cfg.head_dim * 2, "bytes of KV per token (fp16)")   # 64  (128 with full MHA)

Running it prints True and 64 bytes per token (128 without grouped-query attention). The generated tokens are [30, 30, 30, ...]: with random weights the model has nothing meaningful to say, which is fine; the test checks the mechanism. When you load trained weights (next lesson), the same code produces text.

This "two implementations must agree" pattern is the most important testing habit in engine work: a slow, obviously correct reference, and a fast one checked against it. Real projects do this with per-layer tensor comparisons against a framework like PyTorch.

#Where the time goes (and why it matters later)

Count the work for one decode step in this toy model: every weight matrix is multiplied once by a vector, so the work is 2 × parameters FLOPs and every weight is read once. In a real model those reads dominate time: this is the memory-bound regime (roofline). Everything in the next lessons (quantisation, vectorised kernels, batching) attacks that fact.

#Exercises

Add a debug function that prints the shape of every intermediate (for example q has n_head × head_dim entries). Verify that k has n_kv × head_dim.
Set n_kv = 4 (full multi-head). Does cache equality still hold? What happens to bytes per token?
Make generate sample with temperature=0.8 and a fixed seed. Is it reproducible? Why?
Count the multiply-adds in one decode step for this config. (Hint: sum of matrix sizes.)
Intentionally break RoPE (skip it for keys). Does the cache-vs-recompute test still pass? What does that tell you about what the test can and cannot detect?

Answers (sketch): (2) Equality holds; bytes double to 128. (3) Yes; the random generator is seeded, and the model is deterministic. (4) Per layer wq 64×16 + wk 32×16 + wv 32×16 + wo 16×64 + gate 24×16 + up 24×16 + down 16×24 = 1,024+512+512+1,024+384+384+384 = 4,224; two layers 8,448; plus output head 40×16 = 640: ≈ 9,088 multiply-adds. (5) It still passes, because both paths share the same (wrong) code: the test catches cache bugs, not modelling bugs. Check against a reference implementation to catch those.

#Next

Loading real weights from a file: GGUF and mmap.