最後に、何日もドキュメントを読んだり、スレッドインデックスをデバッグしたりした後、MI300XのDSCでフラッシュアテンションをゼロから実装することができました。 最初のバージョン(オレンジ)は、元のフラッシュアテンション紙の基本的なスカラーバージョンです。 2番目のバージョン(緑)は同じアルゴリズムですが、マトリックスコア(AMDテンソルコア)を使用しており、ご覧のとおり、これはスカラーコアよりも*大幅に*高速です。 行列コアを使用して、Sij = Qi @ Kj ^ TとPij @ Vjの両方を計算しました。 AMDマトリックスコアのいくつかの「落とし穴」: - これらは波面ごとに動作し、AMD では波面は 64 スレッドであるため、現在の波の ID とその波内のスレッドの ID の両方を追跡する必要があります。 - マトリックス・コアのコア動作は4x1の外積であるため、出力レイアウトはレジスタ内でシャッフルされるため、並べ替えステップが必要です。 - (私の知る限り)マトリックスコアのhipcc組み込み関数はどこにも文書化されていません。AMDの例がたくさんあるリポジトリがありますが、それ以外はLLVMコードベースをgrepする必要があります。 私は時々コードを磨き、おそらくAMDのフラッシュアテンションについてより詳細な投稿を書くでしょう。 ああ、ところで、これを可能にしてくれた@HotAisleに声をかけてください!