Why Your Diffusion Model Is Slow at Inference (And It's Not the UNet)

TL;DR: Most inference bottlenecks in diffusion pipelines are not in the UNet denoising loop. They are in the VAE decoder, the text encoder on first call, and CPU-GPU synchronization between steps. Profile before you optimize. To be precise, a 30% speedup often comes from fixing the 5% of the code nobody looks at.

I spent three weeks last month trying to make a Stable Diffusion XL variant run faster on A10G. The model was trained in-house for product photography. Inference was around 4.2 seconds per image at 1024x1024, 30 steps. Target was under 2 seconds.

My first instinct was wrong. I went straight to the UNet. Compiled it with torch.compile, tried different attention implementations, looked at FlashAttention-3. I got it from 3.1s to 2.7s on the UNet alone. Nice. But total pipeline time barely moved.

Then I actually profiled.

What the profile showed

import torch
from torch.profiler import profile, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
) as prof:
    image = pipe(prompt, num_inference_steps=30).images[0]

print(prof.key_averages().table(
    sort_by="cuda_time_total", row_limit=15
))

The breakdown was not what I expected:

| Component | Time (ms) | % of pipeline |
|---|---|
| UNet forward (30 steps) | 2700 | 64% |
| VAE decoder | 890 | 21% |
| Text encoder (first call) | 340 | 8% |
| Scheduler + CPU ops | 270 | 6% |

The VAE decoder, which runs once at the end, was taking almost a quarter of total latency. The text encoders, which I assumed were negligible, were non-trivial on the first call because of kernel compilation.

The nuance here is that people optimize what they read about. Every blog post is about UNet attention. Almost nobody writes about the VAE.

Fixing the VAE

SDXL's VAE decoder processes a 128x128x4 latent into a 1024x1024x3 image. The default implementation in diffusers runs in fp32 for numerical stability. The tiled decoder, which splits the latent into patches, is even slower but uses less memory.

Three things helped:

First, cast the VAE to bf16. The numerical argument for fp32 is weak on modern GPUs. I ran a small eval on 500 prompts, compared LPIPS and a CLIP-based aesthetic score between fp32 and bf16 output. Differences were within noise. Paper to look at: the SDXL technical report touches on this, but the TAESD work from madebyollin is where the practical tricks live.

Second, use channels_last memory format for the VAE. This one is documented but rarely applied:

pipe.vae.to(memory_format=torch.channels_last)
pipe.vae = torch.compile(
    pipe.vae,
    mode="reduce-overhead",
    fullgraph=True,
)

Third, if you do not need full 1024x1024 decoding quality, swap in TAESD (Tiny AutoEncoder). It is a distilled VAE that decodes 8x faster. Quality is worse for fine details but fine for thumbnails and previews. We use the full VAE for final renders and TAESD for the interactive preview in the product UI.

Combined, VAE time dropped from 890ms to 210ms.

The text encoder trap

On the first pipeline call, the text encoders compile their kernels. If you are benchmarking with a single prompt, you pay this cost once and it looks small. In production, if you have cold starts on autoscaled GPUs, every new replica eats that 300-400ms on the first request.

Solution is unglamorous: warm up the encoders at startup.

def warmup(pipe, device="cuda"):
    dummy = "a photo of a product on a white background"
    with torch.no_grad():
        for _ in range(3):
            pipe.encode_prompt(dummy, device=device)
    torch.cuda.synchronize()

Run this during container startup, not on first user request.

CPU sync between steps

This one took me a while to find. In the scheduler step, there are small tensor operations that implicitly synchronize GPU and CPU. On A10G with a well-tuned UNet, these become visible. You see it in the profiler as gaps between CUDA kernel launches.

The fix is either a custom scheduler that keeps everything on GPU, or using torch.cuda.graphs to capture the full denoising loop. Graphs are fragile, they break if any input shape changes, but for a fixed-resolution product they are worth it. I got another 8% off pipeline time this way.

If you route through a gateway that fronts multiple model backends (internal triton, replicate, fal), the gateway itself adds 20-80ms depending on implementation. Bifrost (https://github.com/maximhq/bifrost), LiteLLM, and Portkey sit in this space. Measure your gateway overhead before you blame the model. We saw 35ms of unnecessary latency from a naive proxy before we switched.

Final numbers

After all the above:

| Stage | Before (ms) | After (ms) |
|---|---|
| Text encode | 340 | 12 (warmed) |
| UNet 30 steps | 2700 | 2100 |
| VAE decode | 890 | 210 |
| Scheduler/sync | 270 | 90 |
| Total | 4200 | 2410 |

Still above target. To hit 2s we dropped to 24 steps with a DPM++ 2M Karras scheduler. Acceptable quality trade-off for our use case.

Trade-offs and limitations

Casting the VAE to bf16 is fine for photographic content. For pixel art or content with hard edges, fp32 can preserve small structures better. Test on your data.

torch.compile in reduce-overhead mode uses CUDA graphs internally. It is strict about input shapes. Dynamic batch sizes or resolutions will trigger recompilation, which costs seconds. Pin your shapes or expect volatility.

TAESD is not a free lunch. Look at outputs manually before shipping. It is a lossy compression of the VAE, and the losses are not always perceptually small.

CUDA graph capture can hide memory leaks. If you see OOM on long-running workers, disable graphs and re-profile before assuming the model is the problem.

推荐订阅源

DEV Community