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
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 fordtype). - 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 whensinkis True.