Skip to main content

Mojo function

softmax_2_pass

softmax_2_pass[simd_width: Int, buffer_size: Dim, dtype: DType](output: NDBuffer[dtype, 1, origin, DimList(buffer_size)], input: NDBuffer[dtype, 1, origin, DimList(buffer_size)])

Performs an unbatched softmax on an input tensor using the two-pass online algorithm.

The unbatched two-pass online softmax is described in "Online normalizer calculation for softmax" (https://arxiv.org/abs/1805.02867) and "A full-stack search technique for domain optimized deep learning accelerators" (https://dl.acm.org/doi/abs/10.1145/3503222.3507767) and is defined as:

procedure SoftmaxUnbatched(InputInput)
runningMax = -∞
runningSum = 0
STAGE 1:
for i = 0 to N do
newMax = max(runningMax, Input[i])
runningSum = runningSum*exp(runningMax-newMax) + exp(Input[i]-newMax)
runningMax = newMax
end for
for i = 0 to N do
Output[i] = exp(Input[i] - runningMax) / runningSum
end for
procedure SoftmaxUnbatched(InputInput)
runningMax = -∞
runningSum = 0
STAGE 1:
for i = 0 to N do
newMax = max(runningMax, Input[i])
runningSum = runningSum*exp(runningMax-newMax) + exp(Input[i]-newMax)
runningMax = newMax
end for
for i = 0 to N do
Output[i] = exp(Input[i] - runningMax) / runningSum
end for

Parameters:

  • simd_width (Int): The simd_width to use in vectorization.
  • buffer_size (Dim): The size of the input and output buffers.
  • dtype (DType): The dtype of the input and output buffers.

Args:

  • output (NDBuffer[dtype, 1, origin, DimList(buffer_size)]): The output buffer in which to store the softmax values.
  • input (NDBuffer[dtype, 1, origin, DimList(buffer_size)]): The input buffer used to compute the softmax.

Was this page helpful?