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

device_sampling_from_prob

def device_sampling_from_prob[vec_size: Int, block_size: Int, dtype: DType, deterministic: Bool = False](i: Int, d: Int, low: Float32, u: Float32, prob_vec: SIMD[DType.float32, vec_size], aggregate: Float32, sampled_id_sram: Pointer[Int, address_space=AddressSpace.SHARED, _safe=False]) -> Tuple[Float32, Int]

Device-level sampling from probability distribution with atomic operations.

Parameters:

vec_size (Int): Number of elements each thread loads per vectorized access.
block_size (Int): Number of threads per block.
dtype (DType): Element type of the probability distribution.
deterministic (Bool): If True, use deterministic sampling (defaults to False).

Args:

i (Int): Chunk iteration index used to compute element offsets within the row.
d (Int): Total number of elements in the row (vocabulary size).
low (Float32): Lower-bound threshold; only probabilities greater than this value participate in sampling.
u (Float32): Target cumulative probability to locate, scaled by the remaining mass q.
prob_vec (SIMD[DType.float32, vec_size]): Vector of probabilities for the current chunk.
aggregate (Float32): Running sum of filtered probabilities from previously processed chunks.
sampled_id_sram (Pointer[Int, address_space=AddressSpace.SHARED, _safe=False]): Shared-memory slot holding the sampled index, updated via atomic minimum.

Returns:

Tuple[Float32, Int]: Tuple of (new_aggregate, thread_local_max_valid_idx). The caller is responsible for reducing max_valid_idx across the block after all chunks are processed.