8月1日、HyperAI主催の第9回Meet AI Compilerテクニカルサロンが北京で開催されました。このイベントでは、AIコンパイル技術の最新の進歩に焦点が当てられ、業界や研究機関の複数の専門家がプログラミング言語、演算子開発、コンパイル最適化、推論実行に関する知見を共有し、高水準言語表現からハードウェア実行に至るまでのAIコンパイラの協調的な進化を紹介しました。

で、テンセントの高性能コンピューティングエンジニアである朱洪林氏は、「FalconGEMM：低複雑度行列乗算でハードウェアの限界を超える」と題したプレゼンテーションで、低複雑度行列乗算のためのアルゴリズムと演算子の最適化手法について、自身のチームによる研究成果を発表した。

cuBLASなどの成熟した演算子ライブラリによって行列乗算性能がハードウェアの限界に近づき、従来のカーネルレベルの最適化の余地がますます狭まっているという課題に直面し、チームはアルゴリズムの複雑さから再び出発しました。StrassenやAlphaTensorなどの低複雑度行列乗算アルゴリズムをベースに、統一されたLCMA（低複雑度行列アルゴリズム）フレームワークを構築しました。QDSL、演算子融合、パーシステントカーネル、きめ細かいスケジューリング、コストモデルを組み合わせることで、「乗算回数の削減」という理論的な利点をGPU上での実際のパフォーマンス向上へと転換しました。

NVIDIA H20のFP16およびBF16テストにおいて、FalconGEMMは多数の行列形状でcuBLASを上回り、ピーク時には約10%～16%の性能向上を実現しながら、言語モデルベンチマークにおける標準的な行列乗算とほぼ同等の数値精度を維持しています。

HyperAIは、共有されたコンテンツを元の意味を変えることなく編集・要約しました。

WeChat公式アカウント「HyperAI」をフォローし、背景にキーワード「」を入れて返信してください。0801 AIコンパイラ認定講演者のプレゼンテーション用PPTは、「…」をクリックすると入手できます。

ストラッセンの研究を皮切りに、行列乗算の最適化空間を再検討する。

行列乗算は、深層学習における最も重要な基本演算子の1つであり、通常、モデルの計算時間の大部分を占めます。 CUDA、MKL、cuBLASといったソフトウェアスタックは長年にわたって最適化されており、多くのシナリオにおいて、単一のGEMMカーネルのパフォーマンスはすでにハードウェアのピーク性能に非常に近いレベルに達しています。つまり、命令レベル、パイプラインレベル、メモリアクセスレベルといった局所的な最適化だけを続けていくと、改善の余地はますます限られてきています。そこで、チームは古典的なStrassenアルゴリズムに再び注目することにしました。

シュトラッセンアルゴリズムは、1969 年にフォルカー・シュトラッセンによって提案されました。最も基本的な 2×2 行列の乗算の場合、従来の方法では 8 回の乗算が必要ですが、シュトラッセンは入力行列を再構成することで 7 回の乗算のみで済み、その後、追加の加算と減算演算を使用して最終結果を復元します。これは、乗算の計算量を 1/8 に削減することに相当します。オペランドがスカラーのみの場合は、このトレードオフはメリットがありませんが、オペランドが部分行列になると、行列の加算 (O(N²)) と行列の乗算 (O(N³)) の複雑さの差が顕著になります。これにより、「行列の乗算を1回減らし、行列の加算を数回増やす」という行為に、実用的な価値が生まれ始める。

ストラッセンアルゴリズムを再帰的に使用すれば、乗算回数をさらに削減できます。例えば、従来4×4ブロック行列の乗算には64回のブロック乗算が必要でしたが、ストラッセンアルゴリズムを2層使用すると49回で済みます。ただし、再帰層の数を増やすと、加算、データ構成、メモリアクセスのオーバーヘッドも増加します。そのため、実際のシステムでは、計算量の削減と追加のオーバーヘッドとのバランスを取るために、限られた数のレイヤーしか使用しないことが多い。

2022年、DeepMindのAlphaTensorはこのアルゴリズム空間をさらに拡張しました。行列乗算をテンソル分解問題に変換し、強化学習を用いてより低いランクの分解手法を探索することで、古典的なStrassenアルゴリズムに加えて、M、N、Kの形状が異なる多数の低複雑度行列乗算アルゴリズムが存在する可能性を示しました。

しかし、アルゴリズムの発見から工学的応用へと移行する際には、依然として実際的な問題が存在する。低複雑度アルゴリズムごとに個別の手書きGPUカーネルが必要になるとすれば、開発および保守コストは明らかに高くなりすぎるだろう。この問題を解決するため、チームはこれらのアルゴリズムをLCMA（低複雑度行列アルゴリズム）として抽象化しました。LCMAは、入力行列内のどのサブブロックを事前に結合する必要があるか、実際に何回の行列乗算が行われるか、そして中間結果が最終的にどのように出力行列に結合されるかを統一的に記述します。対応する実装は、Codegenを使用して自動的に生成されます。

したがって、「シュトラッセンカーネルをどのように実装するか」という疑問が生じる。焦点は「高いパフォーマンスを維持しながら、さまざまな低複雑度行列アルゴリズムをサポートできる統一的なフレームワークを構築する方法」に移った。

同時に、低複雑度アルゴリズムは数値精度の問題にも対処しなければなりません。ストラッセン乗算は標準的な行列乗算と代数的に等価ですが、浮動小数点演算は結合法則に厳密に従うわけではなく、計算順序の変更によって追加の丸め誤差が生じる可能性があります。したがって、LCMAは性能を追求しつつ、低精度計算における誤差伝播を制御する必要もあります。

統一されたアルゴリズム記述が確立されたので、次のステップは適切なGPU実装を見つけることでした。チームはCUDA、Triton、TiLang、QDSLを試しました。CUDAは最も強力なハードウェア制御を提供しましたが、多数の異なるLCMAアルゴリズムに直面した場合、レジスタ、共有メモリ、中間加算構造すべてに個別の調整が必要となり、拡張と保守のコストが高くなりました。

Tritonは基本的なStrassenシナリオではCUDAに近いパフォーマンスを実現しますが、アルゴリズムがより大きなブロック構造に拡張されると、複数の中間計算にわたってレジスタバッファを正確に再利用する必要があり、Tritonは追加のスピルが発生しやすくなります。TiLangはレジスタと共有メモリの制御においてより柔軟性がありますが、チームテストでのパフォーマンスは依然としてTritonより約51～101 TP3T低くなっています。基本的なStrassenでは、理論上のゲインはわずか12.51 TP3Tであるため、この損失はアルゴリズムのゲインを損なうほど大きくなっています。

ファイナル、チームはFalconGEMMの主要な実装バックエンドとしてQDSLを選択した。 QDSLはCUDAに近い開発粒度を提供すると同時に、コード生成機能と組み込みPTXのサポートも備えています。これにより、既存の高性能実装の移行が容易になり、さまざまなLCMA記述に基づくバッチコード生成にも適しているため、その後の統合やカスタマイズされた最適化において、より高い柔軟性が得られます。

LCMAからFalconGEMMへ、アルゴリズムの性能向上をGPUに展開する

最も単純なStrassen GPU実装は、いくつかのステップに分けられます。まず、サブマトリックスAとBを組み合わせて7つの新しい入力ペアを生成し、次に7つのバッチGEMMを実行し、最後に7つの中間結果セットを最終マトリックスCに結合します。通常のGEMMと比較すると、計算負荷の高いマトリックス乗算部分は、元の7/8にすぎません。したがって、前処理と後処理に費やす追加時間が、削減された計算負荷の8分の1未満であれば、全体としてメリットを得られる可能性がある。

チームはまずNVIDIA H20でテストを行った。H20はメモリ帯域幅が広く、ピーク時の計算負荷が比較的低いため、「計算負荷を軽減する代わりにデータ処理を増やす」というこのアプローチに適している。行列サイズが約2048³以上になると、基本的な実装で既に安定した性能向上が見られた。しかし、行列サイズが小さくなると、入力組み合わせ、中間結果の書き戻し、出力組み合わせの割合が急速に増加し、節約できた計算負荷がすぐに消費されてしまう。

そのため、その後の最適化の焦点はGEMM自体から中間メモリへのアクセスへと移った。最も直接的なアプローチは演算子融合である。中間結果をグローバルメモリに繰り返し書き戻すのではなく、できるだけチップ上に保持するようにしてください。しかし、入力A/Bを直接GEMMに結合させる方法は、同じサブブロックが複数のSMで使用される可能性があり、重複したロードや合計が発生しやすいため、適切ではありません。これに対し、バッチ処理されたGEMMとCombine Hを後処理で融合させる方がより実現可能です。

真の課題は、ストラッセンの7つの中間結果が、最終的な4つの出力サブマトリックスにそれぞれ異なる方法で寄与するという点にある。Hを並列ユニットとして使用する場合、複数のSMが同時に同じCに書き戻す可能性があり、深刻なアトミック競合が発生する。一方、Cを並列ユニットとして使用する場合、Hの一部が異なるSMによって繰り返し計算されることになる。どちらのアプローチも、乗算回数を減らすという利点を打ち消してしまう。

チームは最終的に、ストラッセンの中間結果に基づいてタスクを整理することを断念した。代わりに、行列の空間座標に基づいてタスクをグループ化した。7つのバッチGEMM内の同じ位置にある7つの乗算タイルを1つのグループにまとめ、同じSM上で実行した。こうすることで、グループが計算を完了した後、結果をチップ上の最終Cに直接蓄積することができ、中間結果をグローバルメモリに書き戻す必要がなくなり、SM間の重大な書き込み競合を回避できた。

この融合アプローチは、Strassenによって導入された追加のメモリアクセスを大幅に削減しますが、グループの粒度が大きくなることで負荷の不均衡が生じます。例えば、4096³行列の乗算では、粗粒度スケジューリングによって約211 TP3Tの追加的な波の無駄が発生し、Strassen自体の計算削減量である12.51 TP3Tを超える場合もあります。

この目的を達成するために、チームはStream-Kからヒントを得て、必要に応じてグループを2つのSMに分割して実行できるようにした。スケジューリング層は引き続きグループを基本単位として使用しますが、実際の実行はさらにタイルレベルまで細分化できるため、グループレベルのデータ再利用の利点を維持しながら、テールSMのアイドル時間を短縮し、ハードウェア利用率を向上させることができます。

しかし、負荷分散の問題を解決した後、新たな問題が発生しました。それはL2キャッシュのスラッシングです。グループが分割された後、異なる種類の中間乗算が同じウェーブ内で混在する可能性があり、アクセスされるデータは独立していたため、L2ヒット率が大幅に低下しました。同時に、GEMMはすでにテンソルコアを多用しており、メモリアクセス負荷がほぼ最大負荷に達すると、H2Oは電力制限に達しました。実際のテストでは、コア周波数が約1.8GHzから1.6GHzに低下し、計算性能が低下したため、フュージョンの利点の一部が再び相殺されました。

L2キャッシュのスラッシングに対処するため、チームは分割グループの順序をさらに調整し、同じタイプの中間結果が同じウェーブ内で処理されるようにしました。混合処理はごく一部のテールウェーブでのみ発生します。これにより、きめ細かなスケジューリングによって提供される負荷分散を維持しつつ、L2データの局所性を向上させ、最終的に大幅な周波数低下の問題を解消することができました。

言及する価値があるのは、これらのスケジューリング最適化の重要な基盤となるのは、パーシステントカーネルです。通常のカーネルではタスクアドバイザー（CTA）はブロックの処理が完了すると終了しますが、パーシステントカーネルではCTAがストリーミングサービス（SM）上に長期間常駐し、後続のタスクを継続的に受け入れることができます。これにより、開発者はグループやタイルの実行順序をより柔軟に制御し、オンチップリソースの再利用を促進できます。タスクの分割、スケジューリングの再編成、キャッシュの最適化なども、同じカーネル内で実行可能です。

キャッシュの並べ替えからコストモデル、ピークパフォーマンスの向上まで 10%-16%

マージ、負荷分散、キャッシュ再配置の後、FalconGEMMはLCMAの計算上の利点をより多くの形状で活用できるようになります。ただし、LCMAが通常のGEMMよりも常に優れているわけではありません。その本質は変わらない。余分なデータ処理を犠牲にして、乗算計算の回数を減らすということだ。元のGEMMが既にメモリアクセスによって制限されている場合、計算量をさらに削減しても十分なメリットは得られません。低複雑度アルゴリズムがより有利になるのは、計算密度が高い場合のみです。

したがって、チームはさらに、LCMAを使用するタイミングと、さまざまなLCMAの中からどれを選択するかを決定するために、Rooflineのようなコストモデルを設計した。目標は実行時間を正確に予測することではなく、「適切なアルゴリズムを選択する」ことであるため、このモデルは主にさまざまな方式の計算コストとメモリアクセスコストを分析し、対象GPUの計算能力と帯域幅を組み合わせることで、その計算/メモリアクセスのボトルネック範囲を推定します。

このモデルでは、低複雑度アルゴリズムにおける乗算回数の削減が計算効率の向上につながり、追加のデータ組み合わせとメモリへの繰り返しアクセスが新たなメモリオーバーヘッドとなります。前述の融合最適化により、このメモリアクセスコストはさらに削減されます。したがって、FalconGEMMは、異なるM、N、Kの形状に基づいて、従来のGEMMと様々なLCMAスキームの利点の境界を判断し、より適切な実装を自動的に選択することができます。

QDSLのコード生成機能を活用することで、フレームワーク全体は最終的に比較的完全な実行フローを形成します。まず、LCMA記述に基づいて対応する融合永続カーネルが生成され、融合によって中間メモリへのアクセスが削減されます。次に、コストモデルが特定の形状に適した行列乗算アルゴリズムを選択します。最後に、QDSLがターゲットコードを自動的に生成してコンパイルします。このようにして、LCMAはもはや固定されたStrassen実装ではなく、ワークロードに基づいて動的に選択できるアルゴリズム空間を形成します。

性能テストは主にNVIDIA H20で実施されました。結果によると、さまざまな低精度行列乗算シナリオでは、FalconGEMMは、多数の形状においてcuBLASを凌駕し、ピーク性能は約10%～16%向上する。グループ分割とキャッシュ再配置が完了すると、大きな形状は低複雑度アルゴリズムによってもたらされる計算上の利点を安定して得ることができ、小さな形状のパフォーマンスも向上し、L2キャッシュジッタによって引き起こされるパワーウォールによる周波数低下が回避されます。

コストモデル選択の結果に基づくと、このモデルはLCMAリターン閾値を超えるほとんどの形状に対してより優れたパフォーマンスを発揮する実装を選択することができた。これは、「アルゴリズム選択＋融合カーネル」アプローチが、計算密度が異なる行列乗算シナリオを効果的にカバーできることを示している。

性能に加えて、数値精度もFalconGEMMが検証しなければならない重要な課題です。チームは初期の低精度実験で大きな誤差に遭遇しましたが、これは主に浮動小数点加算の順序変更が原因でした。例えば、A+C+B−Cは代数的にはA+Bと等しいですが、有限精度浮動小数点計算では必ずしも厳密に等しいとは限りません。

さらなる分析の結果、真に重大な誤差は、FP32累積演算自体ではなく、主に低精度への変換に起因することが明らかになった。一般的なFP16/BF16入力シナリオでは、行列乗算は通常、まずFP32累積演算を使用して実行され、その後低精度に変換される。中間結果が頻繁に変換される場合、FP32仮数部の情報は継続的に破棄される。

融合アプローチは、実際にはこの問題を軽減します。WGMMAの出力はFP32精度を維持するのに対し、FalconGEMMは最終的なC値をFP32を使用してチップ上で直接結合および累積し、計算完了後にターゲット精度にキャストし直します。複数の独立したカーネル間で低精度の中間結果を繰り返し書き戻す方法と比較して、この方法は1回以上の精度変換を削減し、計算順序の変更によって生じる誤差をFP32の下位ビットに多く残すことができます。

言語モデルのベンチマークにおいて、FalconGEMMと標準的な行列乗算を用いて得られた最終スコアは、ごくわずかな違いしかなく、ほぼ同じであった。これは、現在の実装ではモデルの精度が大幅に低下しないことを示している。

次の段階では、チームは2つの方向で開発を進める予定です。1つ目は、バッチグループとKループの順序を調整してA/B結合をさらに統合し、中間入力結果のメモリアクセスをさらに削減すること。2つ目は、LCMAをアテンションに拡張することです。フラッシュアテンションは計算量対メモリ量の比率が高く、低複雑度行列アルゴリズムをブロック分解とパイプラインとさらに組み合わせることができれば、新たなパフォーマンスの可能性も秘めています。

StrassenやAlphaTensorからLCMAやFalconGEMMに至るまで、この研究の意義は、既に高度に最適化されたGEMMを数パーセント高速化するだけにとどまりません。それは、カーネル自体が既にハードウェアの限界に近づいている場合、パフォーマンス最適化は命令やパイプラインをさらに深く掘り下げるだけでなく、アルゴリズムの複雑さにおける新たな改善の余地を見出すことができるという、別のアプローチを提示しています。そして、コンパイル、融合、スケジューリングを通じて、理論上の計算削減を真の実行時効率の向上へと転換できるのです。