Fused-Softmax shifts an AI attention bottleneck onto faster GPU hardware

Tianxing Wu, Juuso Korhonen and Jesus Carabano Bravo describe the method in a preprint submitted to arXiv on November 6, 2026. The supplied record does not establish peer-reviewed publication. Their central claim is that attention can benefit more from low-precision hardware if softmax is redesigned alongside the matrix multiplications surrounding it.

Attention lets a model weigh relationships between elements of its input. It first multiplies query and key matrices to produce scores, converts those scores into weights using softmax, then multiplies the weights by a value matrix. In video generation, longer sequences representing more frames or higher resolutions make this work particularly demanding.

Modern GPUs have specialized matrix cores that handle large matrix multiplications efficiently. But softmax largely uses vector-processing instructions, whose throughput has not grown at the same pace. Reducing the precision of matrix operands makes those multiplications faster without necessarily accelerating the intervening softmax. The supposedly smaller calculation can therefore hold up the pipeline.

Fused-Softmax moves much of that intervening work onto matrix cores. One change concerns numerical stability. Conventional implementations repeatedly update a reference maximum and rescale accumulated results when it grows. The new method provisionally freezes that reference, allowing its subtraction to be incorporated into a matrix multiplication and avoiding routine rescaling. If later scores exceed a safety threshold, a recovery procedure updates the reference and corrects the accumulated quantities.

A second change replaces exponential evaluation with the Schraudolph approximation, which exploits how floating-point numbers encode exponents and fractional values. Its linear transformation is folded into the score matrix multiplication. A two-stage conversion then constructs an intermediate BF16 value and converts it into the lower-precision weights needed for the next multiplication. Finally, the method calculates the normalization sum through another matrix multiplication, using the same quantized weights as the output calculation.

The approach requires no model retraining, but it is not exact attention. The researchers separated its numerical effects from those of quantization using synthetic inputs across ten random seeds. Against conventional BF16 attention, the complete pipeline produced relative output errors of 5.69% for the MXFP8 configuration, 6.46% for MXFP6 and 20.53% for MXFP4. Most of that error already existed when the matrix operands were quantized: adding Fused-Softmax increased it by no more than 0.3 percentage points. That distinction matters—the softmax changes add relatively little error, but aggressive quantization still carries a substantial cost.

On a single AMD Instinct MI355X, the three attention implementations ran approximately 2.1, 2.8 and 3.2 times as fast as the BF16 baseline at the reported long-sequence test shape. These are attention-kernel gains, not equivalent improvements across a complete model. The configuration names also simplify mixed formats: the MXFP4 implementation retains six-bit attention weights.

To test practical consequences, the team generated videos with Wan2.2-T2V-A14B using 248 prompts, 720p resolution, 81 frames and 20 denoising steps. Quality experiments changed only self-attention, retaining BF16 for other model operations. They also changed the scheduler’s flow-shift setting from the benchmark reference’s default to the model developer’s recommended value, because the default degraded even baseline videos.

Pure MXFP4 produced visible changes, including a missing tongue in a grooming cat and altered scenery around a motorcycle. A hybrid schedule used MXFP6 for the first and last three denoising steps and MXFP4 for the fourteen between them. Its aggregate VBench score was 72.01, compared with 71.83 for BF16 and 71.13 for pure MXFP4. These automatic scores do not establish identical videos or a meaningful quality improvement; separate frame-comparison metrics showed differences from the baseline.

For the motorcycle prompt, the hybrid reduced generation time from 173.3 to 84.4 seconds in an optimized model configuration, a 2.05-fold speedup. Both sides of that comparison used the same additional optimizations, including other quantized matrix operations and parallel video decoding. Timing covered one prompt, with the slowest of five measured runs reported after warm-up.

The limits extend beyond video quality. Comparisons with competing methods used published results from different GPUs, rather than a controlled same-hardware experiment. The overflow threshold is heuristic, and the fastest four-bit implementation has an extreme-score range where its shortcut can miss a recovery trigger, although the authors observed no such failures in evaluated calls. Tests on two other video models were illustrative rather than equally extensive benchmark evaluations.

The authors report that the kernels ship in AMD’s open-source AITER library. The broader result is a hardware-design lesson: smaller numbers alone are insufficient when another stage becomes the bottleneck. Reassigning that stage to faster execution units can unlock gains, provided its numerical compromises are checked against the application.