GPU kernel optimization: Softmax — Part 1

I recently dove back into CUDA programming with a focus on fused GPU kernels: a powerful optimization technique used in deep learning frameworks.

One of the most common examples is the Softmax function, which converts raw scores into probabilities. It’s typically used in the final layer of neural networks and plays a key role in the attention mechanism behind large language models.

Softmax formula: subtracting max(x) is required for numerical stability

Softmax is also a great starting point to explore GPU optimizations. While the basic implementation is simple, it involves two reduction operations (max and sum) over each input row. If not handled carefully, this can lead to redundant memory access and poor parallelism.

In this article, I’ll walk through my experiments with the Triton Softmax tutorial and how it compares to CUDA implementations, mainly inspired by SzymonOzog’s repo. Along the way, I learned a lot about the trade-offs between high-level abstraction and low-level control.

The result? A custom Softmax kernel with better performance on my slightly aging GeForce RTX 2070. You can find the full implementation on my GitHub repo.

Benchmark of softmax implementations (higher is better)

🐢 A Naive Softmax: Clean, Simple… and Slow

Before diving into optimized GPU kernels, let’s look at a basic implementation of Softmax in PyTorch. This version is taken from the Triton documentation:

import torch

def naive_softmax(x):
    """Compute row-wise softmax of X using native pytorch

    We subtract the maximum element in order to avoid overflows.
    Softmax is invariant to this shift.
    """
    x_max = x.max(dim=1)[0]                 # First reduction
    z = x - x_max[:, None]                  # Needed for computation stability
    numerator = torch.exp(z)
    denominator = numerator.sum(dim=1)      # Second reduction
    ret = numerator / denominator[:, None]
    return ret

x = torch.randn(3, 4, device='cuda')
# [[-0.0660, -0.0489, -0.3908, -1.2046],
#  [ 0.1957,  0.7052, -0.8617, -0.2857],
#  [-1.9676,  2.2595,  0.9407, -0.0621]]

naive_softmax(x)  # Each line is transformed to probabilities adding up to 1
# [[0.3268, 0.3324, 0.2362, 0.1047],
#  [0.2755, 0.4586, 0.0957, 0.1702],
#  [0.0106, 0.7245, 0.1938, 0.0711]]

Note: Softmax is a vector operation, but it’s common to apply it row-wise on a matrix (each row is treated independently).

This implementation is easy to read, but it has a major drawback: Each high-level PyTorch operation launches a separate GPU kernel, which means repeated memory transfers, synchronization overhead, and no chance to fuse operations for speed.

This is where lower-level languages like Triton and CUDA shine.

🔱 Enter Triton: Fusing for Speed, Pythonically

Triton is a domain-specific language designed to make writing fused GPU kernels easy, without diving too deep into CUDA. If you’ve used torch.compile() on a model recently, you’ve already used Triton under the hood.

But Triton isn’t just for automatic graph compilation, it also lets you manually write custom GPU kernels with fine-grained control, using a Python-like syntax.

Here’s a simplified Softmax implementation from the Triton tutorial:

@triton.jit
def kernel_fused_softmax(
    in_ptr, in_row_stride: int,
    out_ptr, out_row_stride: int,
    n_rows: int, n_cols: int,
    block_size: tl.constexpr,
):
    row_idx = tl.program_id(0)  # Index of the row to process

    # Parallelization is abstracted behind vectorized operations
    # So we have to load a whole row of the input
    col_offsets = tl.arange(0, block_size)  # We expect block_size >= n_cols
    mask = col_offsets < n_cols
    in_row_ptrs = compute_row_ptrs(
        in_ptr, in_row_stride, row_idx, col_offsets
    )
    row = tl.load(in_row_ptrs, mask=mask, other=float("-inf"))

    # 1. Max reduction over a row
    row_max = tl.max(row, axis=0)

    # 2. Stable exponential of a row
    numerator = tl.exp(row - row_max)

    # 3. Sum reduction of the exponentiated row
    denominator = tl.sum(numerator, axis=0)

    # 4. Compute and store the result of softmax
    res = numerator / denominator
    out_row_ptrs = compute_row_ptrs(
        out_ptr, out_row_stride, row_idx, col_offsets
    )
    tl.store(out_row_ptrs, res, mask=mask)

This kernel expects an input tensor of shape: (n_rows, n_cols). Each block (program in Triton terms) handles a full row, processing n_cols elements. block_size is typically set as the next power of 2 ≥ n_cols for efficient reductions.

One thing that immediately caught my attention was this: As the number of columns grows, block_size grows too, eventually hitting hardware limits (like the 1024-thread max on my RTX 2070). How does Triton handle this behind the scenes? What happens when your tensor gets wider than the hardware allows? We’ll explore that in the next article!

How Reductions Work Under the Hood in Triton

The Triton language provides handy abstractions around reduction operations such as tl.max and tl.sum. but it doesn’t expose much detail on how they’re actually implemented.

Luckily, I came across a great article by Fei Kong, which provided insight into Triton’s lowering strategy:

*tl.sum() is lowered to a intra-warp reduction with warp shuffle instruction first, followed by inter-warp reduction across warps within shared memory*

Let’s break down what it means!

Part 1: Dividing the Work

Reduction operations are trickier to parallelize than element-wise ops. That’s because they involve data dependencies: the output depends on multiple inputs being combined in order. However, operations like max and sum are associative:

In this case, we can apply a first round of operator on values 2 by 2 in parallel, then synchronize the results and continue, halving the number of parallel operations to execute every time:

Comparison of maximum computation algorithms

Instead of having N independent steps, we now only need log(N) synchronization rounds. That’s a big performance win!

Part 2: Thread Communication

Parallel reduction only works if threads can communicate efficiently — and that’s where GPU memory hierarchies come in.

Threads within a warp (typically 32 threads) can exchange values very quickly using special warp shuffle instructions.

Each block consists of multiple warps (e.g. 32 warps × 32 threads = 1024 threads per block on my RTX 2070). Once each warp has reduced its own subset of data, it shares its intermediate result via shared memory.

Schema of a GPU reduction algorithm execution

This hierarchical reduction strategy (warp-local → block-wide) is what powers Triton's fast tl.sum() and tl.max().

🔍 For more implementation detail, I recommend this explanation.

Now that we understand how triton implements these reductions, we can start our own CUDA implementation!

🔨 Rebuilding Softmax in CUDA

To mirror Triton’s fused approach, I built a simplified CUDA kernel that performs row-wise Softmax. You’ll find the full implementation of the reduction helpers in my repo.

template <typename T>
__global__ void
softmax_kernel_simplified(T* input, T* output, int n_cols, int n_rows) {
    const int num_warps = CEIL_DIV(blockDim.y, WARP_SIZE);
    __shared__ float reduction[WARP_SIZE];

    // Index of the row to process (block wise)
    const int row_index = blockIdx.x*blockDim.x + threadIdx.x;

    // Contrary to Triton that uses vectorized operations, CUDA is working
    // at the thread level. Here each thread is processing a single row value
    // indexed using `col_index`.
    const int col_index = threadIdx.y;

    // Check boundaries
    if (col_index >= n_cols || row_index >= n_rows)
        return;

    // 1. Max reduction over a row
    float row_max = block_wide_max(
        num_warps, col_index, input[row_index * n_cols + col_index], reduction
    );

    // 2. Stable exponential of a row
    float numerator = __expf(input[row_index * n_cols + col_index] - row_max);

    // 3. Sum reduction of the exponentiated row
    float denominator = block_wide_sum(
        num_warps, col_index, numerator, reduction
    );

    // 4. Compute and store the result of softmax
    output[row_index * n_cols + col_index] = numerator / denominator;
}

Similar to earlier, the input tensor has the shape: (n_rows, n_cols). The kernel is launched with one block per row. Each block is processing an entire row (n_colselements) so the block size is pinned to n_cols . For example, with an input of shape (128, 256): we run 128 blocks of 256 threads (8 warps of 32 threads).

This effectively means that the implementation only works for input tensors with a width smaller than MAX_BLOCK_SIZE, which is 1024 on my hardware.

Quadrupling the effort

To go further, I implemented a simple optimization: using CUDA’s built-in float4 type to process four elements per thread. This approach boosts memory throughput and cuts down the total number of threads needed per row.

I implemented some utils functions to manipulate float4 values. Here’s the core of the change:

auto in_vector = reinterpret_cast<const float4*>(&input[row_index * n_cols]);

// 1. Max reduction over a row
float max_val = max_float4(in_vector[col_index]);
max_val = block_wide_max(num_warps, col_index, max_val, reduction);

// 2. Stable exponential of a row
auto numerator = stable_exp_float4(in_vector[col_index], max_val);

// 3. Sum reduction of the exponentiated row
float denominator = sum_float4(numerator);
denominator = block_wide_sum(num_warps, col_index, denominator, reduction);

// 4. Compute and store the result of softmax
auto out_vector = reinterpret_cast<float4*>(&output[row_index * n_cols]);
out_vector[col_index] = divide_float4(numerator, denominator);

This change bumps our input width limit from 1024 to 4096 on the same hardware!

📊 Wrapping Up: Torch vs. Triton vs. CUDA

To evaluate our work, I benchmarked three Softmax implementations — PyTorch (baseline), Triton (from the tutorial), and my own CUDA kernel — across increasing tensor widths.

Here’s what the results tell us:

• 🟠 Triton leads the pack, especially as tensor width grows. Its performance benefits from highly optimized launch heuristics and smart scheduling.
• 🟢 My CUDA implementation tracks closely behind Triton, showing that we’re on the right path with our fused design and warp-level reductions.

This deep dive helped me appreciate both Triton’s high-level abstractions and CUDA’s low-level power. In Part 2, we’ll tackle larger inputs by exploring how Triton’s handling it, and eventually surpassing the Triton tutorial performances.

Thank you for reading!