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 function

TopKMaskLogitsKernel

def TopKMaskLogitsKernel[block_size: Int, vec_size: Int, dtype: DType, out_idx_type: DType, LogitsLayoutType: TensorLayout, logits_origin: ImmOrigin, MaskedLogitsLayoutType: TensorLayout, masked_logits_origin: MutOrigin](logits: TileTensor[dtype, LogitsLayoutType, logits_origin], masked_logits: TileTensor[dtype, MaskedLogitsLayoutType, masked_logits_origin], top_k_arr: Optional[Pointer[Scalar[out_idx_type], MutUntrackedOrigin, _safe=False]], top_k_val: Int, d: Int)

Masks logits to retain only the top-k values per row, setting the rest to negative infinity.

Each block processes one row. A ternary search over the logit range finds the k-th largest value (the pivot), then every element greater than the pivot is kept while the rest are replaced with the dtype's minimum value.

Parameters:

block_size (Int): Number of threads per block.
vec_size (Int): Number of elements each thread loads per vectorized access.
dtype (DType): Element type of the logits and masked_logits tiles.
out_idx_type (DType): Index type used for per-row top-k override values.
LogitsLayoutType (TensorLayout): Memory layout of the input logits tile.
logits_origin (ImmOrigin): Origin tag for the immutable input logits tile.
MaskedLogitsLayoutType (TensorLayout): Memory layout of the output masked_logits tile.
masked_logits_origin (MutOrigin): Origin tag for the mutable output masked_logits tile.

Args:

logits (TileTensor[dtype, LogitsLayoutType, logits_origin]): Input logits tile [batch_size, d].
masked_logits (TileTensor[dtype, MaskedLogitsLayoutType, masked_logits_origin]): Output buffer for masked logits, same shape as logits.
top_k_arr (Optional[Pointer[Scalar[out_idx_type], MutUntrackedOrigin, _safe=False]]): Optional per-row top-k values; overrides top_k_val when present.
top_k_val (Int): Default number of largest logits to retain per row.
d (Int): Row length (vocabulary size).