Presisi dan kompilasi

Presisi rendah dan kompilasi adalah dua cara untuk mempercepat inferensi Diffusers. Muat bobot dalam bfloat16 atau float16, lalu kompilasi denoiser (UNet atau transformer) dengan torch.compile atau kompilasi regional.

Tipe data model

Presisi dan tipe data dari bobot model memengaruhi kecepatan inferensi karena presisi yang lebih tinggi membutuhkan lebih banyak memori untuk dimuat dan lebih banyak waktu untuk melakukan komputasi. Diffusers memuat bobot model dalam float32 ketika Anda mengabaikan dtype, sehingga mengubah tipe data adalah cara sederhana untuk mendapatkan inferensi yang lebih cepat.

import torch
from diffusers import StableDiffusionXLPipeline

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

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

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"

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

import torch
from diffusers import StableDiffusionXLPipeline

torch.backends.cuda.matmul.allow_tf32 = True

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

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

Perhatian hasil kali titik berskala

Scaled dot product attention (SDPA) adalah perhatian default pada PyTorch 2.0 dan versi setelahnya, melalui AttnProcessor2_0 atau backend native dari dispatcher perhatian. PyTorch memilih kernel SDPA untuk perangkat keras Anda. Untuk FlashAttention, SageAttention, xFormers, kernel Hub, dan backend lainnya, lihat Attention backends.

torch.compile

torch.compile mempercepat inferensi dengan mengompilasi kode dan operasi PyTorch menjadi kernel yang dioptimalkan. Anda biasanya mengompilasi denoiser yang mendominasi waktu eksekusi, yaitu pipeline.unet atau pipeline.transformer, dan terkadang juga VAE.

Aktifkan pengaturan compiler berikut untuk kecepatan maksimal (lihat daftar lengkap untuk opsi lainnya).

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

Muat dan kompilasi denoiser serta VAE. Gunakan pipeline.unet pada pipeline UNet seperti Stable Diffusion XL, atau pipeline.transformer pada pipeline Flux dan gaya DiT. Ada beberapa mode berbeda yang dapat Anda pilih, tetapi "max-autotune" mencari kernel tercepat dan menggunakan CUDA graphs. CUDA graphs mengurangi overhead dengan meluncurkan beberapa operasi GPU melalui satu operasi CPU.

Mengubah tata letak memori menjadi channels_last juga dapat mempercepat model yang banyak menggunakan konvolusi seperti UNets dan VAEs. Lakukan benchmark terlebih dahulu, karena beberapa model justru berjalan lebih lambat.

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]

Kompilasi berjalan lambat pada percobaan pertama, terutama dengan "max-autotune", karena compiler mencari kernel dan kemudian terus menggunakan objek pipeline terkompilasi yang sama (lihat Compile Time Caching dan Caching Configuration untuk cara mengonfigurasi perilaku caching). Memanggil pipeline terkompilasi pada ukuran gambar yang berbeda memicu kompilasi ulang.

Kompilasi bentuk dinamis

torch.compile melacak bentuk input dan kondisi, dan jika berbeda, model akan dikompilasi ulang. Misalnya, jika model dikompilasi pada gambar dengan resolusi 1024x1024 dan digunakan pada gambar dengan resolusi berbeda, kompilasi ulang akan dipicu.

Untuk menghindari kompilasi ulang, tambahkan dynamic=True untuk mencoba menghasilkan kernel yang lebih dinamis sehingga kompilasi ulang dapat dihindari ketika kondisi berubah.

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

Menentukan use_duck_shape=False memberi instruksi kepada compiler apakah harus menggunakan variabel simbolik yang sama untuk merepresentasikan ukuran input yang sama. Untuk detail lebih lanjut, lihat komentar ini.

Tidak semua model mungkin mendapat manfaat dari kompilasi dinamis secara bawaan dan mungkin memerlukan perubahan. Lihat PR ini yang meningkatkan implementasi AuraFlowPipeline agar dapat memanfaatkan kompilasi dinamis.

Silakan buka issue jika kompilasi dinamis tidak berfungsi sesuai harapan untuk model Diffusers.

Kompilasi regional

Kompilasi regional memangkas latensi cold-start dengan hanya mengompilasi blok kecil yang sering diulang dari sebuah model - biasanya satu lapisan transformer - dan memungkinkan penggunaan ulang artefak terkompilasi untuk setiap kemunculan berikutnya. Untuk banyak arsitektur difusi, hal ini memberikan peningkatan kecepatan runtime yang setara dengan kompilasi grafik penuh dan dapat memangkas waktu kompilasi hingga 8–10x.

Gunakan metode compile_repeated_blocks(), sebuah helper yang membungkus torch.compile, pada denoiser. Panggil pada pipeline.unet atau pipeline.transformer, tergantung pada pipeline.

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)

Untuk mengaktifkan kompilasi regional untuk model baru, atur _repeated_blocks ke nama kelas dari blok yang diulang (string). Untuk UNet SDXL, gunakan BasicTransformerBlock. Denoiser Transformer mencantumkan nama kelas blok berulang mereka sendiri dengan cara yang sama.

_repeated_blocks = ["BasicTransformerBlock"]

Ada juga metode compile_regions di Accelerate yang secara otomatis memilih blok kandidat dalam model untuk dikompilasi. Sisa graf dikompilasi secara terpisah. Ini berguna untuk eksperimen cepat karena tidak ada banyak opsi untuk Anda mengatur blok mana yang akan dikompilasi atau menyesuaikan flag kompilasi.

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() sengaja dibuat eksplisit. Daftarkan blok yang akan diulang dalam _repeated_blocks dan helper hanya akan mengompilasi blok-blok tersebut. Hal ini menawarkan perilaku yang dapat diprediksi dan memudahkan penalaran tentang penggunaan ulang cache dalam satu baris kode.

Graph breaks

Setel fullgraph=True agar torch.compile menghasilkan error saat terjadi graph break, alih-alih secara diam-diam memecah graph, yang akan mengurangi percepatan.

Sinkronisasi GPU

Setiap prediksi denoiser, pipeline memanggil metode step dari scheduler, yang mengindeks sigmas. Menjaga sigmas di GPU dapat memaksa sinkronisasi CPU↔GPU pada setiap langkah. Biaya ini mudah terlewat sampai denoiser dikompilasi. Lebih baik biarkan sigmas tetap di CPU (scheduler Diffusers seperti Euler sudah melakukannya).

Benchmarks

Lihat dataset diffusers/benchmarks untuk melihat data latensi inferensi dan penggunaan memori untuk pipeline yang dikompilasi.

Repositori diffusers-torchao juga berisi hasil benchmarking untuk versi terkompilasi dari Flux dan CogVideoX.

Kernels

Kernels adalah pustaka untuk membangun, mendistribusikan, dan memuat kernel komputasi yang dioptimalkan di Hub. Pustaka ini mendukung kernel attention dan kernel CUDA kustom untuk operasi seperti RMSNorm, GEGLU, RoPE, dan AdaLN.

Panduan Integrasi Pipeline Diffusers menunjukkan cara mengintegrasikan kernel dengan skill add cuda-kernels. Skill ini memungkinkan agen, seperti Claude atau Codex, untuk menulis kernel kustom yang ditargetkan ke model tertentu dan perangkat keras Anda. Postingan Custom kernels for all from Codex and Claude memiliki detail lebih lanjut.

Sebagai contoh, kernel RMSNorm kustom (yang dihasilkan oleh skill add cuda-kernels) dengan torch.compile mempercepat generasi LTX-Video sebesar 1,43x pada H100.

Kuantisasi dinamis

Kuantisasi aktivasi int8 dinamis tersedia di torchao. Gunakan TorchAoConfig dengan konfigurasi seperti Int8DynamicActivationInt8WeightConfig, atau API quantize_ dari torchao. Jangan gunakan helper apply_dynamic_quant yang lebih lama.

Matriks proyeksi yang digabungkan

~DiffusionPipeline.fuse_qkv_projections bersifat eksperimental.

Sebuah input diproyeksikan ke tiga subruang, yang direpresentasikan oleh matriks proyeksi Q, K, dan V, dalam sebuah blok perhatian. Proyeksi-proyeksi ini biasanya dihitung secara terpisah, tetapi Anda dapat menggabungkannya secara horizontal menjadi satu matriks tunggal dan melakukan proyeksinya dalam satu langkah. Hal ini meningkatkan ukuran perkalian matriks dari proyeksi input dan juga meningkatkan dampak dari kuantisasi.

pipeline.fuse_qkv_projections()

Langkah selanjutnya

• Baca artikel blog Presenting Flux Fast: Making Flux go brrr on H100s untuk mempelajari lebih lanjut tentang cara Anda menggabungkan semua optimasi ini dengan TorchInductor dan AOTInductor untuk mendapatkan percepatan sekitar 2,5x menggunakan resep dari flux-fast.

  Resep-resep ini mendukung perangkat keras AMD dan Flux.1 Kontext Dev.

• Lihat posting blog torch.compile and Diffusers: A Hands-On Guide to Peak Performance untuk memaksimalkan kinerja dengan torch.compile untuk model difusi.