NVIDIA Blackwellで実現したFlashAttention-4が計算とメモリボトルネックを突破
NVIDIAのBlackwellアーキテクチャ上で動作する最新のFlashAttention-4(FA4)が、計算とメモリのボトルネックを大幅に克服した。Transformerモデルは、大規模言語モデル(LLM)の基盤となる技術で、入力全体を一度に処理する「自己注意機構(self-attention)」により、長距離の文脈依存を捉えることが可能。しかし、この機構は計算とメモリの複雑さが二次関数的になるため、長文処理ではメモリ不足が深刻な課題となっていた。 FlashAttentionは、標準的なアテンションと同等の結果を、より効率的に計算するアルゴリズム。その最新版であるFlashAttention-4は、NVIDIA Blackwell(HGX B200)向けにハードウェア・ソフトウェア共同設計され、ピーク性能1,605 TFLOPS/sを達成。理論最大性能の71%を活用し、従来のcuDNNやTriton比で最大1.3倍の速度向上を実現。特に逆伝播(backward pass)では、専用のオンチップメモリ「TMEM」を活用して中間データを直接保存。これにより、共有メモリ(SMEM)のトラフィックを大幅に削減し、計算ユニットの空き時間を最小化。 Blackwellの特徴である「非対称スケーリング」(計算性能は倍増だがメモリ帯域はほぼ変化なし)に対応するため、FA4は以下を実現: ・MUFU(指数演算用ユニット)のソフトウェアエミュレーションによる高速化 ・128×128の大きなタイル処理を可能にしたスケジューリング最適化 ・非行列演算資源の制限を補うための演算最小化 ・非同期Tensor Coreのパイプライン設計で、計算とメモリ操作の重なりを最大化 これらの工夫により、FA4は32,768トークンの長文処理において、前バージョン比で前向きパスで最大3.6倍、逆伝播で3.15倍の速度向上を達成。SGLangやvLLMといった推論フレームワークでも活用可能で、NVIDIA cuDNN 9.14にも統合された。 FA4は、AIモデルの長文処理能力を飛躍的に高め、マルチGPU・マルチノード環境でのスケーラビリティを強化。今後のLLM開発における性能と効率の向上に大きく貢献すると期待されている。
