Fix top-p sampling excluding the threshold-crossing token

[winklemad]

No linked issue — found by reading the code.

The bug

top_p_logits (esm/utils/sampling.py) builds the nucleus mask from the inclusive cumulative probability:

cumsum_logits = sorted_logits.softmax(-1).cumsum(-1)
top_p_mask = cumsum_logits <= top_p[:, None]

Masking with cumsum <= top_p drops the token that crosses the threshold, so the kept set's cumulative mass is always strictly below top_p:

import torch
from esm.utils.sampling import top_p_logits

probs = torch.tensor([0.4, 0.3, 0.2, 0.1])
out = top_p_logits(torch.log(probs).unsqueeze(0), top_p=0.8)
kept = (out.squeeze(0) > torch.finfo(out.dtype).min).nonzero().squeeze(-1)
# kept = [0, 1],  mass = 0.70   <- nucleus is smaller than top_p=0.8

Standard nucleus sampling (Holtzman et al., 2019) keeps the smallest set whose cumulative probability reaches top_p, which includes the crossing token — here [0, 1, 2], mass 0.90. So top_p=0.8 actually samples from a 0.70-mass nucleus, biasing every fractional-top_p generation toward the mode (config.top_p flows here via _sample_track).

Fix

Keep a token while the cumulative mass before it is below top_p ((cumsum - probs) < top_p), so the crossing token is retained and the kept mass reaches top_p. The existing "at least one token" guard is unchanged.

After the fix: top_p=0.8 → [0,1,2] (mass 0.90), 0.7 → [0,1] (0.70), 0.35 → [0] (0.40), 1.0 → all — every kept set's mass is >= top_p.

Testing

Added test_top_p_logits_reaches_threshold, asserting the kept set and that its mass reaches top_p for several thresholds. It fails before this change and passes after; esm/utils/sampling_test.py passes and ruff check / ruff format are clean on the changed lines.