Abstract

A language model turns a sequence of token ids into a distribution over the next token, and does it by repeating one block many times. This article follows that computation end to end: tokenization, the embedding and position lookup, the nine internal steps of a prefill block, the final projection to logits, and the sampling loop that turns them back into text. Every stage is given with its tensor shape, so the reader can track a batch through the model without guessing. It then sets out the two facts the rest of the series builds on — that a decode step reuses keys and values rather than recomputing them, and that attention is the only place in the block where tokens interact. A closing section maps the notation to code and counts where the parameters actually go, which is not where the diagrams suggest.

1. Introduction

Large Language Models (LLMs) are (mostly) autoregressive models: they predict the next token given the tokens so far, and repeat this loop until stopping. This post focuses on inference: the overall flow, what vectors get computed, and how the attention mechanism combines information across tokens. We also focus on the what data is generated in each phase of an LLC and how they are used in other phases of an inference process.


2. Inference: the end-to-end loop

At a high level, inference is:

  1. Tokenize the prompt into token IDs.
  2. Embed token IDs into vectors (plus positional information).
  3. Run NlayersN_{\mathrm{layers}} Transformer blocks (attention + MLP).
  4. Produce logits for the next token; convert to probabilities.
  5. Select the next token (argmax or sampling), append it, and repeat.

The loop runs until an end condition: max tokens, an EOS token, or a stopping rule.

2.1 Overview

Prompt text flows through tokenization, embeddings, repeated Transformer blocks, and next-token sampling before being displayed.
Figure 1. End-to-end LLM inference. Each Transformer block applies causal self-attention and a per-token MLP before the final hidden state produces the next-token logits.

3. Tokenize and Positions embedding

Here, there are LL tokens in each prompt, each embedded in a DD-dimensional space. A batch contains BB such prompts.

3.1 Tokenization (text → IDs)

Modern LLMs typically use a subword tokenizer (BPE, Unigram, WordPiece, etc.). The prompt becomes a sequence of token IDs:

(t1,t2,…,tL),ti∈{1,…,∥V∥} (t_1, t_2, \dots, t_L), \quad t_i \in \{1, \dots, \|V\|\}

where ∥V∥\|V\| is vocabulary size.

Practical notes:

  • Tokenization defines the model’s “alphabet.” It impacts context length, efficiency, and what counts as a single step of generation.
  • Special tokens (BOS/EOS, system delimiters, etc.) are often inserted by the serving stack.

3.2 Turning token IDs into vectors

The prompt is now a sequence of integers. The next two steps turn each integer into a DD-dimensional vector, then add position information to it, producing the block input X0X_0.

3.3 Embedding lookup

The model stores an embedding table E∈R∥V∥×DE \in \mathbb{R}^{\|V\| \times D}.

For each token ID tit_i, we take a row of this table:

xi=E[ti]∈RD x_i = E[t_i] \in \mathbb{R}^D

Stacking all tokens gives:

X(b)=[x1Tx2T⋮xLT]∈RL×D,X∈RB×L×D X^{(b)} = \begin{bmatrix} x_1^T \\ x_2^T \\ \vdots \\ x_L^T \end{bmatrix} \in \mathbb{R}^{L \times D}, \qquad X \in \mathbb{R}^{B \times L \times D}

3.4 Positional information

Transformers need position information because attention alone is permutation-invariant. Common approaches:

  • Absolute positional embeddings: X←X+PX \leftarrow X + P with P∈RB×L×DP \in \mathbb{R}^{B \times L \times D}.
  • Rotary / relative position methods (e.g., RoPE): position is injected into the attention computation (often into QQ and KK) rather than added to XX.

Either way, the result is an initial tensor X0∈RB×L×DX_0 \in \mathbb{R}^{B \times L \times D}.

4. Inside one Transformer block (MHA + FFN)

A transformer block is the key block for the success of LLMs. At a high-level, it mixes information across tokens (via attention) and then applies a nonlinear transformation to each token vector (via the MLP/FFN). The mixing step allows the model to build up contextual understanding, due to the presence of positional encoding. While the MLP step allows it to create new features and representations that are not just linear combinations of the input.

4.1 What a Transformer block computes (vectors everywhere)

Each block maps Xℓ↦Xℓ+1X_\ell \mapsto X_{\ell+1} using (i) multi-head self-attention and (ii) an MLP, typically with residual connections and layer normalization. Here, ℓ∈{1,…,Nlayers}\ell \in \{1, \ldots, N_{\mathrm{layers}}\} indexes the block number, and Xℓ∈RB×L×DX_\ell \in \mathbb{R}^{B \times L \times D} is the hidden-state tensor at that layer.

The exact ordering varies (pre-norm vs post-norm), but the core computations are consistent.

4.2 Complete prefill chain: from XX to context-mixed vectors

Let:

  • BB: batch size
  • LL: prompt sequence length
  • DD: model hidden dimension
  • HH: number of attention heads
  • dh=D/Hd_h = D / H: dimension per head
  • DffD_{\mathrm{ff}}: FFN intermediate dimension
  • LN⁡\operatorname{LN}: layer normalization

Step 1. Input hidden states. The block receives one DD-dimensional vector for every token position.

X∈RB×L×D X \in \mathbb{R}^{B \times L \times D}
⇓normalize before attention\Downarrow\quad\text{normalize before attention}

Step 2. Pre-attention normalization. Normalize each token vector before forming its attention projections.

X~=LN⁡(X),X~∈RB×L×D \widetilde{X} = \operatorname{LN}(X), \quad \widetilde{X} \in \mathbb{R}^{B \times L \times D}
⇓QKV projections\Downarrow\quad\text{QKV projections}

Step 3. Query, key, and value projections. Learned linear maps create three views of each normalized token.

Q=X~WQ,K=X~WK,V=X~WVWQ,WK,WV∈RD×DQ,K,V∈RB×L×D \begin{aligned} Q &= \widetilde{X}W_Q, \quad K = \widetilde{X}W_K, \quad V = \widetilde{X}W_V \\ W_Q, W_K, W_V &\in \mathbb{R}^{D \times D} \\ Q, K, V &\in \mathbb{R}^{B \times L \times D} \end{aligned}
⇓split into H heads\Downarrow\quad\text{split into H heads}

Step 4. Split into attention heads. Reshape each projection into HH independent dhd_h-dimensional subspaces.

Q,K,V∈RB×H×L×dhdh=DH \begin{aligned} Q, K, V &\in \mathbb{R}^{B \times H \times L \times d_h} \\ d_h &= \frac{D}{H} \end{aligned}
⇓for every head h\Downarrow\quad\text{for every head h}

Step 5. Causal attention. Each head scores only current and earlier positions, then forms a weighted sum of values.

Sh=QhKhTdh+Mcausal,Sh∈RB×L×LPh=softmax⁡(Sh),Ph∈RB×L×LAh=PhVh,Ah∈RB×L×dh \begin{aligned} S_h &= \frac{Q_h K_h^T}{\sqrt{d_h}} + M_{\mathrm{causal}}, \quad S_h \in \mathbb{R}^{B \times L \times L} \\ P_h &= \operatorname{softmax}(S_h), \quad P_h \in \mathbb{R}^{B \times L \times L} \\ A_h &= P_h V_h, \quad A_h \in \mathbb{R}^{B \times L \times d_h} \end{aligned}
⇓concatenate heads and project\Downarrow\quad\text{concatenate heads and project}

Step 6. Merge the heads. Concatenate the per-head results and project them back to the model dimension.

A=Concat⁡(A1,…,AH),A∈RB×L×DO=AWO,WO∈RD×D,O∈RB×L×D \begin{aligned} A &= \operatorname{Concat}(A_1, \ldots, A_H), \quad A \in \mathbb{R}^{B \times L \times D} \\ O &= AW_O, \quad W_O \in \mathbb{R}^{D \times D}, \quad O \in \mathbb{R}^{B \times L \times D} \end{aligned}
⇓add attention residual\Downarrow\quad\text{add attention residual}

Step 7. Attention residual. Add the attention update to the original, unnormalized block input.

Y=X+O,Y∈RB×L×D Y = X + O, \quad Y \in \mathbb{R}^{B \times L \times D}
⇓normalize before FFN\Downarrow\quad\text{normalize before FFN}

Step 8. Pre-FFN normalization. Normalize the residual state before the position-wise feed-forward network.

Y~=LN⁡(Y),Y~∈RB×L×D \widetilde{Y} = \operatorname{LN}(Y), \quad \widetilde{Y} \in \mathbb{R}^{B \times L \times D}
⇓apply FFN and add residual\Downarrow\quad\text{apply FFN and add residual}

Step 9. Feed-forward update. Transform each normalized token independently, then add that update to the residual state.

F1=ϕ(Y~W1+b1),F1∈RB×L×DffF2=F1W2+b2,F2∈RB×L×DXout=Y+F2,Xout∈RB×L×D \begin{aligned} F_1 &= \phi(\widetilde{Y}W_1 + b_1), \quad F_1 \in \mathbb{R}^{B \times L \times D_{\mathrm{ff}}} \\ F_2 &= F_1W_2 + b_2, \quad F_2 \in \mathbb{R}^{B \times L \times D} \\ X_{\mathrm{out}} &= Y + F_2, \quad X_{\mathrm{out}} \in \mathbb{R}^{B \times L \times D} \end{aligned}

This is a pre-norm block: layer normalization occurs before attention and before the FFN. Some architectures use a different normalization placement or RMSNorm instead.

Why multiple heads? A single attention head can only learn one “type” of relationship between positions (e.g., syntactic adjacency). With HH heads, the model can attend to different aspects simultaneously — one head might track subject–verb agreement, another might focus on nearby tokens, yet another on coreference.

Multi-head attention repeats the attention computation for every head in every transformer layer. For example, Llama 2 7B has 32 layers with 32 attention heads per layer: at each layer, those 32 heads run in parallel on the same token sequence. Conceptually, head hh has its own learned WQ(h)W_Q^{(h)}, WK(h)W_K^{(h)}, and WV(h)W_V^{(h)}, so it forms different queries, keys, values, and attention weights from the same input.

During decoding, each head also produces its own key and value history. The layer’s KV cache keeps these separately along its head dimension: Kcache,Vcache∈RB×H×t×dhK_{\mathrm{cache}}, V_{\mathrm{cache}} \in \mathbb{R}^{B \times H \times t \times d_h}. Thus, the key/value vectors for one head are not shared with another head; the serving system stores the HH distinct per-head histories and reads the appropriate one when computing the next token.

For a numerical walkthrough of causal attention, see the worked example in §6.

4.3 The causal mask: why it matters

In autoregressive generation the model must not peek at future tokens — token ii can only attend to positions 1,…,i1, \dots, i. The causal mask enforces this by adding −∞-\infty to illegal positions before softmax, which drives those attention weights to zero.

Why is this critical?

  • Without the mask, the model could “cheat” during training by reading the token it is supposed to predict, making the learned weights useless.
  • At inference the future tokens don’t exist yet, so the mask matches reality — but it must be present during training so the model learns correct causal distributions.

The residual connection then adds the original input back: X′=X+MHA(X)X' = X + \text{MHA}(X), so information is never lost. The FFN runs next (§4.4).

4.4 MLP / FFN: per-token nonlinearity

In many write-ups, the per-token MLP is called the FFN layer. It is conceptually simple but does a lot of the model’s nonlinear work. Unlike attention, it never mixes tokens: it transforms each position’s vector independently.

Typical shape choices use an expansion factor Dff≫DD_{\mathrm{ff}} \gg D (often Dff≈4DD_{\mathrm{ff}} \approx 4D):

FFN⁡(X)=W2 σ(XW1+b1)+b2 \operatorname{FFN}(X) = W_2\,\sigma(X W_1 + b_1) + b_2

with W1∈RD×DffW_1 \in \mathbb{R}^{D \times D_{\mathrm{ff}}} and W2∈RDff×DW_2 \in \mathbb{R}^{D_{\mathrm{ff}} \times D}, where σ\sigma is often GELU or a gated variant (e.g., SwiGLU). This increases representational capacity beyond linear mixing.

Key points for inference:

  • The FFN is position-wise: it does not mix tokens (no L×LL \times L interaction). It transforms each token vector independently.
  • The compute is dominated by dense GEMMs (B×L×DB \times L \times D by D×DffD \times D_{\mathrm{ff}}, then B×L×DffB \times L \times D_{\mathrm{ff}} by Dff×DD_{\mathrm{ff}} \times D), which GPUs handle efficiently.
  • Gated variants (e.g., SwiGLU) change the exact formula but keep the same idea: expand D→DffD \to D_{\mathrm{ff}}, apply a nonlinearity/gate, then project back to DD.
  • Because the work is dense and regular rather than data-dependent, the FFN usually accounts for the larger share of both parameters and FLOPs in a block.

5. From final vectors to the next token

After NlayersN_{\mathrm{layers}} blocks, we have Xout∈RB×L×DX_{\mathrm{out}} \in \mathbb{R}^{B \times L \times D}. To predict the next token for one request, we use its last-position vector xL∈RDx_L \in \mathbb{R}^D.

Compute logits over the vocabulary:

z=xLWU+b,WU∈RD×∥V∥ z = x_L W_U + b, \quad W_U \in \mathbb{R}^{D \times \|V\|}

Many models tie weights so that WU≈ETW_U \approx E^T.

Convert logits to probabilities with temperature TT:

p(y=v∣context)=softmax⁡(zT)v p(y = v \mid \text{context}) = \operatorname{softmax}\left(\frac{z}{T}\right)_v

Temperature rescales every logit before softmax. At T=1T=1, the model uses its original distribution; T<1T<1 sharpens it toward high-probability tokens, making output more predictable, while T>1T>1 flattens it and makes lower-probability tokens more likely. It is useful when the application needs to trade deterministic behavior for variety without changing the model’s weights.

Then choose the next token:

  • Greedy decoding: choose arg⁡max⁡vp(v)\arg\max_v p(v). This is deterministic and useful for reproducible outputs, but can sound repetitive because it always chooses the locally most likely token.
  • Top-kk sampling: retain only the kk most likely tokens, renormalize their probabilities, and sample from that set. This removes an arbitrarily long tail of implausible tokens.
  • Nucleus (top-pp) sampling: retain the smallest set of tokens whose cumulative probability is at least pp, then renormalize and sample. Unlike top-kk, the candidate-set size adapts: it grows when the model is uncertain and shrinks when it is confident.
  • Repetition penalties: lower the score of tokens that have already appeared in the generated context before sampling. They help prevent loops and copied phrases, but an overly strong penalty can make necessary names or technical terms less likely.

Generation also needs stopping controls: a maximum-token limit bounds runtime and output size; an EOS token lets the model signal completion; and a serving application can define stop sequences such as a chat delimiter to stop at a protocol boundary.

This yields a new token ID tL+1t_{L+1}, appended to the context, and the loop repeats.

That paragraph is a complete specification, and it is short enough to translate directly into code. §9.8 does exactly that, and shows how temperature, top-kk, and top-pp slot in between the logits and the sampled token.


6. Worked example: prefill and two-token decoding

This deliberately small model has one attention head with D=dh=2D=d_h=2 and identity projections, so Q=K=V=XQ=K=V=X. It generates two tokens after the prompt I like: first tea, then ..

Use these token vectors:

xI=[1,0],xlike=[0,1],xtea=[1,1] x_{I} = [1, 0],\quad x_{like} = [0, 1],\quad x_{tea} = [1, 1]

For the toy output projection, use the following vocabulary vectors:

etea=[0,1],e.=[1,0.4],ecoffee=[0.2,0.5] e_{tea} = [0, 1],\quad e_{.} = [1, 0.4],\quad e_{coffee} = [0.2, 0.5]

6.1 Prefill: process the prompt at once

Prefill runs the full prompt through the model in parallel. For I like, the input is:

Xprompt=[1001] X_{\mathrm{prompt}} = \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}

The causal mask prevents I from seeing like, while like can see both positions. With Q=K=V=XpromptQ=K=V=X_{\mathrm{prompt}}, the masked attention scores and probabilities are:

Sprefill=12[1−∞01],Pprefill≈[1.000.000.330.67] S_{\mathrm{prefill}} = \frac{1}{\sqrt{2}} \begin{bmatrix} 1 & -\infty \\ 0 & 1 \end{bmatrix}, \qquad P_{\mathrm{prefill}} \approx \begin{bmatrix} 1.00 & 0.00 \\ 0.33 & 0.67 \end{bmatrix}

The final prompt position produces:

hlike=0.33 [1,0]+0.67 [0,1]≈[0.33,0.67] h_{like} = 0.33\,[1, 0] + 0.67\,[0, 1] \approx [0.33, 0.67]

At the same time, prefill stores the prompt keys and values for this layer:

Kcache=Vcache=[1001] K_{\mathrm{cache}} = V_{\mathrm{cache}} = \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}

Use hlikeh_{like} to score the next token:

ztea=0.67,z.≈0.60,zcoffee≈0.40 z_{tea} = 0.67,\quad z_{.} \approx 0.60,\quad z_{coffee} \approx 0.40

After softmax, tea has the highest probability, so it is the first generated token. Notice that the first generated token is selected from the last prefill output; it does not require reprocessing the prompt.

6.2 Decode: append tea and generate the second token

The next decode step processes only the newly generated token tea. Its query, key, and value are q3=k3=v3=[1,1]q_3=k_3=v_3=[1,1]. The query attends over the two cached prompt keys plus its new key:

s3=12[1,  1,  2]⟹p3≈[0.248,  0.248,  0.503] s_3 = \frac{1}{\sqrt{2}}[1,\; 1,\; 2] \qquad\Longrightarrow\qquad p_3 \approx [0.248,\; 0.248,\; 0.503]

The new hidden state is the weighted value sum:

htea=0.248 [1,0]+0.248 [0,1]+0.503 [1,1]≈[0.751,  0.751] h_{tea} = 0.248\,[1, 0] + 0.248\,[0, 1] + 0.503\,[1, 1] \approx [0.751,\; 0.751]

Append the new key and value rather than recomputing the prompt’s entries:

Kcache=Vcache=[100111] K_{\mathrm{cache}} = V_{\mathrm{cache}} = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix}

Now the same output projection selects the second generated token:

ztea≈0.75,z.≈1.05,zcoffee≈0.53 z_{tea} \approx 0.75,\quad z_{.} \approx 1.05,\quad z_{coffee} \approx 0.53

The period has the highest probability, so the completed generated suffix is tea .. In a real system, the period would be passed to another one-token decode step unless an EOS token or stop sequence ends generation.


7. What changes during real inference

Real deployments add a few important engineering layers:

  • KV cache: during autoregressive decoding, you don’t recompute KK and VV for all past tokens each step. You cache them per layer and only compute the new token’s projections. What that cache costs, and the attention variants built to shrink it, is the subject of the next article.
  • Batching: serving stacks batch multiple requests/tokens to use GPU efficiently.
  • Precision: inference often uses FP16/BF16, sometimes INT8/FP8 quantization.
  • Decoding policy: sampling settings (temperature, top-pp, repetition penalties) strongly affect outputs without changing model weights.

8. KV caching: decoding as append-only attention

The key observation is that in autoregressive decoding, at step tt you only need one new query (for the new position) but you need keys/values for all positions so far.

For each layer ℓ\ell, maintain cached matrices:

Kcache(ℓ),Vcache(ℓ)∈RB×H×t×dh K^{(\ell)}_{\mathrm{cache}}, V^{(\ell)}_{\mathrm{cache}} \in \mathbb{R}^{B \times H \times t \times d_h}

When a new token arrives, you compute only its projections and append:

kt=xtWK,  vt=xtWV⇒Kcache←[Kcache; kt],  Vcache←[Vcache; vt] k_t = x_t W_K,\; v_t = x_t W_V\quad\Rightarrow\quad K_{\text{cache}} \leftarrow [K_{\text{cache}};\, k_t],\; V_{\text{cache}} \leftarrow [V_{\text{cache}};\, v_t]

Then compute the new query qt=xtWQq_t = x_t W_Q and attend over the cached keys/values:

αt=softmax⁡(qt(Kcache)Tdh+Mt)∈RB×H×1×t,ht=αtVcache∈RB×H×1×dh \alpha_t = \operatorname{softmax}\left(\frac{q_t (K_{\mathrm{cache}})^T}{\sqrt{d_h}} + M_t\right) \in \mathbb{R}^{B \times H \times 1 \times t},\quad h_t = \alpha_t V_{\mathrm{cache}} \in \mathbb{R}^{B \times H \times 1 \times d_h}

Here MtM_t is the (trivial) causal mask for this single query row; in practice you just ensure the new token cannot attend to “future” positions (which do not exist yet).

8.1 ASCII timeline

Prefill (prompt length = L): compute all layers on all tokens

tokens:   [ 1   2   3  ...  L ]
cache K:  [k1  k2  k3  ...  kL]
cache V:  [v1  v2  v3  ...  vL]

Decode step L+1 (generate one token): compute only the new token’s q,k,v

new token:                  [ L+1 ]
append K:   [k1  k2  ...  kL | k(L+1)]
append V:   [v1  v2  ...  vL | v(L+1)]

attention for the new token:
q(L+1)  x  [k1..k(L+1)]^T  -> weights over positions 1..L+1  -> weighted sum of [v1..v(L+1)]

8.2 Why it matters

  • Compute: without caching, each new token would re-run attention over all previous tokens and recompute their K,VK,V repeatedly. With KV caching, you reuse K,VK,V and only do the “new query against all cached keys” work.
  • Asymptotics: prompt processing (“prefill”) is still quadratic in prompt length due to full self-attention, but token-by-token decoding avoids recomputing past projections and is much closer to linear-in-context per generated token.
  • Memory: caching costs memory proportional to context length, number of layers, and number of heads (this is often the main throughput limiter at long context).

That last point is quantitative enough to check. For a 32-layer model with 32 KV heads of dimension 128 in BF16, one token of cache is 512 KB, so an 8K context costs 4.3 GB and a 32-request batch costs 137 GB — more than half of a modern accelerator’s memory, spent on cache rather than weights. The next article derives that formula and shows how grouped-query attention changes its constant.


9. Notation: from these notes to real code

The symbols above are the paper notation. They are unambiguous for math, but no library uses them as attribute names. Before reading a real implementation it helps to have the translation table ready.

Sections 9.1–9.3 are those tables. Sections 9.4–9.9 then put them to work: each module of the block in turn, the parameter budget that follows from it, and finally the sampling loop that turns logits into text.

9.1 Dimensions

NotesnanoGPTHF transformers (Llama)Meaning
BBbatch_size / bbatch_sizesequences in flight
LLblock_size / tseq_len / max_lengthtokens per sequence
DDn_embd / Chidden_sizewidth of one token vector
HHn_headnum_attention_headsattention heads
dhd_hn_embd // n_head / hshead_dimwidth of one head
DffD_{\mathrm{ff}}4 * n_embdintermediate_sizeFFN width
NlayersN_{\mathrm{layers}}n_layernum_hidden_layersnumber of blocks
∥V∥\|V\|vocab_sizevocab_sizevocabulary size

Note the ordering difference: these notes write B×L×DB \times L \times D (batch first), while older Karpathy code writes (B,T,C)(B, T, C) and the original paper’s figures are often (L,B,D)(L, B, D) or (B,L,C)(B, L, C). The axis names matter more than the letters.

9.2 Weights

NotesnanoGPTHF transformers (Llama)vLLM / fused kernels
EEtransformer.wtemodel.embed_tokensembed_tokens
PPtransformer.wpe— (RoPE instead)—
WQW_Qc_attn (Q, K, V fused)q_projqkv_proj (fused)
WKW_Kc_attn (fused)k_projqkv_proj (fused)
WVW_Vc_attn (fused)v_projqkv_proj (fused)
WOW_Oc_projo_projo_proj
W1W_1c_fcup_projgate_up_proj (fused)
W2W_2c_projdown_projdown_proj
σ\sigma (gate)—gate_projgate_up_proj (fused)
LN⁡\operatorname{LN} (pre-attn)ln_1input_layernorminput_layernorm
LN⁡\operatorname{LN} (pre-FFN)ln_2post_attention_layernormpost_attention_layernorm
LN⁡f\operatorname{LN}_fln_fmodel.normfinal_layernorm
WUW_Ulm_headlm_headlm_head

9.3 Tensors inside forward

NotesnanoGPTHF transformers
XXxhidden_states
ShS_hatt (before mask)attn_weights
PhP_hatt (after softmax)attn_weights
AhA_hyattn_output
zzlogitslogits
ppprobsprobs / scores

These tables are only useful against something concrete, so here is the whole of §4.2 as it appears in code, with the notes names kept deliberately. This is the eager, textbook form — every intermediate from steps 3 to 6 exists as a named tensor.

def forward(self, X):
    B, L, D = X.shape

    # step 3: Q = X W_Q,  K = X W_K,  V = X W_V
    Q = self.W_Q(X)
    K = self.W_K(X)
    V = self.W_V(X)

    # step 4: split into H heads -> (B, H, L, d_h)
    Q = Q.view(B, L, self.H, self.d_h).transpose(1, 2)
    K = K.view(B, L, self.H, self.d_h).transpose(1, 2)
    V = V.view(B, L, self.H, self.d_h).transpose(1, 2)

    # step 5: S_h = Q K^T / sqrt(d_h) + M_causal
    #         P_h = softmax(S_h)
    #         A_h = P_h V
    S_h = (Q @ K.transpose(-2, -1)) / math.sqrt(self.d_h)
    M_causal = torch.triu(
        torch.ones(L, L, dtype=torch.bool, device=X.device), diagonal=1
    )
    S_h = S_h.masked_fill(M_causal, float("-inf"))
    P_h = F.softmax(S_h, dim=-1)
    A_h = P_h @ V

    # step 6: concat heads,  O = A W_O
    A = A_h.transpose(1, 2).contiguous().view(B, L, D)
    O = self.W_O(A)
    return O

Note what is absent. There is no qkv_proj, no fused gate, and no attn_weights — the fused names in the tables above appear only once a serving stack or a fused kernel is involved. That gap is the subject of the rest of this series: the KV cache explains why the KK and VV computed here are also stored, PagedAttention explains where they are kept, and FlashAttention explains how the two L×LL \times L intermediates above are removed from memory entirely.

9.4 The MLP (§4.4)

The FFN is the simplest module in the model, and it is worth seeing all of it, because its brevity hides how much of the parameter count it owns:

class MLP(nn.Module):
    def __init__(self, cfg):
        self.W_1 = nn.Linear(cfg.D, cfg.D_ff, bias=cfg.bias)   # D -> D_ff
        self.W_2 = nn.Linear(cfg.D_ff, cfg.D, bias=cfg.bias)   # D_ff -> D

    def forward(self, Y_tilde):
        F_1 = F.gelu(self.W_1(Y_tilde))   # (B, L, D_ff)
        F_2 = self.W_2(F_1)               # (B, L, D)
        return self.dropout(F_2)

Three things to notice. The input is called Y_tilde rather than X, because it arrives already normalized ($\widetilde{Y}$ in §4.2 step 8). Every dimension stays (B, L, …) — nothing here contracts or expands the token axis, which is the code-level statement that the FFN is position-wise. And the only nonlinearity in the entire block is that one gelu call.

9.5 The block (§4.2, steps 2–9)

The block adds no new math. Its whole content is the ordering and the two residual additions:

class Block(nn.Module):
    def forward(self, X):
        X_tilde = self.LN_1(X)      # step 2   normalize before attention
        O = self.attn(X_tilde)      # steps 3-6
        Y = X + O                   # step 7   residual: raw X, not X_tilde

        Y_tilde = self.LN_2(Y)      # step 8   normalize before FFN
        F_2 = self.mlp(Y_tilde)     # step 9
        X_out = Y + F_2             # step 9   residual again
        return X_out

One detail is easy to misread in the formulas and obvious in the code: the residual adds X, the unnormalized block input, not X_tilde. The normalization sits on the branch that feeds attention and is not on the path the signal travels — that is what “pre-norm” means, and it is why the residual stream can accumulate information across many layers without its scale being repeatedly reset.

The other consequence is that X and X_out have the same shape. That is the entire reason NlayersN_{\mathrm{layers}} of these stack with no shape bookkeeping anywhere.

9.6 The Transformer (§3, §5)

Everything above is a component. The top-level module is mostly wiring — but three of its lines are not:

class Transformer(nn.Module):
    def __init__(self, cfg):
        self.E = nn.Embedding(cfg.V, cfg.D)              # |V| x D
        self.P = nn.Parameter(torch.zeros(1, cfg.L, cfg.D))
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.N_layers)])
        self.LN_f = nn.LayerNorm(cfg.D, bias=cfg.bias)
        self.W_U = nn.Linear(cfg.D, cfg.V, bias=False)
        self.W_U.weight = self.E.weight                 # weight tying

    def forward(self, idx, targets=None):
        X = self.E(idx) + self.P[:, :L, :]              # §3: embed + position
        for block in self.blocks:
            X = block(X)
        X_out = self.LN_f(X)                            # §5: final norm
        z = self.W_U(X_out)                             # logits (B, L, V)
        if targets is not None:
            return z, F.cross_entropy(z.view(B*L, V), targets.reshape(B*L))
        return z, None

self.W_U.weight = self.E.weight is the weight tying from §5, and it is a assignment rather than a copy: the two modules now share one tensor, so the embedding table is trained by gradients arriving from both ends of the model. self.P is a plain nn.Parameter rather than a lookup table, because positions are a fixed range 0..L0..L; a RoPE model would delete this line entirely.

Note also that forward returns logits for every position, not just the last one. During training all LL positions are scored at once — that single detail is what makes teacher forcing cheap, and it is why the loss reshapes to (B*L, V).

9.7 Where the parameters actually go

Reading the modules top to bottom suggests attention is the substantial part. Counting them says otherwise. For the small model used in this series (D=256D=256, Dff=1024D_{\mathrm{ff}}=1024, Nlayers=6N_{\mathrm{layers}}=6):

ComponentParametersShare
Attention, all layers (4D24D^2 each)1.57 M33%
MLP, all layers (2⋅D⋅Dff2 \cdot D \cdot D_{\mathrm{ff}} each)3.15 M65%
Embeddings + positions + norms0.09 M2%
Total4.80 M

At the usual expansion factor Dff=4DD_{\mathrm{ff}} = 4D, the MLP holds exactly twice the attention parameters — and the same holds for FLOPs per token. Attention gets the discussion because it is the part that scales quadratically with sequence length and the part that is memory-bound at decode, but at a fixed sequence length the MLP is the larger block. That is why the end-to-end measurement shows attention at only 3.5% of a training step at L=256L=256: the other two-thirds of the block does not care which attention kernel you use.

9.8 Sampling from the model (§5, §8)

Everything so far produces one thing per position: a vector of logits. Turning that into text is a loop, and the loop is short enough to show in full.

@torch.no_grad()
def generate(model, idx, max_new_tokens, temperature=0.8, top_k=50):
    model.eval()
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -model.cfg.L:]          # crop to the context window
        z, _ = model(idx_cond)                    # all positions, logits (B, L, V)
        z = z[:, -1, :] / temperature             # keep the last position only

        if top_k is not None:                     # mask everything below the k-th
            values, _ = torch.topk(z, min(top_k, z.size(-1)))
            z[z < values[:, [-1]]] = float("-inf")

        p = F.softmax(z, dim=-1)                  # §5: logits -> probabilities
        idx_next = torch.multinomial(p, 1)        # §5: sample one token
        idx = torch.cat([idx, idx_next], dim=1)   # append and repeat
    return idx

Four details in those nine lines carry most of the practical content.

z[:, -1, :] is the whole point of the loop. model() returns logits for every position, but only the last one predicts a token that does not exist yet. The earlier rows are a byproduct of scoring the prompt — useful during training, discarded here.

idx[:, -max_len:] is not optional. Position embeddings only cover 0..L0..L, so a sequence grown past the window would index out of range. Cropping keeps generation legal; it also silently forgets the beginning of the conversation, which is the mechanism behind the “context window” behaviour everyone notices in practice.

The crop happens on the input, not on the model. Nothing in the weights changes size. A model trained at one sequence length can be sampled with any prefix length up to that limit, which is why the same checkpoint works for a 20-token prompt and a full-context one.

The in-place mask is safe only because of the preceding line. z[z < ...] = -inf mutates a tensor; here z was already rebound by the division, so it is a fresh tensor and no autograd graph or caller-held value is corrupted. Without that rebinding, the same line would write into the tensor model() returned.

Temperature, top-kk, and top-pp all sit between the logits and multinomial, and they compose without touching each other:

z = z / temperature                 # sharpen (<1) or flatten (>1)
if top_k:                           # keep the k most likely
    kth = torch.topk(z, top_k).values[..., -1:]   # the k-th largest, per row
    z = z.masked_fill(z < kth, float("-inf"))
if top_p:                           # keep the smallest set reaching cumulative p
    sorted_z, order = torch.sort(z, descending=True)
    probs = F.softmax(sorted_z, dim=-1).cumsum(-1)
    cutoff = probs - probs[..., -1:] + 1e-4 > top_p
    sorted_z = sorted_z.masked_fill(cutoff, float("-inf"))
    z = torch.empty_like(z).scatter_(-1, order, sorted_z)
p = F.softmax(z, dim=-1)

Note what none of that does: it never changes the model. Greedy decoding, temperature 0.1, and nucleus sampling at p=0.95p=0.95 are the same weights answering the same question differently. That is why §5 can present them as a policy layer rather than an architectural one, and why two services running the same checkpoint can feel like different models.

9.9 What this loop costs, and why §8 exists

The loop above recomputes the entire prefix on every step. Each call is a full forward pass over the prompt plus everything generated so far, so generating TT tokens costs

∑t=1TO(t2d)=O(T3d) \sum_{t=1}^{T} O(t^2 d) = O(T^3 d)

in attention work alone — cubic in the number of tokens generated. The arithmetic is correct; it is just done over and over.

With the KV cache from §8, each step processes exactly one new token against a stored history:

∑t=1TO(t d)=O(T2d) \sum_{t=1}^{T} O(t\, d) = O(T^2 d)

Quadratic instead of cubic, which is a different class of cost and not a constant-factor tuning. The loop changes in one place:

for _ in range(max_new_tokens):
    z, _ = model(idx_new, cache)      # only the NEW token goes in
    idx_new = sample(z[:, -1, :])     # one row of logits, as before
    cache.append(k_new, v_new)        # §8: extend instead of recompute

And a second change follows from the first: the cost per step now depends on the length of the cache, not on how much work the model does, which is the memory-bound regime measured in the next article. That is why a decode step at 100K context is dominated by reading keys and values rather than by any matrix multiply, and why paged caches and serving kernels exist at all.

One thing this loop still gets wrong for production: it serves one sequence. Batching many requests means each needs its own cache of a different length, which is exactly the ragged and paged problem those two articles address.


10. What’s next

The rest of this series follows one operation at a time from this picture into hardware behaviour, and each article now carries measurements taken for it rather than borrowed from a paper. The claims that have been measured, and where the evidence lives:

ClaimValueArticle
Decode arithmetic intensity1.0 FLOP/byte, 281× below the ridge pointThe KV Cache and Attention Variants
Reservation vs paging, 128K context3 requests vs 3,854PagedAttention
Quadratic intermediates removed17× less peak memory at 8KFlashAttention
Attention speedup vs step speedup10.4× kernel, 10.0× step at 4KFlashAttention
Fused vs conventional accuracy, BF167.5e-03 vs 1.25e-02FlashAttention
Ragged vs padded prefill, batch 3214×FlashInfer
Decode bandwidth, batch 1 → 3223% → 75% of peakFlashInfer
Scheduling a live batchno profile admitted run alongside oneContinuous Batching

Those numbers are hardware-specific and the articles say so; what transfers is the shape of each result. For a complementary view of how these operations are accelerated on Intel CPUs — fused GEMM kernels, blocked weight layouts, flash attention with BRGEMM, and OpenMP parallelism — see Intel TPP for LLM Inference.

References

  1. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention Is All You Need. NeurIPS, 2017.
  2. Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language Models are Unsupervised Multitask Learners. OpenAI, 2019.
  3. Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. RoFormer: Enhanced Transformer with Rotary Position Embedding. 2021.
  4. Angela Fan, Mike Lewis, and Yann Dauphin. Hierarchical Neural Story Generation. ACL, 2018.
  5. Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The Curious Case of Neural Text Degeneration. ICLR, 2020.
  6. Nitish Shirish Keskar, Bryan McCann, Lav R. Varshney, Caiming Xiong, and Richard Socher. CTRL: A Conditional Transformer Language Model for Controllable Generation. 2019.