AMD AIチップ向け計算プログラムが高速に

MoonMath AIチームは、AMDのMI300X GPU向けにbf16 forward attention kernelをリリースしました。これは手書きのアセンブリではなく、HIPで書かれています。このコードはMITライセンスの下でオープンソースとして公開されています。MoonMath.aiチームは、テストされたすべてのshapeにおいて、AMD独自の最適化されたkernelであるAITER v3を上回ると報告しています。ベアメタルアクセスは、AMDのクラウドプロバイダーであるHotAisleから提供されました。Attentionは、すべてのtransformer内部にあるsoftmax(QKᵀ/√d)·Vの結合された演算です。MI300XはAMDのCDNA3データセンターGPUであり、ISAターゲットはgfx942です。このkernelはそのハードウェアでのみ動作します。TL;DR MoonMath.aiは、AMD MI300X向けのbf16 forward attention kernelをHIPで記述し、アセンブリではなく（MITライセンスで）オープンソース化しました。これは、あらゆるshapeとrounding modeにおいて、AMDのAITER v3を上回ります — 幾何平均で1.18倍/1.15倍/1.08倍、最大で1.26倍です。核となるトリックは、1命令のasm wrappersにより、コンパイラがレジスタを割り当てる間にオペコードを選択できることです。速度向上のほとんどはメモリ配置によるものです — KはLDSに、VはL1にホットに、Qとアキュムレータはレジスタに配置されます。実際のSGLang PRでは、Wan2.1ビデオ拡散を品質劣化なしで1.23倍高速化するためにこれを使用しました。Kernelの理解: kernelとは、GPUの多数のコア上で直接実行され、特定の計算（ここではattentionの計算）をハードウェアが許す限り高速に実行するための小さなプログラムです。このkernelは、MI300Xのみでbf16のforward attentionを計算します。入力はBSHDまたはBHSDレイアウトのいずれかを受け入れ、転置は不要です。ヘッドの次元は128に固定されています。cross-attentionを含むあらゆるシーケンス長をサポートします。ただし、いくつかの制限があります。causal mask、GQA、varlen batchingはありません。出力はbf16で、gfx942ハードウェアでのみ排他的に動作します。数値演算は厳密に制御されています。3つのrounding modeすべてがAITERのモードごとの丸めルールと一致します。すべての有限出力はAITERの1 bf16 ULP以内に収まります。NaNおよびInfの処理はビット単位で同一であり、結果は決定論的です。核となるトリック：1命令のasm Wrappers: この核となる技術は、よくあるジレンマを回避します。コンパイラの組み込み関数はコードをきれいに保ちますが、コンパイラに再