Softmax: Stable Row Reductions

Core Concept

Stable Softmax subtracts each row's maximum, computes exponentials and their sum, then normalizes. Its core is two related reductions without large intermediates.

Why It Affects Performance

Softmax has modest FLOPs and is often limited by memory traffic and reduction synchronization. Separate max, exp, sum, and divide kernels repeatedly access global memory.

Example or Pseudocode

x = tl.load(row + offs, mask=mask, other=-float('inf')).to(tl.float32)
m = tl.max(x, axis=0)
z = tl.exp(x - m)
s = tl.sum(z, axis=0)
tl.store(out + offs, z / s, mask=mask)

Profiling and Performance Judgment

Compare eager, torch.compile, Triton, and library implementations using kernel count and DRAM bytes. Test irregular row lengths, extreme logits, and FP16 error.

Five-Minute Exercise

Why not compute exp(x)/sum(exp(x)) directly?