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

softmax_kernel

def softmax_kernel[BLOCK_SIZE: Int, input_fn: def[_dtype: DType, _simd_width: Int, _rank: Int](IndexList[_rank]) capturing thin -> SIMD[_dtype, _simd_width], dtype: DType, sink_type: DType, rank: Int, OutputLayoutType: TensorLayout, output_origin: MutOrigin, OutputStorage: TensorStorage, SinkWeightsLayoutType: TensorLayout, accum_type: DType = get_accum_type[dtype](), *, sink: Bool = False, logsoftmax: Bool = False](shape: IndexList[rank], output: TileTensor[dtype, OutputLayoutType, output_origin, Storage=OutputStorage], sink_weights: TileTensor[sink_type, SinkWeightsLayoutType, ImmutAnyOrigin])

GPU kernel implementing the three-pass softmax with optional sink-attention and logsoftmax variants.

Each block reduces one row: step 1 finds the row max (optionally clamped with a per-head sink weight), step 2 computes exp(x - max) and the row sum, and step 3 normalizes (and applies log when logsoftmax is set).

Parameters:

BLOCK_SIZE (Int): The number of threads per block.
input_fn (def[_dtype: DType, _simd_width: Int, _rank: Int](IndexList[_rank]) capturing thin -> SIMD[_dtype, _simd_width]): The elementwise input lambda.
dtype (DType): The dtype of the input and output buffers.
sink_type (DType): The dtype of the sink weights.
rank (Int): The rank of the input and output tensors.
OutputLayoutType (TensorLayout): The layout type of the output tensor.
output_origin (MutOrigin): The origin of the output tensor.
OutputStorage (TensorStorage): The storage type of the output tensor.
SinkWeightsLayoutType (TensorLayout): The layout type of the sink weights tensor.
accum_type (DType): The accumulation dtype (defaults to the accumulation type for dtype).
sink (Bool): Whether to apply sink-attention bias to the row max.
logsoftmax (Bool): Enable to apply elementwise log() to outputs after softmax.

Args:

shape (IndexList[rank]): The shape of the tensor as an IndexList.
output (TileTensor[dtype, OutputLayoutType, output_origin, Storage=OutputStorage]): The output buffer in which to store the softmax values.
sink_weights (TileTensor[sink_type, SinkWeightsLayoutType, ImmutAnyOrigin]): Per-head sink weights used when sink is True.