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 package
sampling
Token-sampling kernels: penalties, top-k/top-p masking and sampling.
topk_fi holds the single-block kernels, which serve every target;
topk_fi_cluster holds thread-block-cluster variants of the hot spec-decode
kernels for NVIDIA SM90+ devices. The dispatchers in this file pick between
them by device, so callers import from this package and never choose a
launch shape themselves.
Modules
-
coop_row: Provides cross-block reductions for row-parallel sampling. -
sampling: Provides token-sampling utilities including logit-penalty application and repetition-penalty kernels. -
topk_fi: Implements a float-integer packed top-K kernel that jointly tracks values and indices in a single register. -
topk_fi_cluster: Cluster-launched top-k/top-p kernels.
Functions
-
topk_topp_masked_probs: Computes per-row top-k/top-p masked softmax. -
topk_topp_sampling_from_prob: Joint top-k + top-p sampling from probability distribution.