Sampling: temperature, top-k, top-p

With logits z_i and temperature T:

p_i(T) = exp(z_i / T) / Σ_j exp(z_j / T)

The ratio of two probabilities is exp((z_a − z_b) / T). For "mat" and "floor" the logit gap is 0.5. The ratio is exp(0.5) = 1.65 at T = 1, and exp(1.0) = 2.72 at T = 0.5. A lower temperature makes every gap larger.

Temperature 0 in DeepInfer. Division by 0 is not possible. SamplingParams replaces any temperature below SAMPLING_EPS = 1e-5 with 1e-5, and also sets top_p = 1.0 and top_k = -1 (deepinfer/sampling_params.py:10, deepinfer/sampling_params.py:76-79). Then a logit gap of only 0.01 becomes 0.01 / 1e-5 = 1,000, and exp(−1,000) is 0 in float32. So the best token gets probability 1 in practice. Exactly equal logits still share the probability.

A negative temperature raises ValueError (deepinfer/sampling_params.py:94-95). The reference multiplies the logits by 1/T (deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:209).

The reference implementation of DeepInfer's kernel does this for each row (deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:121-150):

1. Compute p = softmax(logits / T) over the full vocabulary. The division by T happens just before this function (deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:209).
2. Sort the tokens by p, largest first. Ties keep token order (a stable sort).
3. Top-p: for each token, compute the mass before it: the sum of the probabilities of all better tokens. Remove the token if that mass is greater than top_p (line 139). So the first token always stays.
4. Top-k: rank the tokens that are left. Remove every token with rank ≥ top_k, unless top_k is −1 (line 149).
5. Set the logits of removed tokens to −∞ and take log-softmax. The kept probabilities now add up to 1.

Worked example (panel 2 of the figure, T = 1, top_p = 0.8, top_k = 2):

Top-p alone keeps 3 tokens with mass 0.859, the smallest group that reaches 0.8. With top-k 2, the final group is {mat, floor}.

The order matters. DeepInfer measures top-p on the full distribution, then cuts to k. Some other engines cut to k first and measure top-p on the renormalized top-k probabilities. Example: p = (0.5, 0.3, 0.2), top_k = 2, top_p = 0.55.

• DeepInfer order: mass before the second token is 0.5 ≤ 0.55, so it stays. Result: 2 tokens.
• Top-k first: renormalized (0.625, 0.375). Mass before the second token is 0.625 > 0.55, so it goes. Result: 1 token.

The kernel tests check the DeepInfer rule (deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:366-377). Without its smallest token, the kept prefix has a full-softmax mass of at most top_p. The prefix has at most k tokens.

Limits. top_p must be in (0, 1]; top_k must be −1 or between 1 and the int32 maximum, so top_k = 0 raises ValueError (deepinfer/sampling_params.py:96-99). A top_k larger than the vocabulary keeps all tokens (deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:149).

The draw in DeepInfer does not use a random generator with state. It makes the random number from a hash, so the result depends only on the inputs. The docstring says that the result is "deterministic in (seq_seeds, sample_positions, is_nextn_predict)", and independent of batch size and row order (deepinfer/model_executor/layers/tilelang_ops/random_sampling.py:17-20). The PyTorch reference shows the steps:

def random_sampling_ref(
    probs: torch.Tensor,
    seq_seeds: torch.Tensor,
    sample_positions: torch.Tensor,
    is_nextn_predict: bool = False,
) -> torch.Tensor:
    """Pure-PyTorch sampler: hash (seq_seeds, 3 * sample_positions + case) to a
    uniform and invert the fp64 CDF. For Ascend, this will be replaced later."""
    sample_case = 1 if is_nextn_predict else 0
    seq = (sample_positions.to(torch.int64) * 3 + sample_case) * i64(0x9E3779B97F4A7C15)
    combined = seq_seeds.to(torch.int64) + (seq ^ (seq * (1 << 13)))
    u = _hash_to_uniform(combined).unsqueeze(-1)
    cdf = probs.double().cumsum(dim=-1)
    target = u * cdf[:, -1:]
    return (cdf > target).long().argmax(dim=-1).int()

1. Line 79: mix the token position with a case number: position · 3 + case, times a large odd constant. The case is 0 for a normal sample. Case 1 is the new sample after the main model rejects a draft token, and case 2 is the accept test of a draft (deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:239, deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:245).
2. Line 80: mix the value once more with a shift and an XOR, and add the request seed seq_seed.
3. Line 81: the SplitMix64 hash turns the 64-bit value into a uniform number u in [0, 1) (deepinfer/model_executor/layers/tilelang_ops/random_sampling.py:58-67).
4. Lines 82-84: build the running sum (CDF) of the kept probabilities in float64. Return the first token whose running sum is greater than u × total.

Worked example (panel 3): kept probabilities (0.622, 0.378), CDF (0.622, 1.000). With u = 0.70: 0.622 is not greater than 0.70; 1.000 is. The sampler returns the second token, "floor". The value 0.70 is illustrative; the real u comes from the hash.

Which position? In a prefill step, the position is seqlen_k, the number of tokens in the sequence so far. That is the position of the token that the step makes (csrc/engine/prepare_sample.cpp:198). In a decode step, it is seqlen_k − seqlen_q + 1: the position of the token that the first new row predicts (csrc/engine/prepare_sample.cpp:206). The draft module samples with the tag "mtp", which adds the seed offset 1000000007, so its random numbers differ from those of the main sample (deepinfer/model_executor/models/mtp/dflash.py:329, deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:15, deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:93).

The default seed is random. If you do not set seq_seed, SamplingParams takes 64 random bits from Python's random module (deepinfer/sampling_params.py:85-91). So two runs differ. For repeatable output, the README asks for a fixed seq_seed together with DEEPINFER_BATCH_INVARIANT=1 and DEEPINFER_ENABLE_MTP_VARLEN=0, on NVIDIA only (README.md:28). The same seed alone is not enough: the logits themselves can differ slightly between batches. See Determinism and batch invariance.

Defaults are in the constructor signature (deepinfer/sampling_params.py:47-62). Note the order in the constructor: the checks run before the temperature clamp, so temperature=0, top_k=0 still raises ValueError for top_k (deepinfer/sampling_params.py:75-79).

Engine-wide settings that change sampling:

• DEEPINFER_SAMPLER_FORCE_TOPK (default 0, off; deepinfer/environ.py:150): keep only the N largest logits of each row before all other steps (deepinfer/model_executor/layers/sampler_utils.py:166-167). This is faster, but tokens outside the N best can never be drawn.
• DEEPINFER_DECODE_TOKEN_BLACKLIST (default empty; deepinfer/environ.py:151): token ids whose logits the sampler sets to −∞ (deepinfer/model_executor/layers/sampler.py:168-169).

The fused kernel returns "processed target logprobs" (deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:51-52). These are the log-probabilities after temperature, top-p and top-k. So the logprob of an output token is the log of its final probability, not of the raw model probability.

In the worked example, "floor" was drawn. Its returned logprob is ln 0.378 = −0.973. The raw model probability was 0.264 (ln = −1.332). A removed token, for example "sofa", has logprob −∞.

top_logprobs = n selects the n best entries of the same processed row (deepinfer/model_executor/layers/sampler.py:636-648). With top_k = 2 and top_logprobs = 4, only two entries are finite.