vllm-omni/docs/user_guide/diffusion_acceleration.md

Diffusion Acceleration Overview

vLLM-Omni supports various acceleration methods to speed up diffusion model inference with minimal quality degradation. These include cache methods that intelligently cache intermediate computations to avoid redundant work across diffusion timesteps, parallelism methods that distribute the computation across multiple devices, and quantization methods that reduce memory footprint while preserving accuracy.

Supported Acceleration Methods

vLLM-Omni currently supports two main cache acceleration backends:

1. TeaCache - Hook-based adaptive caching that caches transformer computations when consecutive timesteps are similar
2. Cache-DiT - Library-based acceleration using multiple techniques:
   • DBCache (Dual Block Cache): Caches intermediate transformer block outputs based on residual differences
   • TaylorSeer: Uses Taylor expansion-based forecasting for faster inference
   • SCM (Step Computation Masking): Selectively computes steps based on adaptive masking

Both methods can provide significant speedups (typically 1.5x-2.0x) while maintaining high output quality.

vLLM-Omni also supports quantization methods:

3. Quantization - Reduces DiT linear layers from BF16 to FP8 or Int8, providing ~1.28x speedup with minimal quality loss. Supports per-layer skip for sensitive layers.

vLLM-Omni also supports parallelism methods for diffusion models, including:

1. Ulysses-SP - splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads.

2. Ring-Attention - splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results, keeping the sequence dimension sharded.

3. CFG-Parallel - runs the positive/negative prompts of classifier-free guidance (CFG) on different devices, then merges on a single device to perform the scheduler step.

4. Tensor Parallelism - shards DiT weights across devices to reduce per-GPU memory usage.

5. VAE Patch Parallelism - shards VAE decode/encode spatially across ranks to reduce VAE peak memory (and can speed up VAE decode).

6. HSDP - Hybrid Sharded Data Parallel shards model weights across GPUs to reduce per-GPU memory usage, enabling inference of large models on limited GPU memory.

Quick Comparison

Cache Methods

Method	Configuration	Description	Best For
TeaCache	cache_backend="tea_cache"	Simple, adaptive caching with minimal configuration	Quick setup, balanced speed/quality
Cache-DiT	cache_backend="cache_dit"	Advanced caching with multiple techniques (DBCache, TaylorSeer, SCM)	Maximum acceleration, fine-grained control

Quantization Methods

Method	Configuration	Description	Best For
FP8	quantization="fp8"	FP8 W8A8 on Ada/Hopper, weight-only on older GPUs	Memory reduction, inference speedup
Int8	quantization="int8"	Int8 W8A8	Memory reduction, inference speedup

Supported Models

The following table shows which models are currently supported by each acceleration method:

ImageGen

Model	Model Identifier	TeaCache	Cache-DiT	Ulysses-SP	Ring-Attention	CFG-Parallel	Tensor-Parallel	VAE-Patch-Parallel
LongCat-Image	meituan-longcat/LongCat-Image	❌	✅	❌	❌	✅	✅	❌
LongCat-Image-Edit	meituan-longcat/LongCat-Image-Edit	❌	✅	❌	❌	✅	✅	❌
Ovis-Image	OvisAI/Ovis-Image	❌	✅	❌	❌	✅	❌	❌
Qwen-Image	Qwen/Qwen-Image	✅	✅	✅	✅	✅	✅	✅
Qwen-Image-2512	Qwen/Qwen-Image-2512	✅	✅	✅	✅	✅	✅	✅
Qwen-Image-Edit	Qwen/Qwen-Image-Edit	✅	✅	✅	✅	✅	✅	❌
Qwen-Image-Edit-2509	Qwen/Qwen-Image-Edit-2509	✅	✅	✅	✅	✅	✅	❌
Qwen-Image-Layered	Qwen/Qwen-Image-Layered	❌	✅	✅	✅	✅	✅	❌
Z-Image	Tongyi-MAI/Z-Image-Turbo	✅	✅	❌	❌	❌	✅ (TP=2 only)	✅
Stable-Diffusion3.5	stabilityai/stable-diffusion-3.5	❌	✅	❌	❌	✅	❌	❌
Bagel	ByteDance-Seed/BAGEL-7B-MoT	✅	✅	❌	❌	❌	❌	❌
FLUX.1-dev	black-forest-labs/FLUX.1-dev	❌	✅	❌	❌	✅	✅	❌
NextStep-1.1	stepfun-ai/NextStep-1.1	❌	❌	❌	❌	✅	✅	❌
FLUX.2-klein	black-forest-labs/FLUX.2-klein-4B	❌	✅	❌	❌	❌	✅	❌
FLUX.2-dev	black-forest-labs/FLUX.2-dev	❌	✅	❌	❌	❌	✅	❌

VideoGen

Model	Model Identifier	TeaCache	Cache-DiT	Ulysses-SP	Ring-Attention	CFG-Parallel	HSDP	VAE-Patch-Parallel
Wan2.1-T2V	Wan-AI/Wan2.1-T2V-1.3B-Diffusers	❌	✅	✅	✅	✅	✅	✅
Wan2.1-T2V	Wan-AI/Wan2.1-T2V-14B-Diffusers	❌	✅	✅	✅	✅	✅	✅
Wan2.2	Wan-AI/Wan2.2-T2V-A14B-Diffusers	❌	✅	✅	✅	✅	✅	✅
LTX-2	Lightricks/LTX-2	❌	✅	✅	✅	✅	❌	❌
DreamID-Omni	XuGuo699/DreamID-Omni	❌	❌	❌	❌	✅	❌	❌

Quantization

Model	Model Identifier	FP8	Int8
Qwen-Image	Qwen/Qwen-Image	✅	✅
Qwen-Image-2512	Qwen/Qwen-Image-2512	✅	✅
Z-Image	Tongyi-MAI/Z-Image-Turbo	✅	✅

Performance Benchmarks

The following benchmarks were measured on Qwen/Qwen-Image and Qwen/Qwen-Image-Edit models generating 1024x1024 images with 50 inference steps:

!!! note "Benchmark Disclaimer" These benchmarks are provided for general reference only. The configurations shown use default or common parameter settings and have not been exhaustively optimized for maximum performance. Actual performance may vary based on:

- Specific model and use case
- Hardware configuration
- Careful parameter tuning
- Different inference settings (e.g., number of steps, image resolution)

For optimal performance in your specific scenario, we recommend experimenting with different parameter configurations as described in the detailed guides below.

Model	Cache Backend	Cache Config	Generation Time	Speedup	Notes
Qwen/Qwen-Image	None	None	20.0s	1.0x	Baseline (diffusers)
Qwen/Qwen-Image	TeaCache	rel_l1_thresh=0.2	10.47s	1.91x	Recommended default setting
Qwen/Qwen-Image	Cache-DiT	DBCache + TaylorSeer (Fn=1, Bn=0, W=8, TaylorSeer order=1)	10.8s	1.85x	-
Qwen/Qwen-Image	Cache-DiT	DBCache + TaylorSeer + SCM (Fn=8, Bn=0, W=4, TaylorSeer order=1, SCM fast)	14.0s	1.43x	-
Qwen/Qwen-Image-Edit	None	No acceleration	51.5s	1.0x	Baseline (diffusers)
Qwen/Qwen-Image-Edit	Cache-DiT	Default (Fn=1, Bn=0, W=4, TaylorSeer disabled, SCM disabled)	21.6s	2.38x	-

To measure the parallelism methods, we run benchmarks with Qwen/Qwen-Image model generating images (2048x2048 as long sequence input) with 50 inference steps. The hardware devices are NVIDIA H800 GPUs. sdpa is the attention backends.

Configuration	Ulysses degree	Generation Time	Speedup
Baseline (diffusers)	-	112.5s	1.0x
Ulysses-SP	2	65.2s	1.73x
Ulysses-SP	4	39.6s	2.84x
Ulysses-SP	8	30.8s	3.65x

Quick Start

Using TeaCache

from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

omni = Omni(
    model="Qwen/Qwen-Image",
    cache_backend="tea_cache",
    cache_config={"rel_l1_thresh": 0.2}  # Optional, defaults to 0.2
)

outputs = omni.generate(
    "A cat sitting on a windowsill",
    OmniDiffusionSamplingParams(
        num_inference_steps=50,
    ),
)

Using Cache-DiT

from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

omni = Omni(
    model="Qwen/Qwen-Image",
    cache_backend="cache_dit",
    cache_config={
        "Fn_compute_blocks": 1,
        "Bn_compute_blocks": 0,
        "max_warmup_steps": 8,
        "enable_taylorseer": True,
        "taylorseer_order": 1,
    }
)

outputs = omni.generate(
    "A cat sitting on a windowsill",
    OmniDiffusionSamplingParams(
        num_inference_steps=50,
    ),
)

Using Ulysses-SP

Run text-to-image:

from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
ulysses_degree = 2

omni = Omni(
    model="Qwen/Qwen-Image",
    parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
)

outputs = omni.generate(
    "A cat sitting on a windowsill",
    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
)

Run image-to-image:

from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
ulysses_degree = 2

omni = Omni(
    model="Qwen/Qwen-Image-Edit",
    parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
)

outputs = omni.generate(
    {
        "prompt": "turn this cat to a dog",
        "multi_modal_data": {"image": input_image}
    },
    OmniDiffusionSamplingParams(num_inference_steps=50),
)

Using Ring-Attention

Run text-to-image:

from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
ring_degree = 2

omni = Omni(
    model="Qwen/Qwen-Image",
    parallel_config=DiffusionParallelConfig(ring_degree=2)
)

outputs = omni.generate(
    "A cat sitting on a windowsill",
    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
)

Using Tensor Parallelism

from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig

omni = Omni(
    model="Tongyi-MAI/Z-Image-Turbo",
    parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
)

outputs = omni.generate(
    prompt="a cat reading a book",
    num_inference_steps=9,
    width=512,
    height=512,
)

Using VAE Patch Parallelism

from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig

omni = Omni(
    model="Tongyi-MAI/Z-Image-Turbo",
    parallel_config=DiffusionParallelConfig(
        tensor_parallel_size=2,
        vae_patch_parallel_size=2,
    ),
)

outputs = omni.generate(
    prompt="a cat reading a book",
    num_inference_steps=9,
    width=1024,
    height=1024,
)

Using HSDP

HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.

from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig

omni = Omni(
    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
    parallel_config=DiffusionParallelConfig(
        use_hsdp=True,
        hsdp_replicate_size=1,  # Number of replica groups
        hsdp_shard_size=8,      # Number of GPUs to shard across (or -1 for auto)
    ),
)

outputs = omni.generate(
    "A cat playing piano",
    OmniDiffusionSamplingParams(num_inference_steps=50),
)

Using CFG-Parallel

Run image-to-image:

CFG-Parallel splits the CFG positive/negative branches across GPUs. Use it when you set a non-trivial true_cfg_scale.

from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
cfg_parallel_size = 2

omni = Omni(
    model="Qwen/Qwen-Image-Edit",
    parallel_config=DiffusionParallelConfig(cfg_parallel_size=cfg_parallel_size)
)

outputs = omni.generate(
    {
        "prompt": "turn this cat to a dog",
        "multi_modal_data": {"image": input_image}
    },
    OmniDiffusionSamplingParams(num_inference_steps=50, true_cfg_scale=4.0),
)

Using FP8 Quantization

from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

omni = Omni(
    model="<your-model>",
    quantization="fp8",
)

outputs = omni.generate(
    "A cat sitting on a windowsill",
    OmniDiffusionSamplingParams(num_inference_steps=50),
)

Using Int8 Quantization

from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

omni = Omni(
    model="<your-model>",
    quantization="int8",
)

outputs = omni.generate(
    "A cat sitting on a windowsill",
    OmniDiffusionSamplingParams(num_inference_steps=50),
)

Documentation

For detailed information on each acceleration method:

• TeaCache Guide - Complete TeaCache documentation, configuration options, and best practices
• Cache-DiT Acceleration Guide - Comprehensive Cache-DiT guide covering DBCache, TaylorSeer, SCM, and configuration parameters
• Quantization Guide - Quantization for DiT models with per-layer control
• Tensor Parallelism - Guidance on how to enable TP for diffusion models.
• Sequence Parallelism - Guidance on how to set sequence parallelism with configuration.
• CFG-Parallel - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks.
• VAE Patch Parallelism - Guidance on how to reduce VAE memory via patch/tile parallelism.
• HSDP - Hybrid Sharded Data Parallel for memory-efficient inference of large models.