Softmax

The math, and why the naive version overflows

For a row of scores x, softmax(x)_i = exp(x_i) / Σ_j exp(x_j). The definition is one line, but computed literally it fails on real inputs: exp(89) already overflows a float32, and logits from a large model routinely reach that scale. The fix follows from an invariance: adding any constant c to every score multiplies both the numerator and the denominator by exp(c), so the result is unchanged. Choosing c = −max(x) makes the largest exponent exactly exp(0) = 1, every other exponent at most 1, and underflow of far-away scores harmless, since a score forty units below the max contributes nothing to the sum anyway. Every serious implementation, from NumPy to cuDNN, is the same three reductions: a max, a sum of shifted exponentials, and a divide.

Forward, four ways

The four implementations below compute the same row-wise softmax on a 2-D array. Reading them against each other shows exactly what the GPU versions add: nothing mathematical, only a mapping of the three reductions onto parallel hardware, one program or thread block per row.

import numpy as np

def softmax(x, axis=-1):
    """Numerically stable row-wise softmax.

    Subtracting the row max changes nothing mathematically
    (top and bottom are both scaled by exp(-max)) but keeps
    every exponent in (-inf, 0], so exp never overflows.
    """
    shifted = x - x.max(axis=axis, keepdims=True)
    e = np.exp(shifted)
    return e / e.sum(axis=axis, keepdims=True)

import torch

class Softmax(torch.autograd.Function):
    """Row-wise softmax with a hand-written backward.

    Forward saves y rather than x: the gradient of softmax
    is expressible entirely in terms of its output, which is
    why frameworks never need to keep the logits around.
    """

    @staticmethod
    def forward(ctx, x):
        shifted = x - x.amax(dim=-1, keepdim=True)
        y = shifted.exp()
        y = y / y.sum(dim=-1, keepdim=True)
        ctx.save_for_backward(y)
        return y

    @staticmethod
    def backward(ctx, dy):
        (y,) = ctx.saved_tensors
        # dx_i = y_i * (dy_i - sum_j dy_j y_j)
        return (dy - (dy * y).sum(dim=-1, keepdim=True)) * y

softmax = Softmax.apply

import torch
import triton
import triton.language as tl

@triton.jit
def softmax_kernel(x_ptr, out_ptr, n_cols,
                   x_stride, out_stride,
                   BLOCK: tl.constexpr):
    # One program per row. BLOCK is the next power of two
    # above n_cols, so the whole row fits in registers and
    # each reduction is a single tl.max / tl.sum.
    row = tl.program_id(0)
    offs = tl.arange(0, BLOCK)
    mask = offs < n_cols
    x = tl.load(x_ptr + row * x_stride + offs,
                mask=mask, other=float('-inf'))
    x = x - tl.max(x, axis=0)   # -inf padding drops out here
    num = tl.exp(x)
    den = tl.sum(num, axis=0)   # exp(-inf) = 0, so padding adds 0
    tl.store(out_ptr + row * out_stride + offs,
             num / den, mask=mask)

def softmax(x):
    rows, cols = x.shape
    out = torch.empty_like(x)
    BLOCK = triton.next_power_of_2(cols)
    softmax_kernel[(rows,)](x, out, cols,
                            x.stride(0), out.stride(0),
                            BLOCK=BLOCK)
    return out

// One thread block per row; each of the three reductions
// (max, sum, normalize) is a strided loop over the row
// followed by a shared-memory tree reduction.
__global__ void softmax_kernel(const float* __restrict__ x,
                               float* __restrict__ out,
                               int cols) {
    extern __shared__ float shm[];
    const float* xr  = x   + (long)blockIdx.x * cols;
    float*       outr = out + (long)blockIdx.x * cols;

    // 1. row max
    float m = -INFINITY;
    for (int i = threadIdx.x; i < cols; i += blockDim.x)
        m = fmaxf(m, xr[i]);
    shm[threadIdx.x] = m;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s)
            shm[threadIdx.x] = fmaxf(shm[threadIdx.x],
                                     shm[threadIdx.x + s]);
        __syncthreads();
    }
    m = shm[0];
    __syncthreads();

    // 2. sum of shifted exponentials
    float sum = 0.f;
    for (int i = threadIdx.x; i < cols; i += blockDim.x)
        sum += expf(xr[i] - m);
    shm[threadIdx.x] = sum;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s)
            shm[threadIdx.x] += shm[threadIdx.x + s];
        __syncthreads();
    }
    float inv = 1.f / shm[0];

    // 3. normalize
    for (int i = threadIdx.x; i < cols; i += blockDim.x)
        outr[i] = expf(xr[i] - m) * inv;
}

// launch, with blockDim a power of two:
// softmax_kernel<<<rows, 256, 256 * sizeof(float)>>>(x, out, cols);

The Triton and CUDA versions are numerically identical to the NumPy reference; a test that compares all four to within float32 tolerance on random rows, including rows containing -inf and rows of equal values, is the first thing to write.

Online softmax: the trick FlashAttention is built on

The three-pass structure above reads the row three times, and for attention that is fatal: the row is a whole matrix of scores that never fits in fast memory. The online variant folds the max and the sum into a single left-to-right pass. Keep a running maximum m and a running sum s of exponentials relative to m; when a new element v raises the maximum, the sum accumulated so far is expressed relative to a stale m, so rescale it by exp(m − m_new) before adding the new term. The rescale is exact, not an approximation:

def online_softmax(row):
    m, s = float('-inf'), 0.0
    for v in row:                      # one pass, O(1) state
        m_new = max(m, v)
        s = s * math.exp(m - m_new) + math.exp(v - m_new)
        m = m_new
    return [math.exp(v - m) / s for v in row]

Because the state is just (m, s), the pass can be tiled: process the row in blocks, and merge each block's (m, s) pair with the running pair using the same rescaling. That is precisely what FlashAttention does, streaming tiles of the attention score matrix through on-chip memory and carrying (m, s) per query row, so the full score matrix is never materialized. My flash-attention-cuda repository builds the forward pass from exactly this kernel, step by step from a naive baseline, with each step's speedup measured on an RTX 4090.

The backward pass

Differentiating y_i = exp(x_i) / Σ exp gives the Jacobian ∂y_i/∂x_j = y_i(δ_ij − y_j): a diagonal term minus a rank-one term. Multiplying the incoming gradient dy through it collapses to dx_i = y_i(dy_i − Σ_j dy_jy_j), which is one dot product and one elementwise multiply, no Jacobian ever materialized. Two things are worth noticing. The backward needs only the output y, not the input x, which is why the PyTorch tab saves y. And when softmax is immediately followed by cross-entropy loss, the composite gradient simplifies further to y − onehot(target), which is why every framework fuses the two and why you should never put an explicit softmax before CrossEntropyLoss.

Performance notes

Softmax is memory-bound: three reads and a write per element against a handful of exponentials, so the roofline is bandwidth, not FLOPs. The Triton version holds the row in registers and reads it once, which is why a fused one-pass kernel beats the three-pass CUDA version above on rows that fit in a block, and the CUDA version in turn exists to make the reduction structure explicit. The remaining distance to cuDNN is fusion: production kernels fold the shift, exponential, and normalize into neighboring operations (the attention matmul, the loss) rather than round-tripping through global memory. On the numerics side, expf can be swapped for the faster __expf intrinsic at a small accuracy cost, a trade worth measuring rather than assuming.