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_cluster

def topk_topp_masked_probs_cluster[dtype: DType, block_size: Int = Int(1024), TopKArrLayoutType: TensorLayout = Layout[TypeList[Int64](), TypeList[ComptimeInt[Int(1)]]()], TopPArrLayoutType: TensorLayout = Layout[TypeList[Int64](), TypeList[ComptimeInt[Int(1)]]()], TemperatureLayoutType: TensorLayout = Layout[TypeList[Int64](), TypeList[ComptimeInt[Int(1)]]()], ProbsLayoutType: TensorLayout = Layout[TypeList[Int64, Int64](), TypeList[Int64, ComptimeInt[Int(1)]]()], TopKArrEngine: TensorEngine = DefaultEngine, TopPArrEngine: TensorEngine = DefaultEngine, TemperatureEngine: TensorEngine = DefaultEngine](ctx: DeviceContext, logits: TileTensor[dtype, Engine=logits.Engine, address_space=logits.address_space, linear_idx_type=logits.linear_idx_type], probs: TileTensor[.float32, ProbsLayoutType, MutAnyOrigin], top_k_val: Int, top_p_val: Float32 = 1, top_k_arr: Optional[TileTensor[.int64, TopKArrLayoutType, ImmutAnyOrigin, Engine=TopKArrEngine]] = None, top_p_arr: Optional[TileTensor[.float32, TopPArrLayoutType, ImmutAnyOrigin, Engine=TopPArrEngine]] = None, temperature: Optional[TileTensor[.float32, TemperatureLayoutType, ImmutAnyOrigin, Engine=TemperatureEngine]] = None)

Computes per-row top-k/top-p masked softmax on a cluster device.

See TopKTopPMaskedProbsKernel for what the output means. The package's topk_topp_masked_probs dispatcher routes here on NVIDIA SM90+ devices; a batch that fills the machine or a slice too wide for shared memory still falls back to the single-block kernel at runtime.

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.
• ​TopKArrEngine (TensorEngine): Engine policy of top_k_arr.
• ​TopPArrEngine (TensorEngine): Engine policy of top_p_arr.
• ​TemperatureEngine (TensorEngine): Engine policy of temperature.

Args:

• ​ctx (DeviceContext): Device context.
• ​logits (TileTensor[dtype, Engine=logits.Engine, address_space=logits.address_space, linear_idx_type=logits.linear_idx_type]): Input logits [batch_size, d].
• ​probs (TileTensor[.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[.int64, TopKArrLayoutType, ImmutAnyOrigin, Engine=TopKArrEngine]]): Optional per-row top-k [batch_size].
• ​top_p_arr (Optional[TileTensor[.float32, TopPArrLayoutType, ImmutAnyOrigin, Engine=TopPArrEngine]]): Optional per-row top-p [batch_size].
• ​temperature (Optional[TileTensor[.float32, TemperatureLayoutType, ImmutAnyOrigin, Engine=TemperatureEngine]]): Optional per-row temperature [batch_size]; 0 is clamped.

Raises:

Error: If the tensor shapes disagree.