Image generation: diffusion and kernels

Coming from LLM inference, diffusion models felt alien. In LLM serving, the core loop is autoregressive: generate one token, feed it back, repeat. In diffusion, the core loop is iterative denoising: start with random noise, predict and remove a bit of noise, repeat 20 to 50 times. Each step is a full forward pass through a large neural network. The entire image is refined simultaneously, not pixel by pixel.

This means the optimization landscape is completely different from what I had been working with, and it took me a while to build the right mental model.

How diffusion inference works

The standard Stable Diffusion pipeline (and its successors like SDXL and SD3) has three components:

• Text encoder. A CLIP or T5 model that converts your text prompt into a sequence of embedding vectors. This runs once per generation and is relatively cheap.
• Denoising backbone. A UNet (in SD 1.5 and SDXL) or a DiT (Diffusion Transformer, in SD3 and Flux) that predicts the noise to be removed at each step. This runs once per denoising step, so 20 to 50 times per image. This is where nearly all the compute goes.
• VAE decoder. A variational autoencoder that converts the denoised latent representation into a full-resolution pixel image. Runs once at the end.

# Simplified diffusion inference loop
import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")

# What happens inside pipe():
# 1. Text encoding (once):
#    text_embeds = text_encoder(prompt_tokens)
#
# 2. Start with random noise:
#    latents = torch.randn(1, 4, 128, 128)  # latent space
#
# 3. Iterative denoising (N steps):
#    for t in scheduler.timesteps:   # e.g., 30 steps
#        noise_pred = unet(latents, t, text_embeds)
#        latents = scheduler.step(noise_pred, t, latents)
#
# 4. VAE decode (once):
#    image = vae.decode(latents)

image = pipe("a cat wearing a space helmet", num_inference_steps=30).images[0]

The compute profile

Unlike LLM decode (memory-bound) or embedding inference (compute-bound on small models), diffusion inference is solidly compute-bound. The UNet in SDXL has about 2.6 billion parameters, and you run it 30 times. That is roughly 78 billion parameter-forward-passes per image. The DiT in SD3 Medium is 2 billion parameters, and Flux Dev is 12 billion.

On an A100, a single SDXL image at 1024x1024 with 30 steps takes about 5 to 8 seconds in FP16. The breakdown is roughly:

• Text encoding: ~50 ms (less than 1%)
• UNet denoising (30 steps): ~5 to 7 seconds (90%+)
• VAE decode: ~200 to 400 ms (5 to 8%)

This tells you exactly where to focus optimization: the denoising loop.

Reducing the number of steps

The highest-impact optimization for diffusion is not a kernel trick. It is reducing the number of denoising steps. The scheduler (also called the sampler or solver) controls how many steps are needed to produce a good image:

• DDPM (the original): 1000 steps. Impractical for inference.
• DDIM: 50 to 100 steps. The first practical sampler.
• DPM++ 2M Karras: 20 to 30 steps. The workhorse for most SD deployments.
• LCM (Latent Consistency Models): 4 to 8 steps. Distilled models that trade some quality for dramatic speedup.
• Lightning/Turbo: 1 to 4 steps. Further distilled for near-real-time generation.

Going from 30 steps to 8 steps gives roughly a 3.75x speedup with no other changes. The quality tradeoff depends on the specific model and scheduler combination, but for many production use cases (thumbnails, previews, iterative design), 8-step generation is perfectly acceptable.

Classifier-free guidance and its cost

Most text-to-image models use classifier-free guidance (CFG): at each denoising step, the model runs twice, once with the text conditioning and once without (the "unconditional" pass). The two predictions are blended to steer the output toward the prompt:

# CFG doubles the compute per step
noise_pred_uncond = unet(latents, t, null_text_embeds)  # unconditional
noise_pred_text = unet(latents, t, text_embeds)          # conditioned

# Blend: steer toward the text prediction
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

This means every denoising step costs 2x the compute. For 30 steps with CFG, you are running the UNet 60 times. Optimizations that eliminate or reduce CFG are therefore very valuable:

• CFG distillation. Train the model to internalize the guidance, eliminating the unconditional pass. SDXL Turbo and Lightning models do this.
• Batched CFG. Instead of running the conditional and unconditional passes sequentially, batch them as a single forward pass with batch size 2. This increases memory usage but improves GPU utilization because larger batch sizes have better arithmetic intensity.
• Reduced guidance at later steps. Some practitioners use high CFG for the first few steps (to establish composition) and reduce it for later steps (refinement), saving some unconditional passes.

Attention kernel optimization

The UNet and DiT architectures use self-attention and cross-attention layers, just like LLMs. The same attention kernel optimizations apply:

• FlashAttention. Reduces memory from O(n^2) to O(n) and speeds up the attention computation. For high-resolution generation (1024x1024 and above), the spatial attention sequences are long enough (4096+ tokens for the deepest UNet level) that FlashAttention makes a meaningful difference.
• xformers memory-efficient attention. The xformers library provides memory-efficient attention kernels that are well-tested with diffusion models. Many diffusers pipelines enable this with a single flag.
• Scaled dot-product attention (SDPA). PyTorch 2.0+ includes F.scaled_dot_product_attention, which automatically selects the best available kernel (FlashAttention, memory-efficient, or math fallback). This is now the default in most diffusion libraries.

# Enable memory-efficient attention in diffusers
pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to("cuda")

# Option 1: xformers (if installed)
pipe.enable_xformers_memory_efficient_attention()

# Option 2: PyTorch SDPA (default in recent diffusers)
# Automatically used when xformers is not explicitly enabled
# Uses FlashAttention-2 when available

Compilation and graph optimization

Because the denoising backbone runs the same computation 20 to 50 times with different inputs, it is an ideal candidate for graph compilation:

• torch.compile. Compiling the UNet with torch.compile(mode="reduce-overhead") fuses operations and eliminates Python overhead. First-step latency increases (compilation cost), but subsequent steps are 10 to 30% faster. Over 30 steps, the amortized benefit is substantial.
• TensorRT. NVIDIA's TensorRT can compile the UNet into a highly optimized engine with layer fusion, kernel auto-tuning, and FP16/INT8 precision. The diffusers library has TensorRT integration, and NVIDIA provides optimized UNet engines as part of their TensorRT demos.
• ONNX Runtime. Export the model to ONNX and run with ORT's CUDA execution provider. This gives cross-platform compatibility and reasonable optimization without the complexity of TensorRT.

In practice, torch.compile is the easiest win for diffusion inference. The compilation cache means you pay the compile cost once and reuse it across all subsequent images.

VAE decoding: the forgotten bottleneck

The VAE decoder converts the small latent representation (128x128 for SDXL at 1024x1024 output) into full-resolution pixels. It is a single forward pass, but for high-resolution outputs it can take 200 to 500 ms. Two optimizations help:

• Tiled VAE decoding. Decode the latent in overlapping tiles instead of all at once. This dramatically reduces peak memory (important for high-resolution generation) and can be slightly faster due to better cache utilization. The diffusers library supports this via pipe.enable_vae_tiling().
• FP16 VAE. Many pipelines run the VAE in FP32 by default for numerical stability, but FP16 works fine for most models and halves the VAE memory and compute.

Batching and throughput

Diffusion batching is straightforward compared to LLMs. All images in a batch go through the same number of denoising steps (assuming the same scheduler settings), so there is no variable-length padding issue. You simply stack the latents and text embeddings along the batch dimension.

The constraint is memory. A single SDXL generation at 1024x1024 in FP16 uses about 8 to 10 GB of GPU memory. An A100 with 80 GB can batch roughly 6 to 8 images simultaneously. For throughput-optimized serving, you want to fill the GPU with a batch and run all steps together.

Diffusion inference inverts the LLM optimization playbook. In LLMs, you fight memory bandwidth during decode. In diffusion, you fight compute during denoising. In LLMs, more steps means more tokens (good). In diffusion, fewer steps means faster images (also good). The only constant is this: measure first, then optimize the actual bottleneck.

This wraps up the core modalities: embeddings, ASR, TTS, and now image generation. Next we will look at video generation, where context parallelism becomes essential because the spatial and temporal dimensions of video push even the largest GPUs to their limits.