The Forward Pass

In the last chapter the engine printed Paris, then Berlin, then Rome, and you read the loop that drove it. One call inside that loop was still a closed box: model(x, cache=cache), the forward pass. That single call is what turned The capital of France is into the token Paris.

Two things happen here. You read forge’s real forward pass, top to bottom, until every line is one you could have written. Then you do the thing that separates an engine from a demo: you confirm the forward pass computes the right logits, by checking its output against a second, independent implementation of the same model.

A forward pass is one trip through the model. Token ids go in; scores come out. Given the tokens so far, it produces — for every position, and for each of the 151,936 words the model knows — a score. Those scores are the logits. The decode loop from the last chapter reads the scores at the final position and picks the next token. Everything the model has learned is brought to bear in this one trip.

2.1 One trip through the model

Here is the whole forward pass, from forge’s model.py.

    def forward(self, in_idx: torch.Tensor, cache: KVCache | None = None) -> torch.Tensor:
        _, t = in_idx.shape
        x = self.tok_emb(in_idx)

        if cache is not None:
            pos_start = self.current_pos
            pos_end = pos_start + t
            self.current_pos = pos_end
        else:
            pos_start, pos_end = 0, t

        mask = causal_mask(t, pos_start, pos_end, in_idx.device)

        for layer, block in enumerate(self.trf_blocks):
            past = cache.get(layer) if cache is not None else None
            x, entry = block(x, mask, self.cos, self.sin, pos_start, past)
            if cache is not None:
                cache.update(layer, entry)

        return self.out_head(self.final_norm(x).to(self.config.dtype))

Listing 2.1 · forge’s forward pass. Ids in at the top, logits out at the bottom, with a stack of identical blocks in between.

Read it with cache set to None for now — that is the branch this chapter is about. The cached path, which lets the engine avoid redoing work on every token, is the next chapter; skip its two if cache is not None arms on this read.

Follow the shape of x down the function. It arrives as token ids, (batch, tokens). The first line, self.tok_emb(in_idx), looks each id up in a table and replaces it with a learned vector, giving (batch, tokens, emb_dim) — every token is now a point in the model’s space instead of a bare integer. That shape then holds steady all the way down. causal_mask builds a mask so each position’s scores can depend only on itself and the tokens before it; the model is never allowed to read ahead. The for loop sends x through each transformer block in turn, and every block hands back an x of the exact same shape it was given. Only the final line changes the shape: self.out_head(...) projects each position from emb_dim back out to one score per word, (batch, tokens, vocab_size). Those are the logits.

So the pass has three movements: turn ids into vectors, push those vectors through a stack of identical blocks, then read the vectors back out as scores over the vocabulary. The blocks are where the model’s knowledge lives, and there are twenty-eight of them in the 0.6B model.

2.2 Inside one block

Every block does the same two things, in the same order.

    def forward(
        self,
        x: torch.Tensor,
        mask: torch.Tensor | None,
        cos: torch.Tensor,
        sin: torch.Tensor,
        start_pos: int,
        past: KVEntry | None,
    ) -> tuple[torch.Tensor, KVEntry]:
        attended, entry = self.att(self.norm1(x), mask, cos, sin, start_pos, past)
        x = x + attended
        return x + self.ff(self.norm2(x)), entry

Listing 2.2 · one transformer block: attention, then a feed-forward network, each added back onto the input.

The block runs x through attention, then adds the result back onto x. It runs that sum through a feed-forward network, and adds that back on too. Those two x + ... lines are residual connections: each sublayer proposes an adjustment, and the block adds the adjustment rather than replacing x with it, so information has a clear path straight down the stack. Notice that x is normalized — self.norm1, self.norm2 — before it enters each sublayer, never after. That ordering is fixed by the weights the model was trained with.

Attention is the step where a position gathers from the positions before it; the feed-forward network then refines each position on its own. Inside those two calls are the four pieces this book builds next, one chapter each:

• RMSNorm — the self.norm1 and self.norm2 that steady x before each sublayer.
• Rotary position embeddings — how attention is told where each token sits in the sequence.
• Grouped-query attention — the sideways reach itself, sharing key and value projections across heads to save memory.
• The feed-forward network — the per-position refinement, self.ff.

The entry the block returns is the key/value state for the cache, which the next chapter picks up. For this chapter, the block is a box that takes an x and returns a better x of the same shape.

2.3 Correct, not just running

The forward pass runs and produces a tensor of logits. That is not the same as producing the right logits. A single transposed matrix, a normalization applied after a sublayer instead of before, a position embedding rotated the wrong way — any of these still returns a perfectly shaped (batch, tokens, vocab_size) tensor. The engine would run, and the model would generate fluent, confident, wrong text. Shape tells you the plumbing connects; it says nothing about whether the water is clean.

So you need a check the forward pass cannot fake. forge’s answer is parity: take a second implementation of the same architecture, written independently, load the same weights into both, feed both the same token ids, and require the two sets of logits to agree to floating-point tolerance. If any line of forge’s forward pass computes the wrong thing, its logits drift away from the reference’s, and the check fails.

def test_uncached_logits_match_at_1e5(pair: tuple[Qwen3, Any]) -> None:
    mine, theirs = pair
    ids = token_ids(2, 9, seed=7)
    with torch.no_grad():
        torch.testing.assert_close(mine(ids), theirs(ids), rtol=1e-5, atol=1e-5)

Listing 2.3 · forge’s parity gate. pair holds forge’s model and an independent reference carrying the same weights; assert_close requires the two logit tensors to agree to a relative 1e-5, with a matching absolute floor.

The weights here are small and random — a tiny two-layer model, not the 0.6B download. That is deliberate: correctness is a property of the arithmetic, not of any particular checkpoint, and random weights exercise every line of the forward pass in a millisecond. forge’s test asserts the match silently; the companion driver for this chapter runs the same comparison with only forge installed and prints the largest disagreement between the two:

The gap is not zero, and it should not be. The two implementations do the same math in a slightly different order — forge scales the attention scores after the matmul, the reference scales the queries before it — and floating-point addition is not perfectly associative, so their last digits disagree. 5.96e-07 is the size of that rounding noise: about seventeen times smaller than the 1e-5 gate, far below any real error. A genuine mistake in the forward pass does not hide down there. It moves the gap up to the size of the logits themselves.

The forward pass now holds up to a check that can fail. It is still built from four boxes — normalization, position embeddings, attention, the feed-forward network — and the next chapters open them in the order a block reaches for them, starting with RMSNorm, the normalization that steadies x before each sublayer. The key/value cache and the sampler follow, and by the end of Part I every call in the loop from Chapter 1 runs on code you have read line by line.