Learning Triton One Kernel at a Time: Softmax

The softmax function converts a vector of raw scores into a probability distribution, and is a crucial component in attention mechanisms and language modeling. To prevent numerical overflow, a common technique is to subtract the maximum value of the input vector from every element before exponentiation. A naive implementation requires three passes over the data, but a more efficient 'online softmax' algorithm fuses the max and sum steps into a single iterative pass, reducing memory reads. For integration with frameworks like PyTorch, the backward pass must be implemented manually by deriving the softmax gradient. The goal is to build an efficient softmax kernel in Triton that is compatible with PyTorch's autograd.