Zyphra Luncurkan ZYFA: FP8 FlashAttention Buat AMD MI355X, Cepat 1.76x!

Kamu yang ngurusin LLM di skala besar pasti tahu betapa pentingnya attention kernel yang cepat dan akurat. Nah, Zyphra baru-baru ini ngeluarin ZYFA (Zyphra FlashAttention on AMD), sekumpulan kernel FP8 dan BF16 prefill attention yang dikhususkan buat AMD Instinct MI355X. Yang bikin ini menarik: kodenya open-source di bawah lisensi Apache 2.0, jadi bebas dipakai dan dimodifikasi. Di paper teknis mereka (arXiv 2611.07496), ZYFA menunjukkan speedup 1.16x sampai 1.76x dibanding AITER — library kernel default AMD — saat dijalankan di aktivasi model asli, bukan cuma benchmark sintetis.

Tapi sebelum kita terlalu excited, perlu dicatat: ZYFA cuma jalan di MI355X (arsitektur CDNA4, gfx950), butuh ROCm 7.2.3 dan PyTorch 2.12. Jadi kalau kamu masih pakai MI300X atau GPU NVIDIA, kernel ini nggak bakal jalan. Interface-nya drop-in compatible dengan flash_attn_func, jadi integrasi ke codebase existing nggak perlu rewrite besar-besaran. ZYFA ini sebenarnya mesin attention di balik Zyphra Optimization Engine (ZOE) yang nge-power Zyphra Cloud, jadi ini bukan cuma eksperimen lab — udah dipakai production.

Kenapa FP8 Attention Susah di MI355X

Ini bagian teknis yang lucu: instruksi matrix FP8 di MI355X jalan 2x lebih cepat dari BF16 peak. Tapi softmax di antaranya? Tetap jalan di unit vektor dengan kecepatan sama untuk kedua presisi. Bayangin: untuk 32 queries vs 128 keys di head dim 128, kerja matrix FP8 butuh 1.024 siklus, tapi softmax FP8 exact butuh 896 siklus vektor. Artinya, overlap harus near-perfect biar unit matrix nggak idle. Thread trace Zyphra nunjukin AITER nganggur 31,7% siklus di pipe matrix. Masalahnya, MI355X nggak track completion instruksi matrix di hardware, jadi urutan instruksi mesti dibikin manual biar overlap terjadi. Inilah kenapa kernel FP8 sering underperform di real world.

Tiga Trik ZYFA Biar Cepat dan Akurat

ZYFA pakai tiga teknik utama. Pertama, Lean P: ganti operasi eksponensial di softmax cuma 1 multiply-add + 1 konversi integer. Pola bit e4m3 tumbuh 8 tiap kali nilai dobel, jadi kode 2^x kira-kira 8x + konstanta. Hasilnya: siklus vektor per blok turun dari 896 jadi 384 — jauh di bawah 1.024 siklus matrix. Error tetap dalam 7,5% dari probabilitas asli (exact rounding: 5,9%). Kedua, two-wave ping-pong schedule: 2 wave berbagi tiap SIMD. Satu nge-issue instruksi matrix, yang lain hitung probabilitas, terus swap. Dengan Lean P, pipe matrix busy 86% siklus steady-state, naik dari 60% pakai exact P. Versi HIP butuh 43-46% lebih banyak siklus, jadi ZYFA emit assembly tangan lewat generator Python.

Ketiga, per-block scaling: di Mistral-7B layer 20, 91% pasangan (row, blok 64-key) duduk lebih dari 8 binade di bawah row max, di mana e4m3 jadi subnormal atau nol. ZYFA skala tiap blok pakai pangkat 2; instruksi matrix scaled CDNA4 buang faktor itu gratis. Tambah 1 MFMA extra buat jumlahkan probabilitas yang sudah dibulatkan (+6% kerja matrix), biar bobot totalnya tetap 1. Kombinasi ketiga ini bikin ZYFA ngejar throughput teoretis dengan lebih dekat.

Angka Nyata: Cepat tapi Clock Turun

Di benchmark Gaussian 32K tokens head-dim-128, ZYFA FP8 ngecap 2.672 TFLOP/s vs AITER 1.944 TFLOP/s — 53% dari peak datasheet FP8 (AITER cuma 39%). Di shape MLA (192/128), ZYFA 2.464 vs 1.237, rasio 1,99x. Tapi di aktivasi model asli, speedup lebih kecil: 1,16x (Mistral-7B), 1,17x (Qwen3-8B), 1,76x (DeepSeek-V2-Lite). Kenapa? Clock MI355X bergantung pada berapa bit operand yang flip. AITER biarin probabilitas kecil jadi nol/subnormal — switch bit lebih sedikit, clock lebih tinggi (2.108 MHz). ZYFA pakai Lean P jaga nilai kecil, jadi clock turun ke 1.628 MHz. Total siklus ZYFA tetap lebih sedikit, tapi gain clock AITER makan sebagian keuntungan. Tim Zyphra saran: benchmark pakai aktivasi asli di sustained clock, jangan cuma lihat peak teoretis.

Akurasi Jauh Lebih Baik, Bukan Cuma Cepat

Ini yang bikin ZYFA menarik buat production: error FP8-nya 62-78% lebih rendah dari AITER di head dim 128. Lawan fp64 attention, error ZYFA 0,22x AITER di Mistral-7B dan 0,38x di Qwen3-8B. Di full Mistral-7B, KL divergence dari fp64 cuma 3,4×10⁻³ (ZYFA) vs 5,9×10⁻³ (AITER). Output ZYFA bitwise identical di 7.200 launch — deterministik banget. Mayoritas error FP8 full-model sebenarnya dari quantisasi Q, K, V. Rotasi Hadamard di Q dan K kira-kira setengahin error input itu, cuma nambah ~0,1% waktu attention. ZYFA juga duduk dekat Composable Kernel dan Triton soal akurasi.

Tapi ada batasan keras: ZYFA cuma causal prefill forward. Nggak ada decode, nggak ada paged KV, cuma MI355X. Bandingin sama FlashAttention-3/4 (NVIDIA H100/B200) atau AITER (MI300X-MI355X, support decode/MLA/paged), scope ZYFA memang sempit. Keunggulannya: Apache 2.0, FP8 accuracy claim paling transparan (62-78% lower error), dan share of peak 53% FP8 — paling tinggi yang dilaporkan buat AMD. Kalau kamu tim ML infra yang pakai MI355X cluster dan butuh prefill cepat + akurat, ZYFA worth a shot. Cek paper, blog Zyphra, dan GitHub repo-nya buat detail lengkap.