Fused softmax makes quantized attention faster on AMD GPUs
Fused-Softmax for efficient quantized attention · arXiv:2611.09225
In plain words
Attention helps a model decide which parts of its input matter. On GPUs, its softmax step turns scores into weights, but that step can slow down the fast matrix calculations around it.
Fused-Softmax Attention is a training-free method that moves most softmax work onto matrix cores and corrects rare numerical overshoots.
On one long-sequence workload on AMD’s MI355X, its quantized kernels ran faster than BF16 attention; Wan2.2 video tests found that a mixed-precision schedule kept quality scores close to BF16.
The research
For long attention calculations, faster low-precision matrix hardware helps only if the softmax between matrix multiplications can keep up.
This work changes the expectation that quantizing attention alone will unlock that hardware: on its tested AMD GPU, the redesigned softmax enables higher kernel speed, while video tests show why the most aggressive precision needs care.
Fused-Softmax Attention shifts softmax work onto matrix cores and uses a rollback gate to enable faster quantized attention on AMD GPUs.
The method’s normal-input comparison measured the complete pipeline against conventional softmax on the same quantized operands.
Across three formats, it added at most 0.3 percentage points of global relative error beyond operand quantization, while the MI355X kernels ran faster than BF16 attention at the tested shape.
In Wan2.2, a hybrid precision schedule scored similarly to BF16 on VBench, an automatic evaluation that averages six video-quality dimensions. Pure MXFP4 scored lower, so the result supports that schedule in this setup rather than every quantized schedule.
The kernel speed figures come from one MI355X workload shape, and the video-quality comparison covers Wan2.2 in a specific setup. VBench scores each video on six quality dimensions, not for similarity to BF16; the latency result is for one prompt.
What the evidence shows
- The softmax bottleneck
Fused-Softmax Attention is a training-free quantized attention method for AMD CDNA 4 GPUs. It targets a hardware mismatch: attention alternates matrix multiplications with softmax calculations, while low-precision matrix operations have become faster than vector operations.
That can leave softmax as the bottleneck, preventing attention from using the matrix hardware’s potential speed.
- Move work, keep correction
The method moves most softmax work from vector units to matrix cores. It folds maximum subtraction and an approximate exponential into the score multiplication, uses quantized probabilities to produce the attention output and its normalization sum, and checks for overshoots that require correction.
The design aims to reduce the slow vector work while preserving a recovery path for numerical failures.
- Error and speed results
The error comparison separates the added cost of the redesigned softmax from error already caused by quantizing attention operands; on normal inputs, that added cost was small.
At the tested long-sequence shape, faster kernels show that matrix-core use can translate to speed on the MI355X, though this one workload does not establish speed on other hardware or shapes.
- A video-generation test
The video evaluation asks whether kernel changes matter in generation. VBench automatically scores each video on six quality dimensions and averages them on a 0–100 scale; in this Wan2.2 setup, the hybrid schedule scored 72.01 versus 71.83 for BF16 attention.
This indicates similar evaluator scores across 248 prompts, with only self-attention changed and other model operations kept in BF16.
- The quality boundary
The precision schedule matters: pure MXFP4 attention scored 71.13 on VBench versus 71.83 for BF16 attention in the same Wan2.2 evaluation, and example videos showed visible changes. The hybrid schedule recovered the visual appearance in those examples and had a similar VBench score to BF16.
The tests support a bounded conclusion about Wan2.2; they do not show that every quantized schedule preserves video quality.
Deep
- Original paper authors
- Tianxing Wu, Juuso Korhonen, Jesus Carabano Bravo
- Submitted to arXiv
- License
- Not specified by source