02 — From-Scratch Inference (Round 2 Build)

“Round 1 taught you to train. Round 2 teaches you to serve. The two skills share vocabulary and share nothing else.”

This is the meat of Phase 1. Time budget: 3–5 weeks of evenings.

The goal: a single Python file (~500–800 LOC) that loads a real HuggingFace model from safetensors and produces byte-identical output to transformers.generate(do_sample=False). When this runs green, you will understand — physically — what every LLM inference server is doing.

---

The Target Model

Pick ONE of these and stick with it:

Model	Params	Attention	Why it’s a good target
meta-llama/Llama-3.2-1B	1.24 B	GQA (32 Q, 8 KV)	Small, canonical, modern arch, easy fp16 on any GPU
Qwen/Qwen2.5-1.5B	1.54 B	GQA (12 Q, 2 KV)	Same family as bigger Qwens; strong benchmark posture
HuggingFaceTB/SmolLM2-1.7B	1.71 B	GQA	Small, high-quality; a pleasure to inspect

All three use RoPE and GQA — exactly the mechanisms you must implement. Pick Llama-3.2-1B unless you have a strong reason otherwise; the Llama family is the most widely-documented reference.

---

Section 1 — The Safetensors Format

Spec: https://github.com/huggingface/safetensors

Every modern model checkpoint on HuggingFace ships as safetensors (increasingly, never pickle). You should be able to parse the format yourself — no library. Do this before you use safetensors.torch.load_file.

The layout, in bytes

[ 8 bytes:  u64 header_size, little-endian ]
[ header_size bytes: UTF-8 JSON metadata    ]
[ remainder: raw tensor bytes, tightly packed ]

The JSON metadata is a dict of tensor-name → {dtype, shape, data_offsets}. data_offsets is [start, end] inclusive-exclusive within the tensor-bytes region.

Do this drill (30 minutes)

import json, struct, numpy as np

def parse_safetensors(path):
    with open(path, 'rb') as f:
        header_size = struct.unpack('<Q', f.read(8))[0]
        header = json.loads(f.read(header_size).decode('utf-8'))
        tensors = {}
        for name, info in header.items():
            if name == '__metadata__':
                continue
            shape = info['shape']
            dtype = info['dtype']         # e.g., 'F16', 'BF16', 'F32'
            start, end = info['data_offsets']
            f.seek(8 + header_size + start)
            raw = f.read(end - start)
            # dtype → numpy dtype mapping (bf16 needs special handling)
            tensors[name] = (raw, dtype, shape)
        return tensors

Acceptance: read Llama-3.2-1B’s safetensors file, dump the tensor-name → shape mapping, and match it against model.state_dict().keys() from HF. Zero mismatches.

Why this matters: in Phase 4 you’ll want to mmap huge model files and stream them into VRAM. Understanding the format now is one afternoon of investment that pays for years.

---

Section 2 — The Tokenizer

Use the library, but understand BPE

Use tokenizers (HuggingFace’s Rust-backed library) for the real work — do not reimplement production BPE for a real model. But do implement toy BPE in ~100 lines of Python on a small corpus. This is the exact assignment Stanford CS336 Assignment 1 walks you through.

The BPE algorithm, in six lines of pseudocode

Initialize vocabulary = all bytes (256 tokens).
Repeat until vocab_size reached:
    Count all adjacent pairs in the corpus.
    Find the most frequent pair (a, b).
    Add a new token "ab" to the vocab.
    Replace all (a, b) with "ab" in the corpus.

Modern tokenizers add: byte-level BPE (Llama, GPT), regex pre-tokenization (GPT-2’s famous r"""'s|'t|...""" regex), special tokens (BOS, EOS, PAD, tool tokens).

The round-trip test

from tokenizers import Tokenizer
tok = Tokenizer.from_pretrained('meta-llama/Llama-3.2-1B')
text = "The quick brown fox jumps over the lazy dog."
ids = tok.encode(text).ids
back = tok.decode(ids)
assert back == text

This MUST be lossless for well-formed input. If your inference stack decodes anything wrong, this is the first place to check.

The Karpathy tokenizer video

“Let’s build the GPT Tokenizer” (~2h): https://www.youtube.com/watch?v=zduSFxRajkE

Watch this. Then implement your toy BPE from scratch on a small corpus (e.g., TinyShakespeare). Compare vocab merges to what tokenizers produces. Look at the actual Llama-3 tokenizer.json — wc -l gives you ~128k vocab entries; open it in an editor and see the actual merges. This concreteness sticks.

---

Section 3 — The Model, Piece by Piece

Build these classes in a single model.py. All shapes documented in comments — wrong shape is bug #1 in every from-scratch implementation.

3.1 Config

@dataclass
class ModelConfig:
    vocab_size: int
    hidden_size: int          # aka d_model
    intermediate_size: int    # SwiGLU MLP hidden dim
    num_hidden_layers: int    # transformer blocks
    num_attention_heads: int  # Q heads
    num_key_value_heads: int  # KV heads (GQA)
    head_dim: int             # usually hidden_size // num_attention_heads
    max_position_embeddings: int
    rope_theta: float         # RoPE base frequency (10000 or 500000 for Llama-3)
    rms_norm_eps: float
    tie_word_embeddings: bool

Load it from HF’s config.json. Print it. Every field should have meaning to you.

3.2 RMSNorm

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps
    def forward(self, x):
        # x: (..., dim)
        variance = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)
        return self.weight * x

Why RMSNorm and not LayerNorm? RMSNorm drops the mean-subtraction (Zhang & Sennrich, 2019: https://arxiv.org/abs/<phone_number_or_numberic_id_or_random_id_10>). Slightly faster, no quality hit. Llama, Qwen, Mistral all use it. Note there’s no bias term either — the recentering-invariance argument justifies both simplifications.

3.3 RoPE (Rotary Positional Embedding)

Paper: Su et al., “RoFormer” (2021): https://arxiv.org/abs/<phone_number_or_numberic_id_or_random_id_11>

The math, briefly: for a position m and a pair of dimensions (2i, 2i+1), rotate (x_{2i}, x_{2i+1}) by angle m * θ_i where θ_i = base^(-2i/d). In matrix form for one 2-D pair:

[ cos(mθ)  -sin(mθ) ] [ x_{2i}   ]
[ sin(mθ)   cos(mθ) ] [ x_{2i+1} ]

Applied to Q and K before attention. Not to V. The property that matters: <RoPE(q, m), RoPE(k, n)> depends only on m - n (relative position). This is why RoPE generalizes better than absolute learned positional embeddings.

Standard implementation (Llama-style):

def precompute_rope_cache(head_dim, max_seq_len, base=10000.0, device='cuda'):
    # freqs: [head_dim/2]
    freqs = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    # positions: [max_seq_len]
    t = torch.arange(max_seq_len, device=device).float()
    # outer product: [max_seq_len, head_dim/2]
    freqs = torch.outer(t, freqs)
    # cos and sin caches
    cos = freqs.cos()  # [max_seq_len, head_dim/2]
    sin = freqs.sin()
    return cos, sin

def apply_rope(x, cos, sin):
    # x: [batch, seq, n_heads, head_dim]
    # split last dim into pairs [head_dim/2, 2] and interpret as (real, imag)
    x1 = x[..., 0::2]   # even dims
    x2 = x[..., 1::2]   # odd dims
    # rotate
    rotated1 = x1 * cos - x2 * sin
    rotated2 = x1 * sin + x2 * cos
    # interleave back
    return torch.stack([rotated1, rotated2], dim=-1).flatten(-2)

The trap: HF’s Llama uses a different dimension pairing (contiguous halves, not interleaved). Look at transformers/models/llama/modeling_llama.py apply_rotary_pos_emb — you’ll see it splits x[..., :d/2] and x[..., d/2:] and rotates them as (a, b) → (a*cos - b*sin, b*cos + a*sin). Use HF’s convention if you want token-identical output.

3.4 Grouped-Query Attention (GQA)

Paper: Ainslie et al., “GQA” (2023): https://arxiv.org/abs/<phone_number_or_numberic_id_or_random_id_12>

The idea: instead of H query heads each with its own K, V (MHA) or all sharing one K, V (MQA), use G KV heads where each is shared by H/G query heads. Llama-3-8B: 32 Q, 8 KV, ratio 4. Llama-3-70B: 64 Q, 8 KV, ratio 8. KV cache shrinks by H/G× — the whole reason GQA exists.

class GroupedQueryAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_heads = config.num_attention_heads
        self.n_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.n_rep = self.n_heads // self.n_kv_heads   # broadcast factor

        self.wq = nn.Linear(config.hidden_size, self.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(config.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(self.n_heads * self.head_dim, config.hidden_size, bias=False)

    def forward(self, x, kv_cache, position_ids, cos, sin):
        B, T, C = x.shape
        q = self.wq(x).view(B, T, self.n_heads,    self.head_dim)
        k = self.wk(x).view(B, T, self.n_kv_heads, self.head_dim)
        v = self.wv(x).view(B, T, self.n_kv_heads, self.head_dim)

        q = apply_rope(q, cos[position_ids], sin[position_ids])
        k = apply_rope(k, cos[position_ids], sin[position_ids])

        # Append to KV cache
        k, v = kv_cache.append(k, v)   # returns full k, v history

        # Repeat KV to match Q head count
        k = repeat_kv(k, self.n_rep)   # (B, T_full, n_heads, head_dim)
        v = repeat_kv(v, self.n_rep)

        # Transpose to (B, n_heads, T, head_dim) for matmul
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # Scaled dot-product attention with causal mask
        scale = 1.0 / math.sqrt(self.head_dim)
        scores = torch.matmul(q, k.transpose(-2, -1)) * scale
        # Causal mask: q of position i attends to k of positions 0..i
        # (During decode T=1, mask is trivial; during prefill it's a triangle.)
        scores = apply_causal_mask(scores, position_ids)
        attn = torch.softmax(scores, dim=-1, dtype=torch.float32).to(q.dtype)
        out = torch.matmul(attn, v)   # (B, n_heads, T, head_dim)

        out = out.transpose(1, 2).contiguous().view(B, T, self.n_heads * self.head_dim)
        return self.wo(out)

def repeat_kv(x, n_rep):
    # x: (B, T, n_kv_heads, head_dim) -> (B, T, n_kv_heads * n_rep, head_dim)
    if n_rep == 1:
        return x
    B, T, H, D = x.shape
    return x[:, :, :, None, :].expand(B, T, H, n_rep, D).reshape(B, T, H * n_rep, D)

The compute-KV-savings drill: For Llama-3-8B (32 Q heads, 8 KV heads, head_dim 128), full MHA would use 32 KV heads at inference. GQA uses 8. KV cache memory is exactly 8/32 = 1/4 of the equivalent MHA. Compute exactly: 2 * 32 * 8 * 128 * 2 bytes/token = 131,072 bytes/token. Full MHA equivalent would be 2 * 32 * 32 * 128 * 2 = 524,288. Ratio 4×.

3.5 SwiGLU MLP

Paper: Shazeer, “GLU Variants Improve Transformer” (2020): https://arxiv.org/abs/<phone_number_or_numberic_id_or_random_id_13>

class SwiGLU(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.gate = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up   = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

Note: 3 linear layers, not 2 (as in vanilla FFN). Parameter count is 3× hidden×intermediate, not 2×. Llama-3 intermediate size is chosen so total MLP params ≈ 4× hidden² × [Llama-3’s factor]. Details in the arithmetic notebook.

3.6 The Transformer Block

class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.attn      = GroupedQueryAttention(config)
        self.mlp_norm  = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.mlp       = SwiGLU(config)

    def forward(self, x, kv_cache, position_ids, cos, sin):
        x = x + self.attn(self.attn_norm(x), kv_cache, position_ids, cos, sin)
        x = x + self.mlp(self.mlp_norm(x))
        return x

Pre-norm (Llama-style), not post-norm (original transformer). Pre-norm is more numerically stable for deep networks.

3.7 The Full Model

class Llama(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.embed = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_hidden_layers)])
        self.norm  = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        if config.tie_word_embeddings:
            self.lm_head.weight = self.embed.weight
        cos, sin = precompute_rope_cache(config.head_dim, config.max_position_embeddings, config.rope_theta)
        self.register_buffer('cos', cos)
        self.register_buffer('sin', sin)

    def forward(self, input_ids, kv_caches, position_ids):
        x = self.embed(input_ids)
        for layer, kv in zip(self.layers, kv_caches):
            x = layer(x, kv, position_ids, self.cos, self.sin)
        x = self.norm(x)
        return self.lm_head(x)   # (B, T, vocab)

---

Section 4 — The KV Cache

Why it exists

During autoregressive decoding, at step t we compute Q, K, V for position t, then attend to K/V of positions 0..t. If we don’t cache K and V, each step re-runs the full transformer over the entire history — O(t) work at every step — for total O(N²) generation cost. With a cache, each step is O(t) memory access but only new-token compute → total O(N) work per token, dominated by memory bandwidth.

Implementation

Allocate contiguous tensors up-front (avoids re-allocation each step):

class KVCache:
    def __init__(self, max_batch, max_seq, n_kv_heads, head_dim, dtype=torch.float16, device='cuda'):
        self.k = torch.zeros((max_batch, max_seq, n_kv_heads, head_dim), dtype=dtype, device=device)
        self.v = torch.zeros((max_batch, max_seq, n_kv_heads, head_dim), dtype=dtype, device=device)
        self.length = 0

    def append(self, k_new, v_new):
        # k_new: (B, T_new, n_kv_heads, head_dim)
        B, T_new, H, D = k_new.shape
        self.k[:B, self.length:self.length+T_new] = k_new
        self.v[:B, self.length:self.length+T_new] = v_new
        self.length += T_new
        return self.k[:B, :self.length], self.v[:B, :self.length]

The naive-vs-cached measurement

Write two versions of generate():

1. Naive: re-run the model on the entire sequence [prompt..., generated_so_far] for every new token. Discard everything but the last logit.

2. Cached: run the model on [new_token] only, with the cache handling the history.

Measure wall-clock time to generate 128 tokens from a 32-token prompt with both. On a 1B model on a 3090 you should see roughly:

• Naive: quadratic slowdown; last token takes ~5–10× the first token’s time.

• Cached: flat per-token latency after the first (prefill) token.

Plot both. The plot is a required artifact of Phase 1.

---

Section 5 — The Sampling Ladder

Implement each sampler as a LogitsProcessor callable that takes and returns a (B, vocab) tensor of logits. Chain them.

class TemperatureProcessor:
    def __init__(self, T): self.T = T
    def __call__(self, logits): return logits / self.T

class TopKProcessor:
    def __init__(self, k): self.k = k
    def __call__(self, logits):
        v, _ = torch.topk(logits, self.k, dim=-1)
        threshold = v[..., -1, None]
        return torch.where(logits < threshold, torch.full_like(logits, -float('inf')), logits)

class TopPProcessor:
    def __init__(self, p): self.p = p
    def __call__(self, logits):
        sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
        probs = torch.softmax(sorted_logits, dim=-1)
        cumprobs = torch.cumsum(probs, dim=-1)
        mask = cumprobs > self.p
        mask[..., 1:] = mask[..., :-1].clone()
        mask[..., 0] = False
        sorted_logits.masked_fill_(mask, -float('inf'))
        return sorted_logits.scatter(-1, sorted_idx, sorted_logits)

class RepetitionPenaltyProcessor:
    def __init__(self, penalty, seen_ids): self.penalty = penalty; self.seen = seen_ids
    def __call__(self, logits):
        for i, ids in enumerate(self.seen):
            for tid in ids:
                logits[i, tid] = logits[i, tid] / self.penalty if logits[i, tid] > 0 else logits[i, tid] * self.penalty
        return logits

def sample(logits, processors):
    for p in processors:
        logits = p(logits)
    probs = torch.softmax(logits, dim=-1)
    return torch.multinomial(probs, 1)

Greedy verification: the whole payoff of Phase 1 is that with sample = argmax, your output matches HF exactly. If it doesn’t, the bug is in one of: RoPE convention, causal mask, RMSNorm epsilon location, GQA broadcasting, or safetensors weight-name mapping. In that order of frequency.

---

Section 6 — Weight Loading (The Boring, Critical Part)

HF names weights like model.layers.0.self_attn.q_proj.weight. You name them layers.0.attn.wq.weight. You need a mapping:

HF_TO_MINE = {
    'model.embed_tokens.weight':         'embed.weight',
    'model.layers.{i}.input_layernorm.weight':               'layers.{i}.attn_norm.weight',
    'model.layers.{i}.self_attn.q_proj.weight':              'layers.{i}.attn.wq.weight',
    'model.layers.{i}.self_attn.k_proj.weight':              'layers.{i}.attn.wk.weight',
    'model.layers.{i}.self_attn.v_proj.weight':              'layers.{i}.attn.wv.weight',
    'model.layers.{i}.self_attn.o_proj.weight':              'layers.{i}.attn.wo.weight',
    'model.layers.{i}.post_attention_layernorm.weight':      'layers.{i}.mlp_norm.weight',
    'model.layers.{i}.mlp.gate_proj.weight':                 'layers.{i}.mlp.gate.weight',
    'model.layers.{i}.mlp.up_proj.weight':                   'layers.{i}.mlp.up.weight',
    'model.layers.{i}.mlp.down_proj.weight':                 'layers.{i}.mlp.down.weight',
    'model.norm.weight':                                     'norm.weight',
    'lm_head.weight':                                        'lm_head.weight',   # or tied to embed
}

Write a load_from_hf(model, safetensors_path) function that loads and shape-checks every tensor. The first time a shape mismatch happens, you’ll want to have this catcher in place.

---

Section 7 — Long-Context and RoPE Extensions (Awareness)

Read but don’t implement yet:

• Position Interpolation (PI): rescale positions by L_new / L_original before RoPE. Fast to apply, some quality loss. Chen et al. 2023: https://arxiv.org/abs/<phone_number_or_numberic_id_or_random_id_14>

• NTK-aware scaling: scale the RoPE base frequency instead of positions. Preserves high-frequency information better.

• YaRN (Yet another RoPE extensioN): piecewise treatment of high vs low frequencies + attention temperature scaling. Peng et al., ICLR 2024: https://arxiv.org/abs/<phone_number_or_numberic_id_or_random_id_15>. Repo: https://github.com/jquesnelle/yarn

• LongRoPE (Microsoft, 2024): search-based non-uniform interpolation, scales to 2M+ tokens.

• Llama-3.1’s approach: YaRN-style, scaling factor 8× to reach 131k from 8k base. See config.json fields: rope_scaling.factor, rope_scaling.low_freq_factor, rope_scaling.high_freq_factor, rope_scaling.original_max_position_embeddings.

Best 2025 explainer covering APE → RoPE → PI → NTK → YaRN evolution: https://amaarora.github.io/posts/2025-09-21-rope-context-extension.html

You will implement long-context RoPE in Phase 4 when you serve at 128k contexts. For now, know what these are so rope_scaling in config.json doesn’t confuse you.

---

Section 8 — The Verification Ritual

A checklist to run before declaring Phase 1 exit:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import my_impl

tok = AutoTokenizer.from_pretrained('meta-llama/Llama-3.2-1B')
hf_model = AutoModelForCausalLM.from_pretrained('meta-llama/Llama-3.2-1B', torch_dtype=torch.float16).cuda()
my_model = my_impl.load('meta-llama/Llama-3.2-1B').cuda()

prompt = "The three most important numbers in inference engineering are"
ids = tok(prompt, return_tensors='pt').input_ids.cuda()

# 1. Token-identical greedy output for 50 tokens.
hf_out = hf_model.generate(ids, max_new_tokens=50, do_sample=False)
my_out = my_impl.generate(my_model, ids, max_new_tokens=50, do_sample=False)
assert torch.equal(hf_out, my_out), (hf_out, my_out)

# 2. Logit-close: for the last token of the prompt, logits within 1e-2 relative in fp16.
with torch.no_grad():
    hf_logits = hf_model(ids).logits[:, -1, :]
    my_logits = my_impl.forward_once(my_model, ids)[:, -1, :]
diff = (hf_logits - my_logits).abs().max().item()
assert diff < 1e-1, diff   # fp16 accumulates ~1e-2 to 1e-1 error on real models

# 3. Perplexity on a real corpus (WikiText-2 test set) within 5% of HF.
ppl_hf = evaluate_perplexity(hf_model, tok, corpus)
ppl_my = evaluate_perplexity(my_model, tok, corpus)
assert abs(ppl_hf - ppl_my) / ppl_hf < 0.05

When all three pass, you own inference. When they don’t, use these diagnostic strategies (they will find every bug):

1. Compare hidden state after each layer. Instrument layer_0 output, layer_1 output, etc., and diff against HF’s output_hidden_states=True. The first layer that diverges is where your bug is.

2. Compare Q, K, V, and attention scores individually. RoPE convention mismatch shows up in Q/K but not V. Causal mask bugs show in scores.

3. Bit-check for the classic: sinusoidal RoPE pair convention (interleaved vs half-half).

---

Time Estimate

Task	Hours
Safetensors parser	2
Toy BPE tokenizer	4
RMSNorm, SwiGLU, config parsing	3
RoPE + verification against HF	4–6 (this bites more than expected)
GQA attention + repeat_kv	4
Full model assembly + weight loader	4
Naive generate + cached generate + comparison plot	4
Sampling ladder	3
Verification ritual + debugging	6–12 (the bug-hunt is the education)
Total	~40–50 hours — 3–5 weeks of evenings

When the verification ritual is green, commit. Push. Tweet the perplexity graph. Move to Phase 2.