The safe softmax alters the naive softmax (how softmax is typically conceptualized) by first finding first maximum element, then scaling all elements by it so as to reduce the chance of overflow and underflow. The safe softmax uses more memory accesses. Online softmax reduces the number memory accesses used by soft softmax back to the same number used in the naive softmax.
Can you remember the 3 implementations?
Online softmax is a key ingredient of Flash Attention.
