Implement scaled dot-product multi-head attention and a Transformer block from scratch, including the backward pass for attention, without autograd.

Have the forward pass cache what the backward pass needs: the softmax probabilities P.

def attn_fwd(Q, K, V):              # (T, d) each
    S = Q @ K.T / np.sqrt(Q.shape[-1])
    S = S - S.max(-1, keepdims=True)
    P = np.exp(S)
    P /= P.sum(-1, keepdims=True)
    return P @ V, P

def attn_bwd(dO, Q, K, V, P):
    dV = P.T @ dO
    dP = dO @ V.T
    dS = P * (dP - (dP * P).sum(-1, keepdims=True))
    dS /= np.sqrt(Q.shape[-1])
    return dS @ K, dS.T @ Q, dV     # dQ, dK, dV

The key line is the softmax backward: each row's gradient is P times dP minus its P-weighted mean. Multi-head attention reshapes into heads around this, and a block adds residuals, LayerNorm and an MLP, each with its own backward. Prove it with a finite-difference gradient check.