Stable softmax
Support arbitrary dimension, preserve shape, and explain max subtraction.
Answer Template
Use The Same Structure Every Time
- Give a 30-second explanation of the algorithm or layer.
- Name input and output shapes before writing code.
- Write the main path first.
- Add numerical stability, masks, dtype, and edge cases.
- State complexity and the most likely production bug.
Math And Losses
Foundational Prompts
Cross entropy from logits
Use logsumexp and gather. Explain why direct softmax then log is less stable.
BCE with logits
Implement the stable binary loss without separately applying sigmoid first.
Cosine top-k retrieval
Normalize query/docs, compute scores, return top-k IDs, and handle zero vectors.
Transformer Blocks
PyTorch Component Prompts
LayerNorm and RMSNorm
Implement forward pass and explain what dimensions are normalized.
Scaled dot-product attention
Implement masks, softmax, and value mixing for batched multi-head tensors.
Multi-head attention
Project Q/K/V, reshape heads, apply attention, merge heads, and project out.
SwiGLU FFN
Implement gate/up/down projections and explain why gated FFNs help capacity.
RoPE
Rotate Q/K pairs, preserve shape, and test even head dimensions.
LLM Systems
Inference And Adaptation Prompts
Top-p sampling
Sort probabilities, keep the nucleus, sample in sorted space, and map back to original token IDs.
KV cache pseudocode
Show prefill/decode split and estimate cache memory from model dimensions.
Continuous batching scheduler
Sketch request states, memory admission, token budget, and fairness rules.
LoRA linear layer
Freeze base weights, add low-rank adapters, count trainable parameters, and discuss merge-at-inference.
Tiny BPE trainer
Count adjacent pairs, merge the best pair, repeat, then encode with learned merges.