vllm-omni/docs/user_guide/diffusion/offloader/layerwise_offload.md

Layerwise Offloading

Layerwise, or blockwise, offloading keeps one transformer block on the accelerator and prefetches the next block while the current block computes. It is best suited to compute-heavy video DiTs whose block execution can hide host-to-device transfers.

Execution flow

Each block has a pre-forward and post-forward hook. Parameters are consolidated in pinned host tensors and rematerialized for execution on a dedicated copy stream.

Block	Pre-forward hook	Forward	Post-forward hook
block 0	Prefetch block 1	Compute block 0	Free block 0
block 1	Prefetch block 2	Compute block 1	Free block 1
...	...	...	...
last block	Prefetch block 0	Compute last block	Free last block

Encoders, VAE modules, and non-block DiT modules such as embeddings and norms remain device resident.

Usage

from vllm_omni import Omni

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

vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
  --omni --enable-layerwise-offload

Model integration

Transformer classes declare containers of executable blocks:

class WanTransformer3DModel(nn.Module):
    _layerwise_offload_blocks_attrs = ["blocks"]


class Flux2Transformer2DModel(nn.Module):
    _layerwise_offload_blocks_attrs = [
        "transformer_blocks",
        "single_transformer_blocks",
    ]

See the layerwise design for the discovery and hook invariants. OffloadPlan is a separate declarative topology path for distributed layerwise offload.

Limitations

• Single device only.
• Setup consolidates and pins block parameters, increasing cold-start time.
• Performance depends on block compute time and host-to-device bandwidth; lightweight blocks may not hide transfers.