Debug a Small PyTorch Transformer with Four Bugs, Then Add a KV Cache

This machine learning coding round has two parts. First, a small decoder-only transformer language model and its training step contain **four bugs**; find and fix all of them. Second, add a key-value (KV) cache so that text generation does not recompute attention over the whole sequence for every new token. The original code was not reported. The PyTorch code below is an illustrative stand-in of the same kind; it contains exactly four bugs. ```python import math import torch import torch.nn as nn import torch.nn.functional as F class SinusoidalPositions(nn.Module): def __init__(self, max_len, d_model): super().__init__() pe = torch.zeros(max_len, d_model) position = torch.arange(max_len).unsqueeze(1).float() div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div) pe[:, 1::2] = torch.sin(position * div) self.register_buffer("pe", pe) def forward(self, x): # x: (batch, seq, d_model) return x + self.pe[: x.size(1)] class CausalSelfAttention(nn.Module): def __init__(self, d_model, n_heads, max_len): super().__init__() assert d_model % n_heads == 0 self.n_heads, self.d_head = n_heads, d_model // n_heads self.qkv = nn.Linear(d_model, 3 * d_model) self.proj = nn.Linear(d_model, d_model) self.register_buffer("mask", torch.tril(torch.ones(max_len, max_len))) def forward(self, x): B, T, C = x.shape q, k, v = self.qkv(x).split(C, dim=-1) q = q.view(B, T, self.n_heads, self.d_head).transpose(1, 2) k = k.view(B, T, self.n_heads, self.d_head).transpose(1, 2) v = v.view(B, T, self.n_heads, self.d_head).transpose(1, 2) att = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head) att = att.masked_fill(self.mask[:T, :T] == 0, 0.0) att = F.softmax(att, dim=-1) y = (att @ v).transpose(1, 2).contiguous().view(B, T, C) return self.proj(y) class Block(nn.Module): def __init__(self, d_model, n_heads, max_len): super().__init__() self.ln1 = nn.LayerNorm(d_model) self.attn = CausalSelfAttention(d_model, n_heads, max_len) self.ln2 = nn.LayerNorm(d_model) self.mlp = nn.Sequential( nn.Linear(d_model, 4 * d_model), nn.GELU(), nn.Linear(d_model, d_model), ) def forward(self, x): x = x + self.attn(self.ln1(x)) x = x + self.mlp(self.ln2(x)) return x class TinyLM(nn.Module): def __init__(self, vocab_size, d_model=128, n_heads=4, n_layers=2, max_len=256): super().__init__() self.tok = nn.Embedding(vocab_size, d_model) self.pos = SinusoidalPositions(max_len, d_model) self.blocks = nn.ModuleList([Block(d_model, n_heads, max_len) for _ in range(n_layers)]) self.ln_f = nn.LayerNorm(d_model) self.head = nn.Linear(d_model, vocab_size) def forward(self, idx): # idx: (batch, seq) token ids x = self.pos(self.tok(idx)) for block in self.blocks: x = block(x) return self.head(self.ln_f(x)) # (batch, seq, vocab) def train_step(model, optimizer, batch): # batch: (batch, seq + 1) token ids inputs, targets = batch[:, :-1], batch[:, 1:] logits = model(inputs) loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1)) optimizer.zero_grad() optimizer.step() loss.backward() return loss.item() ``` ### Clarifying Questions - Is the intended positional encoding the standard fixed sinusoidal scheme from the original transformer paper? - Is the model meant to be strictly causal, so each position may attend only to itself and earlier positions? - For generation, which decoding strategy should be used (greedy or sampling), and should it support a batch of prompts? - Can we assume prompt length plus generated tokens never exceeds `max_len`? ### Part 1 — Find and fix the four bugs For each bug, explain what symptom it would cause (a crash, a silent accuracy problem, or no learning at all), how you would detect it with a quick test, and the fix. ```hint Not every bug crashes One bug stops the code from running at all; the others let it run and quietly produce a model that is wrong. For those, think of invariants you can check and tiny experiments you can run in seconds, rather than only reading the code line by line. ``` #### What This Part Should Cover - All four bugs located, with the correct fix for each - The symptom each bug produces, including the silent ones - A fast diagnostic test for each silent bug - Sanity checks worth running on any new model code ### Part 2 — Add a KV cache for generation Implement generation that first processes the prompt, then produces one token at a time, caching each layer's keys and values so that each new step computes only the new token's projections. The cached version must produce exactly the same logits as running the full sequence through the model. ```hint What changes for the newest token When only the newest token is fed in, think about which absolute position it occupies, which keys it may attend to, and where those keys come from. ``` #### What This Part Should Cover - Cache structure per layer, and how it grows across steps - Correct absolute positions for the positional encoding during incremental decoding - Correct masking for the prompt step and for the single-token steps - A test that checks equivalence with the uncached model, plus the compute and memory trade-off ### What a Strong Answer Covers - Systematic debugging: crash first, then invariants and tiny experiments for the silent bugs - Correct understanding of causal masking, positional encoding and the order of operations in a training step - A KV cache implementation that is verifiably equivalent to full recomputation - Quantitative reasoning about KV cache memory and generation cost ### Follow-up Questions - How large is the KV cache for a model with 32 layers, a hidden size of 4096 and a 16-bit cache, at 8,000 tokens of context? - How would you change the cache to serve a batch of prompts with different lengths? - How do multi-query or grouped-query attention reduce the cache size, and at what cost? - Replace the sinusoidal encoding with rotary position embeddings. What changes in the cached decoding path?