What does “normal random sampling” mean in LLM text generation?

Uncategorized
1

I am working on research involving the numerical behavior of LLM token sampling, and I need to reproduce the ordinary stochastic next-token sampling process as faithfully as possible.

At each decoding step, I obtain the model logits and convert them to probabilities:

logits = model(...).logits[:, -1, :]
probs = torch.softmax(logits, dim=-1)

I can then sample the next token in at least two ways.

Using PyTorch:

next_token = torch.multinomial(probs, num_samples=1)

or by explicitly implementing inverse-transform sampling:

cdf = torch.cumsum(probs, dim=-1)
u = torch.rand(1, device=probs.device)
next_token = torch.searchsorted(cdf, u)

Mathematically, both are intended to sample from the same categorical distribution:

$$
X \sim \mathrm{Categorical}(p).
$$

However, my question is specifically about the concrete finite-precision implementation, not only the mathematical distribution.

With a large LLM vocabulary and float32 probabilities, I have observed cases where the explicitly accumulated CDF contains adjacent equal values:

cdf[i] == cdf[i - 1]

even though:

probs[i] > 0

In that case, the interval width for token i becomes zero in this particular inverse-CDF implementation.

I have also generated tokens using:

torch.multinomial(probs, 1)

and observed behavior that is not always reproduced by my explicit FP32 CDF sampler.

I checked the PyTorch documentation, but it appears to describe the intended sampling distribution rather than guarantee that torch.multinomial(probs, 1) is implemented as the same serial cumsum + uniform random number + searchsorted procedure shown above.

So my questions are:

  1. For torch.multinomial(probs, 1), what sampling algorithm does PyTorch actually use?

    • Does it use inverse-transform/CDF sampling?
    • Can the implementation differ between CPU and CUDA?
    • Can it differ depending on dtype or num_samples?

  2. Is this implementation:

cdf = torch.cumsum(probs.float(), dim=-1)
u = torch.rand(1, device=probs.device)
sample = torch.searchsorted(cdf, u)

expected to have the same finite-precision sampling behavior as:

torch.multinomial(probs, 1)

or are they only equivalent in exact arithmetic?

  1. In PyTorch-based LLM generation, when someone refers to “random sampling” or “sampling from the model distribution,” is torch.multinomial(probs, 1) the usual concrete implementation, or is the phrase intended only at the abstract distribution level?

  2. If I want to reproduce the ordinary stochastic sampling behavior of a PyTorch LLM implementation, should I use the framework’s native sampler (torch.multinomial) rather than assume that “random sampling” means a particular inverse-CDF implementation?