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
topk
Provides top-K selection kernels using warp- and block-level reductions for CPU and GPU.
Structs
-
TopK_2: Tracks the single best (value, index) pair per thread during top-K reductions. -
TopKHeap: Fixed-capacity register heap for per-thread top-M tracking.
Functions
-
apply_gumbel_noise_kernel: Adds Gumbel(0,1) noise to logits for sampling via the Gumbel-max trick. -
fused_token_sampling_cpu: Generalized implementation of the Top K algorithm with sampling. Returns the sampled index from the innermost dimension of the input tensor for each row/subvolume. -
fused_token_sampling_gpu: Top K algorithm with fused sampling. Returns the sampled indices from the Top-K of the innermost dimension of the input tensor for each row/subvolume. -
gumbel_sampling_fused_gpu: Fused Gumbel sampling: applies Gumbel(0,1) noise and selects the argmax in a single GPU kernel launch (no intermediate noised-logits HBM buffer). -
gumbel_sampling_gpu: Gumbel sampling using the Gumbel-max trick for categorical distributions. -
top_k: Implementation of the Top K algorithm. Returns the top or bottom K elements and their index along a specified axis. -
top_k_shape_impl: Compute the output shape of a top/bottom k operation. -
topk_gpu: Generalized implementation of the Top K algorithm with/without sampling. Returns the sampled index from the innermost dimension of the input tensor for each row/subvolume or the top K values and indices across the tensor.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!