Writing a Fused Softmax

A naively written softmax reads a row of floats from HBM three times: once for the max, once for the exponentials, once for the normalisation. On an A100, HBM bandwidth is roughly 2 TB/s while the chip can execute hundreds of TFLOP/s. For softmax - which does almost no arithmetic - those three round-trips are the entire cost. Fusing them into a single kernel pass is not an optimisation nicety; it is most of the work.

Why Softmax Is Memory-Bound by Default

Softmax along a row of length N is:

y_i = exp(x_i - max(x)) / sum_j(exp(x_j - max(x)))

The subtraction of max(x) is numerically important: without it, exp(x_i) overflows for large logits. But it creates a dependency: you cannot compute exp(x_i - max(x)) until you have seen the entire row to find max(x). A straightforward implementation therefore runs three separate GPU kernels:

1. Reduce over the row to compute m = max(x).
2. Compute e_i = exp(x_i - m) and write a temporary tensor to HBM.
3. Reduce again to compute s = sum(e_i), then divide.

Each kernel reads or writes the full row. For a matrix of shape (M, N) this costs roughly 5MN + 2M element reads plus 3MN + 2M element writes (quoting the Triton tutorial's accounting). A fused kernel that keeps values in registers and shared memory across reductions costs MN reads and MN writes. The ratio is the speedup: around 4x in practice.

The Triton fused-softmax tutorial benchmarks this directly and finds the Triton implementation is roughly 4x faster than torch.jit.script-ed PyTorch and competitive with (or faster than) torch.softmax for large enough rows.

The Online Softmax Algorithm

The core algorithmic insight is the online softmax trick, which lets you compute max and sum in a single streaming pass rather than two:

# Streaming over blocks of the row
m = -inf
d = 0.0
for each block b of x:
    m_new = max(m, max(b))
    d = d * exp(m - m_new) + sum(exp(b - m_new))
    m = m_new

After the pass, m holds the row maximum and d holds the correct normalisation denominator sum(exp(x_i - m)). A second pass writes the outputs:

for each block b of x:
    y_b = exp(b - m) / d

This two-pass-over-HBM version is already twice as good as the naive three-pass. In a Triton or CUDA kernel you can often load a block, compute the online update, and emit the output all in one pass for rows that fit in shared memory, reducing HBM traffic to one read and one write per element.

The correctness of the rescaling factor exp(m - m_new) is worth checking once. When you encounter a new block with a larger local max m_new, previously accumulated exponentials exp(x_i - m) are each too large by exactly exp(m - m_new). Multiplying the running sum d by that factor corrects them without re-reading the old data.

A Triton Implementation Skeleton

Triton makes this pattern clean to write. A minimal kernel looks like:

import triton
import triton.language as tl

@triton.jit
def fused_softmax_kernel(
    X, Y,
    stride_xm, stride_xn,
    M, N,
    BLOCK_SIZE: tl.constexpr,
):
    row = tl.program_id(0)          # one program per row
    row_start = row * stride_xm
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < N

    # Load the row (padded with -inf for reduction safety)
    x = tl.load(X + row_start + cols, mask=mask, other=-float('inf'))

    # Single-pass: max reduction, then sum of exp, then normalise
    x_max = tl.max(x, axis=0)
    x = x - x_max               # subtract for numerical stability
    numerator = tl.exp(x)
    denom = tl.sum(numerator, axis=0)
    y = numerator / denom

    tl.store(Y + row_start + cols, y, mask=mask)

Several points to note:

• tl.constexpr on BLOCK_SIZE allows the compiler to unroll loops and allocate registers statically.
• The single tl.load reads all N columns for this row into registers. If N exceeds BLOCK_SIZE, you need a loop over multiple loads with the online rescaling logic above.
• tl.max and tl.sum are warp-wide reductions that stay entirely on-chip.
• No intermediate tensor is written to HBM; numerator lives in registers.

One program handles exactly one row. Launch with grid = (M,). This gives clean memory access patterns: all threads in a warp load contiguous columns, satisfying coalescing requirements for HBM.

What the Compiler Does vs. What You Do

torch.compile with the default Inductor backend will often fuse a softmax written as plain PyTorch automatically, using its own codegen. For standard cases this is adequate. You write a Triton kernel manually when:

• The row length is non-standard (e.g., vocabulary size of 128,256 in an LLM).
• You need to fuse softmax with a preceding or following operation (attention scoring, temperature scaling, top-k masking) to avoid extra HBM writes.
• You need bfloat16 or FP8 with custom handling of the accumulator precision.

FlashAttention (arXiv:2205.14135) is the canonical example of taking this further: it fuses the entire QK^T matmul, the row-wise softmax, and the V matmul into one tiled kernel, eliminating the N x N attention matrix from HBM entirely. The online softmax trick is the key ingredient that makes this possible without sacrificing numerical correctness.

When It Falls Down

Row too wide for SRAM. If N exceeds what fits in shared memory (or the register file), the kernel must loop over tiles and apply the online rescaling. The loop introduces synchronisation overhead, and the benefit shrinks for very wide rows if the number of tiles is large enough to require multiple HBM reads anyway.

Small batch, short rows. For rows of length 32 or 64, the arithmetic intensity is so low that even a fused kernel cannot saturate the GPU. The dominant cost shifts to kernel launch overhead and warp scheduling. At this point torch.softmax on CPU may actually be faster.

Mixed precision edge cases. Accumulating exp in float16 overflows for inputs above about 11.0. The correct approach is to accumulate the sum in float32 even when the input and output are float16/bfloat16. A naive fused kernel that keeps everything in fp16 will silently produce NaN or Inf for large logit ranges (common in un-normalised attention scores at the start of training).

Auto-tuning over-fitting. Triton's @triton.autotune searches over BLOCK_SIZE configs at first run. If the benchmark rows are a different length from production rows, the chosen config may perform worse than the default PyTorch path. Always autotune on shapes representative of your actual workload.

Gradient pass is not free. The backward of softmax requires the forward output y. If the fused kernel does not retain y in a way compatible with autograd, PyTorch will recompute or save large intermediates. Profile with torch.cuda.memory_snapshot() to check.

Further Reading

• Triton fused softmax tutorial - the canonical hands-on walkthrough, with benchmarks showing the 4x speedup over Torch JIT.
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (arXiv:2205.14135) - Dao et al.; shows how the online softmax trick scales to fusing an entire attention layer.
• NVIDIA Deep Learning Performance Guide: GPU Background - authoritative treatment of arithmetic intensity, the roofline model, and HBM vs. on-chip memory hierarchies.