For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

topk_softmax_sample_kernel

def topk_softmax_sample_kernel[block_size: Int, vec_size: Int, dtype: DType, out_idx_type: DType, LogitsLayoutType: TensorLayout, logits_origin: ImmOrigin, SampledLayoutType: TensorLayout, sampled_origin: MutOrigin](logits: TileTensor[dtype, LogitsLayoutType, logits_origin], sampled_indices: TileTensor[out_idx_type, SampledLayoutType, sampled_origin], top_k_arr: Optional[Pointer[Scalar[out_idx_type], MutUntrackedOrigin, _safe=False]], top_k_val: Int, temperature_val: Float32, temperature: Optional[Pointer[Float32, MutUntrackedOrigin, _safe=False]], seed_val: UInt64, seed: Optional[Pointer[UInt64, MutUntrackedOrigin, _safe=False]], d: Int)

Samples a token index from the top-k logits using softmax probabilities in a single kernel.

Each block processes one row. The kernel finds the k-th largest logit via ternary search, computes softmax over the top-k elements cached in shared memory, then draws a single categorical sample on thread 0.

Parameters:

block_size (Int): Number of threads per block.
vec_size (Int): Number of elements each thread loads per vectorized access.
dtype (DType): Element type of the logits tile.
out_idx_type (DType): Index type used for the sampled output indices.
LogitsLayoutType (TensorLayout): Memory layout of the input logits tile.
logits_origin (ImmOrigin): Origin tag for the immutable input logits tile.
SampledLayoutType (TensorLayout): Memory layout of the output sampled_indices tile.
sampled_origin (MutOrigin): Origin tag for the mutable output sampled_indices tile.

Args:

logits (TileTensor[dtype, LogitsLayoutType, logits_origin]): Input logits tile [batch_size, d].
sampled_indices (TileTensor[out_idx_type, SampledLayoutType, sampled_origin]): Output buffer for sampled token indices [batch_size].
top_k_arr (Optional[Pointer[Scalar[out_idx_type], MutUntrackedOrigin, _safe=False]]): Optional per-row top-k values; overrides top_k_val when present.
top_k_val (Int): Default number of largest logits to consider per row.
temperature_val (Float32): Default softmax temperature scaling factor.
temperature (Optional[Pointer[Float32, MutUntrackedOrigin, _safe=False]]): Optional per-row temperature values; overrides temperature_val.
seed_val (UInt64): Default random seed for the generator.
seed (Optional[Pointer[UInt64, MutUntrackedOrigin, _safe=False]]): Optional per-row seed values; overrides seed_val.
d (Int): Row length (vocabulary size).