diffusers/pr_14867/en/optimization/fp16.md

Precision and compilation

Lower precision and compilation are two ways to speed up Diffusers inference. Load weights in bfloat16 or float16, then compile the denoiser (the UNet or the transformer) with torch.compile or regional compilation.

Model data type

The precision and data type of the model weights affect inference speed because a higher precision requires more memory to load and more time to perform the computations. Diffusers loads model weights in float32 when you omit dtype, so changing the data type is a simple way to quickly get faster inference.

bfloat16 is similar to float16 but it is more robust to numerical errors. Hardware support for bfloat16 varies, but most modern GPUs are capable of supporting bfloat16.

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]

float16 is similar to bfloat16 but may be more prone to numerical errors.

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]

TensorFloat-32 (tf32) mode is supported on NVIDIA Ampere and newer GPUs and it computes the convolution and matrix multiplication operations in tf32. Storage and other operations are kept in float32. It speeds up operations that run in float32, so it helps most when the pipeline or parts of it stay in float32.

PyTorch only enables tf32 mode for convolutions by default and you'll need to explicitly enable it for matrix multiplications.

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]

Refer to the mixed precision training docs for more details.

Scaled dot product attention

Scaled dot product attention (SDPA) is the default attention on PyTorch 2.0 and later, through AttnProcessor2_0 or the attention dispatcher's native backend. PyTorch picks an SDPA kernel for your hardware. For FlashAttention, SageAttention, xFormers, Hub kernels, and other backends, see Attention backends.

torch.compile

torch.compile accelerates inference by compiling PyTorch code and operations into optimized kernels. You typically compile the denoiser that dominates runtime, either pipeline.unet or pipeline.transformer, and sometimes the VAE as well.

Enable the following compiler settings for maximum speed (refer to the full list for more options).

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

Load and compile the denoiser and VAE. Use pipeline.unet on UNet pipelines such as Stable Diffusion XL, or pipeline.transformer on Flux and DiT-style pipelines. There are several different modes you can choose from, but "max-autotune" searches for the fastest kernels and uses CUDA graphs. CUDA graphs reduce the overhead by launching multiple GPU operations through a single CPU operation.

Changing the memory layout to channels_last can also speed up convolution-heavy models like UNets and VAEs. Benchmark it first, because some models run slower.

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]

Compilation is slow the first time, especially with "max-autotune", because the compiler searches for kernels and keeps using the same compiled pipeline object afterward (see Compile Time Caching and Caching Configuration for how to configure the caching behavior). Calling the compiled pipeline on a different image size triggers recompilation.

Dynamic shape compilation

torch.compile keeps track of input shapes and conditions, and if these are different, it recompiles the model. For example, if a model is compiled on a 1024x1024 resolution image and used on an image with a different resolution, it triggers recompilation.

To avoid recompilation, add dynamic=True to try and generate a more dynamic kernel to avoid recompilation when conditions change.

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

Specifying use_duck_shape=False instructs the compiler if it should use the same symbolic variable to represent input sizes that are the same. For more details, check out this comment.

Not all models may benefit from dynamic compilation out of the box and may require changes. Refer to this PR that improved the AuraFlowPipeline implementation to benefit from dynamic compilation.

Feel free to open an issue if dynamic compilation doesn't work as expected for a Diffusers model.

Regional compilation

Regional compilation trims cold-start latency by only compiling the small and frequently-repeated block(s) of a model - typically a transformer layer - and enables reusing compiled artifacts for every subsequent occurrence. For many diffusion architectures, this delivers comparable runtime speedups to full-graph compilation and can cut compile time by up to 8–10x.

Use the compile_repeated_blocks() method, a helper that wraps torch.compile, on the denoiser. Call it on pipeline.unet or pipeline.transformer, depending on the 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)

To enable regional compilation for a new model, set _repeated_blocks to the class names of the repeated blocks (strings). For SDXL's UNet, use BasicTransformerBlock. Transformer denoisers list their own repeated block class names the same way.

_repeated_blocks = ["BasicTransformerBlock"]

There is also a compile_regions method in Accelerate that automatically selects candidate blocks in a model to compile. The remaining graph is compiled separately. This is useful for quick experiments because there aren't as many options for you to set which blocks to compile or adjust compilation flags.

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() is intentionally explicit. List the blocks to repeat in _repeated_blocks and the helper only compiles those blocks. It offers predictable behavior and easy reasoning about cache reuse in one line of code.

Graph breaks

Set fullgraph=True so torch.compile raises an error on a graph break instead of silently splitting the graph, which reduces the speedup.

GPU sync

After each denoiser prediction, the pipeline calls the scheduler's step method, which indexes sigmas. Keeping sigmas on the GPU can force a CPU↔GPU sync on every step. That cost is easy to miss until the denoiser is compiled. Prefer leaving sigmas on the CPU (Diffusers schedulers such as Euler already do this).

Benchmarks

Refer to the diffusers/benchmarks dataset to see inference latency and memory usage data for compiled pipelines.

The diffusers-torchao repository also contains benchmarking results for compiled versions of Flux and CogVideoX.

Kernels

Kernels is a library for building, distributing, and loading optimized compute kernels on the Hub. It supports attention kernels and custom CUDA kernels for operations like RMSNorm, GEGLU, RoPE, and AdaLN.

The Diffusers Pipeline Integration guide shows how to integrate a kernel with the add cuda-kernels skill. This skill enables an agent, like Claude or Codex, to write custom kernels targeted towards a specific model and your hardware. The Custom kernels for all from Codex and Claude post has more detail.

For example, a custom RMSNorm kernel (generated by the add cuda-kernels skill) with torch.compile speeds up LTX-Video generation 1.43x on an H100.