Zyphra ZYFA: 1.76x FP8 FlashAttention Speedup on AMD MI355X

Attention computation on AMD Instinct MI355X has long relied on AMD's own AITER library, and Zyphra's newly released open-source project ZYFA offers a more competitive alternative. This set of FP8 and BF16 prefill attention kernels optimized for MI355X achieves 1.16 to 1.76 times the speed of AITER on real model activations, while reducing error by 62% to 78% in the head dimension 128 scenario. For teams focused on inference performance in the AMD ecosystem, this is a new variable worth evaluating.

Why prefill attention deserves dedicated optimization

Large model inference is divided into two stages: prefill and decode, and their performance bottlenecks are completely different. The decode stage is limited by memory bandwidth and is usually memory-bound; the prefill stage must process the entire input sequence at once, has high computational density, and is compute-bound. The FlashAttention series of work focuses on IO and computational efficiency in prefill scenarios, and FP8 further unleashes the throughput potential of matrix multiplication.

As AMD's high-end accelerator card for AI training and inference, MI355X supports low-precision formats at the hardware level, but whether the software stack can deliver this compute power depends on the quality of kernel implementation. ZYFA's entry point is clear: it does not build a general framework, but pushes prefill attention, a single high-frequency hotspot, to the limit. This strategy of breakthrough at a single point often proves more effective than broad, all-encompassing solutions in cost-sensitive inference deployment scenarios.

The core idea of FlashAttention is IO-aware: the bandwidth of GPU on-chip SRAM (shared memory/registers) is far higher than that of HBM memory bandwidth, while standard attention implementations need to repeatedly write and read Q, K, V and the intermediate attention score matrix to and from HBM. FlashAttention uses tiling to split Q/K/V into small blocks loaded into on-chip memory for computation, and uses the online softmax technique to accumulate results block by block without saving the complete score matrix, thereby reducing memory read/write volume from O(N²) to nearly O(N). For prefill, a long-sequence and computationally intensive scenario, reducing HBM round trips often improves end-to-end speed more than reducing floating-point operations, which is also the common starting point for all prefill attention kernel optimization.

ZYFA's performance and accuracy data

According to results published by Zyphra, ZYFA's FP8 kernel achieves a 1.16 to 1.76 times speedup compared with AMD AITER. The existence of a speedup range indicates that the benefit is related to specific sequence length, batch configuration, and head dimension, rather than being a fixed multiplier.

More noteworthy is the accuracy performance. Under the head dimension 128 setting, ZYFA reduces error by 62% to 78% compared with AITER. What makes people hesitate most about low-precision attention is numerical stability. FP8 has limited mantissa bits, and the dynamic range of attention scores is large, so obvious precision loss can occur with the slightest misstep. Being able to push error down to this level means ZYFA has made targeted designs in quantization strategy or accumulation paths, making it not just "faster" but "fast enough to use."

Note that these numbers come from real model activations rather than synthetic benchmarks, so their reference value is higher than pure microbenchmarks. However, the public information does not disclose the specific model scale, sequence length distribution, and testing framework, so horizontal comparison still requires reproduction and verification on your own workload.

FP8 usually refers to two 8-bit floating-point formats, E4M3 and E5M2: E4M3 uses a 4-bit exponent and 3-bit mantissa, with a smaller dynamic range but higher precision, suitable for weights and activations; E5M2 uses a 5-bit exponent and 2-bit mantissa, with a larger range but lower precision, and is often used for gradients. Compared with BF16 (8-bit exponent, 7-bit mantissa), FP8 has only about 3 mantissa bits, meaning the relative representation error for the same value is magnified by about 16 times. In attention computation, the dynamic range of QK^T scores is also particularly wide, so FP8 attention generally requires per-tensor or per-head scaling factors to map values into a representable range, and dequantization or block-wise scaling before softmax. Accumulation is usually still performed in FP32; if low precision is also used for accumulation, error will accumulate rapidly. ZYFA's ability to significantly reduce error under head dimension 128 likely comes from finer design in scaling granularity, accumulation paths, or softmax handling.

The significance of the BF16 path and open-source strategy

ZYFA also provides a BF16 kernel, which is easily overshadowed by the spotlight on FP8 but is important for practical deployment. Not all models and all layers are suitable for directly using FP8, and BF16 remains a safe baseline in mixed-precision deployment. A project covering both precisions means teams can make trade-offs layer by layer within the same kernel system, rather than stitching together multiple libraries.

Choosing open source directly targets the pain points of AMD's software ecosystem. Compared with NVIDIA's CUDA ecosystem, AMD's ROCm still lags in the richness of optimization libraries for core operators such as attention, and the supplement of high-quality third-party kernels can significantly lower the migration threshold. By open-sourcing ZYFA, Zyphra both demonstrates its own engineering capability and accumulates reusable infrastructure for the AMD inference stack.

AITER is AMD's officially maintained AI Tensor Engine for ROCm, providing a collection of high-performance operators including attention, GEMM, MoE, and others, and is one of the default backends for running inference frameworks such as vLLM and SGLang on the MI300/MI350 series. Because it directly calls AMD low-level assembly and the Composable Kernel (CK) library, it has long been regarded as the performance baseline for attention on ROCm. By choosing AITER as the comparison target, ZYFA is in effect making incremental optimizations on top of the official baseline; this also explains why its speedup range is not exaggerated—AITER itself is no longer an unoptimized naive implementation, and achieving a 1.16 to 1.76 times improvement on top of it shows that there is still room for scheduling or quantization gains. When evaluating, using AITER as the comparison baseline is more meaningful than comparing with the native PyTorch implementation.

Several things to confirm before deployment

When evaluating ZYFA, the speedup range (1.16 to 1.76 times) should be regarded as an upper-bound reference rather than an expected value. In real business scenarios, sequence length distributions often have obvious long tails, and when short sequences account for a high proportion, kernel launch overhead will dilute the gains. It is recommended to first use sampled inputs from production traffic for A/B comparison, focusing on P99 latency rather than the average.

On accuracy, the 62% to 78% error reduction was measured under head dimension 128, and performance for other head dimensions needs separate verification. For numerically sensitive tasks such as mathematical reasoning and long-context retrieval, it is recommended to conduct spot checks on output consistency with a small amount of traffic to confirm that precision changes do not propagate to final results.

Ecosystem maturity is also a practical factor. The documentation, issue response, and version compatibility cadence of an open-source kernel determine whether it can enter a long-term maintained technology stack. If the team's dependence on the AMD platform is long-term, incorporating ZYFA into the benchmark system as a regular comparison item is more valuable than a one-time evaluation.

Conclusion

The value of ZYFA does not lie in any single speedup number, but in the fact that it proves there is still room to be tapped in FP8 attention on AMD MI355X, and that this room can be filled by community projects rather than only by vendor libraries. For teams currently evaluating AMD inference solutions, this is a low-cost, high-information test subject: running one round of your own workload will tell you whether it is worth entering the production path.

One-line Summary

Zyphra open-sources ZYFA, providing FP8/BF16 prefill attention kernels for AMD MI355X that are 1.16-1.76 times faster than AITER.

Paragraph Summary

Zyphra released the open-source project ZYFA, providing FP8 and BF16 prefill attention kernels for AMD Instinct MI355X. On real model activations, its FP8 kernel is 1.16 to 1.76 times faster than AMD's official AITER library, and under the head dimension 128 scenario it reduces error by 62% to 78%, balancing speed and numerical stability. The article analyzes why the prefill stage deserves dedicated optimization due to its high computational density, points out that ZYFA adopts a single-point breakthrough strategy and also provides a BF16 path to support layer-wise trade-offs in mixed-precision deployment. The open-source strategy directly addresses the shortcomings of the AMD ROCm ecosystem in the richness of core operator libraries. The article reminds readers that the speedup range should be regarded as an upper-bound reference, that A/B verification using production traffic sampling is needed before deployment, and that attention should be paid to accuracy performance under other head dimensions and the project's long-term maintenance cadence.