AI Engineer Exhaustive
HardPlusLLaMA

Full Forward Pass

LLaMA

Hard

Complete Llama 3 model forward pass

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

The Llama 3 language model is a decoder-only Transformer that maps a token-id tensor to next-token logits. Section 3.2 and Figure 1 of the herd paper describe a dense stack trained for next-token prediction; Table 3 fixes depth, width, head counts, SwiGLU, a 128K-token vocabulary, and RoPE with (\theta=500{,}000). The paper does not write the inference control flow. That lives in Transformer.forward in the official repo: embed, slice the RoPE table, build a causal mask that accounts for the KV cache, run every TransformerBlock, apply a final RMSNorm, then project to the vocabulary.

This note is the full pass — embeddings through logits — not sampling. Llama.generate in llama/generation.py calls this forward once per decode step; temperature / top-(p) sit outside the model.

How it works

Inputs. tokens has shape (B \times S) (integer ids). start_pos is the cache index of the first new token (0 on a prefilling call that owns the whole prompt).

Embed. A vocab-parallel embedding (E \in \mathbb{R}^{V \times d}) (tok_embeddings) produces

h_0 = E[\texttt{tokens}] \in \mathbb{R}^{B \times S \times d}.

(V=128{,}000) and (d) is 4096 / 8192 / 16384 (Table 3). The paper’s tokenizer story (tiktoken 100K plus 28K multilingual tokens) affects how strings become ids, not this matmul.

RoPE slice. At construction the model precomputes freqs_cis with precompute_freqs_cis(head_dim, 2\cdot\texttt{max_seq_len}, rope_theta). The forward takes the rows

\texttt{freqs_cis}[\texttt{start_pos}:\texttt{start_pos}+S]

so rotary angles match absolute positions in the cache, not the local index (0\ldots S-1).

Mask. If (S=1) the mask is None (a single new query can see the whole cache). If (S>1) the code builds an (S \times S) upper-triangular (-\infty) matrix (torch.triu(..., diagonal=1)) and left-pads start_pos zeros:

M = \bigl[\, 0_{S \times \texttt{start_pos}} \;\big|\; \mathrm{triu}_{-\infty}(S)\bigr] \in \mathbb{R}^{S \times (\texttt{start_pos}+S)}.

Scores are scores + M inside attention, so query (i) cannot see keys after cache index (\texttt{start_pos}+i). This is the inference causal mask. The paper’s packed-document mask (no attention across document boundaries) is a pre-training detail and is not in this file.

Stack. For each of (L) blocks ((L \in {32,80,126})):

h_{\ell} = \mathrm{Block}_{\ell}(h_{\ell-1},\; \texttt{start_pos},\; \texttt{freqs_cis},\; M).

Each block is pre-norm GQA + SwiGLU as in the block note. Attention writes the new keys/values into cache_k / cache_v at [0:B, start_pos:start_pos+S] and reads [0:B, 0:start_pos+S].

Unembed.

\mathrm{logits} = W_{\mathrm{out}}\,\mathrm{RMSNorm}(h_L) \in \mathbb{R}^{B \times S \times V},

cast to float32. (W_{\mathrm{out}}) is a bias-free column-parallel linear (output). There is no tied embedding in the official module: tok_embeddings and output are separate.

The public generation.py Llama.build still asserts max_seq_len <= 8192. The July 2024 paper describes a 128K continued-pretrain; that length is a training/recipe fact, not what the April-era inference wrapper advertises.

Official code

llama/model.py — Transformer.__init__ (embed, ModuleList of blocks, norm, output, freqs_cis) and Transformer.forward (the sequence above). Decode orchestration is llama/generation.py (Llama.generate passes tokens[:, prev_pos:cur_pos] and prev_pos). Hyperparameters come from checkpoint params.json, not from ModelArgs defaults alone.

Watch-outs

Sources