Zyphraは次をリリースしました: ZYFA(Zyphra FlashAttention on AMD):AMD Instinct MI355X向けのFP8およびBF16プリフィル注意機構カーネルのセットです。 AMD Instinct MI355X GPU。ZYFAは長コンテキストプリフィル向けのZyphraのFP8注意機構エンジンであり、コードは Apache 2.0の下でオープンソース。実際のモデルのアクティベーションにおいて、そのFP8カーネルは1.16〜1.76倍高速に動作します AITERより1.16〜1.76倍高速に動作します。ヘッド次元128では、ZYFAのFP8の誤差はAITERより62〜78%低くなっています。
TL;DR
- フレームワーク: カーネルライブラリです。FP8 e4m3およびBF16での因果的フォワード(プリフィル)注意機構、ヘッド次元128およびMLA形状(192/128)に対応。
- 動作環境: AMD Instinct MI355Xのみ(CDNA4、gfx950)、ROCm 7.2.3、PyTorch 2.12。ドロップインの
flash_attn_funcインターフェース。 - パフォーマンス: 実際のアクティベーションでAITERのFP8スループットの1.16〜1.76倍。BF16はほぼ同等(1.01〜1.02倍)。
- 最高性能: 32KトークンでFP8 2,672 TFLOP/s、FP8データシートピークの53%(AITER:39%)。
- 最低性能: Mistral-7Bのアクティベーションではわずか1.16倍。FP8時のクロックはAITERの2,108 MHzに対し1,628 MHzに低下。
ZYFAとは何か?
ZYFAは、Zyphraによるコードとともに技術 論文 の形で公開されています。Zyphra Cloudを動かしている Zyphra Optimization Engine(ZOE)の注意機構エンジンを支えています。BF16カーネルはHIPで書かれています。FP8カーネルはPythonジェネレータが出力する手動スケジュールのアセンブリです。
なぜMI355XではFP8注意機構の性能が低いのか?
FP8行列命令はBF16ピークの2倍の速度で動作します。しかしFP8が高速化できるのはQKᵀとPVの2つの行列積のみです。その間のsoftmaxはベクトルユニット上で、両方の精度において同じ速度で実行されます。
128個のキーに対する32個のクエリでは、FP8行列演算に1,024サイクルかかります。正確なFP8ソフトマックスには896ベクトルサイクルかかります。そのため、行列ユニットを稼働させ続けるにはカーネルはほぼ完璧なオーバーラップを必要とします。Zyphraのスレッドトレースによると、AITERの行列パイプはサイクルの31.7%がアイドル状態でした。
MI355Xは行列命令の完了をハードウェアで追跡しないため、そのオーバーラップは命令の順序によって実現する必要があります。
ZYFAはどのように動作するのか?
ZYFAは3つの技法を使います:
Lean P:指数関数の代わりに1回の乗算加算
e4m3のビットパターンは、整数として読むと、値が2倍になるごとに8ずつ増加します。したがって2^xのコードはおよそ8xに定数を加えたものになります。Lean Pは各FP8確率を1回の乗算加算と1回の整数変換で書き込みます。これにより、ブロックあたりのソフトマックスのベクトルサイクルが896から384に削減されます。
2ウェーブのピンポンスケジュール
2つのウェーブが各SIMDを共有します。一方が行列命令を発行している間に、もう一方が確率を計算し、その後入れ替わります。Lean Pにより、定常状態のサイクルのうち行列パイプの稼働率は、正確なPでの60%から86%に向上します。
ブロックごとのスケーリングと丸められた行和
確率の各64キーブロックは、丸めの前に2のべき乗でスケーリングされます。CDNA4のスケーリング付き行列命令は、追加コストなしでこの係数を除去します。ZYFAはさらに、ブロックあたり1回の追加行列命令で計算される丸められた確率の総和で割り算を行います。
ZYFAはどのくらい速いのか?
実アクティベーションでの16KにおけるFP8の速度向上と持続クロックは次のとおりです 1.16倍(Mistral-7B)、1.17倍(Qwen3-8B)、1.76倍(DeepSeek-V2-Lite)。32KでのGaussianベンチマークでは、head-dim-128のFP8カーネルはAITERの1,944に対して2,672 TFLOP/sに達します。MLA形状では1,237に対して2,464に達し、1.99倍の比率です。
キャプチャしたMistral-7BとQwen3-8Bのレイヤーでは、ZYFAは1.94および1.91 TFLOP/Jを達成します。AITERは1.60および1.56です。
実モデルでの向上はベンチマークでの向上より小さいのはなぜか?
MI355Xのクロックは、オペランドの何ビットが反転するかに依存します。AITERはより多くの小さな確率がゼロまたは非正規化数になることを許します。それらは反転するビットが少ないため、クロックはより高く動作します。Mistral-7Bのレイヤーでは、ZYFAは1,628 MHzで動作するのに対し、AITERは2,108 MHzで動作します。それでもZYFAは全体としてより少ないサイクルで済みます。Zyphraチームは、持続クロックでの実アクティベーションによるベンチマークを推奨しています。
ZYFAのFP8アテンションの精度はどのくらいか?
fp64アテンションと比較して、ZYFAのFP8誤差はMistral-7BでAITERの0.22倍、Qwen3-8Bで0.38倍です。Composable KernelやTritonとは数パーセント以内の差に収まっています。Mistral-7B全体では、fp64からのKLダイバージェンスはZYFAで3.4×10⁻³、AITERで5.9×10⁻³です。ZYFAの出力は7,200回の実行全体でビット単位で同一でした。
モデル全体のFP8誤差の大部分は、Q、K、Vの量子化に由来します。QとKのHadamard回転は、その入力誤差をおよそ半分にします。Mistral-7Bでは、アテンション時間に約0.1%が追加されます。
ZYFAは他のアテンションカーネルとどう比較されますか？
| 特徴 | ZYFA FP8 | AITER FP8 | FlashAttention-3 | FlashAttention-4 |
|---|---|---|---|---|
| 開発元 | Zyphra | AMD | Dao-AILabほか | Dao-AILabほか |
| 対象GPU | AMD MI355X | AMD MI300XからMI355X | NVIDIA H100/H800 | NVIDIA B200(およびHopper) |
| 主打ち出しスループット | 2,672 TFLOP/s FP8 | 1,944 TFLOP/s FP8(Zyphra測定) | 約1.2 PFLOP/s FP8 | 1,613 TFLOP/s BF16 |
| ピークに対する割合 | 53% FP8 | 39% FP8(Zyphra測定) | 非公表(FP8) | 71% BF16 |
| Softmaxの手法 | Lean P、指数関数なし | ハードウェア指数関数 | 非同期Tensor Coreによるオーバーラップ | ソフトウェアでエミュレートした指数関数 |
| FP8精度の主張 | AITER比62〜78%低い誤差 | ベースライン | ベースラインのFP8より2.6倍低い誤差 | 非公開 |
| 範囲 | 因果付きプリフィルのフォワードのみ | プリフィル、デコード、MLA、ページング対応 | フォワードとバックワード、FP8フォワード | フォワードとバックワード |
| ライセンス | Apache 2.0 | MIT | BSD-3-Clause | BSD-3-Clause |
ベンダー間比較の数値は異なる入力とクロックに基づいています。ZyphraはNVIDIA GPUを測定していません。
主なポイント
- MI355X上のFP8アテンションは行列ではなくソフトマックスがボトルネックです。
- Lean Pは指数関数を1回の乗算加算に置き換えます。
- 実モデルでの改善は1.16〜1.76倍で、ベンチマークが示す値より小さくなります。
- ヘッド次元128で、FP8の誤差はAITER比62〜78%減少します。
- プリフィルのみ：デコード、ページドKVは非対応、MI355X専用。
ぜひご覧ください 論文、 Zyphraのブログ記事 、そして GitHubリポジトリ。また、私たちの Twitter のフォロー、 150k+ML SubReddit への参加や、ご購読もお忘れなく ニュースレター。ちょっと待ってください!あなたはTelegramを使っていますか? 現在、Telegramでも参加できるようになりました。
GitHubリポジトリ、Hugging Faceページ、製品リリース、ウェビナーなどのプロモーションで当社との提携をお考えですか? ぜひご連絡ください
本記事 Zyphra ZYFA: AMD MI355X向けFP8 FlashAttention、最大1.76倍の高速化 は最初に MarkTechPostで公開されました.