精度とコンパイル

低精度化とコンパイルは、Diffusersの推論を高速化する2つの方法です。bfloat16またはfloat16で重みを読み込み、その後、torch.compileまたはリージョナルコンパイルを使用して、デノイザー（UNetまたはtransformer）をコンパイルします。

モデルのデータ型

モデルの重みの精度とデータ型は推論速度に影響します。精度が高いほど、読み込みに必要なメモリ量が増え、計算に要する時間も長くなるためです。Diffusers は dtype を省略した場合、モデルの重みを float32 で読み込むため、データ型を変更することは、推論を素早く高速化する簡単な方法です。

bfloat16

float16

TensorFloat-32

スケーリングドットプロダクトアテンション

Scaled dot product attention (SDPA)は、PyTorch 2.0以降のデフォルトのattentionです。AttnProcessor2_0またはattention dispatcherのnativeバックエンドを通じて利用できます。PyTorchは、お使いのハードウェアに最適なSDPAカーネルを選択します。FlashAttention、SageAttention、xFormers、Hubカーネル、その他のバックエンドについては、Attention backendsを参照してください。

torch.compile

torch.compile は、PyTorch のコードと操作を最適化されたカーネルにコンパイルすることで推論を高速化します。通常、実行時間を支配するデノイザー、つまり pipeline.unet または pipeline.transformer をコンパイルし、場合によっては VAE もコンパイルします。

最大速度を実現するために、以下のコンパイラ設定を有効にしてください（その他のオプションについては、完全なリストを参照してください）。

import torch
from diffusers import StableDiffusionXLPipeline

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True

デノイザーとVAEを読み込み、コンパイルします。Stable Diffusion XLなどのUNetパイプラインにはpipeline.unetを、FluxやDiTスタイルのパイプラインにはpipeline.transformerを使用します。選択できるモードはいくつかありますが、"max-autotune"は最速のカーネルを検索し、CUDAグラフを使用します。CUDAグラフは、単一のCPU操作で複数のGPU操作を発行することでオーバーヘッドを削減します。

メモリレイアウトをchannels_lastに変更すると、UNetやVAEのような畳み込みが中心のモデルの高速化にもつながります。ただし、一部のモデルでは実行速度が遅くなる場合があるため、事前にベンチマークを実行してください。

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet.to(memory_format=torch.channels_last)
pipeline.unet = torch.compile(
    pipeline.unet, mode="max-autotune", fullgraph=True
)
pipeline.vae.to(memory_format=torch.channels_last)
pipeline.vae.decode = torch.compile(
    pipeline.vae.decode,
    mode="max-autotune",
    fullgraph=True
)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

初回コンパイルは、特に"max-autotune"を使用する場合、コンパイラがカーネルを検索するため遅くなります。その後、同じコンパイル済みパイプラインオブジェクトが再利用されます（キャッシュ動作の設定方法については、コンパイル時間のキャッシュおよびキャッシュ設定を参照してください）。異なる画像サイズでコンパイル済みパイプラインを呼び出すと、再コンパイルがトリガーされます。

動的形状コンパイル

torch.compile は入力形状と条件をトラッキングし、これらが異なる場合はモデルを再コンパイルします。例えば、1024x1024 の解像度の画像でモデルをコンパイルし、異なる解像度の画像で使用した場合、再コンパイルがトリガーされます。

再コンパイルを回避するために、dynamic=True を追加し、条件が変更された際に再コンパイルを回避するためのより動的なカーネルの生成を試みてください。

+ torch.fx.experimental._config.use_duck_shape = False
+ pipeline.unet = torch.compile(
    pipeline.unet, fullgraph=True, dynamic=True
)

use_duck_shape=False を指定すると、コンパイラに、同じである入力サイズを表すのに同じシンボリック変数を使用すべきかどうかを指示します。詳細については、このコメントを参照してください。

すべてのモデルが、そのままの状態で動的コンパイルの恩恵を受けられるとは限りません。変更が必要な場合があります。AuraFlowPipelineの実装を改善し、動的コンパイルの恩恵を受けられるようにしたこのPRを参照してください。

Diffusersモデルで動的コンパイルが期待通りに動作しない場合は、お気軽にissueを作成してください。

地域別コンパイル

Regional compilationは、モデルの小さく頻繁に繰り返されるブロック（通常はtransformer層）のみをコンパイルすることで、コールドスタートのレイテンシを削減します。これにより、以降のすべての出現箇所に対してコンパイル済み成果物を再利用できます。多くの拡散アーキテクチャにおいて、これはグラフ全体のコンパイルと同等のランタイム高速化を実現し、コンパイル時間を最大で8〜10分の1に短縮できます。

デノイザーに、torch.compileをラップするヘルパーであるcompile_repeated_blocks()メソッドを使用します。パイプラインに応じて、pipeline.unetまたはpipeline.transformerに対して呼び出してください。

import torch
from diffusers import StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet.compile_repeated_blocks(fullgraph=True)

新しいモデルでリージョナルコンパイルを有効にするには、_repeated_blocks に繰り返しブロックのクラス名（文字列）を設定します。SDXL の UNet には、BasicTransformerBlock を使用します。Transformer デノイザーも、同様に自身の繰り返しブロックのクラス名を列挙します。

_repeated_blocks = ["BasicTransformerBlock"]

Accelerate には、モデル内でコンパイルする候補ブロックを自動的に選択する compile_regions メソッドもあります。残りのグラフは個別にコンパイルされます。コンパイル対象のブロックを指定したり、コンパイルフラグを調整したりするオプションが少ないため、素早い実験に便利です。

pip install -U accelerate

import torch
from diffusers import StableDiffusionXLPipeline
from accelerate.utils import compile_regions

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet = compile_regions(pipeline.unet, mode="reduce-overhead", fullgraph=True)

compile_repeated_blocks() は意図的に明示的な設計です。_repeated_blocks に反復するブロックを列挙し、ヘルパーはそれらのブロックのみをコンパイルします。これにより、1 行のコードでキャッシュ再利用について予測可能な動作と容易な推論が可能になります。

グラフのブレイク

fullgraph=True を設定し、グラフの分割が黙って行われるのではなく、グラフのブレイク時に torch.compile がエラーを送出するようにします。グラフの分割はスピードアップを低下させるためです。

GPU同期

各デノイザーの予測後、パイプラインはスケジューラーのstepメソッドを呼び出し、sigmasにインデックスを指定します。sigmasをGPU上に保持すると、各ステップでCPU↔GPUの同期が強制される可能性があります。このコストは、デノイザーがコンパイルされるまで見逃されがちです。sigmasをCPU上に保持しておくことをお勧めします（EulerなどのDiffusersスケジューラーはすでにこれを行っています）。

ベンチマーク

コンパイル済みパイプラインの推論レイテンシとメモリ使用量のデータを確認するには、diffusers/benchmarks データセットを参照してください。

diffusers-torchao リポジトリには、Flux と CogVideoX のコンパイル済みバージョンに関するベンチマーク結果も含まれています。

カーネル

Kernels は、Hub 上で最適化された計算カーネルを構築・配布・読み込むためのライブラリです。attention カーネルや、RMSNorm、GEGLU、RoPE、AdaLN などの演算用のカスタム CUDA カーネルに対応しています。

Diffusers Pipeline Integrationガイドでは、add cuda-kernelsスキルを使用してカーネルを統合する方法を示しています。このスキルにより、ClaudeやCodexなどのエージェントが、特定のモデルやあなたのハードウェア向けのカスタムカーネルを作成できるようになります。詳細については、Codex and Claudeによるすべてのカスタムカーネルの記事を参照してください。

例えば、add cuda-kernels スキルによって生成されたカスタム RMSNorm カーネルに torch.compile を組み合わせることで、H100 上で LTX-Video の生成を 1.43 倍高速化できます。

動的量子化

動的 int8 アクティベーション量子化は torchao で利用できます。TorchAoConfig を Int8DynamicActivationInt8WeightConfig のような設定と組み合わせて使用するか、torchao の quantize_ API を使用してください。旧式の apply_dynamic_quant ヘルパーは使用しないでください。

融合された射影行列

~DiffusionPipeline.fuse_qkv_projections は実験的な機能です。

Q、K、Vの射影行列によって表される3つの部分空間に、アテンションブロック内で入力が射影されます。これらの射影は通常、個別に計算されますが、これらを水平方向に結合して単一の行列にし、1ステップで射影を実行することもできます。これにより、入力射影の行列積のサイズが増大し、量子化の影響も向上します。

pipeline.fuse_qkv_projections()

次のステップ

Presenting Flux Fast: Making Flux go brrr on H100s のブログ記事を読み、TorchInductor と AOTInductor をこれらの最適化すべてと組み合わせることで、flux-fast のレシピを使用して約2.5倍の高速化を実現する方法について詳しく学びましょう。

これらのレシピは、AMD ハードウェアおよび Flux.1 Kontext Dev をサポートしています。
拡散モデルでtorch.compileによるパフォーマンスの最大化については、torch.compile と Diffusers: ピークパフォーマンスを実現する実践ガイドというブログ記事をご参照ください。

Update on GitHub