AMD Composable Kernelライブラリ：わずか数行のコードでAIアプリケーション向けの効率的な融合カーネルを実現

グラフ最適化は、AIモデルの学習と推論にかかる時間とリソースを削減する上で重要な役割を果たします。グラフ最適化の最も重要な機能の1つは、テンソル演算子の様々な組み合わせを融合する機会を特定することであり、メモリ割り当てとトラフィックを削減することで計算効率を向上させることができることです。しかし、様々なテンソル演算子を融合した高性能なバックエンドカーネルを高速に提供できるソリューションは大きな課題でもあります。これまでは、このようなソリューションがなかったため、AIアプリケーションでは、実際のハードウェア上で限られたグラフ最適化セットしか実現できませんでした。

AMD Composable Kernel（CK）ライブラリは、AMDの現在および将来世代のGPU向けに、これらの課題に対処することを目的としています[2]。CKは汎用のHIP C++言語を使用し、完全にオープンソース化されており、以下の目標を念頭に置いて作成されています：
・パフォーマンスと生産性： CKの中核にあるのは、慎重に定義され、高度に最適化され、再利用可能なプリミティブのセットです。これらのプリミティブは、CKのすべての従来型および融合型演算子の実装に使用されます。これらのプリミティブを再利用することで、バックエンドの開発サイクルを短縮すると同時に、高いパフォーマンスを確実に提供します。
・現在の問題には熟達し、将来の問題には柔軟： CKは、すべてのAI演算子に対して自己完結型のバックエンド・ソリューションを提供することを目指しています。これにより、外部ライブラリを必要とせず、バックエンド全体をCK内で実装できるため、高度な融合型演算が可能になります。CKの再利用可能なプリミティブは、現在一般的なAIモデル（マシンビジョン、自然言語処理など）に対応しています。将来、新しい問題が出現した場合には、新しいプリミティブを定義することができます。
・シンプルかつ、AIシステムの専門家にとって強力： CKのすべてのAI演算子は、HIP C++のテンプレートとして実装されています。AIシステム開発者は、数行のコードでC++テンプレートを様々な方法でインスタンス化することにより、データ型、演算、メモリレイアウト、オペランドの数など、演算子の属性をカスタマイズすることができます。
・親しみやすいC++インターフェース： CKは、AIアクセラレーションの性能限界を押し広げる上で重要な役割を果たすHPCアルゴリズムの専門家にとって、魅力的で生産的であることを目指しています。CKの再利用可能なプリミティブとその上に構築された演算子は、中間表現（IR）ではなく、すべてHIP C++ APIとして実装されています。HPCアルゴリズムの専門家は、アルゴリズムに付随するカスタマイズされたコンパイラー・パスを構築することなく、より優れたアルゴリズムを開発するだけで、機能を追加し、性能を最適化することができます。
・ポータブル： CKバックエンドのグラフ最適化は、現在および将来の世代のGPUで実行され、最終的にはCPUのように他のアーキテクチャに対してもポータブルになります。

コアコンセプト

CKは、バックエンド開発者の生産性を向上させるために、2つのコンセプトを採用しています：
・1つが我々が “テンソル座標変換 “と呼ぶ革新的な技法を用いた、AIテンソル演算子のアルゴリズム複雑性の軽減です。我々は、テンソル座標変換プリミティブのセットを一般化し定義しました。このプリミティブは、（畳み込み、グループ化ノルム削減、Depth2Spaceなどの）複雑な演算子を、（GEMM、2次元削減、テンソル転送などの）より単純な演算子として再解釈するために用いられます。この手法により、より単純な演算子のために書かれたアルゴリズムを、その複雑な演算子のすべてに再利用することができるようになりました。
・もう１つがタイルベースのプログラミングモデルです。融合演算子の開発は、対応する従来の演算子をそれぞれ “小さな断片 “に分解し、それらを融合演算子に結合すると考えることができます。それぞれの “小さな断片 “は、元のテンソルの小さなタイル上で動作するテンソル演算子であり、我々はこれを “タイルテンソル演算子 “と呼んでいます。CKは高度に最適化されたタイルテンソル演算子プリミティブのセットを含み、CKのすべてのAI演算子はこれらのプリミティブの上に実装されています。現在、これらのプリミティブには、”tile GEMM”、”tile reduction”、”tile tensor transfer “が含まれ、それぞれのプリミティブには、ブロック、ワープ、スレッドの実装があります。

テンソル座標変換プリミティブとタイルテンソル演算プリミティブはCKの再利用可能なプリミティブを形成しています。

コード構造

現在のCKライブラリは、下から上に向かって4つのレイヤーで構成されています。 「テンプレート化されたタイル演算子”、”テンプレート化されたカーネルとインヴォーカー”、”インスタンス化されたカーネルとインヴォーカー”、そして “クライアントAPI” [3]です。

・AIシステムの専門家：「従来のAI演算子と融合AI演算子のための高性能なカーネルをすぐに使いたい」このような人に最適なのは、”クライアントAPI “と “インスタンス化されたカーネルとインヴォーカー “レイヤーで、このクライアントの例[4]のように、GEMM + Add + Add + FastGeLUの融合演算子は、事前にインスタンス化され、コンパイル済みのオブジェクトとしてユーザーに提供されます。

・AIシステムの専門家：「私は、オープンソースのMLフレームワークのために、最先端のモデルレベルのグラフ最適化を行っています。高度に最適化された融合カーネルを幅広く提供する、自己完結型のソリューションが必要です。また、これらのカーネルをカスタマイズする必要があるので、ブラックボックス・ソリューションの “take it or leave it “状況は私のニーズに合いません。」このような人に最適なオプションは、この例[5]のような「テンプレート化されたカーネルとインヴォーカー」レイヤーです。

・HPCアルゴリズムの専門家：「私のチームでは、常に進化し続けるAIモデルのパフォーマンスが重要なバックエンドの自社開発を行っています。当社にはHPCアルゴリズムの専門家がいますが、生産性を向上させ、複数世代のハードウェアで性能の移植性を得るために、ベンダーが提供する最適化されたソースコードを再利用し、修正したいと考えています。そして、ソースコードやワークロードを共有することなく、これを実現したいのです」。このような人に最適なのは、このコード[6]のように、「Templated Tile Operator」レイヤーで、「Templated Tile Operator」を使用してGEMMパイプラインを記述することが必要です。

AITemplate + CKによるエンドツーエンドの推論

Meta社のAITemplate[7]（AIT）は、AMDとNVIDIA GPU用の別々のアクセラレーションバックエンドを持つ統合推論システムです。AITemplateは、AMD GPUのバックエンドとしてCKを利用しており、CKの “Templated Kernel and Invoker “レイヤーと相互作用します。

AITemplateとCKのアプローチは、AMD Instinct MI250上で、いくつかの重要なモデルの最先端の推論性能を確立しました。これらのモデルでサポートされているCKの高度な融合演算子のほとんどは、AITemplateチームのビジョンによって推進されており、多くの融合カーネルはCKとAITemplateチームによって共同設計されています。

ここでは、AMD Instinct™ MI250アクセラレータと代替製品[8]におけるエンドツーエンドのモデル性能を比較します。AMD Instinct MI250上で動作するベンチマークモデルはすべて、AITemplate [9] + CK [10]を使用しています。

ResNet-50

以下のベンチマーク結果は、AMD Instinct MI250上のAIT+CKと、NVIDIA A100-PCIe-40GBおよびA100-DGX-80GB上のTensorRT v8.5.0.12 [11]（TRT）のResNet-50パフォーマンス結果を示しています。その結果、MI250上のAIT+CKは、NVIDIA A100-PCIe-40GB上のTRTよりも最大1.08倍高速化できることがわかりました。

BERT

CKを使用して実装されたアテンション層用のバッチGEMM + Softmax + GEMMバック・トゥ・バック融合演算子テンプレートは、中間結果のための演算ユニットとHBM間のデータトラフィックを完全に除去します。このテンプレートを使用することで、帯域幅に縛られていたアテンション・レイヤーの多くのワークロードがコンピュート・バインドになり、GPUコンピュートのより効率的な利用が可能になります。このアルゴリズムは、FlashAttention[12]に大きくインスパイアされており、オリジナルの FlashAttention と比較して、メモリトラフィックを大幅に削減する改良が施されています。

以下に示すベンチマーク結果は、AMD Instinct MI250 GPU上のAIT+CKと、NVIDIA A100-PCIe-40GBおよびA100-DGX-80GB上のFasterTransformer v5.1.1バグフィックス[13]（FT）のBERTベースモデル（非ケース）の性能結果を比較したものです。シーケンス長が4096の場合、A100-PCIe-40GBおよびA100-DGX-80GBでは、バッチサイズ32でFTがメモリ不足になります。したがって、シーケンス4096については、A100-PCIe-40GBおよびA100-DGX-80GBでFTがサポートするバッチサイズ16の結果のみを示しています。この結果から、MI250上のAIT+CKは、NVIDIA A100-PCIe-40GB上のFTに対して最大3.28倍、NVIDIA A100-DGX-80GB上のFTに対して最大2.91倍のスピードアップを実現できることがわかります。

ヴィジョン・トランスフォーマー（VIT）

以下のベンチマーク結果は、Instinct MI250上のAIT+CKと、NVIDIA A100-PCIe-40GBおよびA100-DGX-80GB上のTensorRT v8.5.0.12（TRT）のVision Transformerベースモデル（画像サイズ224×224）のパフォーマンスを示しています。その結果、Instinct MI250上のAIT+CKは、NVIDIA A100-PCIe-40GB上のTRTよりも最大1.8倍高速化し、NVIDIA A100-DGX-80GB上のTRTよりも最大1.4倍高速化することができました。

安定した拡散

安定した拡散エンドツーエンド

以下のベンチマーク結果は、バッチサイズ1、2、4、6を使用したAMD Instinct MI250におけるAIT+CKのエンドツーエンドの安定拡散性能結果を示しています。バッチサイズ1では、MI250のシングルGCDのみが使用され、バッチ2、4、6では両方のGCDが使用されることに注意してください。

Stable DiffusionにおけるUNet

この記事を書いている時点では、TensorRTを使ってStable Diffusionをエンドツーエンドで実行する方法について、公開されている情報はありません。しかし、この記事Making stable diffusion 25% faster using TensorRT [14]は、TensorRTを使ってStable DiffusionのUNet部分を高速化する方法を示しています。UNetはStable Diffusionにおいて最も重要で時間のかかる部分であるため、UNetの性能はStable Diffusionの性能を大きく反映します。

以下のベンチマーク結果は、AMD Instinct MI250上のAITemplate + CKと、NVIDIA A100-PCIe-40GおよびA100-DGX-80GB上のTensorRT v8.5.0.12（TRT）のUNetパフォーマンス結果です。その結果、AMD Instinct MI250上のAITemplate + CKは、NVIDIA A100-PCIe-40G上のTRTに対して最大2.45倍のスピードアップを、NVIDIA A100-DGX-80GB上のTRTに対して最大2.03倍のスピードアップを実現できることがわかりました。

エンドツーエンドの最適化に関する最終的な考察

AIモデルのエンドツーエンドのパフォーマンス最適化には、現在2つの一般的なアプローチが存在しています。

それは、グラフの最適化からバックエンドまで、すべてをブラックボックスで処理するベンダー提供の専用ソフトウェアです。私たちは、このアプローチのブラックボックス化と限られたカスタマイズ可能性という性質が、AIフレームワークとバックエンドの間に人為的な境界を作り出し、最適化できる可能性のあるものを制限していると考えています。

汎用言語のない特定用途向けAIチップのために、コンパイラーベースのアプローチが開発されています。このアプローチにより、ベンダーはAIに特化したバックエンドプログラミング言語用のコンパイラを構築することができるのです。

コンパイラーベースのアプローチは、すでに汎用言語を持つアーキテクチャのためにも開発されています。モデルレベルのグラフ最適化から、中間表現（IR）を使ったラストワンマイルのカーネル生成までのプロセス全体を効率化することもあり、また、ハードウェアのプログラミングを簡単にすることを目的としたドメイン固有言語（DSL）を実装することもあります。このような努力の課題は、アーキテクチャやワークロードに特化したHPCアルゴリズムの必要性が取り除かれるわけではなく、アーキテクチャやワークロードに特化したコンパイラー・パスを書く必要性に変わるだけであり、AIアクセラレーションの性能限界を押し広げる上で重要な役割を果たすほとんどのHPCアルゴリズム開発者が貢献することが難しくなることです。

グラフ最適化フレームワークがカーネル融合の機会を特定し、柔軟で効率的なバックエンドライブラリを使用してカーネルを構築します。性能結果は、このアプローチがGPU上で優れた性能を発揮し、既存のアプローチでは達成困難な俊敏性を提供することを示しています。優れたライブラリ設計により、本アプローチはAIシステム開発者のニーズに十分に応え、同時にHPCアルゴリズム開発者の貢献も期待できるでしょう。さらに、バックエンド言語の汎用性は、このアプローチが新しいAI問題やアーキテクチャの変更により迅速に適応できる可能性が高いことを意味しています。

CKコードを取得
https://github.com/ROCmSoftwarePlatform/composable_kernel

CK の開発には、Chao Liu、Jing Zhang、Letao Qin、Qianfeng Zhang、Liang Huang、Shaojie Wang、Anthony Chang、Chunyu Lai、Illia Silin、Adam Osewski、Poyen Chen、Rosty Geyyer  Hanwen Chang [15]、Tejash Shah [15]、Xiaoyan Zhou [15]、Jianfeng Yan [15]、Dan Yao、Guangzhao Lu、Raman Jana、Jehandad Khan、Wen-Heng (Jack) Chung、Austin Kerbow、Stanislav Mekhanoshin、Vang Thao、Jeff Byrnes、Chaitanya Sankisなど、AMD の複数のチームが貢献しています。

こちらの記事はAMD本社のブログ記事を機械翻訳したものです。詳しくは元記事をご覧ください。

脚注：

1. Chao LiuはAMDのPMTSソフトウェア開発エンジニア。Jing ZhangはAMDのSMTSソフトウェア開発エンジニア。彼らの投稿は彼ら自身の意見であり、AMDの立場、戦略、意見を代表するものではありません。第三者のサイトへのリンクは便宜上のものであり、明示的に記載されていない限り、AMDはそのようなリンク先のサイトのコンテンツについて責任を負わず、いかなる保証も意味するものではありません。GD-5
2. CPU用CKは開発初期段階。
3. 現時点ではC++ API、Python APIを計画中。
4. GEMM＋Add＋Add＋FastGeLU融合演算子のCK「クライアントAPI」の例 https://github.com/ROCmSoftwarePlatform/composable_kernel/blob/685860c2a9483c9e909d2f8bfb95056672491…
5. GEMM + Add + Add + FastGeLU融合演算子のCK “Templated Kernel and Invoker “の使用例https://github.com/ROCmSoftwarePlatform/composable_kernel/blob/685860c2a9483c9e909d2f8bfb95056672491…
6. GEMMパイプラインを記述するためにCK “Templated Tile Operator “プリミティブを使用する例 https://github.com/ROCmSoftwarePlatform/composable_kernel/blob/685860c2a9483c9e909d2f8bfb95056672491…
7. MetaのAITemplate GitHubリポジトリ. https://github.com/facebookincubator/AITemplate
8. MI200-71： AMD Instinct MI250 OAM（128 GB HBM2e）560WGPU×4基、AMD Infinity Fabric™テクノロジー搭載のAMD EPYC 7713 64コア・プロセッサー・サーバー上で動作するROCm™5.3搭載のAITemplate https://github.com/ROCmSoftwarePlatform/AITemplate、コミットf940d9b） + Composable Kernel https://github.com/ROCmSoftwarePlatform/composable_kernel、コミット40942b9）を使用したAMD MLSE 10.23.22によるテスト対AMD Instinct MI250 OAM（128 GB HBM2e）560WGPU×4基、AMD Infinity Fabric™テクノロジー搭載のAMD EPYC 7713 64コア・プロセッサー・サーバー上で動作するROCm™5.3搭載のAMD MLSE 10.23.22によるテスト。TensorRT v8.5.0.12およびFasterTransformer (v5.1.1バグフィックス)、CUDA® 11.8、Nvidia A100-PCIe-40GB (250W) GPU x 4、TensorRT v8.0.12およびFasterTransformer (v5.1.1バグフィックス)。 5.0.12およびFasterTransformer (v5.1.1バグフィックス)、CUDA® 11.8搭載、2xAMD EPYC 7742 64コア・プロセッサー・サーバー、8xNVIDIA A100 SXM 80GB (400W) GPU。サーバーメーカーによって構成が異なるため、結果が異なる場合があります。性能は、最新のドライバや最適化の使用などの要因によって異なる場合があります。
9. https://github.com/ROCmSoftwarePlatform/AITemplate/tree/f940d9b7ac8b976fba127e2c269dc5b368f30e4e
10. https://github.com/ROCmSoftwarePlatform/composable_kernel/tree/40942b909801dd721769834fc61ad201b5795…
11. TensorRT GitHubリポジトリ。https://github.com/NVIDIA/TensorRT
12. FlashAttention： IOを意識した高速でメモリ効率の高い正確なアテンション。https://arxiv.org/abs/2205.14135。
13. FasterTransformer GitHubリポジトリ。https://github.com/NVIDIA/FasterTransformer
14. TensorRTを使用して、安定した拡散を25%高速化する。https://www.photoroom.com/tech/stable-diffusion-25-percent-faster-and-save-seconds/
15. AMD在籍中