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

Mojo module

topk_fi

Implements a float-integer packed top-K kernel that jointly tracks values and indices in a single register.

Structs

ValueCount: A struct that holds a value and a count, used for block reductions.

apply_min_p_mask_kernel: Zero out probabilities below the per-row min_p threshold.
device_sampling_from_prob: Device-level sampling from probability distribution with atomic operations.
get_min_max_value: Compute the minimum and maximum values from input data using block reduction.
topk_mask_logits: Masks logits to keep only the top-k largest values per row.
topk_sampling_from_prob: Top-K sampling from probability distribution.
topk_softmax_sample: Samples token indices from top-K logits using softmax probabilities.
topk_softmax_sample_kernel:
topk_topp_sampling_from_prob: Joint top-k + top-p sampling from probability distribution.
TopKMaskLogitsKernel:
TopKSamplingFromProbKernel: Kernel for top-k sampling from probability distribution.
TopKTopPSamplingFromProbKernel: Kernel for joint top-k + top-p sampling from probability distribution.