Zyphra releases open-source FP8 attention kernels for AMD MI355X

Zyphra has released ZYFA, an open-source library of FP8 and BF16 prefill attention kernels for AMD Instinct MI355X GPUs. The Apache 2.0 project targets causal forward attention for long-context workloads, with support for FP8 e4m3, BF16, head dimension 128 and MLA shapes of 192/128. It runs on MI355X hardware with ROCm 7.2.3 and PyTorch 2.12, and provides a drop-in flash_attn_func interface.

On real model activations, Zyphra reports FP8 speedups of 1.16x to 1.76x over AITER, AMD’s default kernel library. BF16 performance is approximately comparable, at 1.01x to 1.02x. At 32K tokens, the head-dimension-128 kernel reached 2,672 TFLOP/s, or 53% of the FP8 datasheet peak, compared with 1,944 TFLOP/s for AITER.

At head dimension 128, ZYFA’s FP8 error was 62-78% lower than AITER’s in the reported tests. The implementation replaces the usual softmax exponential with a multiply-add approximation, uses a two-wave scheduling scheme and applies per-block scaling. Results vary by workload: measured gains were only 1.16x on Mistral-7B activations, while the implementation supports prefill only and is limited to MI355X.