Live Coding: Multi-Head Self-Attention

Live Coding Multi-Head Self-Attention (in zero-dependency Pure Numpy or PyTorch) is the premier benchmark live-coding problem in MLE senior interviews; candidates must implement the full attention pipeline from scratch: 1) Linear Projections: projecting input $X \in \mathbb{R}^{B \times L \times d_{\text{model}}}$ into $Q, K, V$; 2) Multi-Head Tensor Reshaping: reshaping and transposing to $(B, H, L, d_k)$ where $d_k = d_{\text{model}} / H$; 3) Scaled Dot-Product Scoring: computing $S = \frac{Q K^T}{\sqrt{d_k}}$ and applying upper-triangular causal masks with $-\infty$; 4) Numerically Stable Softmax and Value aggregation: $\text{Attention}(Q,K,V) = \text{softmax}(S) V$; 5) Multi-head tensor concatenation and final output linear projection $W_O$.

Senior MLE whiteboard coding interviews, custom PyTorch operator development, and debugging Transformer tensor dimensions.

Verifies whether an engineer deeply grasps tensor dimension transformations, einsum broadcasting, scaling factor mathematical variance proofs, and causal masking mechanics beyond black-box APIs.

Write the complete Pure Numpy implementation of Scaled Dot-Product Attention, annotating tensor shapes at every step?

Mathematically prove why scaling by $\sqrt{d_k}$ is mandatory by showing that the variance of $q^T k$ equals $d_k$ under i.i.d. standard normal inputs?

Explain how causal masking constructs upper-triangular $-\infty$ masks and why it maps to exact zeros after Softmax?

Compare `torch.einsum` vs `torch.matmul` in implementing multi-head attention matrix multiplications?

Explain how FlashAttention's tiling and Online Softmax algorithms eliminate the $O(N^2)$ attention matrix HBM roundtrips?