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
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]) -> 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 massq. - 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]): 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.