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_gpu

def topk_gpu[dtype: DType, out_idx_type: DType, //, sampling: Bool = True, largest: Bool = True, KLayoutType: TensorLayout = Layout[*(), *()], TemperatureLayoutType: TensorLayout = Layout[*(), *()], TopPLayoutType: TensorLayout = Layout[*(), *()], MinPLayoutType: TensorLayout = Layout[*(), *()], SeedLayoutType: TensorLayout = Layout[*(), *()], KStorageType: TensorStorage = PointerStorage, TemperatureStorageType: TensorStorage = PointerStorage, TopPStorageType: TensorStorage = PointerStorage, MinPStorageType: TensorStorage = PointerStorage, SeedStorageType: TensorStorage = PointerStorage](ctx: DeviceContext, max_k: Int, input: TileTensor[dtype, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type], out_vals: TileTensor[dtype, Storage=out_vals.Storage, address_space=out_vals.address_space, linear_idx_type=out_vals.linear_idx_type], out_idxs: TileTensor[out_idx_type, Storage=out_idxs.Storage, address_space=out_idxs.address_space, linear_idx_type=out_idxs.linear_idx_type], block_size: Optional[Int] = None, num_blocks_per_input: Optional[Int] = None, k: Optional[TileTensor[DType.int64, KLayoutType, ImmutAnyOrigin, Storage=KStorageType]] = None, temperature: Optional[TileTensor[DType.float32, TemperatureLayoutType, ImmutAnyOrigin, Storage=TemperatureStorageType]] = None, top_p: Optional[TileTensor[DType.float32, TopPLayoutType, ImmutAnyOrigin, Storage=TopPStorageType]] = None, min_p: Optional[TileTensor[DType.float32, MinPLayoutType, ImmutAnyOrigin, Storage=MinPStorageType]] = None, seed: Optional[TileTensor[DType.uint64, SeedLayoutType, ImmutAnyOrigin, Storage=SeedStorageType]] = None, valid: Optional[Pointer[Int8, MutAnyOrigin]] = None)

Generalized implementation of the Top K algorithm with/without sampling. Returns the sampled index from the innermost dimension of the input tensor for each row/subvolume or the top K values and indices across the tensor.

Parameters:

• ​dtype (DType): DType - The data dtype of the input tensor.
• ​out_idx_type (DType): DType - The data dtype of the output indices (default == DType.int).
• ​sampling (Bool): Bool - Whether to return token samples from topK dist (default is True).
• ​largest (Bool): Bool - Whether to find the maximum or minimum value.
• ​KLayoutType (TensorLayout): Layout type of the k buffer.
• ​TemperatureLayoutType (TensorLayout): Layout type of the temperature buffer.
• ​TopPLayoutType (TensorLayout): Layout type of the top_p buffer.
• ​MinPLayoutType (TensorLayout): Layout type of the min_p buffer.
• ​SeedLayoutType (TensorLayout): Layout type of the seed buffer.
• ​KStorageType (TensorStorage): Storage type of the k buffer.
• ​TemperatureStorageType (TensorStorage): Storage type of the temperature buffer.
• ​TopPStorageType (TensorStorage): Storage type of the top_p buffer.
• ​MinPStorageType (TensorStorage): Storage type of the min_p buffer.
• ​SeedStorageType (TensorStorage): Storage type of the seed buffer.

Args:

• ​ctx (DeviceContext): DeviceContext The context for GPU execution.
• ​max_k (Int): Int Largest number of top elements to keep for each batch element.
• ​input (TileTensor[dtype, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type]): TileTensor[dtype] Input tensor as a device TileTensor.
• ​out_vals (TileTensor[dtype, Storage=out_vals.Storage, address_space=out_vals.address_space, linear_idx_type=out_vals.linear_idx_type]): TileTensor[dtype] Output buffer on device for the K largest values.
• ​out_idxs (TileTensor[out_idx_type, Storage=out_idxs.Storage, address_space=out_idxs.address_space, linear_idx_type=out_idxs.linear_idx_type]): TileTensor[DType.int] Output buffer on device for the indices of the K largest values, or sampled token indices. Last dimension is 1 if sampling is True, otherwise K.
• ​block_size (Optional[Int]): Int The number of threads per block (default is 256 from TRT and empirical testing).
• ​num_blocks_per_input (Optional[Int]): Optional[Int] Number of blocks per input (default computed from input size and block size). This is the equivalent of "BLOCKS_PER_BEAM" in TRT-LLM kernel allowing for much larger batch sizes through packing several elements per thread in the first stage.
• ​k (Optional[TileTensor[DType.int64, KLayoutType, ImmutAnyOrigin, Storage=KStorageType]]): Optional TileTensor[DType.int64] Device buffer of top elements to keep for each batch element.
• ​temperature (Optional[TileTensor[DType.float32, TemperatureLayoutType, ImmutAnyOrigin, Storage=TemperatureStorageType]]): The temperature based scaling.
• ​top_p (Optional[TileTensor[DType.float32, TopPLayoutType, ImmutAnyOrigin, Storage=TopPStorageType]]): Only use the tokens whose cumulative probability exceeds this threshold.
• ​min_p (Optional[TileTensor[DType.float32, MinPLayoutType, ImmutAnyOrigin, Storage=MinPStorageType]]): Per-row min-p threshold. Tokens with probability below min_p * max_prob are excluded from sampling.
• ​seed (Optional[TileTensor[DType.uint64, SeedLayoutType, ImmutAnyOrigin, Storage=SeedStorageType]]): The seed to use for the random number generator.
• ​valid (Optional[Pointer[Int8, MutAnyOrigin]]): Optional per-row validity flags (1 = valid, 0 = invalid). The sampling kernel writes 0 for rows where no finite logit was found (e.g. all-NaN rows); rows default to 1 via memset.