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 module
toppminp_gpu
Provides GPU implementations of top-p (nucleus) and min-p sampling for autoregressive token generation.
comptime values
DEBUG_FILE
comptime DEBUG_FILE = False
SEED
comptime SEED = 42
Structs
-
DoubleBuffer: Holds two GPU buffers and alternates between them for double-buffered radix sort passes.
Functions
-
min_p_sampling_gpu: GPU implementation of Min-P sampling for token selection. This function applies temperature scaling, softmax, a radix sort, and then samples tokens based on the calculated probability threshold (Min-P). -
normalize: Normalizes a bfloat16 value to an unsigned 16-bit integer for radix sort by flipping the sign bit for positive values and fully inverting negative values. -
normalize_u32: Returns a uint32 value unchanged since unsigned integers already sort correctly in radix sort. -
radix_sort_pairs_kernel: Radix pair sort kernel for (default) descending order. -
run_radix_sort_pairs_gpu: Runs a multi-pass radix sort on key/index pairs across batches on the GPU using double buffering. -
top_p_sampling_gpu: GPU implementation of Top-P sampling for token selection. This function applies temperature scaling, softmax, a radix sort, and then samples tokens based on the cumulative probability mass (Top-P). -
topk_wrapper: Copy ofKernels/mojo/nn/topk.mojo:_topk_stage1with the addition of max_vals and p_threshold arguments to determine if sorting is needed for top-p/min-p sampling. -
topp_minp_sampling_kernel: Top P-Min P sampling kernel.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!