VFA: グローバル最大値の事前計算によるFlash Attentionにおけるベクトル演算の負荷軽減
arXiv cs.LG / 2026/4/15
📰 ニュースIdeas & Deep AnalysisModels & Research
要点
- FlashAttentionのスタイルによるオンラインsoftmaxはメモリ効率に優れていますが、行max/行和（rowmax/rowsum）やリスケーリングの連鎖といった非行列積（non-matmul）の削減・更新手順が、ベクトル/SIMDのボトルネックとなり得て、高いアクセラレータスループットではレイテンシを支配してしまうことがあります。
- 本論文では、最大値更新の実行コストを下げるVector Relieved Flash Attention（VFA）を提案します。具体的には、maxをキー・ブロックの近似で初期化し、ブロック走査の順序を入れ替えてmaxを早期に安定化させ、後続ブロックでは最大値を「凍結」しつつもオンラインsoftmaxの構造は維持します。
- この考え方を、ブロックスパースのスキップ（例：BLASST）と統合することでVector Relieved Sparse Attention（VSA）へ拡張し、ブロック数とブロック単位のオーバーヘッドの両方を削減します。
- VFA/VSAは、FA4.0の更新ステージで用いられる条件付きリスケール操作を回避し、MMLUやMATH500といったベンチマークで、C16V32ベースラインに対する速度向上を報告しつつ、性能低下は回避しています。
- 結果として、小さなバリアント（例：C8V32、C4V32、C4V16）では現代のハードウェア上で約2×の速度向上に到達し、指数（exponent）容量をより大きくできる将来のアーキテクチャでは最大で約6×の改善が期待されます。
