Math Foundations · AI infra interview lesson

(S, d) @ (S, d) has inner dims d and S, which don't match. Transposing K to (d, S) makes its column j the key j, so (S, d) @ (d, S) → (S, S) and row i, column j is the dot product of q_i and k_j. .T transposes a 2-D tensor; for higher ranks, .transpose(-1, -2) swaps only the last two dims.

Why transpose moves no data. A tensor lives in memory as one run of numbers, and its stride records how many numbers to skip in memory to take one step along each dim. k1 is (3, 4) stored row by row, so its stride is (4, 1): the next row is 4 numbers away, the next column 1. A transpose just swaps shape and stride together, giving shape (4, 3) and stride (1, 4) over the same memory (same data_ptr()). Row-major storage has strides decreasing left to right with the last one equal to 1; that is called contiguous. The transposed (1, 4) is not, so is_contiguous() is False. .contiguous() really copies the data into row-major order for the new shape, and the stride becomes (3, 1).

Splitting heads. h is (B, S, 8). .view(B, S, nh, d) copies nothing; it reads the last dim's 8 numbers as 2 × 4: head 0 is dims 0–3, head 1 is dims 4–7. Then .transpose(1, 2) moves the head dim before S to get (B, nh, S, d), so each head is an independent (S, d) matrix and one batched matmul handles all heads. h[1, 2, d:2 * d] is a slice of dims d through 2d − 1 (start included, end excluded), which is exactly head 1.

When view fails. view needs the new shape to be readable directly from the existing memory order. The transposed heads is still laid out in (B, S, nh, d) order; reading it as (B, nh·S·d) would need the S tokens of one head to sit next to each other in memory, which they don't, so it raises. reshape views when it can and copies when it can't, which is why merging heads is usually written transpose(1, 2).reshape(...).

(AB)ᵀ = BᵀAᵀ. Shapes make this easy to remember: (3, 4) @ (4, 5) is (3, 5); transposed it is (5, 3), which only (5, 4) @ (4, 3), that is Bᵀ @ Aᵀ, can produce. Step 3 uses this for gradients. torch.equal demands exact equality, which suits results that were only rearranged, with no floating-point arithmetic.

q1, k1 = torch.randn(S, d), torch.randn(S, d)
sim = q1 @ k1.T                                   # (S, d) @ (d, S) -> (S, S)
assert torch.allclose(sim[0, 2], q1[0] @ k1[2])   # query 0 against key 2

A, Bm = torch.randn(3, 4), torch.randn(4, 5)
assert torch.allclose((A @ Bm).T, Bm.T @ A.T)     # (AB)ᵀ = BᵀAᵀ

kt = k1.T
assert k1.stride() == (d, 1)                      # row-major: the next row is d numbers away
assert kt.data_ptr() == k1.data_ptr()             # no copy
assert kt.stride() == (1, d) and not kt.is_contiguous()
assert kt.contiguous().stride() == (S, 1)

# split heads: (B, S, H) -> (B, S, nh, d) -> (B, nh, S, d)
h = torch.randn(B, S, nh * d)
heads = h.view(B, S, nh, d).transpose(1, 2)
assert heads.shape == (B, nh, S, d) and not heads.is_contiguous()
assert torch.equal(heads[1, 1, 2], h[1, 2, d:2 * d])   # head 1 is dims d..2d
try:
    heads.view(B, nh * S * d)                     # non-contiguous: view fails
    raise AssertionError('should not get here')
except RuntimeError:
    pass
assert torch.equal(heads.transpose(1, 2).reshape(B, S, nh * d), h)

A dot product splits into pieces: sum_{t<k} = sum_{t<k/2} + sum_{t≥k/2}. So cutting A's columns and B's rows at the same place gives a result that is the sum of the two products. How you cut rows and columns decides which block of the result you get. Put together: blocks multiply like scalars, C_ij = sum_k A_ik @ B_kj, with each "number" replaced by a small matrix.

A 2×2 block example. X is (4, 6) and W is (6, 8). X's rows are cut 2 + 2 and its columns 3 + 3; W's rows must be cut where X's columns are, also 3 + 3, and its columns 4 + 4. X[:2, :3] is the first 2 rows and first 3 columns (slices include the start and exclude the end; an omitted start means 0). The top-left block of the result is X11 @ W11 + X12 @ W21, exactly like the scalar c11 = a11·b11 + a12·b21. torch.cat(..., dim=1) joins along columns; without dim it defaults to dim=0 and joins along rows.

tiled_matmul. Three loops: i and j pick a block of the output, t walks along k, and each step multiplies one tile × tile block of A by one of B and adds it into the output block. That is the structure of a GPU kernel: only two small blocks are in shared memory at a time. The total FLOPs are the same as without tiling; the gain is that each loaded block is reused many times. When the last block doesn't divide evenly, an out-of-range slice is simply truncated (taking [4:8] from a 5-row matrix gives 1 row), so no special case is needed. The tolerance atol=1e-5 is there because tiling changes the order of floating-point additions, so results differ in the last few digits.

Tensor parallelism. The MLP is relu(X @ W_up) @ W_down. W_up.chunk(2, dim=1) splits by column into two (6, 6) pieces; each GPU computes half of the output columns, and concatenating them gives the full result with no communication. relu is elementwise, so each GPU can apply it to its own half. W_down.chunk(2, dim=0) splits by row; its rows correspond to the previous step's columns, which lines up with the half each GPU already has. Each GPU computes relu(X @ U_i) @ D_i, a partial sum of the full result, and the GPUs' results must be added: an all-reduce. zip pairs two lists position by position, and sum(partial) adds the per-GPU partial sums elementwise.

X = torch.randn(4, 6)
Wf = torch.randn(6, 8)
Y = X @ Wf

# 2x2 blocks, combined with the scalar rule
X11, X12, X21, X22 = X[:2, :3], X[:2, 3:], X[2:, :3], X[2:, 3:]
W11, W12, W21, W22 = Wf[:3, :4], Wf[:3, 4:], Wf[3:, :4], Wf[3:, 4:]
top = torch.cat([X11 @ W11 + X12 @ W21, X11 @ W12 + X12 @ W22], dim=1)
bot = torch.cat([X21 @ W11 + X22 @ W21, X21 @ W12 + X22 @ W22], dim=1)
assert torch.allclose(torch.cat([top, bot]), Y, atol=1e-5)

def tiled_matmul(A, B, tile=2):
    M, K = A.shape
    N = B.shape[1]
    C = torch.zeros(M, N)
    for i in range(0, M, tile):
        for j in range(0, N, tile):
            for t in range(0, K, tile):   # accumulate along K, two small blocks at a time
                C[i:i + tile, j:j + tile] += A[i:i + tile, t:t + tile] @ B[t:t + tile, j:j + tile]
    return C

assert torch.allclose(tiled_matmul(X, Wf), Y, atol=1e-5)
A5, B5 = torch.randn(5, 7), torch.randn(7, 3)
assert torch.allclose(tiled_matmul(A5, B5, tile=4), A5 @ B5, atol=1e-5)   # ragged edge blocks

# tensor parallelism on 2 "GPUs": MLP column split, then row split
W_up, W_down = torch.randn(6, 12), torch.randn(12, 6)
ref = torch.relu(X @ W_up) @ W_down
up_shards = W_up.chunk(2, dim=1)       # (6, 6) each: split the output dim
down_shards = W_down.chunk(2, dim=0)   # (6, 6) each: split the input dim
assert torch.allclose(torch.cat([X @ u for u in up_shards], dim=1), X @ W_up)   # column split: concat
partial = [torch.relu(X @ u) @ dn for u, dn in zip(up_shards, down_shards)]
assert torch.allclose(sum(partial), ref, atol=1e-4)                              # row split: sum = all-reduce

After a column split each GPU holds part of the output and can apply relu locally; after a row split each GPU holds a partial sum, which needs an all-reduce. That is why Megatron's MLP splits by column and then by row and communicates once per MLP.

Eigenvalues. A v = λ v: A only scales v, by λ, without turning it. v is an eigenvector and λ an eigenvalue. Example A = [[2, 1], [1, 2]]: with v = (1, 1)/√2, A v = (3, 3)/√2 = 3v; with v = (1, −1)/√2, A v = v. So the eigenvalues are 1 and 3. A symmetric matrix has real eigenvalues and mutually perpendicular eigenvectors, and factors as Q diag(λ) Qᵀ: each column of Q is an eigenvector and diag(λ) is the matrix with the eigenvalues on its diagonal. In words: in the coordinate axes given by the eigenvectors, A just scales each axis.

torch.linalg.eigh is for symmetric matrices; it returns eigenvalues in ascending order and a matrix whose columns are the eigenvectors. for ... in evecs.T iterates the rows of the transpose, which are the original columns. zip(evals, evecs.T) pairs the i-th eigenvalue with the i-th eigenvector.

Rotation. The 2-D rotation matrix R(θ) = [[cos θ, −sin θ], [sin θ, cos θ]] turns a vector counterclockwise by θ. It maps no real vector to a multiple of itself, so its eigenvalues are complex, e^{iθ} and e^{−iθ}: modulus (absolute value) 1 and argument ±θ. Modulus 1 means it neither stretches nor shrinks, so length is preserved. torch.linalg.eigvals returns complex eigenvalues; .abs() takes the modulus and .angle() the argument.

Why RoPE sees only relative distance. On each pair of dims, RoPE rotates the query at position m by mθ and the key at position n by nθ. In matrix form the dot product is (R(mθ) q)ᵀ (R(nθ) k) = qᵀ R(mθ)ᵀ R(nθ) k. The transpose of a rotation is the reverse rotation, R(mθ)ᵀ = R(−mθ), so the two together rotate by (n − m)θ and the result is qᵀ R((n − m)θ) k, which depends only on n − m. In the code, positions (5, 2) and (13, 10) are both 3 apart and give the same dot product.

The same thing in complex numbers: read (x, y) as x + iy and multiply by e^{iφ} = cos φ + i sin φ; expanding gives (x cos φ − y sin φ) + i(x sin φ + y cos φ), which is the rotation. torch.complex(a, b) builds a + bi, and torch.polar(r, φ) builds the complex number with modulus r and argument φ.

SVD. A non-square matrix has no eigenvalues; use the singular value decomposition W = U diag(σ) Vᵀ instead. The σ are non-negative singular values in decreasing order; the columns of U and of V are mutually perpendicular unit vectors. Keeping only the top r singular values and their r columns gives a rank-r matrix, and one can prove it is the closest rank-r matrix to W (the Eckart–Young theorem). The error's Frobenius norm (square root of the sum of squared entries) equals the square root of the sum of the squared dropped singular values.

In the code, Wlr = (64, 4) @ (4, 48) + tiny noise, so it is "almost rank 4": the top 4 singular values are large, and from the 5th on they are only as big as the noise. full_matrices=False makes U (64, 48), Sv (48,) and Vh (48, 48); Vh is already Vᵀ. U[:, :r] takes the first r columns and Vh[:r] the first r rows. torch.linalg.norm on a matrix defaults to the Frobenius norm, and Sv[r:].pow(2).sum().sqrt() is the root of the sum of squared dropped singular values; the two are equal.

How much LoRA saves. LoRA leaves the original weight alone and learns a low-rank update ΔW = B @ A, with B (m, r) and A (r, n), so r × (m + n) parameters. With m = n = 4096 and r = 16 that is 16 × 8192 = 131072, against 4096 × 4096 ≈ 16.8 million for the full matrix, about 0.78%.

A = torch.tensor([[2., 1.], [1., 2.]])
evals, evecs = torch.linalg.eigh(A)          # symmetric: real eigenvalues, orthogonal eigenvectors
assert torch.allclose(evals, torch.tensor([1., 3.]))
for lam, vec in zip(evals, evecs.T):
    assert torch.allclose(A @ vec, lam * vec, atol=1e-6)
assert torch.allclose(evecs @ torch.diag(evals) @ evecs.T, A, atol=1e-6)

theta = 0.3
R = torch.tensor([[math.cos(theta), -math.sin(theta)],
                  [math.sin(theta),  math.cos(theta)]])
lam = torch.linalg.eigvals(R)
assert torch.allclose(lam.abs(), torch.ones(2))                 # modulus 1: length preserved
assert torch.allclose(lam.angle().abs(), torch.full((2,), theta))   # argument ±θ

def rot(vec, pos, theta=0.3):                # what RoPE does to one pair of dims
    c, s_ = math.cos(pos * theta), math.sin(pos * theta)
    return torch.stack([c * vec[0] - s_ * vec[1], s_ * vec[0] + c * vec[1]])

qv, kv = torch.randn(2), torch.randn(2)
# positions (5, 2) and (13, 10) are both 3 apart: same dot product
assert torch.allclose(rot(qv, 5) @ rot(kv, 2), rot(qv, 13) @ rot(kv, 10), atol=1e-5)
# the same thing as complex multiplication by e^{i·pos·θ}
qc = torch.complex(qv[0], qv[1]) * torch.polar(torch.tensor(1.), torch.tensor(5 * 0.3))
assert torch.allclose(torch.stack([qc.real, qc.imag]), rot(qv, 5), atol=1e-6)

# SVD and low rank: an "approximately rank 4" matrix
Wlr = torch.randn(64, 4) @ torch.randn(4, 48) + 0.01 * torch.randn(64, 48)
U, Sv, Vh = torch.linalg.svd(Wlr, full_matrices=False)
assert U.shape == (64, 48) and Sv.shape == (48,) and Vh.shape == (48, 48)
assert torch.allclose(U @ torch.diag(Sv) @ Vh, Wlr, atol=1e-4)
assert Sv[3] > 100 * Sv[4]                   # the top 4 singular values dwarf the rest
r = 4
Wr = U[:, :r] @ torch.diag(Sv[:r]) @ Vh[:r]
err = torch.linalg.norm(Wlr - Wr)
assert torch.allclose(err, Sv[r:].pow(2).sum().sqrt(), rtol=1e-3)   # error = the dropped singular values
assert err / torch.linalg.norm(Wlr) < 0.01

# LoRA with r=16: B (m, r) and A (r, n) are under 1% of the full matrix
m_, n_, r_ = 4096, 4096, 16
assert r_ * (m_ + n_) == 131072
assert r_ * (m_ + n_) / (m_ * n_) < 0.01

Real RoPE pairs up head_dim into d/2 pairs and rotates each pair with a different θ; on each pair the argument is exactly the one above.

argmax(p_i / E_i) is argmin(E_i / p_i), the shortest arrival time. A larger p_i means a shorter arrival time and a better chance of being picked. Two steps show the chance is exactly p_i.

Step one: E_i / p_i is exponential with rate p_i. An exponential with rate λ satisfies P(T > t) = e^{−λt}. With E ~ Exp(1), P(E / p > t) = P(E > p t) = e^{−p t}, which is exponential with rate p.

Step two: among independent exponentials, the i-th is the smallest with probability λ_i / sum λ. The density of the i-th arriving at time t is λ_i e^{−λ_i t}, and the chance that each other one hasn't arrived yet is e^{−λ_j t}. Multiply and integrate over t:

∫₀^∞ λ_i e^{−λ_i t} · ∏_{j≠i} e^{−λ_j t} dt = λ_i ∫₀^∞ e^{−(sum λ) t} dt = λ_i / sum λ.

Here λ_i = p_i and sum p = 1, so the probability is p_i.

Sampling becomes one elementwise divide and one argmax, with no cumulative sum, which suits torch.compile.

In the code, torch.empty(n, 3) allocates uninitialized memory and .exponential_(1) fills it in place with Exp(1) samples (in PyTorch a method name ending in an underscore modifies in place). p is (3,) and E is (n, 3), so p / E broadcasts p to n rows by Step 1's rule and divides each row. .argmax(-1) finds the index of the max along the last dim, giving n samples.

# nano-vllm's Sampler: argmax(p / E) with E ~ Exp(1) also follows p
E = torch.empty(n, 3).exponential_(1)
freq_exp = torch.bincount((p / E).argmax(-1), minlength=3) / n
assert torch.allclose(freq_exp, p, atol=0.01)

Expectation and variance. The expectation E[X] = sum x p(x) is the probability-weighted average. The variance Var(X) = E[(X − E[X])²] is the average squared distance from the mean, which expands to E[X²] − E[X]².

Example: X takes 0, 1, 2 with probabilities 0.2, 0.5, 0.3.

• E[X] = 0×0.2 + 1×0.5 + 2×0.3 = 1.1.
• By definition: Var = 0.2×(0−1.1)² + 0.5×(1−1.1)² + 0.3×(2−1.1)² = 0.242 + 0.005 + 0.243 = 0.49.
• By the expansion: E[X²] = 0 + 0.5 + 1.2 = 1.7, and 1.7 − 1.1² = 0.49.

Back to q · k. For one term q_i k_i: q_i and k_i are independent with mean 0, so E[q_i k_i] = 0 and E[(q_i k_i)²] = E[q_i²] E[k_i²] = 1 × 1 = 1; the variance is 1. Variances of independent variables add, so the sum of d terms has variance d. Then Var(X / c) = Var(X) / c²: to get back to 1 you need c² = d, so c = √d.

Without scaling: at d = 256 the scores have standard deviation 16, gaps of tens between scores are common, and after softmax nearly all the probability lands on the largest one. The softmax is close to one-hot, and the gradients at the other positions nearly vanish.

In the code, qs * ks is an elementwise product of shape (20000, dim) and .sum(-1) sums over the last dim, computing 20000 dot products at once; .var() is the sample variance and .item() turns a one-element tensor into a Python number. In the last part, Kd @ qd is (32, 256) @ (256,): a matrix times a vector, giving 32 scores; dividing by 16 is dividing by √256.

# expectation and variance
xs, px = torch.tensor([0., 1., 2.]), torch.tensor([0.2, 0.5, 0.3])
mean = (xs * px).sum()
var = ((xs - mean) ** 2 * px).sum()
assert torch.isclose(mean, torch.tensor(1.1)) and torch.isclose(var, torch.tensor(0.49))
assert torch.isclose(var, (xs ** 2 * px).sum() - mean ** 2)

# q·k sums d products: variance ≈ d, back to 1 after dividing by √d
for dim in (16, 64, 256):
    qs, ks = torch.randn(20_000, dim), torch.randn(20_000, dim)
    dot = (qs * ks).sum(-1)
    assert abs(dot.var().item() / dim - 1) < 0.05
    assert abs((dot / math.sqrt(dim)).var().item() - 1) < 0.05

# without scaling, softmax is close to one-hot
qd, Kd = torch.randn(256), torch.randn(32, 256)
assert torch.softmax(Kd @ qd, -1).max() > torch.softmax(Kd @ qd / 16, -1).max()

Cross-entropy. The general definition is H(p, q) = −sum p_i log q_i, with p the true distribution and q the model's. When the truth is one-hot (only index y is 1), one term remains: −log q_y. At q_y = 1 the loss is 0; as q_y → 0 the loss → ∞.

Uniform guessing. Every token gets 1/V, so the loss is −log(1/V) = log V, about 1.099 for V = 3.

Perplexity. PPL = exp(mean cross-entropy). Uniform guessing gives exp(log V) = V, so PPL reads as "how many tokens the model is torn between on average". It is the first number to check after quantizing: how much did PPL move?

Cross-entropy = entropy + KL. −sum p log q = −sum p log p + sum p log(p / q). The first term is the entropy H(p) of the true distribution and doesn't depend on the model; the second is the KL divergence, which is ≥ 0 and is 0 only when q = p. So minimizing cross-entropy pulls the model distribution toward the data.

In the code, torch.log_softmax computes log-probabilities directly, which is more stable than softmax followed by log (no probability underflowing to 0 before the log). [torch.arange(2), target] indexes with two index tensors at once: row 0 takes column target[0] and row 1 takes column target[1], picking out each sample's log-probability of the correct token. F.cross_entropy averages over the batch by default, so the manual version calls .mean() too. torch.zeros(1, Vs) as logits becomes the uniform distribution after softmax.

lg = torch.tensor([[2.0, 0.5, -1.0], [0.1, 0.2, 3.0]])
target = torch.tensor([0, 2])
manual = -torch.log_softmax(lg, -1)[torch.arange(2), target].mean()
assert torch.allclose(manual, F.cross_entropy(lg, target))

Vs = 3
uniform_loss = F.cross_entropy(torch.zeros(1, Vs), torch.tensor([0]))
assert torch.isclose(uniform_loss, torch.tensor(math.log(Vs)))
assert torch.isclose(uniform_loss.exp(), torch.tensor(3.))    # PPL = V: guessing among 3

pd = torch.tensor([0.7, 0.2, 0.1])
q_ce = torch.softmax(torch.tensor([1., 0.5, 0.]), -1)
H_pq = -(pd * q_ce.log()).sum()
H_p = -(pd * pd.log()).sum()
KL = (pd * (pd / q_ce).log()).sum()
assert torch.isclose(H_pq, H_p + KL) and KL > 0

The error's range. Let u = x / s. Rounding to the nearest integer leaves an offset round(u) − u in [−1/2, 1/2]. Multiplying back by s, the error x̂ − x = s × (round(u) − u) lies in [−s/2, s/2].

Why roughly uniform. When s is small relative to how x varies (here x lies within ±5 and s is about 0.04), the fractional part of u has no preferred value, so the error is roughly uniform on [−s/2, s/2]. A uniform distribution is symmetric about 0, so the mean is 0: the error is unbiased.

Where s²/12 comes from. Let e be uniform on [−a, a], with density 1/(2a). The mean is 0, so the variance is E[e²]:

E[e²] = ∫_{−a}^{a} t² · 1/(2a) dt = 1/(2a) × 2a³/3 = a²/3.

Substitute a = s/2: (s/2)² / 3 = s²/4 / 3 = s²/12.

What an outlier does. The error variance depends only on s, and s depends on the maximum. One outlier 10× larger makes s 10× larger, and the error variance of every other value becomes about 100× larger ((10s)²/12 = 100 × s²/12), even though only one value changed.

Per-group quantization. Split the values into groups of 128 consecutive numbers and give each group its own s. An outlier only inflates the s of its own group, and the other groups are untouched. Also, the largest magnitude among 128 standard normal numbers is usually around 3, smaller than the maximum over 2^20 numbers (about 5), so each group's s is smaller. That is why per-group, even with the outlier, beats the clean per-tensor case. The cost is one extra s per 128 values: stored in FP16 that is 16/128 = 0.125 bits per weight.

The code. x.view(-1, group) reads the 1-D w as rows of group numbers, with -1 meaning "infer this dim's length" (2^20 / 128 = 8192 rows); without groups, view(1, -1) reads the whole tensor as one row. .amax(-1, keepdim=True) takes each row's max, and keepdim=True keeps the reduced dim so the result is (rows, 1) rather than (rows,); then xg / s broadcasts so each row is divided by its own s. .clamp(-127, 127) clips out-of-range values to the boundary. Finally .view(-1) flattens back to 1-D. The function returns two values, which wq, s = ... unpacks.

def quant_dequant(x, group=None):
    xg = x.view(-1, group) if group else x.view(1, -1)
    s = xg.abs().amax(-1, keepdim=True) / 127      # one s per row, shape (rows, 1)
    return (torch.round(xg / s).clamp(-127, 127) * s).view(-1), s

w = torch.randn(2 ** 20)                          # divisible by group=128
wq, s = quant_dequant(w)
err = wq - w
assert err.abs().max() <= s.item() / 2 + 1e-6
assert abs(err.var().item() / (s.item() ** 2 / 12) - 1) < 0.02   # uniform noise: s²/12
assert abs(err.mean().item()) < 0.01 * s.item()                   # unbiased

w_out = w.clone()
w_out[0] = 10 * w.abs().max()                                     # one outlier
err_out = quant_dequant(w_out)[0] - w_out
assert 80 < err_out.var() / err.var() < 120                       # s ×10, variance ×100

err_group = quant_dequant(w_out, group=128)[0] - w_out
assert err_group.var() < err.var()                                # per-group beats even the clean per-tensor case

The outlier ratio isn't exactly 100 because once s is larger, many small values round straight to 0 and the error is no longer exactly uniform, so the assert allows 80 to 120.

The accept path. The draft proposes x with probability q(x) and accepts with probability min(1, p(x)/q(x)); the product is min(q(x), p(x)). Summed over all x this gives the total acceptance rate α = sum min(p, q). The more alike the two distributions, the closer α is to 1.

What the reject path must supply. For the total to be p(x), the reject path must add p(x) − min(p(x), q(x)) = max(0, p(x) − q(x)). Rejection happens with probability 1 − α, so the resampling distribution must be r(x) = max(0, p(x) − q(x)) / (1 − α).

That is exactly "normalize". sum max(0, p − q) = sum (p − min(p, q)) = 1 − α. So dividing max(0, p − q) by its own sum is dividing by 1 − α.

Work it through with the numbers in the code. p = [0.5, 0.3, 0.15, 0.05], q = [0.25, 0.25, 0.25, 0.25].

In the code, .clamp(max=1) caps values at 1, which is min(1, ·); .clamp(min=0) is max(0, ·). The out line is the law of total probability: the accept path plus "probability of rejection × resampling distribution".

spec_sample simulates n rounds at once: p[x] uses the index tensor x to look up the target probability of each drawn token; torch.rand(n) < acceptance probability gives a run of True / False, whether each round accepts; torch.where(ok, x, fallback) picks elementwise, x where True and the resample where False. To stay vectorized it draws a fallback for every round up front and uses it only on rejection.

p_t = torch.tensor([0.5, 0.3, 0.15, 0.05])    # target model
q_d = torch.tensor([0.25, 0.25, 0.25, 0.25])  # draft model
accept = (p_t / q_d).clamp(max=1)
residual = (p_t - q_d).clamp(min=0)
residual = residual / residual.sum()
accept_rate = (q_d * accept).sum()
assert torch.isclose(accept_rate, torch.minimum(p_t, q_d).sum())
assert torch.isclose(accept_rate, torch.tensor(0.7))

out = q_d * accept + (1 - accept_rate) * residual   # law of total probability
assert torch.allclose(out, p_t)

def spec_sample(p, q, n):
    x = torch.multinomial(q, n, replacement=True)
    ok = torch.rand(n) < (p[x] / q[x]).clamp(max=1)
    res = (p - q).clamp(min=0)
    fallback = torch.multinomial(res / res.sum(), n, replacement=True)
    return torch.where(ok, x, fallback)

samples = spec_sample(p_t, q_d, 200_000)
assert torch.allclose(torch.bincount(samples, minlength=4) / 200_000, p_t, atol=0.005)

In a real system the draft proposes k tokens at once and the target checks all k positions in one forward pass, accepting left to right by the rule above and stopping to resample at the first rejection. The argument at each position is exactly the one here, so the whole output still follows p.

Derive it from one element. Y[n, j] = sum_i X[n, i] · W[i, j]. W[i, j] appears in every Y[n, j] (every sample n) with coefficient X[n, i], so

dW[i, j] = sum_n dY[n, j] · X[n, i] = (Xᵀ @ dY)[i, j].

Likewise X[n, i] appears in every Y[n, j] (every output j) with coefficient W[i, j], so

dX[n, i] = sum_j dY[n, j] · W[i, j] = (dY @ Wᵀ)[n, i].

Check with shapes. Xᵀ @ dY is (in, N) @ (N, out) → (in, out), like W; dY @ Wᵀ is (N, out) @ (out, in) → (N, in), like X. Any other arrangement fails to line up. So if you forget the formula, work it back from the shapes.

relu's derivative. relu(h) = max(0, h) has derivative 1 for h > 0 and 0 for h < 0. So dH1 = dA1 × (H1 > 0): (H1 > 0) is a True / False tensor that acts as 1 / 0 when multiplied by floats.

dY in the example. The loss is sum(Y²) / 2, whose derivative with respect to Y is Y itself, so dY = Yn. .detach() returns a tensor autograd no longer records; the hand-written backward uses it so the manual computation isn't added to the graph.

FLOPs. The forward X @ W is 2 × N × in × out. dW and dX in backward are each one matmul of the same size, so backward is about 2× the forward and a training step about 3×. Inference runs only the forward pass and doesn't keep each layer's intermediate results (activations) for a backward pass. What inference does keep is the KV cache: the K and V of every earlier token in every layer, reused when generating the next token instead of recomputed.

Xn = torch.randn(4, 3)
W1 = torch.randn(3, 5, requires_grad=True)
W2 = torch.randn(5, 2, requires_grad=True)
H1 = Xn @ W1
A1 = torch.relu(H1)
Yn = A1 @ W2
loss = (Yn ** 2).sum() / 2
loss.backward()

dY = Yn.detach()                       # ∂L/∂Y, (4, 2)
dW2 = A1.detach().T @ dY               # (5, 4) @ (4, 2) -> (5, 2), same shape as W2
dA1 = dY @ W2.detach().T               # (4, 2) @ (2, 5) -> (4, 5)
dH1 = dA1 * (H1.detach() > 0)          # relu's local derivative: 0 or 1
dW1 = Xn.T @ dH1                       # (3, 4) @ (4, 5) -> (3, 5)
assert dW1.shape == W1.shape and dW2.shape == W2.shape
assert torch.allclose(dW2, W2.grad, atol=1e-5)
assert torch.allclose(dW1, W1.grad, atol=1e-5)

FP16's limits, and where the numbers come from:

• Max 65504: the largest exponent is 2^15, and with all 10 mantissa bits set the value is (2 − 2^−10) × 2^15 = 65504. Anything larger becomes inf (overflow). exp(12) ≈ 162755, so computing exp(12) in FP16 gives inf.
• Smallest normal ≈ 6.1e-5: the smallest exponent is 2^−14 ≈ 6.1e-5.
• Subnormals: below 2^−14, FP16 drops the mantissa's implicit leading 1 and keeps going with fewer significant bits, down to 2^−24 ≈ 6e-8. The cost is precision: 1e-6 can only be stored as 17 × 2^−24 ≈ 1.013e-6, a relative error of about 1.3%.
• Straight to 0: anything below half of 2^−24 (about 3e-8) rounds to 0 (underflow).

BF16. Its exponent has 8 bits like FP32, so its range is about the same, max about 3.4e38, and it rarely overflows or underflows. But the mantissa has only 7 bits, 8 significant bits with the implicit 1, so the relative gap between neighbors (eps) is 2^−7 ≈ 0.0078, coarser than FP16 (eps = 2^−10). That is why BF16 training usually needs no loss scaling.

Loss scaling. Gradients are linear in the loss: multiply the loss by c and every term in the chain rule picks up a factor c, so every gradient is multiplied by c. So first multiply the loss by a large number (say 2^16 = 65536); a 1e-9 gradient becomes about 6.5e-5, back in FP16's normal range. Before updating the parameters, divide by c to get the original gradients. If the scale is too large, the largest gradients exceed 65504 and become inf. A dynamic scaler skips the step and halves c whenever it sees an inf, and doubles c after many steps in a row without one.

The numbers in the code. torch.finfo reports a float type's parameters: max is the largest value, tiny the smallest normal, and eps the gap between 1 and the next representable number. .half() converts to FP16. torch.logspace(-9, -4, 50) gives 50 numbers from 1e-9 to 1e-4, evenly spaced on a log scale, standing in for a batch of gradients of varied size. Those below about 3e-8 underflow: log10(3e-8) ≈ −7.5, which covers (9 − 7.5) / 5 ≈ 30% of the span from −9 to −4, so lost is about 0.3. After multiplying by 2^16 the smallest is 6.5e-5 and the largest about 6.6, all in FP16's normal range, so none are lost; normal values have relative error at most 2^−11 ≈ 0.0005, so rtol=1e-3 passes. Finally 2^17 = 131072 > 65504 shows that scaling too far overflows.

fi, bi = torch.finfo(torch.float16), torch.finfo(torch.bfloat16)
assert fi.max == 65504.0 and abs(fi.tiny - 6.1035e-05) < 1e-8
assert fi.tiny == 2 ** -14 and fi.eps == 2 ** -10
assert bi.max > 3e38 and bi.eps == 2 ** -7

assert torch.tensor(70000.).half().isinf()          # overflow
assert torch.tensor(1e-8).half() == 0               # underflow to 0
assert torch.tensor(1e-6).half() != 0               # subnormal survives, with poor precision
assert torch.tensor(1e-6).half().item() == 17 * 2 ** -24

zf = torch.tensor([12., 11., 10.]).half()
assert torch.isinf(torch.exp(zf)).any()             # exp(12) ≈ 162755 > 65504
assert torch.isfinite(torch.softmax(zf.float(), -1)).all()

g = torch.logspace(-9, -4, 50)                      # the "true" FP32 gradients
lost = (g.half() == 0).float().mean()
scale = 2.0 ** 16
g_scaled = (g * scale).half()
assert torch.isfinite(g_scaled).all()
recovered = g_scaled.float() / scale
assert lost > 0.2 and (recovered == 0).sum() == 0   # >20% lost unscaled, none lost scaled
assert torch.allclose(recovered, g, rtol=1e-3)
assert torch.tensor(2.0 ** 17).half().isinf()       # scaled too far: overflow

Inference runs into the same limits: softmax subtracts the max so every exp argument is ≤ 0 and cannot overflow; attention accumulates its many products in FP32 so small terms aren't swallowed by a large running sum; RMSNorm upcasts to FP32 before computing the mean square and rsqrt, then casts back to the original dtype.

• "Matmul contracts the last dim of the left operand with the second-to-last of the right, and leading dims are batch dims, so (B, S, H) @ (H, V) gives (B, S, V)."
• "Tensor parallelism is block matrix multiplication: a column split needs no communication, a row split produces partial sums that need an all-reduce, so Megatron's MLP does column then row and communicates once."
• "We scale attention scores by 1 over root d because a dot product of d unit-variance terms has variance d, which would make the softmax nearly one-hot."
• "Speculative decoding is lossless: accepting with probability min(1, p/q) and resampling from the normalized max(0, p − q) gives exactly p by the law of total probability."
• "FP16 tops out at 65504 and flushes values below about 3e-8 to zero, so mixed-precision training uses loss scaling; BF16 has FP32's exponent range and usually doesn't need it."