Mojo function

apply_min_p_mask_kernel

apply_min_p_mask_kernel[dtype: DType, block_size: Int](probs: UnsafePointer[Scalar[dtype], MutExternalOrigin], min_p_arr: UnsafePointer[Float32, ImmutExternalOrigin], 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.

Args:

• ​probs (UnsafePointer[Scalar[dtype], MutExternalOrigin]): Probability buffer [batch_size * d], modified in-place.
• ​min_p_arr (UnsafePointer[Float32, ImmutExternalOrigin]): Per-row min_p values [batch_size].
• ​d (Int): Vocabulary size (row length).

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit