ZYFA: FP8 FlashAttention Kernels for AMD MI355X, Open Source

Posted on by
Contents show
ZYFA

The fastest attention kernels in the world run on NVIDIA hardware, and ZYFA is a direct answer to that imbalance: NVIDIA ships mature FP8 attention libraries, and AMD does not. Zyphra has released ZYFA, a set of open-source FP8 and BF16 prefill attention kernels for AMD's Instinct MI355X that closes a meaningful slice of that gap, and ZYFA does it with hand-written assembly. On activations captured from real models, ZYFA's FP8 kernels run 1.16 to 1.76 times faster than AITER, AMD's default kernel library, while adding 62 to 78 percent less numerical error at the standard head dimension.

The release ships under Apache 2.0 with a technical paper, and it is not a research curiosity. ZYFA is the attention engine of Zyphra's Optimization Engine, which powers Zyphra Cloud, meaning the ZYFA kernels serve production traffic today. The BF16 kernels are written in HIP, AMD's C++ kernel language; the FP8 kernels are hand-scheduled gfx950 assembly emitted by a Python generator.

This article explains why FP8 attention underperforms on AMD hardware, the three ZYFA techniques that fix it, what the benchmarks establish and where they overstate, and why the methodology note about GPU clocks may be the most useful part of the release.

Why FP8 attention is softmax-bound on MI355X

The puzzle ZYFA solves starts with an arithmetic disappointment: FP8 matrix instructions on the MI355X run at twice the BF16 peak, yet FP8 attention was not delivering twice the speed.

The reason is the softmax sitting between the two matrix products. FP8 only accelerates the QK transposition and the PV products; the softmax between them runs on the vector unit at the same speed in either precision. For a typical block of 32 queries against 128 keys, the FP8 matrix work takes 1,024 cycles while the exact FP8 softmax takes 896 vector cycles, nearly a full second pipeline's worth of work hiding behind the fast one. The kernel needs near-perfect overlap to keep the matrix units busy, and Zyphra's thread traces showed AITER's matrix pipe idle for 31.7 percent of cycles.

A further complication is that the MI355X does not track matrix instruction completion in hardware, so the overlap cannot be requested; it must be constructed by the order of instructions themselves.

This diagnosis reframes the problem. The bottleneck is not the matrix hardware, which is already twice as fast, but the elementwise work between matrix operations. Any fix has to attack the softmax, and that is where ZYFA's three techniques go.

Why this release matters for the AMD software story

The kernel itself matters to a narrow audience, but the gap it fills explains a structural problem in AI infrastructure.

Every major modeling breakthrough of the past three years, flash attention, FP8 training, fused kernels, has shipped for NVIDIA first, because CUDA has the compilers, the documented instruction semantics and the installed base that justify hand-written kernels. AMD's ROCm stack has improved quickly, but the attention layer specifically lagged: AITER, AMD's inference kernel library, selects kernels that leave the matrix pipe idle for a third of the time on the very workload FP8 exists to accelerate.

That gap has real consequences. Procurement decisions weigh software maturity alongside hardware specifications, and the answer for years has been to buy NVIDIA because the kernels exist. Every serious open-source kernel effort on AMD hardware shifts that calculation at the margin, and ZYFA is the strongest data point yet at the attention layer, which is the layer where long-context inference lives.

It also matters who shipped it. A model company writing hand-scheduled assembly for a hardware vendor's chip, because it needed the speed for its own cloud, is the open-source pattern working as intended: the incentive existed, the result was published rather than hoarded, and the license lets every MI355X owner benefit.

The three techniques, in plain terms

Each of ZYFA's three techniques removes cycles from the vector pipe, and together they raise ZYFA's matrix-pipe occupancy to 86 percent of steady-state cycles from 60 percent with the exact softmax.

Lean P: the exponential becomes one multiply-add. The insight is arithmetic: read as an integer, the 8-bit code of a positive FP8 number grows by exactly 8 each time its value doubles, so the code is approximately 8 times the base-2 logarithm plus a constant. Since a softmax probability is 2 raised to the score minus the row maximum, the kernel can write the FP8 code of each probability with one multiply-add and one integer conversion instead of an exponential and a conversion. Softmax vector cycles per block drop from 896 to 384. The cost is small and measured: lean P raises the error of one attention call to 1.22 to 1.36 times the exact version and shifts KL divergence by 1.7 percent relative, which the paper treats as a fair trade and documents honestly.

Two-wave ping-pong scheduling. Two waves share each SIMD unit, one issuing matrix instructions while the other computes probabilities, then swapping. Because the hardware does not track matrix completion, the assembly schedule emitted by the Python generator constructs the overlap manually. This is hand-scheduled assembly at a time when most of the industry has accepted compiler-generated kernels.

Per-block scaling with a rounded row sum. In real activations, most probabilities sit many binades below the row maximum, where FP8 values become subnormal or underflow to zero. ZYFA multiplies each 64-key block of probabilities by a power of two before rounding, and the MI355X's scaled matrix instruction divides the same factor back out in the PV product at no cost. The denominator is computed on the matrix unit with one extra instruction per block, so the output remains a true weighted average whose weights sum to one, which is also why ZYFA's error is lower than AITER's.

The cycle arithmetic, worked once

The softmax-bound diagnosis deserves its numbers laid out, because they are small and decisive.

Take a block of 32 queries against 128 keys, the working unit of the kernel. The FP8 matrix instructions for that block take 1,024 cycles on the matrix pipe. The exact softmax, the exponential and its supporting work, takes 896 cycles on the vector pipe. If the two ran perfectly in parallel, the block would finish in 1,024 cycles and the matrix hardware would never idle. In practice, the hardware does not track matrix completion, so achieving overlap requires the instruction stream itself to alternate correctly, and AITER's assembly left the matrix pipe waiting for 31.7 percent of cycles.

Lean P changes the vector budget to 384 cycles, less than half the matrix time, which creates slack for the ping-pong schedule to absorb. The matrix pipe then runs at 86 percent occupancy, and the two matrix products execute at the doubled FP8 rate almost continuously.

That is the entire optimization in miniature: find the pipe that is starving the fast one, shrink its work arithmetically, and schedule the two to interleave by hand. The techniques are generalizable to any accelerator where cheap matrix instructions outrun elementwise work, which describes most modern AI hardware.

What the numbers establish, and where they overstate

The benchmark story has two layers, and the difference between them is the most instructive part of the release.

On Gaussian inputs at 2K to 32K tokens, ZYFA's FP8 kernels reach 1.26 to 2.02 times AITER's throughput, hitting 2,672 TFLOP/s at 32K tokens with head dimension 128, which is 53 percent of the FP8 datasheet peak against AITER's 39 percent. At the multi-head latent attention shape, the ratio reaches 1.99 times.

On activations captured from real models, Mistral-7B, Qwen3-8B and DeepSeek-V2-Lite, timed at sustained clocks, the gain compresses to 1.16 to 1.76 times, with only 1.16 on Mistral-7B. BF16 kernels roughly tie AITER, which is itself notable: ZYFA's BF16 kernels are new and already at parity.

The compression is explained by a phenomenon the paper documents with unusual care. The MI355X lowers its clock under sustained load depending on how many operand bits flip between instructions. AITER's kernel lets small probabilities underflow to zeros and subnormals, which switch few bits, so its clock rises, on the Mistral layer to 2,108 MHz against ZYFA's 1,628 MHz. ZYFA's per-block scaling keeps values normal and accurate, at the price of clock speed, and it still needs fewer cycles overall. The methodological conclusion is one other kernel papers rarely state: benchmarks on Gaussian inputs overstate the gain, and attention kernels should be timed on captured activations at sustained clocks.

The details that make it production-grade

Several kernel-design decisions beyond the three headline techniques separate a benchmark kernel from one that can serve traffic, and ZYFA's paper works through them.

Work distribution is balanced by pairing causal query tiles first with last, so every job does the same amount of work regardless of position, and a persistent grid of one workgroup per compute unit loops over many jobs instead of exiting after one. The grid is remapped so each of the eight chiplets processes consecutive tile pairs from the same attention head, which lets those jobs share keys and values in the chiplet's own L2 cache rather than competing for bandwidth.

Output rescaling is lazy: the running maximum of a row triggers a rescale only when it grows past a threshold, expressed in base-2 units of the score, rather than on every block. The threshold varies by kernel, 0.125 in FP8 and up to 8 in BF16 at the latent shape.

Determinism is a feature rather than an accident. ZYFA's outputs are identical in every launch, and the log-sum-exp is computed from the rounded probabilities so that results computed over different key ranges merge consistently, which is what chunked prefill requires. These are the properties that turn a fast kernel into one a serving team will actually deploy, and they are the properties most kernel papers leave out.

The ZYFA accuracy work is the sleeper contribution

Speed claims in FP8 usually arrive with an accuracy shrug. ZYFA's accuracy section is more substantial than the kernel's speed story.

Against fp64 attention, ZYFA's FP8 error is 62 to 78 percent lower than AITER's on Mistral-7B and Qwen3-8B, and within a few percent of Composable Kernel and Triton. In full Mistral-7B, KL divergence from fp64 attention is 3.4 times ten to the negative third with ZYFA against 5.9 with AITER. Outputs are bitwise identical across 7,200 launches, where AITER's assembly kernel varies between launches when a GEMM runs in between, a reproducibility property that matters for anyone debugging at scale.

The paper also contributes to the broader FP8 problem rather than only this kernel. In full models, most FP8 error comes from quantizing the inputs Q, K and V rather than from rounding probabilities, and a Hadamard rotation of Q and K before quantization roughly halves that input error while adding only 0.1 to 3.4 percent to ZYFA's attention time. The paper recommends measuring FP8 accuracy as KL divergence from the BF16 model rather than perplexity change, with the striking observation that AITER's kernel lowered Qwen3-8B's perplexity while having the largest KL divergence of all variants, a benchmark failure mode worth knowing about.

Who should actually run it

The drop-in interface makes trying ZYFA cheap for its intended audience, and the intended audience is specific.

Teams serving long-context models on MI355X hardware with ROCm 7.2.3 and PyTorch 2.12 can install from GitHub, import flash_attn_func from zyfa, and replace their prefill attention call with matching arguments. The functions follow flash-attention 2's argument order, support MHA, GQA and MQA, accept per-tensor FP8 descales, and work under torch.compile and in CUDA graphs, with the caveat that each new FP8 shape should run once before graph capture because the first call builds a host-side job table.

The checks before adoption are the ones the paper models. Benchmark on captured activations at sustained clocks, not Gaussian inputs. Measure accuracy as KL divergence from the BF16 model. If deploying FP8, apply the Hadamard rotation to Q and K, which the repository documents with code, and consider key smoothing. Expect the Mistral-7B profile, a 1.16 times gain, unless the workload resembles DeepSeek-V2-Lite's, where 1.76 times was measured.

Teams not on MI355X hardware gain nothing directly today, and should watch the repo rather than integrate it.

Getting started with ZYFA

Installation is two commands from the GitHub repository: clone and pip install with no build isolation, so the kernels compile against the PyTorch already installed. The requirements are an MI355X, ROCm 7.2.3 and a ROCm build of PyTorch 2.12, and the FP8 kernels assemble with the ROCm compiler during the build.

The interface mirrors flash-attention: import flash_attn_func from zyfa, pass q, k and v in the familiar shapes, and set causal to true. FP8 calls take one power-of-two descale per tensor, returned by the bundled fp8_quantize helper, and the kernel reads the descales from GPU memory so a call never waits on the host. Sequences of any length work; lengths that are not multiples of 256 are padded, which yields exact results for causal attention at the cost of copying the inputs.

The limits, stated plainly

The scope restrictions are extensive and honest. ZYFA supports causal forward prefill attention only: no decode, no KV cache paging, no unequal-length batches, no non-causal masks or sliding windows, no dropout, and no backward pass, which is the honest boundary of the first release. It runs on the MI355X only, on ROCm 7.2.3 with PyTorch 2.12. The FP8 kernels cover head dimension 128 and the MLA shape; lean P is not used at the 192 shape, where it fell below the correctness test's signal-to-noise threshold in one configuration. BF16 performance is parity, not victory. And the cross-vendor comparisons in the table come from different inputs and clocks, with Zyphra measuring nothing on NVIDIA hardware.

The audience is therefore narrow and self-selected: teams running long-context prefill on MI355X hardware, which today means organizations that chose AMD deliberately. For them, this is the difference between FP8 attention being unusable and being 1.76 times faster with less error, and it is open source under Apache 2.0 with a drop-in flash_attn_func interface.

For everyone else, the value is the method. The softmax-bound diagnosis, the clock-sensitivity methodology, the Hadamard trick, and the honest Gaussian-versus-activations comparison are transferable to any hardware and any kernel effort. This is the same lesson this site drew from Zyphra's hierarchical attention release: the artifact is useful, and the measurement discipline is the export. It also lands in a week where inference optimization went mainstream, alongside Together Compiler's agent-driven tuning and the serving economics in our inference caching guide, suggesting the marginal gains in inference are now being hunted at the kernel level by teams willing to write assembly.

Questions kernel engineers will ask

Does ZYFA work for training?

No. The kernels cover the causal forward pass only, with no backward pass; inputs that require gradients still run forward, but calling backward raises an error. Training on MI355X in FP8 attention remains out of scope.

Can I use it for decode or with paged KV caches?

Not in this release. Decode, KV caching, paged KV, unequal-length batches, non-causal masks, sliding windows and dropout are all explicitly unsupported. The flash-attention-compatible varlen function requires equal-length sequences.

Is the FP8 accuracy acceptable for production serving?

The paper's evidence is strong for prefill: output error 62 to 78 percent below AITER's, KL divergence lower than the baseline in full models, bitwise-identical outputs across thousands of launches. The model card guidance applies nonetheless: validate on your model, your context lengths and your inputs, and apply the documented Hadamard rotation before judging.

Why assembly instead of a compiler?

The MI355X does not track matrix instruction completion in hardware, so the overlap between matrix and vector work depends entirely on instruction order. A compiler scheduler was not producing it, and the two-wave ping-pong sequence needed hand control. The Python generator emits the assembly, so the kernels remain reproducible and reviewable.

Sources
Subscription Form