Writing a fused softmax in Triton

01 · What eager mode actually does

Four mathematical steps, six memory round trips

Softmax over each row of an N×D matrix — the last step of every attention block, and a component of half the losses in deep learning. As math it's four steps: exponentiate, sum, divide… and one step people forget until NaNs appear: subtract the row maximum first, because exp overflows FP32 just past ~88. With the max subtracted, the largest argument becomes 0, exp(0) = 1, and nothing can overflow. The result is mathematically identical — you're multiplying numerator and denominator by exp(-m).

In PyTorch eager mode those four steps are four separate kernels: max, subtract, exp, then sum-and-divide. Each kernel reads its input from DRAM and writes its output back. For a 4096×4096 activation that's ~67 MB per pass, ~400 MB of total traffic for an operation whose inputs and outputs are 134 MB. The arithmetic is trivial — the traffic is the cost, exactly as day 23 will generalize.

02 · The kernel

One program per row, complete

@triton.jit
def fused_softmax(x_ptr, out_ptr, n_cols, BLOCK: tl.constexpr):
    row = tl.program_id(0)        # one program = one row
    offs = tl.arange(0, BLOCK)   # this row's column offsets
    mask = offs < n_cols

    # load the whole row; masked lanes read -inf, not garbage
    x = tl.load(x_ptr + row * n_cols + offs, mask=mask,
                other=-float('inf'))

    m = tl.max(x, axis=0)      # stability shift (day 16's reduction)
    e = tl.exp(x - m)             # masked lanes: exp(-inf) = 0
    denom = tl.sum(e, axis=0)  # …so they sum to nothing
    y = e / denom

    tl.store(out_ptr + row * n_cols + offs, y, mask=mask)

# launch: one program per row, BLOCK = next power of 2 ≥ n_cols
fused_softmax[(n_rows,)](x, out, n_cols, BLOCK=4096)

Fourteen lines, and every previous day is visible in it. One program per row is day 20's grammar; the other=-inf fill is the masked-load subtlety — padding with 0 would silently corrupt the max, and would survive testing forever; tl.max and tl.sum are day 16's reductions, compiler-generated; and exp(-inf) = 0 is what lets masked lanes slide through the sum untouched. The row (up to BLOCK elements) lives in registers for the entire computation — day 09's budget, spent on purpose.

03 · What we saved

Four kernels, six passes, become one of each

The fused kernel reads its row once and writes it once: 134 MB total where eager mode moved ~400 MB — a 3× cut in traffic, which in a memory-bound kernel (day 11: intensity ~0.5, deep in the amber zone) means roughly 3× wall-clock. It also launches one kernel instead of four, which matters at small sizes where day 13's ~5–10 µs launch tax dominates, and it never allocates the two intermediates — day 23's memory savings showing up early.

04 · The honest caveat

This version has one row in flight per program

Production softmax kernels usually process several rows per program, or split very long rows across programs with a two-pass online scheme — the same rescaling idea day 24's FlashAttention uses. Our single-row version is genuinely correct and genuinely fast for typical widths, and it's the version to internalize; the fancy variants are this one plus day 15-style tiling choices. If a reviewer asks "how many rows per program?", you'll know it's a resource question, not a style question: registers (day 09) versus parallelism (day 06).

Next: day 22 — autotuning: letting Triton search block sizes, warp counts, and pipeline depths instead of guessing.