Chapter 6: Kernel Fusion: Fused Softmax

Softmax appears everywhere — attention, classification heads, mixture weights — and it's the perfect vehicle for learning kernel fusion, the optimization that defines practical Triton work.

A naive softmax in PyTorch runs as a chain of separate kernels: subtract the max, exponentiate, sum, divide. Each kernel reads its input from DRAM and writes its output back to DRAM, so the data makes the slow round-trip multiple times. For a memory-bound operation, that traffic — not the arithmetic — is the entire runtime.

A fused kernel loads each row from DRAM exactly once, performs the whole computation (max → exp → sum → divide) while the row sits in fast on-chip SRAM, and writes the result once. The math is identical; the memory traffic collapses by a factor equal to the number of fused passes. That's the whole game.

Along the way you'll implement softmax as a row-wise reduction (each program owns one row), and you'll make it numerically stable by subtracting the row max before exponentiating — without which large logits overflow to infinity. By the end you'll understand why fusion is the first thing a Triton programmer reaches for.

This chapter covers:

• The Cost of DRAM Round-Trips: why unfused chains are slow
• What Is Kernel Fusion?: one load, all the work, one store
• The Row-Wise Reduction Pattern: a program per row
• Numerically Stable Softmax: subtracting the max
• Fused vs Unfused: measuring the win