Mojo function

TopKSamplingFromProbKernel

TopKSamplingFromProbKernel[ProbsLayoutType: TensorLayout, probs_origin: ImmutOrigin, 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[UnsafePointer[Scalar[out_idx_type], MutExternalOrigin]], top_k_arr: Optional[UnsafePointer[Scalar[out_idx_type], MutExternalOrigin]], 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.

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[UnsafePointer[Scalar[out_idx_type], MutExternalOrigin]]): Optional row indices for batch indexing [batch_size].
top_k_arr (Optional[UnsafePointer[Scalar[out_idx_type], MutExternalOrigin]]): 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.