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_sampling_from_prob

def topk_topp_sampling_from_prob[dtype: DType, out_idx_type: DType, block_size: Int = Int(1024), from_logits: Bool = False, emit_dist: Bool = False, dist_dtype: DType = DType.float32, DistLayoutType: TensorLayout = Layout[*(), *()], TopKArrLayoutType: TensorLayout = Layout[*(), *()], IndicesLayoutType: TensorLayout = Layout[*(), *()], TopPArrLayoutType: TensorLayout = Layout[*(), *()], SeedLayoutType: TensorLayout = Layout[*(), *()], TemperatureLayoutType: TensorLayout = Layout[*(), *()], MinPLayoutType: TensorLayout = Layout[*(), *()], TopKArrStorageType: TensorStorage = PointerStorage, IndicesStorageType: TensorStorage = PointerStorage, TopPArrStorageType: TensorStorage = PointerStorage, SeedStorageType: TensorStorage = PointerStorage, TemperatureStorageType: TensorStorage = PointerStorage, MinPStorageType: TensorStorage = PointerStorage](ctx: DeviceContext, probs: TileTensor[dtype, Storage=probs.Storage, address_space=probs.address_space, linear_idx_type=probs.linear_idx_type], output: TileTensor[out_idx_type, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type], top_k_val: Int, top_p_val: Float32 = 1, deterministic: Bool = False, rng_seed: Optional[TileTensor[DType.uint64, SeedLayoutType, ImmutAnyOrigin, Storage=SeedStorageType]] = None, rng_offset: UInt64 = UInt64(0), indices: Optional[TileTensor[out_idx_type, IndicesLayoutType, ImmutAnyOrigin, Storage=IndicesStorageType]] = None, top_k_arr: Optional[TileTensor[out_idx_type, TopKArrLayoutType, ImmutAnyOrigin, Storage=TopKArrStorageType]] = None, top_p_arr: Optional[TileTensor[DType.float32, TopPArrLayoutType, ImmutAnyOrigin, Storage=TopPArrStorageType]] = None, temperature: Optional[TileTensor[DType.float32, TemperatureLayoutType, ImmutAnyOrigin, Storage=TemperatureStorageType]] = None, min_p: Optional[TileTensor[DType.float32, MinPLayoutType, ImmutAnyOrigin, Storage=MinPStorageType]] = None, out_dist: Optional[TileTensor[dist_dtype, DistLayoutType, MutAnyOrigin]] = None)

Joint top-k + top-p sampling from probability distribution.

Performs stochastic sampling considering only tokens that satisfy both the top-k count constraint AND the top-p nucleus constraint. When top_p_val is 1.0 (default) this behaves identically to topk_sampling_from_prob.

When emit_dist is set, the masked renormalized distribution is written to out_dist as well; see the kernel docstring.

When from_logits is True, probs contains raw logits: softmax with per-row temperature scaling and the optional min-p mask are fused into the sampling kernel, avoiding the [batch_size, d] probability round-trip through global memory and the separate softmax / mask kernel launches.

Parameters:

• ​dtype (DType): Element type of the probs tensor.
• ​out_idx_type (DType): Index type used for the sampled output indices.
• ​block_size (Int): Number of threads per block (defaults to 1024).
• ​from_logits (Bool): If True, probs holds raw logits and softmax with per-row temperature scaling and min-p masking is fused into the kernel (defaults to False).
• ​emit_dist (Bool): If True, also write the masked renormalized distribution to out_dist (defaults to False).
• ​dist_dtype (DType): Element type of out_dist.
• ​DistLayoutType (TensorLayout): Memory layout of the optional out_dist tensor.
• ​TopKArrLayoutType (TensorLayout): Memory layout of the optional top_k_arr tensor.
• ​IndicesLayoutType (TensorLayout): Memory layout of the optional indices tensor.
• ​TopPArrLayoutType (TensorLayout): Memory layout of the optional top_p_arr tensor.
• ​SeedLayoutType (TensorLayout): Memory layout of the optional rng_seed tensor.
• ​TemperatureLayoutType (TensorLayout): Memory layout of the optional temperature tensor.
• ​MinPLayoutType (TensorLayout): Memory layout of the optional min_p tensor.
• ​TopKArrStorageType (TensorStorage): Storage type of the optional top_k_arr tensor.
• ​IndicesStorageType (TensorStorage): Storage type of the optional indices tensor.
• ​TopPArrStorageType (TensorStorage): Storage type of the optional top_p_arr tensor.
• ​SeedStorageType (TensorStorage): Storage type of the optional rng_seed tensor.
• ​TemperatureStorageType (TensorStorage): Storage type of the optional temperature tensor.
• ​MinPStorageType (TensorStorage): Storage type of the optional min_p tensor.

Args:

• ​ctx (DeviceContext): Device context for kernel execution.
• ​probs (TileTensor[dtype, Storage=probs.Storage, address_space=probs.address_space, linear_idx_type=probs.linear_idx_type]): Input probability distribution [batch_size, d], or raw logits when from_logits is True.
• ​output (TileTensor[out_idx_type, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type]): Output sampled indices [batch_size].
• ​top_k_val (Int): Default top-k value (number of top tokens to consider).
• ​top_p_val (Float32): Default top-p value (nucleus probability threshold).
• ​deterministic (Bool): Whether to use deterministic sampling.
• ​rng_seed (Optional[TileTensor[DType.uint64, SeedLayoutType, ImmutAnyOrigin, Storage=SeedStorageType]]): Optional per-row seed tensor [batch_size], indexed by the request's logical row (see indices). If None, defaults to 0.
• ​rng_offset (UInt64): Random offset for Random number generator.
• ​indices (Optional[TileTensor[out_idx_type, IndicesLayoutType, ImmutAnyOrigin, Storage=IndicesStorageType]]): Optional row indices for batch indexing [batch_size].
• ​top_k_arr (Optional[TileTensor[out_idx_type, TopKArrLayoutType, ImmutAnyOrigin, Storage=TopKArrStorageType]]): Optional per-row top-k values [batch_size].
• ​top_p_arr (Optional[TileTensor[DType.float32, TopPArrLayoutType, ImmutAnyOrigin, Storage=TopPArrStorageType]]): Optional per-row top-p values [batch_size].
• ​temperature (Optional[TileTensor[DType.float32, TemperatureLayoutType, ImmutAnyOrigin, Storage=TemperatureStorageType]]): Optional per-row temperature values [batch_size]. Only used when from_logits is True; defaults to 1.0 per row.
• ​min_p (Optional[TileTensor[DType.float32, MinPLayoutType, ImmutAnyOrigin, Storage=MinPStorageType]]): Optional per-row min-p thresholds [batch_size]. Only used when from_logits is True.
• ​out_dist (Optional[TileTensor[dist_dtype, DistLayoutType, MutAnyOrigin]]): Output masked distribution [batch_size, d]. Required when emit_dist is set.

Raises:

Error: If tensor ranks or shapes are invalid.