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

TopKSamplingFromProbKernel

def TopKSamplingFromProbKernel[ProbsLayoutType: TensorLayout, probs_origin: ImmOrigin, OutputLayoutType: TensorLayout, output_origin: MutOrigin, block_size: Int, vec_size: Int, dtype: DType, out_idx_type: DType, deterministic: Bool](probs: TileTensor[dtype, ProbsLayoutType, probs_origin], output: TileTensor[out_idx_type, OutputLayoutType, output_origin], indices: Optional[Pointer[Scalar[out_idx_type], MutUntrackedOrigin, _safe=False]], top_k_arr: Optional[Pointer[Scalar[out_idx_type], MutUntrackedOrigin, _safe=False]], top_k_val: Int, d: Int, rng_seed: UInt64, rng_offset: UInt64)

Kernel for top-k sampling from probability distribution.

This kernel performs top-k sampling by:

Using ternary search to find a pivot threshold.
Rejecting samples iteratively until acceptance criteria is met.
Sampling an index using uniform random numbers from Random generator.

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.

Args:

probs (TileTensor[dtype, ProbsLayoutType, probs_origin]): Input probability distribution [batch_size, d].
output (TileTensor[out_idx_type, OutputLayoutType, output_origin]): Output sampled indices [batch_size].
indices (Optional[Pointer[Scalar[out_idx_type], MutUntrackedOrigin, _safe=False]]): Optional row indices for batch indexing [batch_size].
top_k_arr (Optional[Pointer[Scalar[out_idx_type], MutUntrackedOrigin, _safe=False]]): Optional per-row top_k values [batch_size].
top_k_val (Int): Default top_k value if top_k_arr is null.
d (Int): Vocabulary size.
rng_seed (UInt64): Random seed for Random number generator.
rng_offset (UInt64): Random offset for Random number generator.