Puzzle 16: Softmax Op

fn softmax_gpu_kernel[
    layout: Layout,
    input_size: Int,
    dtype: DType = DType.float32,
](
    output: LayoutTensor[mut=True, dtype, layout],
    input: LayoutTensor[mut=False, dtype, layout],
):
    shared_max = tb[dtype]().row_major[TPB]().shared().alloc()
    shared_sum = tb[dtype]().row_major[TPB]().shared().alloc()
    global_i = block_dim.x * block_idx.x + thread_idx.x
    local_i = thread_idx.x

    var thread_max: Scalar[dtype] = min_finite[dtype]()
    if global_i < input_size:
        thread_max = rebind[Scalar[dtype]](input[global_i])

    shared_max[local_i] = thread_max
    barrier()

    # Parallel reduction to find max
    stride = TPB // 2
    while stride > 0:
        if local_i < stride:
            shared_max[local_i] = max(
                shared_max[local_i], shared_max[local_i + stride]
            )
        barrier()
        stride = stride // 2

    block_max = shared_max[0]

    var exp_val: Scalar[dtype] = 0.0
    if global_i < input_size:
        exp_val = rebind[Scalar[dtype]](exp(input[global_i] - block_max))
        output[global_i] = exp_val

    shared_sum[local_i] = exp_val
    barrier()

    # Parallel reduction for sum
    stride = TPB // 2
    while stride > 0:
        if local_i < stride:
            shared_sum[local_i] += shared_sum[local_i + stride]
        barrier()
        stride = stride // 2

    block_sum = shared_sum[0]

    # Normalize by sum
    if global_i < input_size:
        output[global_i] = output[global_i] / block_sum

Our GPU implementation implements the numerically stable softmax algorithm with highly optimized parallel reduction techniques. Let's dissect the kernel in detail:

Kernel signature and memory management

fn softmax_gpu_kernel[
    layout: Layout,
    input_size: Int,
    dtype: DType = DType.float32,
](
    out: LayoutTensor[mut=True, dtype, layout],
    input: LayoutTensor[mut=False, dtype, layout],
)

The kernel is parameterized with:

• Common layout parameter for both input and output tensors
• Vector size as an Integer parameter
• Configurable data type with float32 as default
• Mutable output tensor for in-place computation
• Non-mutable input tensor (mut=False)

Shared memory allocation

shared_max = tb[dtype]().row_major[TPB]().shared().alloc()
shared_sum = tb[dtype]().row_major[TPB]().shared().alloc()

The kernel allocates two shared memory buffers:

• shared_max: For parallel maximum finding reduction
• shared_sum: For parallel sum computation
• Both use TPB (Threads Per Block = 128) as their size
• Shared memory provides fast access for all threads within a block

Thread indexing

global_i = block_dim.x * block_idx.x + thread_idx.x
local_i = thread_idx.x

Each thread computes:

• global_i: Its global index in the entire computation space
• local_i: Its local index within the current thread block This mapping ensures each thread processes exactly one input element.

Maximum-finding phase

var thread_max: Scalar[dtype] = min_finite[dtype]()
if global_i < input_size:
    thread_max = rebind[Scalar[dtype]](input[global_i])

shared_max[local_i] = thread_max
barrier()

This initializes each thread with:

• The minimum finite value for elements outside the valid range
• The actual input value for threads that map to valid elements
• Storage in shared memory for the reduction process
• A barrier synchronization to ensure all threads complete memory writes

Parallel max reduction

stride = TPB // 2
while stride > 0:
    if local_i < stride:
        shared_max[local_i] = max(shared_max[local_i], shared_max[local_i + stride])
    barrier()
    stride = stride // 2

This implements a parallel tree-reduction pattern:

1. Start with stride = 64 (half of TPB)
2. Each active thread compares two values separated by the stride
3. Store the maximum in the lower index
4. Synchronize all threads with a barrier
5. Halve the stride and repeat
6. After \(\log_2(TPB)\) steps, shared_max[0] contains the global maximum

This logarithmic reduction is significantly faster than a linear scan on large inputs.

Exponentiation with numerical stability

block_max = shared_max[0]

var exp_val: Scalar[dtype] = 0.0
if global_i < input_size:
    exp_val = rebind[Scalar[dtype]](exp(input[global_i] - block_max))
    out[global_i] = exp_val

Each thread:

1. Reads the global maximum from shared memory
2. Subtracts it from its input value before taking the exponential
3. This subtraction is crucial for numerical stability - it prevents overflow
4. The largest exponent becomes \(e^0 = 1\), and all others are \(e^{negative} < 1\)
5. Stores the intermediate result in the output buffer

Parallel sum reduction

shared_sum[local_i] = exp_val
barrier()

stride = TPB // 2
while stride > 0:
    if local_i < stride:
        shared_sum[local_i] += shared_sum[local_i + stride]
    barrier()
    stride = stride // 2

The second reduction phase:

1. Stores all exponential values in shared memory
2. Uses the same tree-based reduction pattern as for max
3. But performs addition instead of maximum comparison
4. After \(\log_2(TPB)\) steps, shared_sum[0] contains the total sum of all exponentials

Final normalization

block_sum = shared_sum[0]

if global_i < input_size:
    out[global_i] = out[global_i] / block_sum

Each thread:

1. Reads the total sum from shared memory
2. Divides its exponential value by this sum
3. Writes the normalized probability to the output buffer
4. This produces a valid probability distribution that sums to 1

Performance characteristics

The implementation has excellent performance characteristics:

• Complexity: \(O(\log n)\) for both max and sum calculations vs \(O(n)\) in a sequential approach
• Memory efficiency: Uses only \(2 \times TPB\) elements of shared memory
• Work efficiency: Each thread performs approximately \(2 \times \log_2(n)\) operations
• Load balancing: Each thread handles the same amount of work
• Synchronization: Uses minimal barriers, only where necessary
• Memory access: Coalesced global memory access pattern for optimal bandwidth

The algorithm is also numerically robust, handling potential overflow/underflow cases by applying the max-subtraction technique that maintains precision across the wide range of values common in neural network activations.

fn softmax_cpu_kernel[
    layout: Layout,
    input_size: Int,
    dtype: DType = DType.float32,
](
    output: LayoutTensor[dtype, layout, MutableAnyOrigin],
    input: LayoutTensor[dtype, layout, MutableAnyOrigin],
):
    var max_val: Scalar[dtype] = min_finite[dtype]()
    for i in range(input_size):
        max_val = max(max_val, rebind[Scalar[dtype]](input[i]))

    var sum_exp: Scalar[dtype] = 0.0
    for i in range(input_size):
        var exp_val = rebind[Scalar[dtype]](exp(input[i] - max_val))
        output[i] = exp_val
        sum_exp += exp_val

    for i in range(input_size):
        output[i] = output[i] / sum_exp

Our CPU implementation provides a sequential fallback that follows the same mathematical approach but is optimized for single-threaded execution. Let's analyze each phase:

1. Maximum Finding:

   var max_val: Scalar[dtype] = min_finite[dtype]()
   for i in range(input_size):
       max_val = max(max_val, rebind[Scalar[dtype]](input[i]))
   

   We initialize with the minimum finite value and perform a linear scan through the array, keeping track of the maximum value encountered. This has \(O(n)\) complexity but works efficiently on CPU where we don’t have many cores to parallelize across.

2. Exponential Computation and Summation:

   var sum_exp: Scalar[dtype] = 0.0
   for i in range(input_size):
       var exp_val = rebind[Scalar[dtype]](exp(input[i] - max_val))
       out[i] = exp_val
       sum_exp += exp_val
   

   We compute \(e^{x_i - max}\) for each element, store the result in the output buffer, and accumulate the sum \(\sum_{j=1}^{n} e^{x_j - max}\) in a single pass. This approach minimizes memory operations compared to using separate loops.

3. Normalization:

   for i in range(input_size):
       out[i] = out[i] / sum_exp
   

   Finally, we normalize each element by dividing by the sum, producing a proper probability distribution according to the softmax formula:

   $$\Large \text{softmax}(x_i) = \frac{e^{x_i - \max(x)}}{\sum_{j=1}^{n} e^{x_j - \max(x)}}$$

The CPU implementation uses the same numerical stability technique (subtracting the maximum) but with sequential operations rather than parallel ones. It’s simpler than the GPU version since it doesn’t need to handle shared memory or thread synchronization, but it’s also less efficient for large inputs.

Both implementations are registered with MAX Graph’s custom operation system through the @compiler.register("softmax") decorator, allowing seamless execution on either device type based on availability.

    with Graph(
        "softmax_graph",
        input_types=[
            TensorType(
                dtype,
                shape=input_tensor.shape,
                device=DeviceRef.from_device(device),
            ),
        ],
        custom_extensions=[mojo_kernels],
    ) as graph:
        input_value = graph.inputs[0]

        # The output shape is the same as the input for softmax
        # Note: the name must match the name used in `@compiler.register("softmax")` in op/softmax.mojo
        output = ops.custom(
            name="softmax",
            values=[input_value],
            out_types=[
                TensorType(
                    dtype=input_value.tensor.dtype,
                    shape=input_value.tensor.shape,
                    device=DeviceRef.from_device(device),
                )
            ],
            parameters={
                "input_size": input_tensor.shape[0],
                "dtype": dtype,
            },
        )[0].tensor
        graph.output(output)

The Python integration creates a seamless bridge between NumPy arrays and our optimized Mojo GPU kernel. The implementation consists of several key components:

1. Graph Setup and Configuration:

   with Graph(
       "softmax_graph",
       input_types=[
           TensorType(
               dtype,
               shape=input_tensor.shape,
               device=DeviceRef.from_device(device),
           ),
       ],
       custom_extensions=[mojo_kernels],
   ) as graph:
   

   This creates a computation graph named “softmax_graph” that:

   • Defines the input tensor type with proper dtype and shape
   • Maps the tensor to the target device (CPU or GPU)
   • Loads our custom Mojo operations from the specified directory
   • The custom_extensions parameter is crucial for linking to our Mojo implementation

2. Custom Operation Configuration:

   output = ops.custom(
       name="softmax",
       values=[input_value],
       out_types=[
           TensorType(
               dtype=input_value.tensor.dtype,
               shape=input_value.tensor.shape,
               device=DeviceRef.from_device(device),
           )
       ],
       parameters={
           "input_size": input_tensor.shape[0],
           "dtype": dtype,
       },
   )[0].tensor
   

   This sets up our custom operation with:

   • Name matching the @compiler.register("softmax") in our Mojo code
   • Input values passed as a list
   • Output type definition matching the input shape and type
   • Parameters required by our kernel, including the vector size and data type
   • We extract the tensor from the first returned element with [0].tensor

3. Graph Output Definition:

   graph.output(output)
   

   This registers our operation’s result as the graph’s output.

The main script includes comprehensive testing that:

• Generates random input data: np.random.randn(INPUT_SIZE).astype(np.float32)
• Calculates expected results with SciPy: scipy_softmax(input_array)
• Verifies numerical accuracy: np.testing.assert_allclose(..., rtol=1e-5)
• Confirms the output is a valid probability distribution: np.sum(result.to_numpy())

This implementation showcases the power of MAX Graph for integrating high-performance Mojo kernels with Python’s scientific computing ecosystem, providing both efficiency and ease of use.