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 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]], top_k_val: Int32, d: Int32)