Zyphra ZYFA: FP8 FlashAttention for AMD MI355X, Up to 1.76x Faster

Zyphra has released ZYFA (Zyphra FlashAttention on AMD), a set of FP8 and BF16 prefill attention kernels for AMD Instinct MI355X GPUs. ZYFA is Zyphra’s FP8 attention engine for long-context prefill, and the code is open source under Apache 2.0. On real model activations, its FP8 kernels run 1.16-1.76x faster than AITER, AMD’s default kernel library. At head dimension 128, ZYFA’s FP8 error is 62-78% lower than AITER’s.

TL;DR

• Framework: It is a kernel library. Causal forward (prefill) attention in FP8 e4m3 and BF16, head dim 128 and MLA shape (192/128).
• Runs on: AMD Instinct MI355X only (CDNA4, gfx950), ROCm 7.2.3, PyTorch 2.12. Drop-in flash_attn_func interface.
• Performance: 1.16-1.76x AITER’s FP8 throughput on real activations; BF16 roughly ties (1.01-1.02x).
• Best: 2,672 TFLOP/s FP8 at 32K tokens, 53% of the FP8 datasheet peak (AITER: 39%).
• Worst: Only 1.16x on Mistral-7B activations; FP8 clock drops to 1,628 MHz vs AITER’s 2,108 MHz.

What is ZYFA?

ZYFA is shipped in the form of a technical paper along with code from Zyphra. It powers the attention engine of the Zyphra Optimization Engine (ZOE), which runs Zyphra Cloud. The BF16 kernels are written in HIP. The FP8 kernels are hand-scheduled assembly emitted by a Python generator.

Why does FP8 attention underperform on MI355X?

FP8 matrix instructions run at 2x the BF16 peak. But FP8 only speeds up the 2 matrix products, QKᵀ and PV. The softmax between them runs on the vector unit at the same speed in both precisions.

For 32 queries against 128 keys, FP8 matrix work takes 1,024 cycles. The exact FP8 softmax takes 896 vector cycles. So the kernel needs near-perfect overlap to keep matrix units busy. Zyphra’s thread traces show AITER’s matrix pipe idle for 31.7% of cycles.

The MI355X does not track matrix instruction completion in hardware, so instruction order must create that overlap.

How does ZYFA work?

ZYFA uses 3 techniques:

Lean P: one multiply-add instead of an exponential

An e4m3 bit pattern, read as an integer, grows by 8 each time the value doubles. So the code of 2^x is roughly 8x plus a constant. Lean P writes each FP8 probability with 1 multiply-add and 1 integer conversion. This cuts softmax vector cycles per block from 896 to 384.

Two-wave ping-pong schedule

2 waves share each SIMD. One issues matrix instructions while the other computes probabilities, then they swap. With lean P, the matrix pipe is busy 86% of steady-state cycles, up from 60% with exact P.

Per-block scaling and a rounded row sum

Each 64-key block of probabilities is scaled by a power of 2 before rounding. CDNA4’s scaled matrix instruction removes the factor at no cost. ZYFA also divides by the sum of the rounded probabilities, computed with 1 extra matrix instruction per block.

‘;});
$(‘zsp’).innerHTML=h;$(‘zsp-note’).textContent=d.note;requestAnimationFrame(function(){setTimeout(function(){$(‘zsp’).querySelectorAll(‘.bar’).forEach(function(b){b.style.width=b.dataset.w+’%’});post();},30)});}
seg(‘sp’,drawSp);
/* 5 accuracy */
var AC=[[‘Mistral-7B’,12.0,64.1,’0.22x’],[‘Qwen3-8B’,9.11,31.7,’0.38x’],[‘DeepSeek-V2-Lite’,8.36,10.2,’0.82x’]];
function drawAc(){var h=””;AC.forEach(function(r){h+=’

‘;});
$(‘zac’).innerHTML=h;requestAnimationFrame(function(){setTimeout(function(){$(‘zac’).querySelectorAll(‘.bar’).forEach(function(b){b.style.width=b.dataset.w+’%’});post();},30)});}
/* resize */
function post(){try{var h=R.offsetHeight+40;parent.postMessage({zyfaHeight:h},’*’);}catch(e){}}
window.addEventListener(‘load’,post);window.addEventListener(‘resize’,post);setTimeout(post,400);
})();