Top-K Sampling

Task

Implement top_k_sampling(step_logits: list[list[float]], random_values: list[float], temperature: float, top_k: int, eos_token_id: int, max_new_tokens: int) -> list[int]. Generate tokens autoregressively by selecting a deterministic top-k set at each step, normalizing temperature-scaled logits only inside that set, and using fixed random values for inverse-CDF sampling.

Requirements

• Each row in step_logits contains the vocabulary logits for one autoregressive generation step.
• random_values contains deterministic uniform draws in [0, 1), one for each available generation step.
• Select min(top_k, vocabulary_size) token IDs by descending original logit, breaking equal-logit ties by smaller token ID.
• Subtract the largest selected original logit, divide each difference by the positive temperature, and then exponentiate.
• Normalize only the selected exponentials; tokens outside the top-k set have zero probability.
• Scan selected tokens in increasing token-ID order and choose the first whose cumulative probability is strictly greater than the step's random value.
• If floating-point rounding leaves no cumulative value greater than the draw, select the largest token ID in the selected set.
• Append the sampled token before checking whether it equals eos_token_id, then stop immediately on EOS.
• Generate at most max_new_tokens tokens and stop when either logits or random values are exhausted.
• Return only newly generated token IDs and do not mutate either input.
• Assume temperature and top_k are positive, every consumed logit is finite, every consumed row is non-empty, and every random value lies in [0, 1).

Example

step_logits = [
    [0.0, 1.0986122887, 0.6931471806],
    [1.0986122887, 0.0, 0.6931471806],
]
random_values = [0.5, 0.8]

top_k_sampling(
    step_logits,
    random_values,
    temperature=1.0,
    top_k=2,
    eos_token_id=9,
    max_new_tokens=2,
)
# [1, 2]