Forem: Elise Moreau

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

Elise Moreau — Mon, 27 Apr 2026 04:48:07 +0000

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.

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

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

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.

Kimi K2.6 Is a Legit Opus 4.7 Replacement

Elise Moreau — Mon, 27 Apr 2026 04:39:40 +0000

For a long time, Opus 4.7 has been the default recommendation when someone asks for a top tier model. It has been reliable, capable, and strong across a wide range of tasks.

After spending real time with Kimi K2.6 and gathering feedback from customers using it in production workflows, I have started to change my mind. It is the first model I feel comfortable recommending as a practical replacement for Opus 4.7.

Not better, but close enough

Kimi K2.6 is not outright better than Opus 4.7. If you are comparing raw performance on difficult reasoning or edge case tasks, Opus still wins.

What matters more in practice is coverage. Kimi K2.6 can handle around 85 percent of the tasks that Opus can, and it does so at a quality level that is good enough for real work. That gap sounds large on paper, but in day to day usage it is surprisingly small.

Most users are not constantly pushing models to their limits. They need something that works consistently across writing, coding, research, and general problem solving. In that context, Kimi K2.6 holds up very well.

Strong features that actually matter

Two areas where Kimi K2.6 stands out are vision and browser use.

Vision is not just a checkbox feature here. It is genuinely useful for workflows that involve screenshots, documents, or UI level debugging. Being able to mix text and visual context smoothly removes a lot of friction.

Browser use is another big win. It handles multi step information gathering better than expected, especially for longer tasks where the model needs to plan, search, and refine results over time.

These features are not always the headline benchmarks, but they have a real impact on productivity.

Surprisingly good at long horizon tasks

One of the more unexpected strengths of Kimi K2.6 is how well it handles longer time horizon work.

I have been slowly replacing parts of my personal workflows with it, including tasks that require multiple steps, iteration, and context retention. It performs more reliably than I expected, and it does not fall apart as quickly over extended interactions.

This makes it useful for things like research threads, content pipelines, and multi step coding tasks.

The size question and what it signals

Kimi K2.6 is a very large model. There is no getting around that.

But its performance raises an interesting point. Frontier models like Opus 4.7 are not necessarily introducing completely new capabilities. Instead, we are seeing strong alternatives that can replicate most of that value.

If a model can deliver 80 to 90 percent of the experience, the remaining gap starts to matter less, especially when other factors come into play.

Limits, cost, and the shift to local

One of the biggest complaints around models like Opus 4.7 is usage limits. As demand increases, constraints become more noticeable.

This is where models like Kimi K2.6 become more attractive. There is growing interest in running models locally or in more controlled environments, where limits are less of a concern.

It feels like the conversation is starting to shift. Instead of chasing the absolute best model, people are looking for models that are good enough, flexible, and easier to integrate into their own systems.

Final thoughts

Kimi K2.6 is not a perfect replacement for Opus 4.7. If you need the absolute best performance on every task, Opus is still ahead.

But for most real world use cases, Kimi K2.6 gets you very close. Close enough that the tradeoffs start to make sense.

That is what makes it interesting. Not that it beats Opus, but that it makes you question whether you still need Opus at all.

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

Elise Moreau — Thu, 23 Apr 2026 14:48:58 +0000

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,
)

The text encoder trap

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.

Diffusion Model Inference in Production: What the Benchmarks Leave Out

Elise Moreau — Thu, 23 Apr 2026 14:31:27 +0000

BIFROST COMMENT

The routing overhead caught us off guard. We were running caption generation through a larger model for every input when 70% of them only needed a fast small model. Adding a gateway with cost-aware routing (we landed on Bifrost for this, though LiteLLM and Portkey do the same thing: https://github.com/maximhq/bifrost) cut LLM spend in our vision pipeline by 38% without touching the heavy-model cases.

Forem: Elise Moreau

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

What the profile showed

Fixing the VAE

The text encoder trap

CPU sync between steps

Final numbers

Trade-offs and limitations

Further reading

Kimi K2.6 Is a Legit Opus 4.7 Replacement

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

What the profile showed

Fixing the VAE

The text encoder trap

CPU sync between steps

Final numbers

Trade-offs and limitations

Further reading

Diffusion Model Inference in Production: What the Benchmarks Leave Out

BIFROST COMMENT