Flash Attention

Flash Attentionとは、人工知能の分野において、トランスフォーマーモデルの核となるアテンション機構の計算効率を劇的に改善するアルゴリズムおよびその実装手法のことです。従来のアテンション計算は、シーケンス長の二乗に比例してメモリ消費量が増加するというボトルネックを抱えていましたが、この手法ではハードウェアのメモリ階層を最適化することでその制限を打破しました。具体的には、GPUの高速なオンチップメモリであるSRAMと、大容量だが低速な高帯域メモリであるHBMの間でのデータ移動の回数を最小限に抑えるタイリング技術を活用しています。これにより、中間結果を巨大なメモリ上に一度書き出すことなく、オンチップ上で効率的に計算を完結させることが可能となります。結果として、メモリ使用量を大幅に削減しつつ、計算速度を飛躍的に向上させることに成功しました。大規模言語モデルの開発現場において、長文のコンテキストを扱う際のハードウェア的な制約を緩和する重要な技術として広く普及しています。