AMD、TeaCacheでRyzen AIのNPU上のSD3/3.5画像生成を高速化

AMDはROCm Blogsで、キャッシュ手法TeaCacheを使ってAMD Ryzen AIのNPU上でStable Diffusion 3および3.5の画像生成を高速化する方法を公開しました。Ryzen AI Software 1.8.0の環境でrun.pyを-O1オプション付きで起動するとTeaCacheが有効になり、再トレーニングやアーキテクチャ変更は不要です。

AMDはCPU、GPUなどの半導体製品を開発する企業で、PC向けにRyzenを提供しています。

TeaCacheの仕組みとRyzen AIへの適用

TeaCacheは、拡散モデルの隣接するタイムステップ間で中間出力の変化が小さい場合に、完全なフォワードパスを省略してキャッシュ済み出力を再利用する手法です。潜在表現の相対L1距離を累積し、しきい値δを超えたら再計算して累積距離を0にリセットします。記事では、再スケーリングを行わないnaive cachingと、ポリノミアル関数fで再スケーリングするrescaled cachingを比較しています。TeaCacheはSD向けの係数を公開していないため、著者らは(1,0)、(2,1,0)とFlux用の既定係数を試しました。

PyTorch実装では、JointTransformerBlocksをCachedTransformerBlocksに置き換えるラッパークラスで、ブロック内部でキャッシュ判断を行います。一方、Ryzen AIの静的なONNXベースパイプラインではTransformerのONNXモデルがブラックボックスのため、AMDは次の2点を追加しました。

• Block 1のnormを担う追加ONNXモデルを、Transformer本体の前に別のONNX Runtime呼び出しとして実行
• TeaCache有効時のみ初期化されるTeaCacheStateで、累積L1距離をCPU側で外部から追跡

Transformerは累積距離にかかわらず最初と最後のタイムステップでは必ず実行します。現在のパイプラインはTransformerベースのSD3/SD3.5のみに対応し、UNetベースの他のStable Diffusionモデルには別のパイプラインが必要です。

実験条件とメモリ使用量

項目	値
プラットフォーム	Strix、Kraken、GPT、GPT2プロセッサー搭載のAMD Ryzen AI
OS	Microsoft Windows
ソフトウェア	Ryzen AI Software 1.8.0
SD3 PyTorch推論ステップ数	20
SD3/SD3.5 ONNXテキスト画像生成ステップ数	40
解像度	1024x1024
SD3/3.5テキスト画像生成 システムメモリ	約9.0～10.5 GB(512×512および1024×1024)
SD3/3.5テキスト画像生成 NPUメモリ	約5.5～6.1 GB
ControlNet システムメモリ	最大約11.9 GB
SD3.0インペインティング(1024×1024)	システム約15.1 GB、NPU約9.1 GB

AMDは32 GB RAMを推奨しています。16 GBでも基本的なテキスト画像生成は動作する可能性がありますが、ControlNetとインペインティングには推奨していません。

試し方

1. Ryzen AIをセットアップする
2. Ryzen AIのインストールに含まれるGenAI-SDパッケージのテストディレクトリへ移動する
3. ryzen-ai-1.8.0 Conda環境を有効化する
4. Ryzen AI Software 1.8.0同梱のONNX Runtime Vitis AI Execution Providerを使い、run.pyを-O1付きで起動してSD3.0またはSD3.5でTeaCacheを有効化する

記事ではSD3.5/SD3.0のテキスト画像生成に加え、SD3.0のControlNet(Canny、Pose、Tile、Depth)の実行例を載せています。ONNXベースのパイプラインでは、TeaCache用の追加ONNXモデルとTeaCacheStateを使用します。