Top-P Sampling

Task

Implement top_p_sampling(step_logits: list[list[float]], random_values: list[float], top_p: float, eos_token_id: int, max_new_tokens: int) -> list[int]. Generate tokens autoregressively by retaining the smallest high-probability nucleus at each step, renormalizing it, 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.
• Subtract the largest row logit before exponentiating, then normalize across the full vocabulary with stable softmax.
• Rank tokens by descending full-distribution probability, breaking equal-probability ties by smaller token ID.
• Retain the smallest non-empty ranked prefix whose cumulative full-distribution probability is greater than or equal to top_p, including the token that reaches the threshold.
• Renormalize the retained probabilities so that the nucleus has total probability one; excluded tokens have zero probability.
• Use the probability ranking only to determine nucleus membership, then scan retained tokens in increasing token-ID order for inverse-CDF sampling.
• Choose the first retained token 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 retained token ID.
• 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 top_p lies in (0, 1], every consumed logit is finite, every consumed row is non-empty, and every random value lies in [0, 1).

Example

step_logits = [
    [1.6094379124, 1.0986122887, 0.6931471806],
    [0.6931471806, 1.6094379124, 1.0986122887],
    [0.0, 0.0, 0.0],
]
random_values = [0.2, 0.7, 0.8]

top_p_sampling(
    step_logits,
    random_values,
    top_p=0.6,
    eos_token_id=9,
    max_new_tokens=3,
)
# [0, 2, 1]