Softmax over large rows is memory-bound: almost no arithmetic, lots of bytes. That makes it a clean case study for the single most useful habit in kernel work — decide what the hardware limit is before you benchmark anything, so you know whether your number is good or just a number.
This post walks through the fused row-softmax in my
Triton LLM inference kernel lab:
the roofline math, the measured results on H100, and the shapes where the
kernel loses to torch.softmax — because knowing when not to use a custom
kernel is half the point of writing one.
One Triton program per row. Each program loads its row into SRAM once, does the numerically stable max-subtract softmax, and writes the row back once:
row = tl.load(input_ptr + offsets, mask=mask, other=-float("inf"))
row = row - tl.max(row, axis=0) # stability: shift by the row max
num = tl.exp(row)
out = num / tl.sum(num, axis=0)
tl.store(output_ptr + offsets, out, mask=mask)
The whole design fits in one sentence: each element crosses the HBM bus exactly twice — one read, one write. Everything else (max, exp, sum, divide) happens in registers/SRAM and is effectively free at these intensities.
Take the largest benchmark shape, 8192 × 4096 in fp16:
2 (read+write) × 8192 × 4096 × 2 B = 134.2 MB134.2 MB ÷ 3.35 TB/s = 40 µsThat 40 µs is the roofline. No scheduling trick, no occupancy tuning, no rewrite in raw CUDA gets a memory-bound kernel below it. The only question a benchmark can answer is how close you get.
| shape (rows × cols) | latency | achieved bandwidth | vs torch.softmax | max abs err |
|---|---|---|---|---|
| 8192 × 4096 | 0.048 ms | 2.8 TB/s (~83% of peak) | 2.78× | 3.8e-6 |
| 4096 × 2048 | 0.024 ms | — | 1.11× | 3.8e-6 |
48 µs measured against a 40 µs roofline — ~83% of theoretical peak. The remaining 17% is the usual tax: launch latency, imperfect coalescing on row boundaries, and DRAM pages that don't stream as cleanly as the back of the envelope assumes. For a readable, unautotuned kernel, I'll take it.
The speedup over torch.softmax is consistent with the traffic argument: a
softmax that isn't fused end-to-end has to touch the row more than twice
(max pass, exp/sum pass, normalize pass), and at 134 MB per pass the extra
trips are the whole game. Correctness is checked against a PyTorch fp32
reference on every run — max error stays at fp16 rounding scale (3.8e-6),
including non-power-of-two column counts.
On rows of ≤1024 columns, this kernel is slower than torch.softmax.
At 4096 × 2048 the win has already shrunk to 1.11×, and below that it
inverts. The reason isn't subtle: a small row is a small amount of actual
work, and the fixed costs — kernel launch, program scheduling — dominate it.
PyTorch's dispatcher is amortizing those costs across a much better-tuned
path than my one-program-per-row layout, which at small sizes launches a lot
of programs that each do almost nothing.
The fused kernel wins exactly where the design says it should: large,
bandwidth-bound reductions. Everywhere else, the right engineering decision
is torch.softmax. A kernel lab that only reports the flattering shapes
isn't a lab; it's marketing.
Attention scores, logit normalization, and sampling all hit softmax-shaped reductions, and in decode-heavy serving the row sizes swing wildly with context length. A serving engine that picks the fused path for long-context prefill and the vendor path for short rows gets the best of both — but you can only build that dispatch rule if you've measured where the crossover is.
Everything here is reproducible with the repo's harness (50 warmup / 200 timed iterations, latency + bandwidth + max-error per shape):
python -m triton_llm_kernel_lab.bench --kernel softmax
Next on the list: Nsight Compute traces for this roofline (the repo's
docs/profiling.md has the workflow), and a flash-decoding split-KV path for
the attention kernel — where the same "count the bytes first" discipline says
the current q_len=1 decode shape is leaving most of the machine idle.