IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

TopKTopPSamplingFromProbKernel

def TopKTopPSamplingFromProbKernel[ProbsLayoutType: TensorLayout, probs_origin: ImmOrigin, OutputLayoutType: TensorLayout, output_origin: MutOrigin, block_size: Int, vec_size: Int, dtype: DType, out_idx_type: DType, deterministic: Bool, from_logits: Bool = False, emit_dist: Bool = False, dist_dtype: DType = DType.float32, ProbsStorageType: TensorStorage = PointerStorage, OutputStorageType: TensorStorage = PointerStorage](probs: TileTensor[dtype, ProbsLayoutType, probs_origin, Storage=ProbsStorageType], output: TileTensor[out_idx_type, OutputLayoutType, output_origin, Storage=OutputStorageType], out_dist: Optional[Pointer[Scalar[dist_dtype], MutAnyOrigin]], indices: Optional[Pointer[Scalar[out_idx_type], ImmutAnyOrigin]], top_k_arr: Optional[Pointer[Scalar[out_idx_type], ImmutAnyOrigin]], top_k_val: Int32, top_p_arr: Optional[Pointer[Float32, ImmutAnyOrigin]], top_p_val: Float32, d: Int32, rng_seed: Optional[Pointer[UInt64, ImmutAnyOrigin]], rng_offset: UInt64, temperature: Optional[Pointer[Float32, ImmutAnyOrigin]], min_p: Optional[Pointer[Float32, ImmutAnyOrigin]])

Kernel for joint top-k + top-p sampling from probability distribution.

Identical to TopKSamplingFromProbKernel but additionally enforces a nucleus (top-p) constraint: a token is accepted only when both the count of tokens above the pivot is less than k AND the cumulative probability of those tokens is less than p.

When top_p_val = 1.0 and top_p_arr is null, this degrades to top-k-only with zero overhead since sum < 1.0 is always true.

When from_logits is True, probs contains raw logits and softmax with per-row temperature scaling is fused into the kernel: every load is transformed to exp((logit - row_max) / temp), the unnormalized softmax value with the row maximum shifted to exactly 1.0. The pivot search over [0, 1] is unchanged; the total unnormalized mass z replaces the normalized distribution's implicit total of 1.0 in the initial CDF budget and scales the top-p threshold. The optional min-p mask is applied inline (in this domain the max "probability" is 1.0, so the mask threshold is simply min_p), matching apply_min_p_mask_kernel semantics in the normalized domain.

When emit_dist is set, the kernel also writes the masked renormalized distribution it drew from to out_dist. Speculative decoding builds its rejection residual from that distribution, and reads the sampled token's own probability back out of it. Requires from_logits, and a non-Apple GPU because the cutoff search uses block collectives.

Parameters:

  • ​ProbsLayoutType (TensorLayout): Memory layout of the input probs tile.
  • ​probs_origin (ImmOrigin): Origin tag for the immutable input probs tile.
  • ​OutputLayoutType (TensorLayout): Memory layout of the output output tile.
  • ​output_origin (MutOrigin): Origin tag for the mutable output output tile.
  • ​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 probs tensor.
  • ​out_idx_type (DType): Index type used for the sampled output indices.
  • ​deterministic (Bool): If True, use deterministic sampling.
  • ​from_logits (Bool): If True, probs holds raw logits and softmax with per-row temperature scaling and min-p masking is fused into the kernel (defaults to False).
  • ​emit_dist (Bool): If True, also write the masked distribution to out_dist (defaults to False).
  • ​dist_dtype (DType): Element type of out_dist.
  • ​ProbsStorageType (TensorStorage): Storage type of the input probs tile.
  • ​OutputStorageType (TensorStorage): Storage type of the output output tile.

Args:

Models

Cookbook

Puzzles

Blog

Community

Contact