Why architecture matters here
Softmax bugs are subtle and catastrophic. A missing max subtraction produces NaN under long context. A missing fp32 accumulation produces drift under large batch. A non-fused kernel bloats memory reads. Each is fixable at the architectural level.
The architecture matters because softmax runs many billions of times in every training and serving step. Small overheads compound into large MFU losses.
With the pieces in mind, you can pick the right kernel and dtype per hardware and workload.
The architecture: every piece explained
The top strip is the algorithm. Logits x are attention scores. Max subtract shifts x by max(x) so exp inputs stay bounded. Exponentiate is elementwise exp on shifted values. Normalize divides by the sum.
The middle row is the numeric details. fp32 accumulate is required for sums even if inputs are bf16/fp16. Logit scaling (1/sqrt(d_k)) is usually baked into the attention op. Online softmax streams by keeping running max and sum, letting kernels process attention in blocks (foundation of Flash Attention). Fused attention combines QK^T + softmax + attention in one kernel to avoid materializing the score matrix.
The lower rows are safety + ops. Log-sum-exp trick computes log-softmax in a stable manner. Testing includes NaN / inf detectors and gradient smoothness checks. Ops covers dtype choice and kernel matrix per hardware.
End-to-end flow
End-to-end: a decoding kernel computes attention. Inputs are bf16. Kernel launches Flash Attention: blocks of Q, K, V loaded to SRAM; online softmax maintains running max and sum in fp32 across blocks; final normalization emits attention output in bf16. Total kernel time 12% less than unfused; memory traffic 5× less. Numerics tests show no divergence from reference softmax. Long-context runs (128k) stable.