ZYFA: Fast and Accurate FP8 Attention on AMD MI355X

wpnews.pro

Zyphra presents ZYFA (Zyphra FlashAttention on AMD) - FP8 and BF16 prefill attention kernels for AMD Instinct MI355X GPUs. ZYFA is part of the attention engine of our Zyphra Optimization Engine (ZOE) on AMD MI355X GPUs, powering our inference cloud and on-prem deployments.

We show that FP8 matrix instructions offer twice the peak throughput of BF16, but attention kernels fall well short of a 2x speedup in practice. On the MI355X at 32K tokens, the FP8 attention kernel of AITER (AMD’s library of inference kernels) reaches 39% of the FP8 peak, while its BF16 kernel reaches 49% of the BF16 peak.

FP8 speeds up only the two matrix multiplications of attention, so the softmax between them becomes the bottleneck. ZYFA’s FP8 kernels shrink the softmax, overlap it with the matrix multiplications, and keep the 8-bit probabilities accurate.

On attention inputs captured from three models and timed at sustained clock speeds (the steady clock under continuous load), they achieve 1.16–1.76× the throughput of AITER’s default kernels (Figure 1), and at head dimension 128 their numerical error is 62–78% lower than that of AITER’s FP8 kernel.

This post summarizes our paper, and the source code is available at github.com/Zyphra/zyfa.

Key Results #

Speed on real activations: On activations captured fromMistral-7B ,Qwen3-8B , andDeepSeek-V2-Lite at 16K tokens and timed at sustained clocks, ZYFA’s FP8 kernels achieve1.16–1.76× the throughput of AITER’s default kernels, and ZYFA’s BF16 kernels achieve 1.01–1.02×1 . The FP8 speedup is smaller than on Gaussian inputs because the clock of the MI355X depends on the data. AITER lets more probabilities underflow to zero, and zeros switch fewer bits, so AITER’s kernel runs at a higher clock.
Speed on kernel benchmarks: On Gaussian inputs at 32K tokens, ZYFA’s FP8 kernels achieve 1.26–1.37× the throughput of AITER’s at head dimension 128 and1.99× at the head shape of DeepSeek’smulti-head latent attention (MLA) . At head dimension 128 the FP8 kernel reaches53% of the FP8 datasheet peak.
Kernel accuracy: At head dimension 128 the error of ZYFA’s FP8 kernels is within a few percent of that of AMD’sComposable Kernel (CK) andTriton kernels and62–78% lower than that of AITER’s.
Full-model accuracy: Most of the FP8 error in a full model comes from quantizing the queries, keys, and values. On Mistral-7B aHadamard rotation of Q and Khalves this error and roughly halves the KL divergence from the BF16 model, at almost no extra time2 .

We describe our methodology and its limitations at the end of this post. In short, all results come from one MI355X node, the kernels cover causal prefill only, and both our kernels and the baselines change quickly, so all results are point-in-time.

What this post covers #

Constraints of FP8 Attention : the vector-pipe cost of the softmax, the software ordering of matrix instructions, and the precision and underflow of FP8 probabilities.
Kernel Design : lean P, the two-wave ping-pong schedule, and per-block scaling with normalization by the rounded row sum.
Throughput andNumerical Accuracy : captured activations, Gaussian inputs, the dependence of the clock on operand values, and FP8 error in full models.
Remaining Performance Gaps and Recommendations ,Future Work , andMethodology and Limitations .
Appendix A : tables of the values plotted in the figures.

Introduction #

Prefill, the forward pass over the prompt before a model generates its first token, spends a growing share of its time in attention as the context grows. FlashAttention keeps the score matrix in on-chip memory instead of high-bandwidth memory (HBM), so the bytes that attention moves to and from HBM grow linearly with the sequence length while its arithmetic grows quadratically. Prefill attention at long contexts is therefore compute-bound, and its speed depends on how well the kernel keeps the GPU’s matrix units busy3.

The MI355X executes 8-bit floating-point (FP8) matrix instructions at twice the peak rate of BF16, so FP8 attention could in principle take half the time of BF16 attention. In practice the gain is much smaller, and AITER’s FP8 kernel reaches a lower share of its potential peak than its BF16 kernel. FP8 also raises a question of accuracy, since attention is a very numerically sensitive block. DeepSeek-V3, for example, trains its linear layers in FP8 but keeps attention in higher precision.

ZYFA consists of causal forward attention kernels in BF16 and FP8 for the MI355X (CDNA4, gfx950). The kernels support head dimension 128 and the head shape of MLA in DeepSeek-V2, with a query-key dimension of 192 and a value dimension of 128.4 Constraints of FP8 Attention on the MI355X describes the three obstacles that limit FP8 attention, and Kernel Design describes the technique in ZYFA’s FP8 kernels that addresses each one. We provide the full analysis in our technical report.

Constraints of FP8 Attention on the MI355X #

Vector-Pipe Cost of the Softmax

FP8 halves the cycles of the two matrix products of attention, QKᵀ and PV, which run on the matrix pipe5. The softmax between them runs on the vector pipe, which executes elementwise arithmetic at the same speed in FP8 and BF16, and an FP8 kernel must also convert every probability to 8 bits. Both pipes execute the instructions of waves, groups of 64 threads that issue instructions together. For one wave computing 32 queries against a block of 128 keys, the vector pipe needs 0.44× the cycles of the matrix pipe in BF16 but 0.94× in FP8 (Figure 2). An FP8 kernel therefore keeps the matrix pipe busy only if the softmax overlaps the matrix products almost perfectly.

Software Ordering of Matrix Instructions

On the MI355X a matrix instruction (MFMA) keeps running while the wave that issued it goes on to other instructions, and the hardware does not track when the MFMA finishes. NVIDIA’s Hopper and Blackwell GPUs track the completion of asynchronous matrix instructions in hardware. On the MI355X the order of the instructions alone must overlap the softmax with the MFMAs and keep every instruction from reading an MFMA result too early. In thread traces of AITER’s FP8 kernel, the matrix pipe is idle for 31.7% of the cycles while the softmax runs.

Precision and Underflow of FP8 Probabilities

The FP8 format of the probabilities, e4m3, keeps all 3 mantissa bits only for values of at least 2⁻⁶, the normal values. Smaller values lose precision, and values below 2⁻¹⁰ underflow to zero. Most probabilities of a long row lie far below the largest one. In layer 20 of Mistral-7B at 16K tokens, the largest probability of a 64-key block lies more than 8 binades (factors of two) below the row maximum for 91% of the pairs of a row and a block. Underflow also biases the output if the row sum adds up the unrounded probabilities, as in AITER’s FP8 kernel, because every probability that underflows to zero still counts in the row sum6.

Kernel Design #

ZYFA’s FP8 kernels address each obstacle with one technique.

Lean P: Linear Approximation of the FP8 Code

A standard FP8 kernel evaluates each probability p = 2ˣ, with x the score minus the row maximum in base-2 units, with an exponential instruction and then rounds it to e4m3. We call this exact P, and it takes 14 vector-pipe cycles per score. Lean P uses the fact that the 8-bit pattern of a positive e4m3 number, read as an integer and called its code, grows by 8 each time the value doubles7. The code of 2ˣ is therefore close to 8x plus a constant, so lean P computes the code with one multiply-add and one conversion to an integer, in 6 cycles per score.

Lean P is an approximation. On the 68 captured layers at head dimension 128 it raises the error of an attention call to 1.22–1.36× that of exact P, which is still 0.26× the error of AITER’s kernel in geometric mean. ZYFA uses lean P by default at head dimension 128 and exact P at the MLA shape.

Two-Wave Ping-Pong Schedule in Assembly

Each SIMD, one of the four execution units of a compute unit, runs two waves of the FP8 kernel. In each phase one wave issues the matrix instructions of its block while the other computes FP8 probabilities, and the two waves swap roles after a barrier (Figure 3). FlashAttention-3 calls this arrangement ping-pong scheduling. Each phase lasts until both waves reach the barrier, so a phase takes as long as the slower of its two waves. With exact P the softmax wave arrives last, and the matrix pipe is busy for 60% of the cycles. With lean P the softmax wave arrives first, and the matrix pipe is busy for 86%.

Because both waves issue from the same SIMD, the overlap depends on the exact placement of every instruction, so a Python generator emits the FP8 kernels as assembly, which fixes the position of every instruction. On captured activations a version of the same kernel written in HIP, AMD’s C++ kernel language, needs 43–46% more cycles, and in the paper’s step-by-step comparison the assembly schedule gives the largest gain of any design step.

Per-Block Scaling and Normalization by the Rounded Row Sum

The FP8 kernels multiply the probabilities of each 64-key block by a power of two so that the largest probability of the block lies close to 1. The scaled FP8 matrix instruction of CDNA4 divides the same power of two back out in the PV product at no extra cost. The kernels also divide by the sum of the rounded probabilities, which one extra MFMA per block computes from the same values that multiply V. Each output is then a weighted average with weights that sum to one, and the extra MFMA adds 6% to the matrix cycles per block (Figure 4).

Throughput #

Throughput on Captured Activations at Sustained Clocks

We time each kernel on the queries, keys, and values that one layer each of Mistral-7B, Qwen3-8B, and DeepSeek-V2-Lite produces at 16K tokens, which we call captured activations. Under sustained load the MI355X lowers its shader clock (SCLK), so each run launches one kernel back to back for seconds and is timed only after the clock has settled. ZYFA’s FP8 kernels achieve 1.16×, 1.17×, and 1.76× the throughput of AITER’s on the three layers (Figure 1). The ratio is largest on DeepSeek-V2-Lite because CK, AITER’s default FP8 kernel at the MLA shape, issues an older matrix instruction with half the throughput per clock. Inside full forward passes at 16K tokens, the median FP8 ratio over all layers is 1.26× on Mistral-7B and 1.25× on Qwen3-8B.

Throughput on Gaussian Inputs

With kernel-only timing on Gaussian inputs, as in typical kernel benchmarks, ZYFA’s FP8 kernels achieve a higher throughput than AITER’s at every sequence length from 1K to 32K tokens, and ZYFA’s BF16 kernels achieve 1.00–1.03× the throughput of AITER’s at 16K and 32K tokens (Figure 5). In BF16 the matrix products take more than twice as many cycles as the softmax, so they limit the kernel time of both libraries. In FP8 with exact P the softmax wave reaches each barrier after the MFMA wave, so the softmax limits the kernel time instead.

Dependence of the Shader Clock on Operand Values

The sustained clock of the MI355X depends on the values in its arithmetic units. Each bit that switches between 0 and 1 draws current, so the clock falls as more operand bits change from one instruction to the next (Figure 6)8. A loop of FP8 MFMAs runs at 2.34 GHz with constant operands and at 2.07 GHz with operands that change every instruction.

On the captured Mistral-7B layer, ZYFA’s FP8 kernel needs fewer cycles than AITER’s but runs at 1,628 MHz, against 2,108 MHz for AITER (Figure 7). AITER writes subnormal values or zeros for the blocks far below the row maximum, and these values switch few bits, while ZYFA’s per-block scaling keeps them normal. On Gaussian inputs, which make attention nearly uniform, both kernels write only normal values, and ZYFA’s speedup on this layer grows from 1.16× to 1.30×. Throughput on random inputs therefore does not reliably predict throughput on real model activations.

Numerical Accuracy #

Error of a Single Attention Call

To isolate the error that a kernel itself adds, we compare its output with fp64 attention on the same quantized inputs. In FP8 at head dimension 128 the errors of ZYFA, CK, and Triton differ by a few percent, while ZYFA’s error is 0.22× and 0.38× that of AITER’s FP8 kernel on Mistral-7B and Qwen3-8B (Figure 8). Most of AITER’s extra error comes from dividing by the sum of the unrounded probabilities. In a full Mistral-7B model, the KL divergence from a model with fp64 attention is 3.4 × 10⁻³ with ZYFA’s FP8 kernel and 5.9 × 10⁻³ with AITER’s.

FP8 Error in Full-Model Inference

Most of the FP8 error in a full model comes from quantizing the inputs (Figure 9). On Mistral-7B, quantizing Q, K, and V adds 8 times the error of quantizing the probabilities with lean P, and Q and K add more error than V. A Hadamard rotation of Q and K over the head dimension leaves QKᵀ unchanged but spreads large channels over all channels, which roughly halves the input error. On Mistral-7B the rotation lowers the KL divergence from the BF16 model from 45.7 to 21.9 times the upper end of the BF16 noise band9, and attention with the rotation takes 1.001× the all-FP8 time. Keeping the most sensitive heads in BF16 is a worse use of time. On Mistral-7B and Qwen3-8B even 90% of heads in BF16 does not reach the noise band, and that mix takes longer than all-BF16 attention.

Remaining Performance Gaps and Recommendations #

At 32K tokens on Gaussian inputs, ZYFA’s FP8 kernel at head dimension 128 reaches 53% of the datasheet peak10. Three gaps remain. The sustained clock stays well below the 2.4 GHz that the datasheet peak assumes, the matrix pipe idles at the transitions between query tiles, and with lean P the MFMA wave also computes the row maximum and the rescale of the output, so the MFMA wave now takes longer than the softmax wave and limits the length of each phase. For kernel developers we recommend the following:

Measure on real activations at sustained clocks , since the clock speed is dependent on input data.
Target the quantization of the inputs , for example with a Hadamard rotation of Q and K, before the quantization of the probabilities.
Reduce the vector work per score before adding matrix throughput. A faster matrix format alone does not help while the softmax is the bottleneck.
Keep packed FP32 instructions out of a softmax that runs next to MFMAs. A packed instruction applies one operation to two values per thread, and a packed FP32 instruction stalls while an MFMA executes on the same SIMD.

Future Work #

We will continue to improve ZYFA with a focus on the following areas:

Closing the remaining gaps above.
Extending lean P to other formats whose code grows by a fixed step each time the value doubles, such as e5m2 and FP4.
Supporting decoding, paged KV caches, and batches of variable length.

The source code is available at github.com/Zyphra/zyfa, and the paper is on arXiv.

Methodology and Limitations #

All experiments run on one node with eight MI355X GPUs, with the default automatic clock policy and a total board power (TBP) limit of 1,400 W, on ROCm 7.2.3. Our baselines are the AITER 0.1.19 default kernel selections.
Kernel-only timing on Gaussian inputs reports the mean device time of 100 launches after 500 warm-up launches. Sustained timing runs 1 s of warm-in launches and then times a window of at least 3 s.
The full-model study covers three models, PG-19 text, and no downstream tasks.
The kernels cover the causal forward pass for equal query and key lengths. They do not cover decoding, paged KV caches, or batches of variable length.
The data-dependent clock may differ with other firmware or power settings, and our kernels and the baselines change quickly, so all results are point-in-time.

Appendix A: Data Tables #

The tables below contain the values plotted in the figures. Find the full results in our technical report.

A1. Captured activations at sustained clocks (Figures 1 and 7)

One layer per model at S = 16K (Mistral-7B layer 20, Qwen3-8B layer 18, DeepSeek-V2-Lite layer 13). The time ratio is AITER’s time divided by ZYFA’s time, so values above 1 mean that ZYFA takes less time. The Gaussian ratio is the same ratio on Gaussian inputs of the same shape.

Model	Precision	ZYFA (TFLOP/s)	AITER (TFLOP/s)	Time ratio	Gaussian ratio
Mistral-7B	FP8	2,537	2,191	1.16	1.30
Qwen3-8B	FP8	2,495	2,130	1.17	1.29
DeepSeek-V2-Lite	FP8	2,337	1,326	1.76	1.84
Mistral-7B	BF16	1,223	1,196	1.02	1.03
Qwen3-8B	BF16	1,202	1,191	1.01	1.02
DeepSeek-V2-Lite	BF16	1,285	1,276	1.01	1.01

Kernel times and sustained shader clocks (SCLK) of the same runs.

| Model | Precision | ZYFA (µs) | AITER (µs) | ZYFA SCLK (MHz) | AITER SCLK (MHz) | 
|---|---|---|---|---|---|

| Mistral-7B | FP8 | 867 | 1,003 | 1,628 | 2,108 | | Qwen3-8B | FP8 | 881 | 1,032 | 1,629 | 2,048 | | DeepSeek-V2-Lite | FP8 | 588 | 1,036 | 1,742 | 2,219 | | Mistral-7B | BF16 | 1,798 | 1,838 | 1,590 | 1,446 | | Qwen3-8B | BF16 | 1,830 | 1,847 | 1,586 | 1,438 | | DeepSeek-V2-Lite | BF16 | 1,069 | 1,077 | 1,449 | 1,403 |

A2. Gaussian kernel benchmarks (Figure 5)

Throughput of ZYFA divided by that of AITER’s default kernel, kernel-only on Gaussian inputs at B = 1. At the MLA shape (query-key dimension 192) the FP8 baseline is CK. The last column gives ZYFA’s and AITER’s throughput at 32K.

Attention	Precision	Head shape	2K	16K	32K	TFLOP/s at 32K (ZYFA / AITER)
MHA	FP8	128	1.44	1.34	1.37	2,672 / 1,944
MHA	FP8	MLA (192)	2.02	1.89	1.99	2,464 / 1,237
GQA	FP8	128	1.42	1.28	1.26	2,563 / 2,027
GQA	FP8	MLA (192)	1.94	1.82	1.88	2,374 / 1,262
MHA	BF16	128	1.05	1.03	1.02	1,265 / 1,236
MHA	BF16	MLA (192)	1.04	1.01	1.00	1,330 / 1,329
GQA	BF16	128	1.07	1.02	1.02	1,310 / 1,283
GQA	BF16	MLA (192)	1.07	1.00	1.00	1,376 / 1,378

A3. Error of one attention call (Figure 8)

Relative RMS error against fp64 attention on the same quantized inputs, mean over the layers of each model. In parentheses, each kernel’s error divided by AITER’s error in the same layer, averaged over layers, so a ratio can differ from the ratio of the two means. AITER’s Triton kernel does not support FP8 at the MLA shape.

|---|---|---|---|---|---|
| Mistral-7B | FP8 | 1.20e-02 (0.22) | 6.41e-02 | 1.16e-02 (0.22) | 1.14e-02 (0.21) | 
| Qwen3-8B | FP8 | 9.11e-03 (0.38) | 3.17e-02 | 8.94e-03 (0.37) | 8.86e-03 (0.37) | 
| DeepSeek-V2-Lite | FP8 | 8.36e-03 (0.82) | 1.02e-02 | 1.02e-02 (1.00) | — | 
| Mistral-7B | BF16 | 1.81e-03 (1.02) | 1.77e-03 | 1.98e-03 (1.12) | 1.81e-03 (1.02) | 
| Qwen3-8B | BF16 | 1.75e-03 (1.01) | 1.73e-03 | 1.94e-03 (1.12) | 1.75e-03 (1.01) | 
| DeepSeek-V2-Lite | BF16 | 1.94e-03 (1.00) | 1.94e-03 | 1.96e-03 (1.01) | 1.77e-03 (0.91) |

A4. Error of one FP8 call by stage (Figure 9)

Relative RMS error of ZYFA’s FP8 kernel, split into input quantization (Q, K, and V), probability quantization (P), and the rest of the kernel’s arithmetic. ZYFA uses lean P on Mistral-7B and Qwen3-8B and exact P at the MLA shape of DeepSeek-V2-Lite.

|---|---|---|---|---|---|
| Mistral-7B | None | 8.3e-02 | 1.0e-02 | 5.0e-03 | 8.3e-02 | 
| Mistral-7B | Hadamard | 4.4e-02 | 1.0e-02 | 4.8e-03 | 4.5e-02 | 
| Qwen3-8B | None | 1.1e-01 | 1.0e-02 | 4.7e-03 | 1.1e-01 | 
| Qwen3-8B | Hadamard | 5.6e-02 | 9.9e-03 | 4.5e-03 | 5.7e-02 | 
| DeepSeek-V2-Lite | None | 9.5e-02 | 8.1e-03 | 4.4e-03 | 9.5e-02 | 
| DeepSeek-V2-Lite | Hadamard | 5.6e-02 | 8.1e-03 | 4.3e-03 | 5.6e-02 |

Footnotes #

1 Mistral-7B-v0.3 layer 20, Qwen3-8B layer 18, and DeepSeek-V2-Lite layer 13, with the queries, keys, and values after the rotary position embedding (post-RoPE) on WikiText-2 text at 16K tokens. The baseline is the kernel that AITER 0.1.19 selects by default, which is Composable Kernel (CK) for FP8 at the MLA shape. ↩

2 The rotation is fused into the pass that quantizes Q, and K and V are quantized as they are written to the KV cache. On Mistral-7B, attention with the rotation takes 1.001× the time of all-FP8 attention. ↩

3 Our Wave 2 inference post measures the rooflines of the MI355X, which give the highest throughput that a kernel can reach at each arithmetic intensity (FLOPs per byte moved to or from HBM), and the peak matrix throughput of each data type. ↩

4 MLA caches one compressed latent vector per token instead of per-head keys and values. Prefill expands the latent vectors into per-head keys and values and runs standard attention with these dimensions. Decoding uses a different form of MLA that ZYFA does not cover. ↩

5 The matrix pipe is the matrix unit of a SIMD, and the vector pipe is its vector unit. Each compute unit of the MI355X has four SIMDs. The matrix pipe executes matrix instructions, and the vector pipe executes elementwise arithmetic such as the exponentials and FP8 conversions of the softmax. ↩

6 Dividing by the unrounded sum treats every probability that underflowed to zero as a key whose value vector is zero, so each output shrinks slightly toward zero. Dividing by the sum of the rounded probabilities instead gives a weighted average, with weights that sum to one, of the values whose probabilities survive rounding. ↩

7 For example, the codes of 1, 2, and 4 are 56, 64, and 72, and the code 60, halfway between 56 and 64, stands for 1.5. ↩ 8 Each bit that switches charges or discharges a small capacitance, so the current that the chip draws grows with the number of bits that change per cycle. Total board power stays below its 1,400 W limit in these runs, and the telemetry points to a fast-acting limit on current rather than on average power. ↩

9 The BF16 noise band is the range of KL divergences no larger than that between two models with equally correct BF16 kernels, ZYFA’s and AITER’s, which round in different orders. ↩

10 The datasheet peak assumes that a matrix instruction issues on every cycle at the 2.4 GHz peak clock. Our Wave 2 inference post compares the datasheet peak of each data type with the throughput that the MI355X reaches in practice. ↩

source & further reading

zyphra.com — original article ZHA: Zyphra Hierarchical Attention