FlashAttention-4、Blackwellのボトルネックに合わせたoverlap重視設計

r/LocalLLaMAで話題になったFlashAttention-4は、NVIDIA Blackwell向けにattention kernelを再設計した取り組みだ。主張の中心は、最新GPUではtensor coreの伸びがshared memory bandwidthやSFUの伸びを上回るため、GEMM単体最適化よりも演算の重ね合わせ(overlap)が支配的になるという点にある。

公開記事では、H100からB200への移行でBF16 tensor throughputが約1.0から2.25 PFLOPsへ増える一方、SFU数とshared memory bandwidthは同じ比率で増えていないと説明する。この前提からFlashAttention-4は次の2点に集中する。

• Forward: softmax exponentialをMMAと重ねる新しいpipelineと、hardware/software混在のexp計算。
• Backward: TMEM配置とBlackwell 2-CTA MMAを使い、shared-memory trafficとatomicコストを削減。

技術記事では、ping-pong tile scheduling、conditional online softmax rescaling、TMEM再利用、DSMEM交換を使ったdQ分解などが具体的に解説される。さらにreduction順序を固定するdeterministic modeも紹介され、ベンチマークではnondeterministic比で約85-90%のthroughputを示した。

性能値としては、B200 BF16で最大1605 TFLOPs/s（約71% utilization）、forwardでcuDNN 9.13比最大1.1-1.3x、Triton比最大2.1-2.7xを報告。記事内ではcuDNNチームとの連携や新しいcuDNN版との比較にも触れている。

実装面ではCuTe-DSL（CUTLASS Python DSL）全面採用も注目点で、重いC++ template中心の開発に比べてcompile timeを約20-30x短縮できたとしている。

当然ながら、これらは著者側公開のbenchmarkであり、実運用ではshape、mask、sequence長に合わせた再評価が必要だ。それでもattentionが主要コストであるLLM学習・推論にとって、有力な最適化方向を示すアップデートと言える。

コミュニティ出典: r/LocalLLaMAスレッド
原文: Together AI FlashAttention-4