For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

TopKTopPSamplingFromProbKernel

def TopKTopPSamplingFromProbKernel[ProbsLayoutType: TensorLayout, probs_origin: ImmOrigin, OutputLayoutType: TensorLayout, output_origin: MutOrigin, block_size: Int, vec_size: Int, dtype: DType, out_idx_type: DType, deterministic: Bool, from_logits: Bool = False](probs: TileTensor[dtype, ProbsLayoutType, probs_origin], output: TileTensor[out_idx_type, OutputLayoutType, output_origin], indices: Optional[Pointer[Scalar[out_idx_type], ImmutAnyOrigin, _safe=False]], top_k_arr: Optional[Pointer[Scalar[out_idx_type], ImmutAnyOrigin, _safe=False]], top_k_val: Int, top_p_arr: Optional[Pointer[Float32, ImmutAnyOrigin, _safe=False]], top_p_val: Float32, d: Int, rng_seed: Optional[Pointer[UInt64, ImmutAnyOrigin, _safe=False]], rng_offset: UInt64, temperature: Optional[Pointer[Float32, ImmutAnyOrigin, _safe=False]], min_p: Optional[Pointer[Float32, ImmutAnyOrigin, _safe=False]])

Kernel for joint top-k + top-p sampling from probability distribution.

Identical to TopKSamplingFromProbKernel but additionally enforces a nucleus (top-p) constraint: a token is accepted only when both the count of tokens above the pivot is less than k AND the cumulative probability of those tokens is less than p.

When top_p_val = 1.0 and top_p_arr is null, this degrades to top-k-only with zero overhead since sum < 1.0 is always true.

When from_logits is True, probs contains raw logits and softmax with per-row temperature scaling is fused into the kernel: every load is transformed to exp((logit - row_max) / temp), the unnormalized softmax value with the row maximum shifted to exactly 1.0. The pivot search over [0, 1] is unchanged; the total unnormalized mass z replaces the normalized distribution's implicit total of 1.0 in the initial CDF budget and scales the top-p threshold. The optional min-p mask is applied inline (in this domain the max "probability" is 1.0, so the mask threshold is simply min_p), matching apply_min_p_mask_kernel semantics in the normalized domain.

Parameters:

ProbsLayoutType (TensorLayout): Memory layout of the input probs tile.
probs_origin (ImmOrigin): Origin tag for the immutable input probs tile.
OutputLayoutType (TensorLayout): Memory layout of the output output tile.
output_origin (MutOrigin): Origin tag for the mutable output output tile.
block_size (Int): Number of threads per block.
vec_size (Int): Number of elements each thread loads per vectorized access.
dtype (DType): Element type of the probs tensor.
out_idx_type (DType): Index type used for the sampled output indices.
deterministic (Bool): If True, use deterministic sampling.
from_logits (Bool): If True, probs holds raw logits and softmax with per-row temperature scaling and min-p masking is fused into the kernel (defaults to False).

Args:

probs (TileTensor[dtype, ProbsLayoutType, probs_origin]): Input probability distribution [batch_size, d], or raw logits when from_logits is True.
output (TileTensor[out_idx_type, OutputLayoutType, output_origin]): Output sampled indices [batch_size].
indices (Optional[Pointer[Scalar[out_idx_type], ImmutAnyOrigin, _safe=False]]): Optional row indices for batch indexing [batch_size].
top_k_arr (Optional[Pointer[Scalar[out_idx_type], ImmutAnyOrigin, _safe=False]]): Optional per-row top_k values [batch_size].
top_k_val (Int): Default top_k value if top_k_arr is null.
top_p_arr (Optional[Pointer[Float32, ImmutAnyOrigin, _safe=False]]): Optional per-row top_p values [batch_size].
top_p_val (Float32): Default top_p value if top_p_arr is null.
d (Int): Vocabulary size.
rng_seed (Optional[Pointer[UInt64, ImmutAnyOrigin, _safe=False]]): Optional per-row seed array [batch_size], indexed by row_idx. If null, defaults to 0.
rng_offset (UInt64): Random offset for Random number generator.
temperature (Optional[Pointer[Float32, ImmutAnyOrigin, _safe=False]]): Optional per-row temperature [batch_size]. Only used when from_logits is True; defaults to 1.0 per row.
min_p (Optional[Pointer[Float32, ImmutAnyOrigin, _safe=False]]): Optional per-row min-p thresholds [batch_size]. Only used when from_logits is True.