GPT-2 Inference Engine — From Scratch

A pure NumPy implementation of GPT-2 inference. No PyTorch in the forward pass — every operation (attention, layer norm, GELU, sampling) is written by hand using only numpy. Weights are loaded once from HuggingFace, then the entire computation graph runs in NumPy.

Built as a learning project to deeply understand how a transformer LLM actually works at the matrix level.

The meaning of life is not a matter of opinion, but a matter of fact.

Output from greedy decoding — GPT-2 small, 124M params, pure NumPy.

---

Architecture

GPT-2 Small specs, derived directly from weight shapes at runtime:

Hyperparameter	Value
vocab_size	50,257
n_embd	768
n_ctx	1,024
n_layer	12
n_head	12
head_dim	64
ffn_inner_dim	3,072
Total params	~124M

---

How It Works

1. Tokenization

Text is converted to token IDs using tiktoken (OpenAI's BPE tokenizer, same one GPT-2 was trained with). Each token maps to an integer between 0–50256.

enc = tiktoken.get_encoding("gpt2")
token_ids = enc.encode("Hello, I am")
# → [15496, 11, 314, 716]

2. Embeddings

Two embedding tables are looked up and summed:

x = wte[token_ids] + wpe[0:T]

• wte (50257 × 768) — token embedding table. Each token ID maps to a learned 768-dim vector encoding its meaning.
• wpe (1024 × 768) — position embedding table. Position 0 adds its own 768-dim vector, position 1 adds another, etc. This is how the model knows token order, since attention has no inherent sense of sequence.

After this step, x is shape (T, 768) — a matrix where each row is a token's initial representation.

3. Transformer Blocks × 12

The core of the model. Each block does two things in sequence, both using pre-norm + residual:

x = x + Attention(LayerNorm(x))   # self-attention sub-layer
x = x + FFN(LayerNorm(x))         # feed-forward sub-layer

The residual (x + ...) means the original signal always flows through — the block only learns a correction to add on top.

Layer Norm

Normalises each token's 768-dim vector to zero mean and unit variance, then applies learned scale w and shift b:

y = (x - mean) / sqrt(var + ε) * w + b

Prevents activations from exploding or vanishing across 12 layers.

Multi-Head Causal Self-Attention

This is where tokens communicate with each other.

Step 1: Project x → Q, K, V via one big linear:
        (T, 768) @ (768, 2304) → (T, 2304)  →  split into 3 × (T, 768)

Step 2: Split each into 12 heads:
        (T, 768) → (12, T, 64)   [12 heads, each 64-dim]

Step 3: For each head, compute attention scores:
        scores = Q @ Kᵀ / √64       shape: (12, T, T)
        Apply causal mask (upper triangle = -∞)
        weights = softmax(scores)    shape: (12, T, T)
        out = weights @ V            shape: (12, T, 64)

Step 4: Merge heads back:
        (12, T, 64) → (T, 768)
        Final linear projection → (T, 768)

The causal mask is critical — it sets all positions where j > i to -∞ before softmax, so token at position i can only attend to tokens 0..i. This enforces left-to-right generation.

The √64 scaling prevents the dot products from growing so large that softmax saturates (all probability mass on one token, gradients vanish).

Why multiple heads? Each head learns to attend to different relationships — one head might track syntax, another coreference, another proximity. The results are concatenated, giving the model a richer joint representation.

Feed-Forward Network (MLP)

Applied independently to each token position after attention:

FFN(x) = GELU(x @ W1 + b1) @ W2 + b2

Expands 768 → 3072 (4×), applies GELU, projects back to 768. This is where most of the model's "knowledge" is believed to be stored — the attention heads route information, the MLP transforms it.

GELU (Gaussian Error Linear Unit) is GPT-2's activation function. Unlike ReLU it has a smooth gradient near zero:

GELU(x) ≈ 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715x³)))

4. Final LayerNorm + LM Head

After all 12 blocks, one final layer norm is applied, then the output is projected to vocabulary logits:

logits = x @ wteᵀ     shape: (T, 50257)

Note: the LM head reuses the token embedding matrix (wte) — this is called weight tying. The intuition is that the embedding that encodes a token's meaning going in is the same geometry used to predict that token coming out. It also saves ~38M parameters.

5. Autoregressive Generation

The model always produces logits for all positions, but we only care about the last position — it's the only one that has seen the entire context.

for each new token:
    logits = forward(full_context)     # (T, 50257)
    next_token = sample(logits[-1])    # last row only
    context.append(next_token)

This means on step 1 we run the full prompt through, on step 2 we run prompt+1 token, on step 3 prompt+2 tokens, and so on. Every step is a full forward pass over the growing sequence — this is the naive approach, and it's the performance bottleneck. See KV Cache below.

6. Sampling Strategies

Three strategies are implemented:

Greedy — always picks the highest-probability token. Deterministic, but tends to be repetitive.

next = argmax(logits)

Temperature — divides logits by T before softmax. T < 1 sharpens the distribution (more confident), T > 1 flattens it (more random). T → 0 converges to greedy.

probs = softmax(logits / temperature)
next = np.random.choice(vocab_size, p=probs)

Top-k — zeroes out all logits outside the top-k before sampling. Prevents the model from ever picking from the long tail of nonsense tokens while still allowing diverse outputs.

keep = top_k indices of logits
sample from softmax(keep / temperature)

---

KV Cache — What It Is and Why We Don't Have It

The problem

In our current implementation, generating 50 tokens from a 10-token prompt means:

Step  1: forward pass on 10 tokens
Step  2: forward pass on 11 tokens
Step  3: forward pass on 12 tokens
...
Step 50: forward pass on 59 tokens

Total compute ∝ 10 + 11 + 12 + ... + 59 = O(T²) work. Every new token forces us to recompute K and V for every previous token in every layer, from scratch, even though they haven't changed.

What KV Cache does

In the attention step, K and V only depend on the input tokens — not on the query. So for any token at position i, its K and V vectors are identical on every generation step. We're recomputing them wastefully.

The KV cache stores them:

# Instead of recomputing for all T tokens each step:
kv_cache[layer] = {
    "k": np.zeros((n_head, T_max, head_dim)),
    "v": np.zeros((n_head, T_max, head_dim)),
}

# Each new step, only compute K, V for the NEW token:
k_new = linear(x_new, c_attn_w)[:, C:2*C]    # just 1 token
v_new = linear(x_new, c_attn_w)[:, 2*C:]

# Append to cache:
kv_cache[layer]["k"][:, step, :] = k_new
kv_cache[layer]["v"][:, step, :] = v_new

# Attention uses full cached K, V but only new Q:
scores = q_new @ kv_cache["k"].transpose(...)  # (n_head, 1, T)

This reduces generation from O(T²) to O(T) — each step is constant cost regardless of context length. At 1024 tokens, that's a ~512× speedup in attention computation.

Memory cost

The tradeoff is memory. For GPT-2 small at max context:

per layer:  2 (K+V) × n_head × n_ctx × head_dim × 4 bytes
          = 2 × 12 × 1024 × 64 × 4
          = ~6 MB per layer

all layers: 12 × 6 MB = ~72 MB

For GPT-4-scale models (96 layers, 128 heads, 10K+ context) this blows up fast — KV cache memory management becomes a major systems problem (paged attention, sliding window attention, etc.).

Why it's not in this implementation

This engine is built for clarity over performance — the goal is to show every matrix operation explicitly. Adding KV cache would require threading a stateful cache dict through the entire call stack and splitting the attention function into "prefill" (process prompt) and "decode" (one token at a time) modes. That's the right next step if you want to extend this into something fast.

---

Setup

pip install numpy tiktoken transformers torch
python inference_engine.py

Weights (~548 MB) are downloaded from HuggingFace on first run and cached locally. No GPU required — runs on CPU only.

---

File Structure

inference_engine.py
│
├── load_weights()          # pull from HuggingFace, convert to NumPy dicts
├── gelu()                  # approximate tanh activation
├── softmax()               # numerically stable
├── layer_norm()            # per-token normalisation with learned scale/bias
├── linear()                # x @ w + b  (HF Conv1D layout: no transpose)
├── attention()             # multi-head causal self-attention
├── mlp()                   # 2-layer FFN with 4× expansion
├── transformer_block()     # pre-norm + residual for attn and mlp
├── gpt2_forward()          # full forward pass, returns (T, 50257) logits
├── greedy_sample()         # argmax
├── temperature_sample()    # softmax with temperature scaling
├── top_k_sample()          # top-k filtering + temperature
├── generate()              # autoregressive loop with streaming output
└── inspect_architecture()  # derive and print all hyperparams from weights

---

Key Implementation Details

Weight layout — HuggingFace GPT-2 uses Conv1D (not nn.Linear), which stores weights as (in_features, out_features). This means x @ w is correct — no transpose needed. Getting this wrong causes an immediate shape mismatch in the QKV projection.

Weight tying — The LM head (logits = x @ wte.T) reuses the input embedding matrix. GPT-2 has ~124M parameters but only stores ~85M unique weights because of this.

Causal mask — Built fresh each forward call as an upper-triangular -∞ matrix. In production implementations this is precomputed and cached.

Numerical stability — Softmax subtracts max(x) before exp() to prevent overflow. Layer norm adds ε=1e-5 to variance before sqrt to prevent division by zero.

---

Possible Extensions

• KV Cache — implement prefill/decode split for O(T) generation
• Batch inference — add a batch dimension to process multiple prompts simultaneously
• Top-p (nucleus) sampling — sample from the smallest set of tokens whose cumulative probability exceeds p
• GPT-2 Medium/Large/XL — the same code works, just pass different n_layer/n_head values
• Beam search — maintain k candidate sequences and pick the highest joint probability

---

References

• GPT-2 Paper — Language Models are Unsupervised Multitask Learners
• Attention Is All You Need
• Andrej Karpathy — nanoGPT
• Jay Mody — picoGPT