Zyphra released ZYFA, an open-source (Apache 2.0) FP8 and BF16 causal prefill attention kernel library for AMD Instinct MI355X (CDNA4, gfx950) running ROCm 7.2.3 and PyTorch 2.12, supporting head dimensions 128 and MLA shapes (192/128) via a drop-in `flash_attn_func` interface. BF16 kernels are written in HIP; FP8 kernels use hand-scheduled assembly generated by a Python tool. On real model activations at sustained clocks, ZYFA’s FP8 throughput exceeds AMD’s AITER library by 1.16× (Mistral-7B), 1.17× (Qwen3-8B), and 1.76× (DeepSeek-V2-Lite), while BF16 roughly matches AITER (1.01–1.02×). Peak FP8 performance reaches 2,672 TFLOP/s at 32K tokens (53% of FP8 datasheet peak vs. AITER’s 39%), and 2,464 TFLOP/s for MLA shape (1.99× AITER). Energy efficiency on captured layers hits 1.94 and 1.91 TFLOP/J versus AITER’s 1.60 and 1.56. ZYFA addresses the softmax bottleneck—where exact FP8 softmax consumes 896 vector cycles versus 1,024 matrix cycles for QKᵀ/PV—via three techniques: Lean P replaces the exponential with one multiply-add and an integer conversion (exploiting e4m3’s log-scale encoding), cutting softmax cycles to 384; a two-wave ping-pong schedule overlaps matrix and vector work, raising matrix pipe utilization to 86% steady-state; per-block power-of-two scaling with a rounded row sum uses CDNA4’s scaled matrix instructions to remove scaling overhead and computes the divisor with one extra matrix instruction per block. Accuracy improves significantly: FP8 error is 62–78% lower than AITER at head dim 128 (0.22× and 0.38× AITER’s error on Mistral-7B and Qwen3-8B), with full-model KL divergence of 3.4×10⁻³ vs. 5.9×10⁻³. Real-model speedups are smaller than Gaussian benchmarks because ZYFA’s fewer subnormals increase bit-flip activity, lowering clock to 1,628 MHz versus AITER’s 2,108 MHz on Mistral-7B, though ZYFA still requires fewer total cycles. Scope is limited to causal prefill forward; no decode, paged KV, or backward pass.
Read original
marktechpost