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
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]], seed: Optional[Pointer[UInt64, ImmutAnyOrigin]])
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]]): Optional per-token temperature scaling. - seed (
Optional[Pointer[UInt64, ImmutAnyOrigin]]): Optional per-token random seed.