FlashAttention-3およびFlashAttention-4の最適化をコンシューマー向けRTX GPUに適用する探査により、データセンター向けの改善がパフォーマンスの向上をもたらさないことが明らかになった。著者は既存のFA-2フォールバックを上回るかどうかを試すため、アテンションカーネルを一から再構築した。

  • FA-3の中核をなす高速なテンソルコア命令(WGMMA)はRTX GPUでは利用不可である。
  • Tensor Memory Accelerator (TMA) は sm_120 に存在するが、転送がボトルネックではないため役立たない。
  • ワープの専門化はパフォーマンスに悪影響を及ぼし、レイテンシを206usから213usに増加させた。
  • FA-4の指数関数近似最適化は関連性がなく、RTX 5090はSFUバウンドではなくテンソルコアバウンドであるためである。

本研究により、コンシューマー向けシリコンにおける通常のアテンションのパフォーマンス上限はFlashAttention-2であり、さらなる向上には精度を犠牲にして低精度化を図る必要があることが結論づけられた。

この分析はprefillおよびcompute-boundのレジームに焦点を当てている;大規模なKVキャッシュに対するデコーディングは別個のmemory-bound課題であり、そこではsplit-KV手法がより関連性が高い。