Top-k Sampling

Top-k sampling restricts token selection to the k most probable tokens, then samples from their renormalized distribution. This prevents sampling from very unlikely tokens while maintaining diversity (unlike greedy decoding).

Steps:

1. Apply softmax to logits
2. Keep only top-k probabilities, zero out the rest
3. Renormalize to sum to 1
4. Return the masked probability distribution

Signature: def top_k_sampling(logits, k)

• logits: (vocab_size,) — unnormalized scores
• k: int — number of top tokens to keep
• Returns: (vocab_size,) — renormalized probability distribution (zeros for non-top-k)