softmax

One program per row, the row in registers, max then exp then sum then divide, one read and one write of HBM. The runnable kernel with its launch, the numerics (subtract the max, accumulate in fp32), the block-size rule and the wide-row limit, and the arithmetic that says the kernel is done at 90% of copy speed.

Four lines of arithmetic with three ways to get it wrong: dividing by a temperature of zero, exponentiating without subtracting the maximum, and applying the filters in an order that changes the result. The implementation, the verification against the analytic distribution, and the overflow that is one logit away.