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_gumbel_noise_kernel

def apply_gumbel_noise_kernel[dtype: DType, OutputLayoutType: TensorLayout, InputLayoutType: TensorLayout, num_sms: Int, num_threads: Int](output: TileTensor[dtype, OutputLayoutType, MutAnyOrigin], input: TileTensor[dtype, InputLayoutType, ImmutAnyOrigin], temperature: Optional[Pointer[Float32, ImmutAnyOrigin, _safe=False]], seed: Optional[Pointer[UInt64, ImmutAnyOrigin, _safe=False]])

Adds Gumbel(0,1) noise to logits for sampling via the Gumbel-max trick.

Parameters:

dtype (DType): Data type of the input and output logit buffers.
OutputLayoutType (TensorLayout): Layout of the output tensor.
InputLayoutType (TensorLayout): Layout of the input tensor.
num_sms (Int): Number of streaming multiprocessors to launch with.
num_threads (Int): Number of threads per block.

Args:

output (TileTensor[dtype, OutputLayoutType, MutAnyOrigin]): Output tensor of noised logits.
input (TileTensor[dtype, InputLayoutType, ImmutAnyOrigin]): Input tensor of logits.
temperature (Optional[Pointer[Float32, ImmutAnyOrigin, _safe=False]]): Optional per-token temperature scaling.
seed (Optional[Pointer[UInt64, ImmutAnyOrigin, _safe=False]]): Optional per-token random seed.