[論文レビュー] Implementing and Optimizing the Scaled Dot-Product Attention on Streaming Dataflow

Gina Sohn, Nathan Zhang+1 more|arXiv (Cornell University)|Apr 25, 2024

Cloud Computing and Resource Management被引用数 4

ひとこと要約

この論文は、スケーリングドットプロダクトアテンション（SDPA）のメモリ最適化実装を、ストリーミングデータフローアーキテクチャ上で提示する。演算の再順序付けによりデータフロー遅延をバランスさせ、O(1)の間欠的メモリでフルスルーレットを達成する。標準実装のO(N)のメモリオーバーヘッドを排除する。FlashAttentionにインspiredされたアルゴリズム的再構築により、データフロー抽象マシン（DAM）上でサイクル正確なシミュレーションにより検証された。

ABSTRACT

Transformer models serve as the backbone of many state-ofthe-art language models, and most use the scaled dot-product attention (SDPA) mechanism to capture relationships between tokens. However, the straightforward implementation of SDPA has quadratic compute and memory complexity with respect to the sequence length. On processor architectures such as GPUs and TPUs, there is a robust body of prior work. However, little work has been performed on non-processor architectures.In this work, we show how the architecture and execution model of Streaming Dataflow Accelerators can help tackle this challenge. We first define abstract hardware that adopts a streaming execution model, and we implement a cycle-accurate simulator of the abstract hardware using the Dataflow Abstract Machine simulation framework. Second, we implement the naive SDPA algorithm on this abstract hardware and show it requires linear (O(N)) intermediate memory. Third, we then modify the naive algorithm, taking inspiration from prior processor-oriented works, by reordering the multiplication and division operations. Finally, we map the modified algorithm to abstract hardware, and confirm that the implementation computes SDPA at full throughput while only using a constant amount (O(1)) of intermediate memory.

研究の動機と目的

• ストリーミングデータフローアーキテクチャのような非プロセッサアーキテクチャにおける、スケーリングドットプロダクトアテンション（SDPA）の二次的メモリ複雑性を解決すること。
• 最小限の間欠的メモリ使用量で、ストリーミングデータフローガード上でフルスルーレットSDPA実行を可能にすること。
• ナードなストリーミング実装で必要なO(N)サイズのFIFOを排除できる、演算のアルゴリズム的再順序付けが可能であることを示すこと。
• サイクル正確なシミュレーションにより、O(1)の間欠的メモリがフルスルーレットSDPAに十分であることを検証すること。

提案手法

• 並列パターンに基づく抽象的なストリーミングデータフローガードモデルを定義する。Map、Reduce、Scan、MemReduceノードを含む。
• データフロー抽象マシン（DAM）フレームワークを用いてサイクル正確なシミュレータを実装し、機能的正しさとメモリ使用量を検証する。
• ナードなSDPAアルゴリズムを抽象ハードウェアにマップし、分岐したデータフローパathにおける不均衡な遅延のためO(N)の間欠的メモリが生じることを確認する。
• 分配法則を用いて、ソフトマックスと行列乗算ステップの演算を再順序付けし、各パス間の計算遅延をバランスさせる。
• 行単位のリダクションを累積最大値と累積合計値の操作に置き換え、デルタ項によるリスケーリングにより数値安定性を維持する。
• 最適化されたアルゴリズムを抽象ハードウェアにマップし、短いFIFO（深さ2）のみを用い、フルスルーレットでのO(1)の間欠的メモリ使用量を確認する。

Figure 1. An abstract diagram of the architecture and execution model for streaming dataflow accelerators

実験結果

リサーチクエスチョン

• RQ1 ストリーミングデータフローアーキテクチャ上で、O(N)未満の間欠的メモリ複雑性でSDPAを実装できるか？
• RQ2 分岐したデータフローパスをバランスさせ、大規模なFIFOを排除するために必要なアルゴリズム的変更は何か？
• RQ3 フルスルーレットと数値安定性を維持しながら、メモリフットプリントをO(1)にまで削減できるか？
• RQ4 ソフトマックスにおける除算と乗算演算の再順序付けは、ストリーミング実行モデルにおけるメモリと遅延にどのように影響するか？

主な発見

• ナードなSDPA実装は、ストリーミングデータフローガード上で、分岐したデータフローパスにおける不均衡な遅延のためO(N)の間欠的メモリを要する。
• ソフトマックスにおける除算演算の再順序付けと累積最大値・累積合計値の使用により、データフローパス間で遅延がバランスされる。
• 最適化されたアルゴリズムは、行単位のリダクションを要素単位のスキャン操作に置き換え、デルタ項によるリスケーリングを用いることで、O(N)深のFIFOの必要性を排除する。
• 最終的な実装は、データフロー抽象マシン上でサイクル正確なシミュレーションにより確認されたように、O(1)の間欠的メモリのみでフルスルーレットを達成する。
• FlashAttentionと同様のスケーリングとリスケーリング技術により、ストリーミング実行モデル内でも数値安定性を維持する。

Figure 2. Implementation of attention using Parallel Patterns. The depth of the short FIFOs is set to 2, and the depth of the Long FIFO is set to $N+2$ . Each node can be mapped to a configuration of a set of compute and memory units in a streaming dataflow hardware.