Breaking the FP8 Attention Bottleneck on AMD MI355X: How Zyphra ZYFA Unleashes Performance

An Unexpected Breakthrough in AMD MI355X FP8 Attention Performance

Zyphra recently released ZYFA (Zyphra FlashAttention on AMD), a set of FP8 and BF16 prefill attention kernels designed for AMD Instinct MI355X GPUs. This open-source project addresses the issue of underutilized performance on MI355X when using FP8 precision for large language model (LLM) inference prefill, a limitation stemming from the hardware's architectural characteristics. ZYFA focuses on optimizing matrix multiplication and Softmax computations within the attention mechanism for long-context scenarios, aiming for more efficient hardware resource utilization.

The release of this kernel library has garnered significant attention because it directly challenges the prevailing perception of MI355X's FP8 performance limitations. Zyphra claims that, on real model activations, ZYFA's FP8 kernels run 1.16x to 1.76x faster than AMD's default AITER library, while also exhibiting 62% to 78% lower FP8 error at a head dimension of 128. This substantial improvement in both performance and accuracy suggests deeper optimization potential for AMD MI355X in LLM acceleration, offering a new, compelling performance benchmark for technical leads relying on this hardware stack for large-scale model deployments.

The Core Bottleneck: Speed Mismatch Between FP8 Matrix Operations and Softmax Vector Units

A key characteristic of AMD Instinct MI355X hardware is that its FP8 matrix instruction peak throughput is theoretically twice that of BF16. This implies that for core computations like matrix multiplication, FP8 can deliver significant speed improvements. However, the attention mechanism doesn't solely consist of matrix multiplications (QKV); it also includes a critical Softmax normalization step. The challenge arises because Softmax operations are primarily executed on vector units, and these vector units do not experience proportional speedups when FP8 precision is used.

This speed mismatch creates a performance bottleneck. Zyphra's analysis indicates that when processing an attention block, such as 32 queries against 128 keys, FP8 matrix work takes approximately 1024 cycles, while an exact FP8 Softmax requires up to 896 vector cycles. This means that even if the matrix units complete their tasks quickly, they must wait for the vector units to finish the Softmax computation, leading to significant idle time for the matrix pipeline. Zyphra's thread traces show that AMD's default AITER library experiences matrix pipeline idleness for as much as 31.7% of cycles during attention processing, resulting in its FP8 attention kernels only achieving 39% of the theoretical FP8 peak throughput. This is the core challenge that ZYFA aims to address.

ZYFA's Technical Solutions: Encoding Tricks, Parallelism, and Fine-Grained Quantization

To address the aforementioned bottleneck, ZYFA employs three core technical strategies. The first is the "Lean P" strategy, which approximates attention probabilities using a single multiply-add and an integer conversion, leveraging the specific encoding properties of FP8 e4m3 floating-point numbers. While traditional exponential operations demand more vector unit cycles, Lean P reduces Softmax vector cycles per block from 896 to 384, significantly lightening the Softmax computation load and better aligning its speed with matrix operations.

Second, ZYFA introduces a "two-wave ping-pong schedule" mechanism. Since MI355X hardware does not track the completion of matrix instructions, Zyphra uses a Python generator to emit hand-scheduled assembly code that allows two thread waves to alternate execution. One wave issues matrix instructions while the other concurrently computes probabilities, then they swap roles. This precise scheduling strategy increases the matrix pipeline's busy percentage from 60% with exact Softmax to 86% with the Lean P strategy, substantially reducing matrix unit idle time.

Finally, ZYFA implements "per-block scaling and a rounded row sum" technique. Before rounding each 64-key block of probabilities to FP8, it scales them by a power of 2, leveraging CDNA4's scaled matrix instruction to cancel this scaling factor at no additional cost. Furthermore, ZYFA uses the sum of the rounded probabilities for normalization, rather than the unrounded sum. This resolves issues where small probabilities rounded to zero would still be counted in the denominator, leadi

Performance Benchmarks: Numeric Evidence from Real Model Activations and Synthetic Tests

ZYFA's innovative techniques translate into measurable performance gains. In single-layer kernel tests using real model activations (16K sequence length), ZYFA's FP8 attention kernels achieved a throughput increase of 1.16x for Mistral-7B, 1.17x for Qwen3-8B, and a remarkable 1.76x for DeepSeek-V2-Lite, all relative to AITER. In contrast, BF16 kernel improvements were more modest, around 1.01x to 1.02x, highlighting ZYFA's specific and effective focus on FP8 optimization.

Beyond real activation data, in kernel-level benchmarks with Gaussian random inputs and a 32K sequence length, ZYFA's FP8 throughput at head dimension 128 reached 2,672 TFLOP/s, representing 53% of its FP8 peak throughput; AITER, conversely, achieved only 1,944 TFLOP/s, with a peak utilization of 39%. In tests with a DeepSeek-style MLA shape (QK dimension 192, V dimension 128), ZYFA demonstrated an even greater advantage, reaching 2,464 TFLOP/s compared to AITER's 1,237 TFLOP/s, a 1.99x ratio. Furthermore, ZYFA's energy efficiency (TFLOP/J) for Mistral-7B and Qwen3-8B layers also surpassed AITER, indicating it is not only faster but also more energy-efficient. Crucially, in a single attention call at head dimension 128, ZYFA's FP8 RMS error relative to fp64 was 62% to 78% lower than AITER's, proving it significantly improves numerical precision while boosting performance.

Performance Considerations: Clock Frequencies and Metric Interpretation

While benchmarks indicate performance gains of up to 1.99x, the acceleration observed on real model activations (1.16-1.76x) is comparatively lower. This discrepancy relates to a crucial characteristic of the MI355X hardware: its sustained clock frequency is influenced by the number of operand bit flips in the data being processed. AITER tends to produce more zeros or subnormals, which involve fewer bit flips, allowing AITER to operate at higher clock frequencies in certain scenarios.

For instance, in the Mistral-7B layer test, ZYFA operated at a clock frequency of 1,628 MHz, while AITER reached 2,108 MHz. Despite ZYFA's lower clock speed, it achieved overall performance advantages by utilizing its cycles more efficiently and completing more meaningful computational work per cycle. This highlights that when evaluating kernel performance, it's essential not just to look at absolute throughput but also to deeply analyze underlying clock frequencies, instruction scheduling efficiency, and computational density. Therefore, while Gaussian random input tests can demonstrate theoretical peak capabilities, real model activation tests better reflect performance in practical applications.

Application Boundaries and Engineering Decisions: When to Adopt ZYFA?

The release of ZYFA offers a clear performance optimization path for AMD Instinct MI355X users, but its applicability and engineering decisions still require careful consideration. Currently, ZYFA kernels primarily cover causal forward prefill attention with equal query and key lengths. It does not yet support advanced features such as the decode stage during inference, paged KV cache, or variable batch sizing. This implies that for end-to-end LLM serving scenarios, ZYFA's performance advantages will need to be combined with other optimizations, such as efficient KV cache management and batching strategies, to be fully realized.

Furthermore, the current results have been published by the Zyphra team, with no mention of independent third-party verification. While its technical mechanisms are clear and data comprehensive, technical teams are advised to conduct their own testing and validation before critical mission deployments. It is important to note that ZYFA's reported error reduction primarily targets single attention kernel operations and does not equate to a direct improvement in full model output accuracy or task performance. In practical applications, a significant portion of FP8 precision errors may stem from the quantization of Q, K, and V inputs themselves, rather than the attention kernel. Therefore, when using ZYFA, a comprehensive evaluation encompassing model quantization strategies and overall accuracy requirements is still necessary. ZYFA is recommended for use in ROCm 7.2.3 and PyTorch 2.12 environments, with explicit support for the CDNA4 (gfx950) archite