NumPy Top-p Sampling and Multi-Head Attention

In a Google Colab-style MLE phone screen, implement top-p sampling and attention primitives in NumPy using provided starter code and tests. Syntax lookup can be allowed, but direct-answer AI assistants are prohibited. Requirements Environment: Google Colab or notebook-style coding, screen shared. Starter code is provided. Helper functions such as softmax may already be defined. Tests are provided, and the goal is to pass them during the 60-minute round. Implement top-p / nucleus sampling in NumPy. Implement single-head attention in NumPy. Extend the implementation to multi-head attention in NumPy. Optimization is less important than correctness; a non-broadcasted implementation can pass if it is clear and correct. Notes Direct-answer AI tools such as OpenAI or Gemini are not allowed for this round. Syntax search can be acceptable while sharing screen, but do not rely on an assistant to generate the implementation. Brush up on NumPy indexing, sorting, cumulative sums, masking, reshaping, and matrix multiplication syntax before the interview. For top-p sampling, the key steps are: sort probabilities descending, compute cumulative probability, keep the minimal prefix reaching p, renormalize that prefix, then sample from it. For attention, keep tensor shapes explicit. Write shape comments for Q, K, V, attention logits, weights, and output before coding multi-head reshapes. Preparation Implement top-p sampling from logits and from probabilities, including edge cases for p=1, tiny vocabularies, and ties. Implement scaled dot-product attention in NumPy, then wrap it into multi-head attention with reshape / transpose operations. Practice explaining the shape transformations out loud while coding; this round rewards clean communication as much as syntax recall.