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

topk_mask_logits

def topk_mask_logits[dtype: DType, out_idx_type: DType, block_size: Int = Int(1024), TopKArrLayoutType: TensorLayout = Layout[*?, *?]](ctx: DeviceContext, logits: TileTensor[dtype, Storage=logits.Storage, address_space=logits.address_space, linear_idx_type=logits.linear_idx_type], masked_logits: TileTensor[dtype, Storage=masked_logits.Storage, address_space=masked_logits.address_space, linear_idx_type=masked_logits.linear_idx_type], top_k_val: Int, top_k_arr: Optional[TileTensor[out_idx_type, TopKArrLayoutType, MutUntrackedOrigin]] = None)

Masks logits to keep only the top-k largest values per row.

Launches TopKMaskLogitsKernel with one block per batch row. Elements below the k-th largest logit are set to the dtype's minimum value so downstream sampling ignores them.

Parameters:

dtype (DType): Element type of the logits and masked_logits tensors.
out_idx_type (DType): Index type used for per-row top-k override values in top_k_arr.
block_size (Int): Number of threads per block (defaults to 1024).
TopKArrLayoutType (TensorLayout): Memory layout of the optional top_k_arr tensor.

Args:

ctx (DeviceContext): Device context for kernel execution.
logits (TileTensor[dtype, Storage=logits.Storage, address_space=logits.address_space, linear_idx_type=logits.linear_idx_type]): Input logits tensor [batch_size, d].
masked_logits (TileTensor[dtype, Storage=masked_logits.Storage, address_space=masked_logits.address_space, linear_idx_type=masked_logits.linear_idx_type]): Output buffer for masked logits, same shape as logits.
top_k_val (Int): Default number of largest logits to retain per row.
top_k_arr (Optional[TileTensor[out_idx_type, TopKArrLayoutType, MutUntrackedOrigin]]): Optional per-row top-k values that override top_k_val.

Raises:

Error: If masked_logits shape does not match logits shape.