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

topk_topp_masked_probs

def topk_topp_masked_probs[dtype: DType, block_size: Int = Int(1024), TopKArrLayoutType: TensorLayout = Layout[*(), *()], TopPArrLayoutType: TensorLayout = Layout[*(), *()], TemperatureLayoutType: TensorLayout = Layout[*(), *()], ProbsLayoutType: TensorLayout = Layout[*(), *()]](ctx: DeviceContext, logits: TileTensor[dtype, Storage=logits.Storage, address_space=logits.address_space, linear_idx_type=logits.linear_idx_type], probs: TileTensor[DType.float32, ProbsLayoutType, MutAnyOrigin], top_k_val: Int, top_p_val: Float32 = 1, top_k_arr: Optional[TileTensor[DType.int64, TopKArrLayoutType, ImmutAnyOrigin]] = None, top_p_arr: Optional[TileTensor[DType.float32, TopPArrLayoutType, ImmutAnyOrigin]] = None, temperature: Optional[TileTensor[DType.float32, TemperatureLayoutType, ImmutAnyOrigin]] = None)

Computes per-row top-k/top-p masked softmax, one block per row.

See TopKTopPMaskedProbsKernel for what the output means.

Parameters:

• ​dtype (DType): Element type of logits.
• ​block_size (Int): Threads per block.
• ​TopKArrLayoutType (TensorLayout): Memory layout of top_k_arr.
• ​TopPArrLayoutType (TensorLayout): Memory layout of top_p_arr.
• ​TemperatureLayoutType (TensorLayout): Memory layout of temperature.
• ​ProbsLayoutType (TensorLayout): Memory layout of probs.

Args:

• ​ctx (DeviceContext): Device context.
• ​logits (TileTensor[dtype, Storage=logits.Storage, address_space=logits.address_space, linear_idx_type=logits.linear_idx_type]): Input logits [batch_size, d].
• ​probs (TileTensor[DType.float32, ProbsLayoutType, MutAnyOrigin]): Output masked renormalized distribution [batch_size, d].
• ​top_k_val (Int): Default top-k; <= 0 or > d keeps every token.
• ​top_p_val (Float32): Default top-p threshold.
• ​top_k_arr (Optional[TileTensor[DType.int64, TopKArrLayoutType, ImmutAnyOrigin]]): Optional per-row top-k [batch_size].
• ​top_p_arr (Optional[TileTensor[DType.float32, TopPArrLayoutType, ImmutAnyOrigin]]): Optional per-row top-p [batch_size].
• ​temperature (Optional[TileTensor[DType.float32, TemperatureLayoutType, ImmutAnyOrigin]]): Optional per-row temperature [batch_size]; 0 is clamped.

Raises:

Error: If the tensor shapes disagree.