In large AI models, the attention mechanism is a key bottleneck. Modern GPUs have very fast matrix multipliers but slower vector arithmetic, and the softmax step in attention relies on vector operations, creating a mismatch.
Fused-Softmax Attention rewrites softmax so its core work, finding the maximum, computing exponentials, and summing probabilities, is done using the fast matrix cores instead of the slow vector units. It approximates the exponential with a fast method and speculatively fixes the maximum to avoid repeated rescaling, with a fallback if the guess is wrong.
On an AMD MI355X GPU, this achieves speeds over three times faster than standard half-precision attention, while adding less error than the quantization of the matrix values themselves. The technique speeds up real video-generation models by about two times without visible quality loss, and is now shipping in AMD's open-source AI library.
