Softmax Kernel

Numerically stable definition

$y_i=\frac{e^{x_i-m}}{{\sum }_j{e}^{x_j-m}},m={\max }_j{x}_j$

Subtracting the max is needed for stability as $e^x$ overflows FP32 above $x\approx 88$, and attention logits can often exceed that.

Pass structure

Version	Passes over row	HBM traffic
Three-pass (max, sum, normalize)	3 reads + 1 write	$4N$ elements
Two-pass (fused max+sum via online)	1 read + 1 read + 1 write	$3N$
One-pass (row fits in registers/shared)	1 read + 1 write	$2N$

A row that fits on-chip should always be one-pass: load once, keep in registers, reduce, write. Lower bound is $2N$, softmax can never do better, so its arithmetic intensity is $O(1)$ and it is always memory-bound.

Online (streaming) softmax

Process the row in blocks, maintaining a running max $m$ and running sum $\ell$. On a new block with local max $m^{\prime }$ and local sum ${\ell }^{\prime }$:

$m_{\text{new}}=\max (m,m^{\prime })$ ${\ell }_{\text{new}}=\ell \cdot e^{m-m_{\text{new}}}+{\ell }^{\prime }\cdot e^{m^{\prime }-m_{\text{new}}}$

The correction factor $e^{m-m_{\text{new}}}$ retroactively rescales everything accumulated so far. This makes max and sum computable in a single pass, and the operator is associative, so it also parallelizes as a tree reduction over $(m,\ell )$ pairs (Reduction Kernels).

# one block of a Triton softmax row
m_new = tl.maximum(m, tl.max(x, 0))
alpha = tl.exp(m - m_new)
l     = l * alpha + tl.sum(tl.exp(x - m_new), 0)
acc   = acc * alpha + ...      # any accumulator must be rescaled too
m     = m_new

Why this matters for attention

FlashAttention never materializes the $N\times N$ score matrix. It walks K/V tiles, and for each tile it needs softmax over scores it has not finished seeing. Online softmax lets it rescale the running output accumulator $O$ by the same $\alpha =e^{m-m_{\text{new}}}$:

$O\leftarrow \alpha O+\tilde{P}{V}_{\text{tile}}$

with the single division by $\ell$ deferred to the end. Memory drops from $O(N^2)$ to $O(N)$, turning attention from HBM-bound to compute-bound.

Implementation notes

• One row per CTA (or per warp for short rows); use __shfl_down_sync for the intra-warp max/sum.
• Fuse whatever surrounds it: mask, scale, dropout, top-k: every fused elementwise op is $2N$ bytes saved (Kernel Fusion).
• Vocabulary-sized softmax (e.g. 128K logits) at decode is a real cost; fuse it with the sampling/top-k step.
• exp2f is cheaper than expf; fold ${\log }_{2}e$ into the scale factor so you compute $2^{(x-m)\cdot s{\log }_{2}e}$.

Tip

Any “reduce then normalize” pattern (softmax, layernorm, RMSNorm, logsumexp) has an online single-pass form. Knowing the rescaling identity is the standard attention-kernel interview question.

FlashAttention · Reduction Kernels · Kernel Fusion · Arithmetic Intensity · Triton Programming Model