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

apply_min_p_mask_kernel

def apply_min_p_mask_kernel[dtype: DType, block_size: Int](probs: Pointer[Scalar[dtype], MutUntrackedOrigin, _safe=False], min_p_arr: Pointer[Float32, ImmUntrackedOrigin, _safe=False], d: Int)

Zero out probabilities below the per-row min_p threshold.

Each block processes one batch row. Threads cooperatively find the row-wise max probability via a block reduction, compute the threshold as min_p * max_prob, and then zero any element below it.

Parameters:

dtype (DType): Element type of the probs buffer.
block_size (Int): Number of threads per block.

Args:

probs (Pointer[Scalar[dtype], MutUntrackedOrigin, _safe=False]): Probability buffer [batch_size * d], modified in-place.
min_p_arr (Pointer[Float32, ImmUntrackedOrigin, _safe=False]): Per-row min_p values [batch_size].
d (Int): Vocabulary size (row length).