AMD folds softmax into the matrix cores: MXFP4 attention hits 4.6 PFLOP/s on MI355X, with up to end-to-end speedup in video generation

Softmax in attention no longer mainly occupies the vector units; most of its work moves into the two matrix multiplies and a format conversion, so low-precision matrix throughput converts directly into attention throughput. The FlashAttention family overlapped matrix multiplies with softmax vector work, but matrix throughput has grown far faster than vector throughput across GPU generations, and quantizing operands widens that gap further, leaving softmax as the bottleneck. On an AMD Instinct MI355X the MXFP8, MXFP6 and MXFP4 kernels reach about 3.0, 4.0 and 4.6 PFLOP/s with speedups over BF16 attention, and the complete pipeline adds at most a few percentage points to the error that quantizing the operands already causes.

The row maximum is frozen and subtracted ahead of time inside the score matrix multiply, the exponential is approximated by an affine map folded into that same multiply, and the row sum also becomes a matrix multiply. The frozen maximum triggers a rollback only when a later tile overflows, and the rollback re-shifts and rescales the accumulators by the difference between the true and speculative maxima, so the steady-state loop needs no per-tile rescaling. The exponential uses the Schraudolph affine approximation followed by a two-stage conversion to quantized probabilities, with block maxima supplying both the E8M0 scales and the overflow gate; on real attention calls from Wan2.2-T2V-A14B rollback is rare, and with the gate disabled every kernel's output is NaN, showing the frozen maximum alone overflows.

On video diffusion models, pure MXFP4 causes visible drift, while a hybrid MXFP6/4 schedule brings the quality loss back close to BF16. Earlier low-bit attention work mostly left softmax unchanged or replaced only the exponential, keeping the row maximum and row sum on the vector units. Across the 248 VBench prompts of Wan2.2-T2V-A14B, MXFP8, MXFP6 and the hybrid schedule match the BF16 overall VBench score, while MXFP4 alone falls noticeably behind; on HunyuanVideo and MiniMax-H3 the hybrid schedule likewise recovers the reference appearance.

Evaluation centers on the AMD Instinct MI355X and the CDNA 4 architecture, accuracy and throughput comparisons use a fixed-shape non-causal forward attention, and video quality is assessed on three diffusion models: Wan2.2-T2V-A14B, HunyuanVideo and MiniMax-H3.

The lag cap on the frozen maximum is a heuristic chosen for the long sequence lengths of diffusion workloads, and accumulation bounds also depend on sequence length and value magnitudes, so they cannot be established from the BF16 exponent range alone; the faster maximum reduction used by MXFP4 guarantees correct gating only over a limited range, beyond which the true overshoot is lost.