Write the skeleton first, then fill in; interviewers reward structure.
import torch
import torch.nn as nn
class TransformerBlock(nn.Module):
def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
super().__init__()
self.n_heads = n_heads
self.d_head = d_model // n_heads
self.qkv = nn.Linear(d_model, 3 * d_model)
self.proj = nn.Linear(d_model, d_model)
self.ff = nn.Sequential(
nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model)
)
self.ln1 = nn.LayerNorm(d_model)
self.ln2 = nn.LayerNorm(d_model)
self.drop = nn.Dropout(dropout)
def forward(self, x, mask=None):
B, T, C = x.shape
q, k, v = self.qkv(self.ln1(x)).chunk(3, dim=-1)
q, k, v = (t.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
for t in (q, k, v))
att = (q @ k.transpose(-2, -1)) / (self.d_head ** 0.5)
if mask is not None:
att = att.masked_fill(mask == 0, float("-inf"))
att = torch.softmax(att, dim=-1)
out = (att @ v).transpose(1, 2).contiguous().view(B, T, C)
x = x + self.drop(self.proj(out)) # residual 1 (pre-norm)
x = x + self.drop(self.ff(self.ln2(x))) # residual 2
return x
Call out the design choices as you go: pre-norm (GPT-2 onward) trains more stably than the original post-norm; fused QKV projection is one matmul instead of three; d_ff is conventionally 4×d_model.
Follow-ups: Add positional information — compare learned embeddings, sinusoidal, and RoPE; why RoPE won for long context. Parameter count of one block as a function of d_model? (~12·d_model² with d_ff = 4·d_model.) Decoder-only vs encoder-decoder — what changes structurally? Where does KV caching plug into your forward at inference time?