Frequency tiers (estimates, relative to interviews targeting this pack's specialty — a ★★★★★ here means "core for these roles", not for generalist ML loops): ★★★★★ expect it in most loops · ★★★★ very common · ★★★ common · ★★ occasional · ★ rare/deep-specialist. Spend ~50% of practice on the top two tiers.
70+ coding problems with full PyTorch solutions and tests
Principal/Senior-Principal LLM/Transformer Interview Prep
Notes
This pack covers the live-coding questions that come up in principal-level Transformer / LLM interviews: tok- enization (BPE / WordPiece / Unigram / byte-level / SentencePiece), attention variants (MHA, MQA, GQA, MLA, sliding-window, FlashAttention online softmax, ring attention, linear attention, sparse / Longformer- style), long-context positional encodings (RoPE, YARN, NTK-aware scaling, ALiBi, positional interpolation), inference (KV cache, paged attention, speculative decoding, beam search, contrastive search, top-p / top-k / temperature / repetition penalty), and MoE (top-k routing, Switch, Expert Choice, load-balance / z-loss, capacity factor). Standard imports across the pack:
import math, numpy as np, torch
import torch.nn as nn, torch.nn.functional as F
from collections import Counter
Problem
Implement a simple pre-tokenizer that lowercases and splits on whitespace and punctuation.
import re
def pre_tokenize(text):
text = text.lower()
return re.findall(r"[a-z0-9]+|[^\sa-z0-9]", text)
Tests
assert pre_tokenize("Hello, World!") == ["hello", ",", "world", "!"]
print("pretok OK")
Problem
Train a BPE tokenizer on a list of words: count adjacent symbol pairs, repeatedly merge the most frequent into a new token.
def get_pairs(word_freqs):
pairs = Counter()
for word, f in word_freqs.items():
sym = word.split()
for i in range(len(sym) - 1):
pairs[(sym[i], sym[i + 1])] += f
return pairs
def merge_pair(word_freqs, pair):
bg = ' '.join(pair); rep = ''.join(pair)
return {w.replace(bg, rep): f for w, f in word_freqs.items()}
def train_bpe(corpus, n_merges):
word_freqs = Counter([' '.join(list(w)) + ' </w>' for w in corpus])
merges = []
for _ in range(n_merges):
pairs = get_pairs(word_freqs)
if not pairs: break
best = pairs.most_common(1)[0][0]
word_freqs = merge_pair(word_freqs, best)
merges.append(best)
return merges
Tests
corpus = ['low'] * 5 + ['lower'] * 2 + ['newest'] * 6 + ['widest'] * 3
merges = train_bpe(corpus, n_merges=10)
assert len(merges) > 0; print("BPE OK", merges[:3])
Problem
Apply learned BPE merges greedily to a single word.
import math
def bpe_encode(word, merges):
sym = list(word) + ['</w>']
rank = {pair: i for i, pair in enumerate(merges)}
while len(sym) > 1:
pairs = list(zip(sym[:-1], sym[1:]))
# pick the pair with the lowest rank (= earliest merged)
best = min(((rank.get(p, math.inf), i) for i, p in enumerate(pairs)),
default=(math.inf, -1))
if best[0] == math.inf: break
i = best[1]
sym = sym[:i] + [sym[i] + sym[i + 1]] + sym[i + 2:]
return sym
Tests
merges = [('e', 's'), ('es', 't'), ('est', '</w>')]
out = bpe_encode("newest", merges)
assert out[-1] == 'est</w>' or 'est</w>' in out
print("bpe encode OK", out)
Problem
Implement byte-level encoding so any UTF-8 string can always be tokenized: map each byte 0≤b < 256 to a unicode character that is unambiguous in token vocabularies.
def bytes_to_unicode():
bs = list(range(33, 127)) + list(range(161, 173)) + list(range(174, 256))
cs = bs[:]; n = 0
for b in range(256):
if b not in bs: bs.append(b); cs.append(256 + n); n += 1
return dict(zip(bs, [chr(c) for c in cs]))
def encode_bytes(text):
table = bytes_to_unicode()
return ''.join(table[b] for b in text.encode('utf-8'))
Tests
out = encode_bytes("Hello!"); assert all(ord(c) > 0 for c in out)
print("byte-level OK", out)
Problem
WordPiece’s training rule: pick the merge that maximises score(a, b) = count(ab)/(count(a) · count(b)).
def wordpiece_best_merge(word_freqs):
pair_freq = Counter(); sym_freq = Counter()
for word, f in word_freqs.items():
sym = word.split()
for s in sym: sym_freq[s] += f
for i in range(len(sym) - 1): pair_freq[(sym[i], sym[i + 1])] += f
best, best_score = None, -1
for (a, b), c in pair_freq.items():
s = c / (sym_freq[a] * sym_freq[b])
if s > best_score: best, best_score = (a, b), s
return best
Tests
wf = Counter({'l o w </w>': 5, 'l o w e r </w>': 2})
print("WordPiece OK", wordpiece_best_merge(wf))
Problem
Given a vocabulary with log-probabilities, find the highest-probability segmentation by Viterbi.
def unigram_segment(text, vocab_logp):
n = len(text); INF = -1e18
dp = [INF] * (n + 1); back = [(-1, "")] * (n + 1); dp[0] = 0.0
for i in range(1, n + 1):
for j in range(i):
piece = text[j:i]
if piece in vocab_logp and dp[j] + vocab_logp[piece] > dp[i]:
dp[i] = dp[j] + vocab_logp[piece]
back[i] = (j, piece)
if dp[n] == INF: return []
out = []; i = n
while i > 0:
j, p = back[i]; out.append(p); i = j
return list(reversed(out))
Tests
vocab = {'a': -1.0, 'ab': -1.5, 'abc': -3.0, 'b': -2.0, 'c': -2.0, 'bc': -3.0}
print("unigram OK", unigram_segment('abc', vocab))
Problem
Map tokens ↔ ids using a fixed vocabulary; handle the unknown token.
class Vocab:
def __init__(self, tokens, unk='<unk>'):
self.unk = unk
self.tok2id = {t: i for i, t in enumerate(tokens)}
if unk not in self.tok2id: self.tok2id[unk] = len(self.tok2id)
self.id2tok = {i: t for t, i in self.tok2id.items()}
def encode(self, toks): return [self.tok2id.get(t, self.tok2id[self.unk]) for t in toks]
def decode(self, ids): return [self.id2tok[i] for i in ids]
Tests
v = Vocab(['a', 'b']); ids = v.encode(['a', 'c'])
assert v.decode(ids) == ['a', '<unk>']; print("vocab OK")
Problem
Compute character coverage of a corpus and which characters fall below a threshold.
def char_coverage(corpus, threshold=0.9999):
c = Counter(''.join(corpus))
total = sum(c.values()); kept = 0; rare = []
for ch, f in c.most_common():
if kept / total < threshold:
kept += f
else:
rare.append(ch)
return kept / total, rare
Tests
cov, rare = char_coverage(['hello', 'helloo'])
assert cov > 0.99; print("coverage OK")
Problem
Pad a list of variable-length token sequences to the same length using a PAD id; return key-padding mask (True at pads).
import torch
def pad_collate(seqs, pad_id=0):
L = max(len(s) for s in seqs)
out = torch.full((len(seqs), L), pad_id, dtype=torch.long)
mask = torch.ones_like(out, dtype=torch.bool)
for i, s in enumerate(seqs):
out[i, :len(s)] = torch.tensor(s); mask[i, :len(s)] = False
return out, mask
Tests
out, mask = pad_collate([[1, 2, 3], [4, 5]], pad_id=-1)
assert out[1, 2].item() == -1 and mask[1, 2].item() == True
print("pad OK")
Problem
Attn(Q, K, V ) = softmax(QK⊤/√dk + M)V with optional additive mask M.
import math
import torch
def sdpa(Q, K, V, mask=None):
d_k = Q.size(-1)
s = (Q @ K.transpose(-2, -1)) / math.sqrt(d_k)
if mask is not None: s = s + mask
return torch.softmax(s, dim=-1) @ V
Tests
Q = K = V = torch.randn(1, 4, 4, 8)
out = sdpa(Q, K, V)
ref = F.scaled_dot_product_attention(Q, K, V)
assert torch.allclose(out, ref, atol=1e-5); print("SDPA OK")
Problem
Project to H heads with shared linear layers, attend, recombine.
import torch.nn as nn
class MHA(nn.Module):
def __init__(self, d, h):
super().__init__()
assert d % h == 0
self.h, self.dk = h, d // h
self.qkv = nn.Linear(d, 3 * d, bias=False); self.proj = nn.Linear(d, d, bias=False)
def forward(self, x, mask=None):
B, T, D = x.shape
q, k, v = self.qkv(x).chunk(3, dim=-1)
q = q.view(B, T, self.h, self.dk).transpose(1, 2)
k = k.view(B, T, self.h, self.dk).transpose(1, 2)
v = v.view(B, T, self.h, self.dk).transpose(1, 2)
out = sdpa(q, k, v, mask).transpose(1, 2).contiguous().view(B, T, D)
return self.proj(out)
Tests
m = MHA(64, 4); y = m(torch.randn(2, 6, 64))
assert y.shape == (2, 6, 64); print("MHA OK")
Problem
Build an additive causal mask of shape (T, T), 0 on / below diagonal, −∞above.
import torch
def causal_mask(T, device='cpu'):
return torch.triu(torch.full((T, T), float('-inf'), device=device), diagonal=1)
Tests
M = causal_mask(4)
assert M[0, 1] == float('-inf') and M[1, 0] == 0; print("causal OK")
Problem
Pre-norm encoder/decoder block: LN → MHA → residual → LN → MLP → residual.
import torch.nn as nn
class TransformerBlock(nn.Module):
def __init__(self, d, h, mlp_ratio=4):
super().__init__()
self.ln1 = nn.LayerNorm(d); self.attn = MHA(d, h)
self.ln2 = nn.LayerNorm(d)
self.mlp = nn.Sequential(nn.Linear(d, mlp_ratio * d), nn.GELU(), nn.Linear(mlp_ratio * d, d))
def forward(self, x, mask=None):
x = x + self.attn(self.ln1(x), mask)
x = x + self.mlp(self.ln2(x))
return x
Tests
blk = TransformerBlock(64, 4); y = blk(torch.randn(2, 6, 64))
assert y.shape == (2, 6, 64); print("block OK")
Problem
Stack L causal Transformer blocks with token embedding + tied output head.
import torch
import torch.nn as nn
class MiniGPT(nn.Module):
def __init__(self, vocab, d=64, h=4, L=2, max_T=64):
super().__init__()
self.emb = nn.Embedding(vocab, d)
self.pos = nn.Embedding(max_T, d)
self.blocks = nn.ModuleList([TransformerBlock(d, h) for _ in range(L)])
self.ln = nn.LayerNorm(d)
def forward(self, ids):
T = ids.size(1)
x = self.emb(ids) + self.pos(torch.arange(T, device=ids.device))
m = causal_mask(T, ids.device)
for b in self.blocks: x = b(x, m)
x = self.ln(x)
return x @ self.emb.weight.T
Tests
g = MiniGPT(50); y = g(torch.randint(0, 50, (1, 8)))
assert y.shape == (1, 8, 50); print("MiniGPT OK")
Problem
Decoder cross-attention: Q comes from the decoder state, K/V from the encoder output.
import torch.nn as nn
import torch.nn.functional as F
class CrossAttn(nn.Module):
def __init__(self, d, h, ctx_d):
super().__init__()
self.h, self.dk = h, d // h
self.q = nn.Linear(d, d); self.k = nn.Linear(ctx_d, d); self.v = nn.Linear(ctx_d, d)
self.proj = nn.Linear(d, d)
def forward(self, x, ctx):
B, T, D = x.shape; M = ctx.size(1)
q = self.q(x).view(B, T, self.h, self.dk).transpose(1, 2)
k = self.k(ctx).view(B, M, self.h, self.dk).transpose(1, 2)
v = self.v(ctx).view(B, M, self.h, self.dk).transpose(1, 2)
out = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).contiguous().view(B, T, D)
return self.proj(out)
Tests
m = CrossAttn(32, 4, 64); y = m(torch.randn(1, 5, 32), torch.randn(1, 8, 64))
assert y.shape == (1, 5, 32); print("crossattn OK")
Problem
H query heads share a single K/V head. Used by PaLM, Falcon for inference speed.
import torch.nn as nn
import torch.nn.functional as F
class MQA(nn.Module):
def __init__(self, d, h):
super().__init__()
self.h, self.dk = h, d // h
self.q = nn.Linear(d, d, bias=False)
self.k = nn.Linear(d, self.dk, bias=False)
self.v = nn.Linear(d, self.dk, bias=False)
self.proj = nn.Linear(d, d, bias=False)
def forward(self, x):
B, T, D = x.shape
q = self.q(x).view(B, T, self.h, self.dk).transpose(1, 2)
k = self.k(x).unsqueeze(1).expand(-1, self.h, -1, -1)
v = self.v(x).unsqueeze(1).expand(-1, self.h, -1, -1)
out = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).contiguous().view(B, T, D)
return self.proj(out)
Tests
m = MQA(64, 4); y = m(torch.randn(1, 5, 64))
assert y.shape == (1, 5, 64); print("MQA OK")
Problem
G KV groups, each shared by H/G query heads. LLaMA-2/3 uses 8 KV with 32–64 query heads.
import torch.nn as nn
import torch.nn.functional as F
class GQA(nn.Module):
def __init__(self, d, h_q, h_kv):
super().__init__()
assert h_q % h_kv == 0
self.h_q, self.h_kv = h_q, h_kv; self.dk = d // h_q
self.q = nn.Linear(d, h_q * self.dk, bias=False)
self.k = nn.Linear(d, h_kv * self.dk, bias=False)
self.v = nn.Linear(d, h_kv * self.dk, bias=False)
self.proj = nn.Linear(d, d, bias=False)
def forward(self, x):
B, T, D = x.shape
q = self.q(x).view(B, T, self.h_q, self.dk).transpose(1, 2)
k = self.k(x).view(B, T, self.h_kv, self.dk).transpose(1, 2)
v = self.v(x).view(B, T, self.h_kv, self.dk).transpose(1, 2)
rep = self.h_q // self.h_kv
k = k.repeat_interleave(rep, dim=1); v = v.repeat_interleave(rep, dim=1)
out = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).contiguous().view(B, T, D)
return self.proj(out)
Tests
m = GQA(64, h_q=8, h_kv=2); y = m(torch.randn(1, 5, 64))
assert y.shape == (1, 5, 64); print("GQA OK")
Problem
Compress K/V into a low-rank latent c ∈ Rr, expand back when computing attention; saves KV-cache memory.
import torch.nn as nn
import torch.nn.functional as F
class MLA(nn.Module):
def __init__(self, d, h, r=64):
super().__init__()
self.h, self.dk = h, d // h
self.q = nn.Linear(d, d, bias=False)
self.kv_lat = nn.Linear(d, r, bias=False)
self.k_up = nn.Linear(r, d, bias=False)
self.v_up = nn.Linear(r, d, bias=False)
self.proj = nn.Linear(d, d, bias=False)
def forward(self, x):
B, T, D = x.shape
c = self.kv_lat(x) # (B, T, r) -- this is what gets cached
k = self.k_up(c).view(B, T, self.h, self.dk).transpose(1, 2)
v = self.v_up(c).view(B, T, self.h, self.dk).transpose(1, 2)
q = self.q(x).view(B, T, self.h, self.dk).transpose(1, 2)
out = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).contiguous().view(B, T, D)
return self.proj(out)
Tests
m = MLA(64, 4, r=8); y = m(torch.randn(1, 5, 64))
assert y.shape == (1, 5, 64); print("MLA OK")
Problem
Each token attends only to itself and the previous w tokens (Mistral’s local attention).
import torch
def sliding_mask(T, w, causal=True):
i = torch.arange(T)[:, None]; j = torch.arange(T)[None, :]
rel = i - j
keep = (rel >= 0) & (rel <= w) if causal else (rel.abs() <= w)
return torch.where(keep, 0.0, float('-inf'))
Tests
m = sliding_mask(5, w=2)
assert m[0, 1] == float('-inf') and m[3, 1] == 0; print("sliding OK")
Problem
Mark a small set of global tokens that attend to all and are attended by all; everyone else uses sliding-window.
def longformer_mask(T, w, global_idx):
m = sliding_mask(T, w, causal=False)
for i in global_idx:
m[i, :] = 0; m[:, i] = 0
return m
Tests
m = longformer_mask(8, 1, global_idx=[0])
assert (m[0] == 0).all(); print("longformer OK")
Problem
Stream K, V in tiles of size b and maintain (m, ℓ, O) state per query so the softmax is computed exactly without materialising the full T ×T matrix.
import math
import torch
def flash_attn_naive(Q, K, V, tile=64):
Tq, d = Q.shape; Tk, _ = K.shape
O = torch.zeros(Tq, d); m = torch.full((Tq,), -float('inf')); l = torch.zeros(Tq)
for j in range(0, Tk, tile):
Kj = K[j:j+tile]; Vj = V[j:j+tile]
S = (Q @ Kj.T) / math.sqrt(d)
m_new = torch.maximum(m, S.max(dim=1).values)
P = torch.exp(S - m_new[:, None])
l_new = torch.exp(m - m_new) * l + P.sum(dim=1)
O = O * (torch.exp(m - m_new) * l / l_new.clamp_min(1e-12))[:, None] \
+ (P @ Vj) / l_new.clamp_min(1e-12)[:, None]
m, l = m_new, l_new
return O
Tests
T, d = 32, 16; Q = torch.randn(T, d); K = torch.randn(T, d); V = torch.randn(T, d)
ref = F.scaled_dot_product_attention(Q.unsqueeze(0), K.unsqueeze(0), V.unsqueeze(0)).squeeze(0)
ours = flash_attn_naive(Q, K, V, tile=8)
assert torch.allclose(ours, ref, atol=1e-4); print("FlashAttn online OK")
Advanced implementation. The full production-shaped version: chunked attention with a running max, rescaled normalizer AND a rescaled weighted-value accumulator — exact attention in O(chunk) working memory. Verified equal to full attention (1e-5) at extreme logit scales.
import math, torch
def chunked_attention(q, K, V, chunk=16):
# single query q: (d,); K: (T, d); V: (T, dv)
m, s = -float("inf"), 0.0
acc = torch.zeros(V.size(1))
for i in range(0, K.size(0), chunk):
scores = K[i:i+chunk] @ q / math.sqrt(q.size(0))
m_new = max(m, scores.max().item())
w = torch.exp(scores - m_new)
s = s * math.exp(m - m_new) + w.sum().item()
acc = acc * math.exp(m - m_new) + w @ V[i:i+chunk] # rescale old accumulator
m = m_new
return acc / s
Problem
Replace softmax with a positive feature map ϕ: attention = ϕ(Q)(ϕ(K)⊤V )/(ϕ(Q)ϕ(K)⊤1). ϕ(x) = ELU(x)+1 keeps it positive (Linear Transformer).
import torch.nn.functional as F
def linear_attn(Q, K, V):
phi = lambda x: F.elu(x) + 1
Qp, Kp = phi(Q), phi(K)
KV = Kp.transpose(-2, -1) @ V
Z = Qp @ Kp.sum(dim=-2, keepdim=True).transpose(-2, -1)
return (Qp @ KV) / Z.clamp_min(1e-9)
Tests
out = linear_attn(torch.randn(1, 8, 16), torch.randn(1, 8, 16), torch.randn(1, 8, 16))
assert out.shape == (1, 8, 16); print("linear attn OK")
Advanced implementation. The causal form needs no (T, T) matrix at all: maintain running sums S = Σφ(kᵢ)vᵢᵀ and z = Σφ(kᵢ) — O(T·d²) time, O(d²) state, and the reason linear attention admits an RNN-style incremental decode. Verified exactly equal to the quadratic masked form (1e-5).
import torch
import torch.nn.functional as F
def linear_attention_causal(q, k, v):
phi = lambda x: F.elu(x) + 1
qf, kf = phi(q), phi(k)
S = torch.einsum("btd,bte->btde", kf, v).cumsum(1) # running sum of k v^T
z = kf.cumsum(1) # running normalizer
num = torch.einsum("btd,btde->bte", qf, S)
den = torch.einsum("btd,btd->bt", qf, z).unsqueeze(-1)
return num / (den + 1e-8)
Problem
Partition K, V across P devices. Each device passes its block in a ring while accumulating softmax state. Skeleton in pure PyTorch.
import math
import torch
def ring_attn(Q, K_blocks, V_blocks):
# K_blocks, V_blocks: lists of length P, each (T_i, d)
Tq, d = Q.shape; O = torch.zeros(Tq, d); m = torch.full((Tq,), -float('inf')); l = torch.zeros(Tq)
for K, V in zip(K_blocks, V_blocks):
S = (Q @ K.T) / math.sqrt(d)
m_new = torch.maximum(m, S.max(dim=1).values)
P_ = torch.exp(S - m_new[:, None])
l_new = torch.exp(m - m_new) * l + P_.sum(dim=1)
O = O * (torch.exp(m - m_new) * l / l_new.clamp_min(1e-12))[:, None] \
+ (P_ @ V) / l_new.clamp_min(1e-12)[:, None]
m, l = m_new, l_new
return O
Tests
T, d = 16, 8; Q = torch.randn(T, d); K = torch.randn(T, d); V = torch.randn(T, d)
ours = ring_attn(Q, [K[:8], K[8:]], [V[:8], V[8:]])
ref = F.scaled_dot_product_attention(Q[None], K[None], V[None]).squeeze(0)
assert torch.allclose(ours, ref, atol=1e-4); print("ring OK")
Problem
PEeven = sin(pos/100002i/d), PEodd = cos(·).
import math
import torch
def sin_pe(L, d):
pe = torch.zeros(L, d); pos = torch.arange(L).unsqueeze(1).float()
div = torch.exp(torch.arange(0, d, 2).float() * (-math.log(10000.0) / d))
pe[:, 0::2] = torch.sin(pos * div); pe[:, 1::2] = torch.cos(pos * div)
return pe
Tests
pe = sin_pe(10, 16); assert pe.shape == (10, 16); print("sin pe OK")
Problem
Rotate pairs of features by θi · pos before attention. Rotation matrix is applied to (x2k, x2k+1) via cos/sin.
import torch
def rope_freqs(dim, max_pos, base=10000.0):
freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
pos = torch.arange(max_pos).float()
angles = torch.outer(pos, freqs)
return torch.cos(angles), torch.sin(angles)
def apply_rope(x, cos, sin):
x1, x2 = x[..., 0::2], x[..., 1::2]
rot = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
return rot.flatten(-2)
Tests
T, D = 4, 8; cos, sin = rope_freqs(D, T)
y = apply_rope(torch.randn(1, T, D), cos, sin); assert y.shape == (1, T, D)
print("RoPE OK")
Problem
For extrapolation to longer sequences, scale the base by α = (s)d/(d−2) where s is the context-length factor.
import torch
def ntk_aware_freqs(dim, max_pos, scale=2.0, base=10000.0):
new_base = base * (scale ** (dim / (dim - 2)))
freqs = 1.0 / (new_base ** (torch.arange(0, dim, 2).float() / dim))
pos = torch.arange(max_pos).float()
a = torch.outer(pos, freqs)
return torch.cos(a), torch.sin(a)
Tests
cos, sin = ntk_aware_freqs(8, 64, scale=4.0)
assert cos.shape == (64, 4); print("NTK OK")
Problem
Stretch positions by factor s: scale the input position index from pos to pos/s before computing RoPE, so a 4k → 16k extension uses s = 4.
import torch
def pi_freqs(dim, max_pos, scale=2.0, base=10000.0):
freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
pos = torch.arange(max_pos).float() / scale
a = torch.outer(pos, freqs)
return torch.cos(a), torch.sin(a)
Tests
cos, sin = pi_freqs(8, 4, scale=2.0)
cos_ref, _ = rope_freqs(8, 8)
assert torch.allclose(cos[1], cos_ref[2], atol=1e-6) # pos 1 with scale 2 = pos 2 base
print("PI OK")
Problem
YARN interpolates linearly for low-frequency components (which need stretching) and uses NTK for high- frequency ones, blending across a transition band.
import math
import torch
def yarn_scale(theta, dim, scale=4.0, beta_lo=1.0, beta_hi=32.0):
# theta_i: per-dim base frequency; high-freq -> NTK; low-freq -> PI
i = torch.arange(0, dim, 2).float()
wave = (2 * math.pi / theta).clamp(min=1.0)
ramp = ((wave / beta_lo).log() / (beta_hi / beta_lo + 1e-9)).clamp(0, 1)
return theta * (1 - ramp) + theta * scale * ramp
Tests
freqs = 1.0 / (10000 ** (torch.arange(0, 8, 2).float() / 8))
out = yarn_scale(freqs, 8, scale=4.0)
assert out.shape == freqs.shape; print("YARN OK")
Problem
Add a per-head bias −mh|i −j| to attention scores; no learned positional embeddings.
import math
import torch
def alibi_slopes(h):
start = 2 ** (-2 ** -(math.log2(h) - 3))
return torch.tensor([start * start ** i for i in range(h)])
def alibi_bias(T, h):
pos = torch.arange(T)
return -alibi_slopes(h).view(h, 1, 1) * (pos[None, :] - pos[:, None]).abs().float()
Tests
b = alibi_bias(8, 8); assert b.shape == (8, 8, 8); print("ALiBi OK")
Problem
Bucket relative positions into ≤K buckets, with logarithmic spacing for large distances.
import math
import torch
def t5_rel_bucket(rel, n_buckets=32, max_dist=128):
n = n_buckets
sign = (rel < 0).long(); rel = rel.abs()
half = n // 2
is_small = rel < half
log_ratio = (rel.float() / half).clamp_min(1).log() / math.log(max_dist / half)
val_large = (half + (log_ratio * (n - half)).long()).clamp_max(n - 1)
return torch.where(is_small, rel, val_large) + sign * 0 # bidirectional bucketing
Tests
buckets = t5_rel_bucket(torch.tensor([0, 1, 16, 64, 200]))
assert buckets.shape == (5,); print("T5 buckets OK", buckets)
Problem
Add a learned Embedding(L, d) on top of token embeddings.
import torch
import torch.nn as nn
class LearnedPE(nn.Module):
def __init__(self, L, d):
super().__init__()
self.pe = nn.Embedding(L, d)
def forward(self, x):
return x + self.pe(torch.arange(x.size(1), device=x.device))
Tests
m = LearnedPE(64, 32); y = m(torch.randn(2, 8, 32))
assert y.shape == (2, 8, 32); print("learned PE OK")
Problem
GELU = xΦ(x). Exact uses erf; GPT-2 uses tanh approx.
import math
import torch
def gelu_exact(x):
return 0.5 * x * (1.0 + torch.erf(x / math.sqrt(2.0)))
def gelu_tanh(x):
c = math.sqrt(2.0 / math.pi)
return 0.5 * x * (1.0 + torch.tanh(c * (x + 0.044715 * x ** 3)))
Tests
x = torch.linspace(-3, 3, 13)
assert torch.allclose(gelu_exact(x), F.gelu(x, approximate='none'), atol=1e-7)
assert torch.allclose(gelu_tanh(x), F.gelu(x, approximate='tanh'), atol=1e-6)
print("GELU OK")
Problem
SwiGLU(x) = (silu(xW1) ⊙xW2)W3. Used by LLaMA and Mistral.
import torch.nn as nn
import torch.nn.functional as F
class SwiGLU(nn.Module):
def __init__(self, d, hidden):
super().__init__()
self.w1 = nn.Linear(d, hidden, bias=False)
self.w2 = nn.Linear(d, hidden, bias=False)
self.w3 = nn.Linear(hidden, d, bias=False)
def forward(self, x):
return self.w3(F.silu(self.w1(x)) * self.w2(x))
Tests
m = SwiGLU(64, 128); y = m(torch.randn(2, 6, 64)); assert y.shape == (2, 6, 64)
print("SwiGLU OK")
Problem
GeGLU(x) = (GELU(xW1) ⊙xW2)W3; ReGLU uses ReLU instead.
import torch.nn as nn
import torch.nn.functional as F
class GeGLU(nn.Module):
def __init__(self, d, h):
super().__init__()
self.fc = nn.Linear(d, 2 * h); self.proj = nn.Linear(h, d)
def forward(self, x):
a, b = self.fc(x).chunk(2, -1)
return self.proj(F.gelu(a) * b)
Tests
m = GeGLU(32, 64); y = m(torch.randn(2, 5, 32)); assert y.shape == (2, 5, 32)
print("GeGLU OK")
Problem
Σ x2 q 1RMSNorm(x) = γ x/i + ε. D
import torch
import torch.nn as nn
class RMSNorm(nn.Module):
def __init__(self, d, eps=1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(d)); self.eps = eps
def forward(self, x):
return self.weight * x / x.pow(2).mean(-1, keepdim=True).add(self.eps).sqrt()
Tests
m = RMSNorm(16); y = m(torch.randn(3, 4, 16))
assert y.shape == (3, 4, 16); print("RMSNorm OK")
Problem
Take arg max over logits at each step until EOS or max length.
import torch
@torch.no_grad()
def greedy_decode(model, prompt_ids, max_new=20, eos=None):
out = list(prompt_ids)
for _ in range(max_new):
logits = model(torch.tensor(out).unsqueeze(0))[:, -1]
nxt = int(logits.argmax(-1).item()); out.append(nxt)
if nxt == eos: break
return out
Tests
class M(nn.Module):
def __init__(self): super().__init__(); self.l = nn.Linear(8, 8)
def forward(self, ids): return self.l(F.one_hot(ids, 8).float())
out = greedy_decode(M(), [0], max_new=3)
assert len(out) == 4; print("greedy OK")
Problem
Apply temperature, then top-k, then top-p filters before sampling.
import torch
def sample_token(logits, T=1.0, top_k=0, top_p=0.0):
logits = logits.clone() / max(T, 1e-8)
if top_k > 0:
v, _ = torch.topk(logits, top_k)
logits[logits < v[..., -1, None]] = -float('inf')
if top_p > 0.0:
sl, idx = logits.sort(descending=True)
cum = torch.softmax(sl, dim=-1).cumsum(dim=-1)
mask = cum > top_p; mask[..., 1:] = mask[..., :-1].clone(); mask[..., 0] = False
sl[mask] = -float('inf')
logits = torch.full_like(logits, -float('inf')).scatter_(-1, idx, sl)
probs = torch.softmax(logits, dim=-1)
return torch.multinomial(probs, 1)
Tests
torch.manual_seed(0)
out = sample_token(torch.tensor([1., 2., 3., 4., 5.]).unsqueeze(0), T=0.001)
assert out.item() == 4; print("sampling OK")
Problem
Decrease the logit of any token already in the context by a factor of θ.
import torch
def repetition_penalty(logits, prev_ids, theta=1.2):
logits = logits.clone()
for tok in set(prev_ids):
logits[..., tok] = torch.where(logits[..., tok] > 0, logits[..., tok] / theta, logits[..., tok] * theta)
return logits
Tests
out = repetition_penalty(torch.tensor([2.0, 2.0, 2.0]).unsqueeze(0), [0])
assert out[0, 0] < 2.0; print("rep penalty OK")
Problem
Standard beam search with cumulative log-probabilities.
import torch
def beam_search(score_fn, init, beam=4, max_len=10, eos=2):
beams = [(init, 0.0)]
for _ in range(max_len):
cands = []
for seq, lp in beams:
if seq[-1] == eos: cands.append((seq, lp)); continue
log_p = score_fn(seq)
v, idx = torch.topk(log_p, beam)
for k in range(beam): cands.append((seq + [int(idx[k])], lp + float(v[k])))
beams = sorted(cands, key=lambda x: x[1], reverse=True)[:beam]
if all(s[-1] == eos for s, _ in beams): break
return beams[0]
Tests
def score(seq):
out = torch.full((5,), -10.0); out[1] = 0.0; out[2] = -1.0
return out
print("beam OK", beam_search(score, [0]))
Problem
CS(v) = (1 −α)p(v|x) −α maxh⟨hv, ht⟩to penalise repeats and encourage diversity.
import torch.nn.functional as F
def contrastive_search(logits, candidate_h, ctx_h, alpha=0.6):
p = F.softmax(logits, dim=-1)
sim = (candidate_h @ ctx_h.transpose(-2, -1)).max(dim=-1).values
return (1 - alpha) * p - alpha * sim
Tests
out = contrastive_search(torch.zeros(4), torch.randn(4, 8), torch.randn(2, 8))
assert out.shape == (4,); print("contrastive OK")
Problem
Append the new token’s K/V to a cache, then attend to the full cached sequence.
import torch
import torch.nn as nn
import torch.nn.functional as F
class CachedAttn(nn.Module):
def __init__(self, d, h):
super().__init__()
self.h, self.dk = h, d // h
self.qkv = nn.Linear(d, 3 * d, bias=False); self.proj = nn.Linear(d, d, bias=False)
def step(self, x, cache):
B, _, D = x.shape
q, k, v = self.qkv(x).chunk(3, -1)
q = q.view(B, 1, self.h, self.dk).transpose(1, 2)
k = k.view(B, 1, self.h, self.dk).transpose(1, 2)
v = v.view(B, 1, self.h, self.dk).transpose(1, 2)
if cache is not None:
k = torch.cat([cache[0], k], dim=2); v = torch.cat([cache[1], v], dim=2)
out = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).contiguous().view(B, 1, D)
return self.proj(out), (k, v)
Tests
m = CachedAttn(32, 4); cache = None
for _ in range(4):
y, cache = m.step(torch.randn(1, 1, 32), cache)
assert cache[0].size(2) == 4; print("KV cache OK")
Problem
Draft model proposes k tokens; target verifies with rejection sampling. Accept the longest prefix that passes the acceptance test; resample after rejection from the residual distribution.
import torch
def speculative_step(draft_probs, target_probs, draft_tokens, rng):
accepted = []
for p_d, p_t, t in zip(draft_probs, target_probs, draft_tokens):
ratio = min(1.0, p_t[t].item() / max(p_d[t].item(), 1e-12))
if rng.random() < ratio:
accepted.append(int(t))
else:
adj = (p_t - p_d).clamp(min=0); adj = adj / adj.sum().clamp_min(1e-12)
return accepted + [int(torch.multinomial(adj, 1))]
return accepted + [int(torch.multinomial(target_probs[-1], 1))]
Tests
import random; rng = random.Random(0)
V = 5; k = 3
dp = [torch.softmax(torch.randn(V), -1) for _ in range(k)]
tp = [torch.softmax(torch.randn(V), -1) for _ in range(k)]
acc = speculative_step(dp, tp, [int(p.argmax()) for p in dp], rng)
assert 1 <= len(acc) <= k + 1; print("spec OK", acc)
Advanced note. The property worth stating unprompted: with greedy verification, the output is provably bit-identical to target-only decoding for any draft model — every emitted token is either a draft token the target agrees with, or the target's own argmax. A bad draft only costs speed, never correctness (verified in this document series with an independently-initialized draft). The sampling-mode generalization preserves the target distribution exactly via rejection sampling: accept with prob min(1, p_target/p_draft), else resample from the clamped residual (p_target − p_draft)₊.
Problem
Maintain a logical → physical block mapping for the KV cache so memory is allocated in fixed blocks of size B (vLLM-style).
class BlockTable:
def __init__(self, block_size=16):
self.block_size = block_size
self.free = list(range(1024)) # pool of physical blocks
self.tables = {} # seq_id -> [phys_blk_ids]
def allocate(self, seq_id, n_tokens):
n_blocks = (n_tokens + self.block_size - 1) // self.block_size
ids = [self.free.pop() for _ in range(n_blocks)]
self.tables[seq_id] = ids
return ids
def append(self, seq_id, n_tokens):
cur = self.tables[seq_id]
used = (sum(self.block_size for _ in cur))
if used < self._slots_used(seq_id, n_tokens):
cur.append(self.free.pop())
return cur
def _slots_used(self, seq_id, n): return n
Tests
bt = BlockTable(block_size=4)
bt.allocate("s", 6); assert len(bt.tables["s"]) == 2
print("paged OK")
Problem
Normalise sequence log-probabilities by Lα to avoid favouring short outputs.
def length_normalize(lp, length, alpha=0.6):
return lp / (length ** alpha)
Tests
assert length_normalize(-10.0, 5, 0.6) > -10.0
print("length norm OK")
Problem
Route each token to the top-k experts (typically k = 1 or 2); combine outputs via softmax weights.
import torch
import torch.nn as nn
import torch.nn.functional as F
class MoE(nn.Module):
def __init__(self, d, n_experts, k=2, hidden=None):
super().__init__()
hidden = hidden or 4 * d
self.gate = nn.Linear(d, n_experts, bias=False)
self.experts = nn.ModuleList([nn.Sequential(nn.Linear(d, hidden), nn.GELU(),
nn.Linear(hidden, d)) for _ in range(n_experts)])
self.k = k
def forward(self, x):
flat = x.view(-1, x.size(-1))
scores = self.gate(flat)
v, idx = scores.topk(self.k, dim=-1)
w = F.softmax(v, dim=-1)
out = torch.zeros_like(flat)
for kk in range(self.k):
for e in range(len(self.experts)):
m = idx[:, kk] == e
if m.any(): out[m] += w[m, kk:kk+1] * self.experts[e](flat[m])
return out.view_as(x), scores, idx
Tests
m = MoE(32, 4, k=2); y, _, _ = m(torch.randn(2, 6, 32))
assert y.shape == (2, 6, 32); print("MoE OK")
Problem
Top-1 routing with auxiliary load-balance loss Laux = N Σi fipi where fi is the fraction routed to expert i and pi is the average gate prob.
import torch
import torch.nn.functional as F
def switch_aux_loss(scores, idx, n_experts):
# scores: (T, N), idx: (T, 1)
p = F.softmax(scores, dim=-1).mean(dim=0) # (N,)
f = torch.zeros(n_experts, device=scores.device)
for e in range(n_experts):
f[e] = (idx[:, 0] == e).float().mean()
return n_experts * (f * p).sum()
Tests
T, N = 16, 4; sc = torch.randn(T, N); idx = sc.argmax(-1, keepdim=True)
loss = switch_aux_loss(sc, idx, N); assert loss > 0
print("switch loss OK", loss.item())
Problem
Cap the number of tokens per expert at cap = ⌈c · T · k/N⌉. Drop overflow tokens or send them through a residual connection.
import math
import torch
def assign_with_capacity(idx, scores, T, N, k, capacity_factor=1.25):
cap = math.ceil(capacity_factor * T * k / N)
counts = torch.zeros(N, dtype=torch.long, device=idx.device)
keep = torch.zeros_like(idx, dtype=torch.bool)
order = scores.gather(-1, idx).argsort(descending=True, dim=0)
for kk in range(k):
for i in order[:, kk].tolist():
e = idx[i, kk].item()
if counts[e] < cap: keep[i, kk] = True; counts[e] += 1
return keep
Tests
T, N, k = 8, 2, 1
idx = torch.randint(0, N, (T, k)); sc = torch.randn(T, N)
keep = assign_with_capacity(idx, sc, T, N, k, capacity_factor=1.0)
assert keep.sum().item() <= T; print("capacity OK")
Problem
Each expert chooses its top-⌈cT/N⌉tokens; this avoids token-level capacity overflow and creates load balance by construction.
import math
def expert_choice(scores, c=1.0):
# scores: (T, N); each expert picks the top-cap tokens
T, N = scores.shape
cap = math.ceil(c * T / N)
expert_topk = scores.transpose(0, 1).topk(cap, dim=1) # values, idx of tokens
return expert_topk.indices # (N, cap)
Tests
out = expert_choice(torch.randn(8, 2)); assert out.shape == (2, 4)
print("expert choice OK")
Problem
Penalise large logits at the router with Lz = Et[log Σi exp si]2.
import torch
def router_z_loss(scores):
return (torch.logsumexp(scores, dim=-1) ** 2).mean()
Tests
loss = router_z_loss(torch.randn(8, 4)); assert loss > 0
print("z-loss OK")
Problem
Approximate the discrete routing decision differentiably: πi = softmax((si + gi)/τ) where gi ∼Gumbel(0, 1).
import torch
import torch.nn.functional as F
def gumbel_route(scores, tau=1.0, hard=True):
g = -torch.log(-torch.log(torch.rand_like(scores) + 1e-20) + 1e-20)
y = F.softmax((scores + g) / tau, dim=-1)
if hard:
i = y.argmax(-1, keepdim=True)
y_hard = torch.zeros_like(y).scatter_(-1, i, 1.0)
y = (y_hard - y).detach() + y
return y
Tests
y = gumbel_route(torch.randn(8, 4))
assert y.shape == (8, 4) and (y.sum(-1) - 1).abs().max() < 1e-4
print("gumbel route OK")
Problem
Compute per-token CE loss while masking PAD tokens.
import torch.nn.functional as F
def lm_loss(logits, targets, pad_id=-100):
return F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=pad_id)
Tests
z = torch.randn(2, 4, 5); t = torch.tensor([[1, -100, 3, 4], [0, 1, -100, -100]])
assert lm_loss(z, t).item() > 0; print("LM loss OK")
Problem
Tie the input embedding and output classifier weights.
import torch.nn as nn
class TiedHeadLM(nn.Module):
def __init__(self, vocab, d):
super().__init__()
self.emb = nn.Embedding(vocab, d)
def forward(self, ids):
h = self.emb(ids)
return h, h @ self.emb.weight.T
Tests
m = TiedHeadLM(50, 16)
_, logits = m(torch.randint(0, 50, (1, 8)))
assert logits.shape == (1, 8, 50); print("tied OK")
Problem
Wrap a sub-module so its forward activations are not stored, but recomputed during backward.
import torch.utils.checkpoint as cp
class CheckpointBlock(nn.Module):
def __init__(self, block): super().__init__(); self.b = block
def forward(self, x): return cp.checkpoint(self.b, x, use_reentrant=False)
Tests
b = nn.Sequential(nn.Linear(8, 8), nn.ReLU(), nn.Linear(8, 8))
m = CheckpointBlock(b); m(torch.randn(2, 8, requires_grad=True)).sum().backward()
print("checkpoint OK")
Problem
Pack variable-length sequences into fixed-size blocks separated by EOS, used for high-throughput LM pretrain- ing.
import torch
def pack_sequences(seqs, block_size, eos=2):
flat = []
for s in seqs:
flat.extend(list(s)); flat.append(eos)
return torch.tensor([flat[i:i+block_size]
for i in range(0, len(flat) - block_size + 1, block_size)])
Tests
out = pack_sequences([[1,2,3],[4,5],[6]], block_size=4, eos=0)
assert out.shape[1] == 4; print("pack OK")
Problem
Linear warmup for Tw steps, then half-cycle cosine decay to ηmin.
import math
def cosine_lr(step, lr_max, warmup, total, lr_min=0.0):
if step < warmup:
return lr_max * step / max(1, warmup)
p = (step - warmup) / max(1, total - warmup)
return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * p))
Tests
assert cosine_lr(0, 1.0, 10, 100) == 0.0
assert cosine_lr(10, 1.0, 10, 100) == 1.0
print("cosine LR OK")
Problem
Standard AdamW: weight decay applied directly to parameters, decoupled from the gradient.
def adamw_step(p, g, m, v, t, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8, wd=1e-2):
m.mul_(b1).add_(g, alpha=1 - b1)
v.mul_(b2).addcmul_(g, g, value=1 - b2)
mhat = m / (1 - b1 ** t); vhat = v / (1 - b2 ** t)
p.data.add_((mhat / (vhat.sqrt() + eps) + wd * p.data), alpha=-lr)
Tests
p = torch.zeros(3); m = torch.zeros(3); v = torch.zeros(3)
g = torch.tensor([1., 2., 3.])
for t in range(1, 11): adamw_step(p, g, m, v, t, lr=0.1)
assert p.abs().min() > 0.05; print("adamw OK")
Problem
Clip a list of gradient tensors to a global L2 norm.
import math
def clip_grad(grads, max_norm):
total = math.sqrt(sum(g.pow(2).sum().item() for g in grads if g is not None))
coef = max_norm / (total + 1e-6)
if coef < 1:
for g in grads:
if g is not None: g.mul_(coef)
return total
Tests
g1 = torch.ones(10) * 100; g2 = torch.ones(10) * 100
clip_grad([g1, g2], max_norm=1.0)
assert (g1.pow(2).sum() + g2.pow(2).sum()).sqrt() < 1.01; print("clip OK")
Problem
Dynamic loss scaling for fp16 training.
import torch
class Scaler:
def __init__(self, init=2.0**16, growth=2.0, backoff=0.5, growth_int=2000):
self.scale, self.growth, self.backoff, self.gi = init, growth, backoff, growth_int
self.steps_no_skip = 0
def step(self, params_grads, optimizer_step):
bad = any(not torch.isfinite(g).all() for _, g in params_grads)
if bad: self.scale *= self.backoff; self.steps_no_skip = 0; return False
for p, g in params_grads: g.div_(self.scale)
optimizer_step(); self.steps_no_skip += 1
if self.steps_no_skip % self.gi == 0: self.scale *= self.growth
return True
Tests
s = Scaler(init=8.0)
ok = s.step([(torch.zeros(3), torch.tensor([float('inf'),0,0]))], lambda: None)
assert ok is False and s.scale == 4.0; print("scaler OK")
Problem
Discretised SSM step: ht = Aht−1 + Bxt, yt = Cht. Implement scan in PyTorch.
import torch
def ssm_scan(x, A, B, C):
# x: (T, d), A,B: (d,)
h = torch.zeros_like(x[0]); ys = []
for t in range(x.size(0)):
h = A * h + B * x[t]
ys.append((C * h))
return torch.stack(ys, 0)
Tests
x = torch.randn(8, 4); A = torch.full((4,), 0.9); B = C = torch.ones(4)
y = ssm_scan(x, A, B, C); assert y.shape == (8, 4); print("SSM OK")
Problem
1D causal convolution: output at t depends only on inputs ≤t, achieved via left-padding.
import torch.nn as nn
import torch.nn.functional as F
class CausalConv1d(nn.Module):
def __init__(self, c, k=3):
super().__init__()
self.k = k
self.conv = nn.Conv1d(c, c, k)
def forward(self, x):
x = F.pad(x, (self.k - 1, 0))
return self.conv(x)
Tests
m = CausalConv1d(16, 3); y = m(torch.randn(1, 16, 8))
assert y.shape == (1, 16, 8); print("causal conv OK")
Notes
Survival tactics for live Transformer / LLM coding:
State the shapes: (B, T, D), (B, H, T, dk). Most bugs come from a wrong permute or view between heads-first and heads-last layout.
Always start with a stable formula: subtract max for softmax, log1p, etc.
Tokenization: BPE / WordPiece / Unigram differ in (i) merge selection rule, (ii) tokenizer training objective, (iii) byte-level vs. char-level fallback. SentencePiece unifies these as a library.
Causal vs. bidirectional: build the additive mask first; never silently rely on a default.
Long context: RoPE base scaling → NTK-aware → Position Interpolation → YARN. ALiBi extrapolates by design but loses attention sharpness.
KV cache: per-layer (K, V ) tensors of shape (B, H, Tcache, dk). With GQA / MQA, only the KV-head dimension differs from MHA. Paged attention adds a logical-to-physical block table to avoid fragmentation.
MoE: top-k routing has 4 components — gate, expert dispatch / combine, capacity, load-balance loss. Switch and Expert Choice differ on who picks whom.
Decoding: greedy / beam (deterministic) vs. top-k / top-p / temperature (stochastic). Always state how repetition penalty / length penalty interact with the others.
Training tricks: weight tying, RMSNorm vs. LayerNorm, SwiGLU vs. GELU MLP, decoupled AdamW, cosine LR + warmup, gradient checkpointing, fp16 / bf16 with grad scaling, sequence packing, document boundary masks.
End every implementation with tests: shape-equality, identity checks (e.g. MHA with Wq = Wk = Wv = I should reduce to a known formula), or a comparison against F.scaled_dot_product_attention for any custom kernel.