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 module

softmax

Provides numerically stable softmax kernels for CPU and GPU, including fused and online variants.

Functions

• ​identity: Returns the input SIMD vector unchanged.
• ​logsoftmax: Computes log-softmax over the given axis using a caller-supplied input lambda.
• ​mul: Returns the element-wise product x * y.
• ​reciprocal: Returns the element-wise reciprocal 1 / x.
• ​reduce_add_simd: This functions adds val to either the scalar value or the vector value depending on the step_simd_width. This is useful when the simd_width varies between iterations as in vectorize.
• ​softmax: Computes softmax over the given axis of input and stores the result in output.
• ​softmax_2_pass: Performs an unbatched softmax on an input tensor using the two-pass online algorithm.
• ​softmax_3_pass: Performs an unbatched softmax on an input tensor using the three-pass algorithm.
• ​softmax_kernel: GPU kernel implementing the three-pass softmax with optional sink-attention and logsoftmax variants.
• ​softmax_with_temperature: GPU softmax with per-row temperature scaling.
• ​sub: Returns the element-wise difference x - y.

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit