vllm-omni/docs/user_guide/diffusion/cpu_offload_diffusion.md

CPU Offloading for Diffusion Models

Overview

vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models:

1. Model-level (Sequential) Offloading: Mutual exclusion between DiT model and encoder - only one is on GPU at a time.
2. Layerwise (Blockwise) Offloading: Keeps only one transformer block on GPU at a time with compute-memory overlap.

Both strategies use pinned memory for faster CPU-GPU transfers. The strategies are mutually exclusive for now - if both are enabled, layerwise takes priority.

Model-level (Sequential) Offloading

How It Works

Model-level offloading implements mutual exclusion between DiT transformer and encoder modules using pre forward hooks:

• When encoders run: DiT transformer is offloaded to CPU
• When DiT runs: Encoders are offloaded to CPU
• VAE: Stays resident on GPU

Before each module's forward pass, the hook automatically moves it to GPU while offloading the other module group to CPU. Transfers use pinned memory for speed.

Usage

Python API:

from vllm_omni import Omni

m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True)

CLI:

vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-cpu-offload

Limitations

• Cold start latency increases
• Adds overhead from CPU-GPU transfers between encoder and denoising phases
• Support single GPU only for now

Layerwise (Blockwise) Offloading

How It Works

Layerwise offloading keeps only one transformer block on GPU at a time.

As each block completes, the next block is prefetched to GPU while the current block is freed. The pre and forward hooks utilized by layerwise offloading apply a separate CUDA stream (copy_stream) to overlap weight transfer with computation, and retain flattened tensors in pinned CPU memory for block parameters re-materialization. Encoders, VAE, and non-block DiT modules (embeddings, norms) always stay on GPU.

Execution Flow:

Block	Pre-forward Hook	Forward	Post-forward Hook
block-0	Prefetch block-1 (async)	Compute block-0	Free block-0
block-1	Prefetch block-2 (async)	Compute block-1	Free block-1
...	...	...	...
block-(n-1)	Prefetch block-0 (async)	Compute block-(n-1)	Free block-(n-1)

Each transformer block has a LayerwiseOffloadHook that prefetches the next block before forward and frees the current block after forward.

Layerwise offloading is primarily recommended for large video generation models where the compute cost per block is high enough to effectively overlap with memory prefetch operations. For example, Wan2.2 T2V and I2V pipelines.

Usage

Python API:

from vllm_omni import Omni

# Text-to-video
m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_layerwise_offload=True)

# Or image-to-video
m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True)

CLI:

# Text-to-video
vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-layerwise-offload

# Or image-to-video
vllm-omni serve diffusion Wan-AI/Wan2.2-I2V-A14B-Diffusers --enable-layerwise-offload

To Support a Model

Models must define the blocks attribute name for layerwise offloading:

class WanTransformer3DModel(nn.Module):
    _layerwise_offload_blocks_attr = "blocks"  # Attribute name containing transformer blocks

    def __init__(self):
        self.blocks = nn.ModuleList([...])  # Transformer blocks

Limitations

• Cold start latency increases because of
  1. components are loaded to CPU first at the very first during initialization,
  2. weight consolidation and pinning
• Performance depends on compute cost and H2D bandwidth as well
• Support single GPU only for now

Implementation Notes

Module Discovery

The offloader automatically discovers pipeline components:

• DiT modules: transformer, transformer_2, dit
• Encoders: text_encoder, text_encoder_2, text_encoder_3, image_encoder
• VAE: vae

Hook System

Both strategies use vLLM-Omni's hook registry system (HookRegistry and ModelHook) to register pre/post forward callbacks on modules, enabling automatic swapping without modifying model code.

Backend Architecture

OffloadBackend (base class)
├── ModelLevelOffloadBackend → uses SequentialOffloadHook
└── LayerWiseOffloadBackend → uses LayerwiseOffloadHook

Factory function get_offload_backend() selects the appropriate backend based on configuration.

Supported Models

Architecture	Example Models	DiT Class	Model-Level Offload	Layerwise Offload	Blocks Attr (Layerwise specific)
Wan22Pipeline	Wan-AI/Wan2.2-T2V-A14B-Diffusers	WanTransformer3DModel	✓	✓	"blocks"
Wan22I2VPipeline	Wan-AI/Wan2.2-I2V-A14B-Diffusers	WanTransformer3DModel	✓	✓	"blocks"
QwenImagePipeline	Qwen/Qwen-Image	QwenImageTransformer2DModel	✓	✓	"transformer_blocks"

Notes:

• Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
• Layerwise Offloading requires DiT class to define _layerwise_offload_blocks_attr pointing to transformer blocks