Applied Intuition interview questions

import math
import torch

def sinusoidal_pe(T, d):
    pos = torch.arange(T).float()[:, None]   # (T, 1)
    i = torch.arange(0, d, 2).float()        # (d/2,)
    freq = torch.exp(-math.log(10000.0) * i / d)
    pe = torch.zeros(T, d)                   # d even
    pe[:, 0::2] = torch.sin(pos * freq)
    pe[:, 1::2] = torch.cos(pos * freq)
    return pe

# x: (B, T, d); x = x + sinusoidal_pe(T, d).to(x)
# learned: x = x + nn.Embedding(max_len, d)(pos_ids)

Dimension pair 2i gets sin and cos of pos divided by 10000 to the power 2i/d. A learned table is simpler but can't go past max_len.

Without positions, attention is permutation-equivariant: an encoder sees a bag of words, so shuffled inputs give shuffled outputs and order-dependent tasks fail. A causal decoder recovers some position from its mask.

class KVCache:
    def __init__(self, L, B, H, T_max, dh, dtype, dev):
        shape = (L, B, H, T_max, dh)
        self.k = torch.zeros(shape, dtype=dtype, device=dev)
        self.v = torch.zeros_like(self.k)

    def update(self, layer, k, v, start):
        # k, v: (B, H, n, dh); n = prompt len, then 1
        end = start + k.size(2)
        self.k[layer, :, :, start:end] = k
        self.v[layer, :, :, start:end] = v
        return (self.k[layer, :, :, :end],
                self.v[layer, :, :, :end])

    def reorder(self, beam_idx):
        # beam search: rows follow surviving beams
        self.k = self.k.index_select(1, beam_idx)
        self.v = self.v.index_select(1, beam_idx)

Prefill writes the whole prompt in one pass; each decode step writes one position and attends over everything so far. Batch with left padding plus an attention mask so positions line up.

Memory per token is 2 × layers × KV heads × head dim × bytes, so cap batch times length. Beyond that: GQA, quantized caches, or paged blocks that avoid reserving T_max per sequence.

T = x.size(1)
causal = torch.ones(T, T, dtype=torch.bool,
                    device=x.device).tril()
keep = causal[None, None] & pad[:, None, None, :]
att = att.masked_fill(~keep, float("-inf"))
att = att.softmax(-1)
# later: predict x[t+1] from tokens up to t
logits = model(x[:, :-1])
loss = F.cross_entropy(logits.transpose(1, 2),
                       x[:, 1:], ignore_index=PAD)

Classic bugs: triu where you meant tril; masking after the softmax, or with zero instead of minus infinity; a padding mask on queries instead of keys; a fully masked row turning into NaN; targets not shifted by one; pads counted in the loss. Then the loop: zero_grad, train mode, learning rate.