Fused Triton softmax

Fusion arithmetic

The unfused PyTorch softmax reads 5MN + 2M elements and writes 3MN + 2M, while a fused kernel reads and writes 2MN. That is the tutorial’s count, and it is the whole argument for writing the kernel: the operation is bandwidth-bound, and the win comes from deleting round trips through DRAM.

One row block per program

@triton.jit
def softmax_kernel(out_ptr, in_ptr, in_stride, out_stride, n_rows, n_cols, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    mask = cols < n_cols
    x = tl.load(in_ptr + row * in_stride + cols, mask=mask, other=-float('inf'))
    x = x - tl.max(x, axis=0)
    num = tl.exp(x)
    den = tl.sum(num, axis=0)
    tl.store(out_ptr + row * out_stride + cols, num / den, mask=mask)

BLOCK is the next power of two at or above n_cols, because a block must have a power-of-two number of elements. The row is a block value: the max, the exponential and the sum all happen on it on chip, and the only write to DRAM is the final store.

Why the padding value is minus infinity

The masked lanes have to lose the maximum. With other=0.0 a row whose values are all negative reports a maximum of zero, every exponential is below one, and the normaliser is wrong in a way that still looks like a probability distribution. Subtracting the max also keeps every exponent at or below zero, which is the same stability argument as the scalar baseline from the first phase.

Things that break

The one-row-one-program shape assumes the row fits in a block, so a very wide row overflows the budget and the fix is a streaming loop with a running max, not a larger constant. tl.exp is the fast approximate exponential, so checks use a tolerance rather than equality. And a benchmark that times the output allocation along with the kernel reports a number that moves with the allocator. These snippets follow Triton 3.9; confirm tl.max axis semantics and the tl.exp implementation against triton.__version__ before trusting a tolerance.