6.5 Image Generation Models

How serving image generation differs from LLM serving, and the optimizations specific to it.

Figure 6.8: Image generation models may accept both text and reference images to create new output images.
Figure 6.8: Image generation models may accept both text and reference images to create new output images.

Working with image and video generation models is entirely different from working with large language models on a few axes.

The first is architecture. While some recent models like HunyuanImage-3.0 more closely resemble LLMs, most image and video generation models are iterative denoisers, not autoregressive token generators. Image generation models are pipelines with multiple small models working together in latent space rather than the uniform decoder architecture of an LLM.

As such, the tooling is different. At the time of publication, SGLang Diffusion and vLLM Omni are brand new. Most image and video generation model inference is implemented lower in the stack, working with PyTorch or TensorRT directly.

The constraints are different too. Image generation models are ten to twenty times smaller than frontier language models, and inference is constrained on compute, not bandwidth.

But perhaps the most significant difference is that image and video generation models offer more direct quality to speed tradeoffs.

Evaluating image model output quality is difficult to do programmatically. Automatic pipelines using vision language models give directional signal at best and may diverge from human preferences. The human eye is mysterious, and most image quality evals work by asking humans to pick among thousands of images to aggregate vibes and preferences into quality benchmarks.

6.5.1 Image Generation Kernel Optimization

When you read a model card for an image generation model from its repository, inference examples generally use the diffusers library with very few optimizations.

In fact, while image generation is theoretically compute bound, you often need to select memory-efficient kernels and use kernel fusion to even reach that bottleneck.

High-performance image model inference uses one of three libraries:

  • SGLang Diffusion: Performant inference engines built for popular image and video generation architectures.
  • TensorRT: High-quality black-box implementations of popular models with NVIDIA’s in-house kernels.
  • PyTorch: Careful kernel selection and fusion yields control, flexibility, and improved high-end performance.

If you want something that works well and you want it now, just use the SGLang Diffusion or TensorRT implementation of a model. But with PyTorch, there’s an opportunity for advanced inference engineers to do deep customization.

The most essential kernel is the attention kernel. Many image generation models use FlashAttention 2 out of the box, but FlashAttention 3 and 4 yield better performance on Hopper and Blackwell GPUs, respectively.

There is a whole barrage of smaller kernels, especially normalization functions like RMSNorm, that are good candidates for fusion to ensure efficient memory usage.

Then, GEMM kernels matter for compute-bound inference. GEMM kernels apply to linear layers, and are generally safe to quantize into 8-bit floating point formats to access two times higher FLOPS on Tensor Cores. Kernels from CuTe, CUTLASS, or DeepGEMM may prove most efficient on a model-by-model basis.

Torch compilation includes automatic kernel fusion with a plugin system for inserting manually selected kernels, and the resulting engine can be cached for faster load times on node startup (which is important because compilation takes several minutes).

Like most high-performance engines, Torch compilation targets the specific GPU model and architecture performing the compilation – if you want to run the model on a B200, do the compilation on a B200.

6.5.2 One Weird Trick for Faster Image Generation

Kernel selection and Torch compilation are all bona fide inference optimization techniques. But the world of inference optimization has fun hacks as well, and here’s one of them.

Figure 6.9: Recall that diffusion is a step-by-step process and that the general outline of the image is established in early steps.
Figure 6.9: Recall that diffusion is a step-by-step process and that the general outline of the image is established in early steps.

Image generation time tracks linearly with step count. That’s why few-step models and latent consistency models are so much faster than full 50-step models. But reducing step count may reduce image quality below an acceptable threshold.

Each pass through the denoising model is run at a batch size of two as each step includes a pass with and without prompt guidance.

As a refresher, the guidance parameter controls how much the prompt-guided image is weighted when combining the two iterations generated on each step. If the guidance is zero, the prompt-guided image does not need to be generated.

After the first few steps, the basic outline of the image is in place, and the rest of the steps are for filling in the details. Thus, prompt adherence is more important in early steps that affect the broad strokes of the image – the model is not going to change its mind on later steps and generate a dog when it is in the middle of generating a cat.

If you turn off guidance partway through the image generation, you save passes through the denoiser without reducing step count. If guidance is skipped for the last 20 steps of a 50-step run, there are only 80 passes through the model instead of 100, and quality generally remains high.