[Refactor] Let diffusion pipelines declare offloadable modules via SupportsModuleOffload (#2427)

Signed-off-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>

.claude/skills/add-diffusion-model/SKILL.md

@@ -282,10 +282,11 @@ For Omni or custom models, create:
 
 Required updates:
 1. `docs/user_guide/diffusion/parallelism_acceleration.md` — parallelism support table
-2. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
-3. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
-4. `examples/offline_inference/xxx/README.md` — offline example docs
-5. `examples/online_serve/xxx/README.md` — online serve docs
+2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table)
+3. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
+4. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
+5. `examples/offline_inference/xxx/README.md` — offline example docs
+6. `examples/online_serve/xxx/README.md` — online serve docs
 
 ### Step 8: Add E2E Tests (Recommended)
 
@@ -512,6 +513,37 @@ After adding parallelism support, update:
 1. `docs/user_guide/diffusion/parallelism_acceleration.md` — add your model to the support table
 2. Record which parallelism methods are supported (USP, Ring, CFG, TP, HSDP, VAE-Patch)
 
+### Step 11: Add CPU Offload Support
+
+Implement `SupportsModuleOffload` on your pipeline class to enable
+`--enable-cpu-offload` and `--enable-layerwise-offload`. The protocol
+declares which submodules the offloader should manage:
+
+```python
+from typing import ClassVar
+from vllm_omni.diffusion.models.interface import SupportsModuleOffload
+
+class YourPipeline(nn.Module, SupportsModuleOffload):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = []  # optional
+```
+
+- `_dit_modules`: denoising submodules (kept on GPU during diffusion loop)
+- `_encoder_modules`: encoder/vision submodules (offloaded to CPU during diffusion loop)
+- `_vae_modules`: VAE(s) (handled by both sequential and layerwise backends)
+- `_resident_modules`: additional modules to pin on GPU during layerwise
+  offloading (e.g. embedders, connectors). Only used by the layerwise
+  backend. Optional — defaults to `[]`.
+
+All attribute names support dotted paths for nested submodules
+(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`).
+
+Pipelines without `SupportsModuleOffload` fall back to scanning
+well-known attribute names (`transformer`, `text_encoder`, `vae`,
+etc.), which fails for non-standard names.
+
 ---
 
 ## Iterative Development Tips

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -36,6 +36,45 @@ m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True)
 vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-cpu-offload
 ```
 
+### To Support a Model
+
+Implement the `SupportsModuleOffload` protocol to declare which
+submodules participate in offloading:
+
+```python
+from typing import ClassVar
+from vllm_omni.diffusion.models.interface import SupportsModuleOffload
+
+class MyPipeline(nn.Module, SupportsModuleOffload):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = []  # optional
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = ...     # DiT — stays on GPU during denoising
+        self.text_encoder = ...    # Encoder — offloaded to CPU during denoising
+        self.vision_model = ...    # Encoder — offloaded to CPU during denoising
+        self.vae = ...             # VAE — always on GPU
+```
+
+- `_dit_modules`: attribute names of denoising submodules (kept on GPU
+  during the diffusion loop).
+- `_encoder_modules`: attribute names of encoder/vision submodules
+  (offloaded to CPU during the diffusion loop).
+- `_vae_modules`: attribute names of VAE(s) (always kept on GPU, not
+  part of the mutual exclusion hooks).
+- `_resident_modules`: attribute names of small submodules that must
+  stay on GPU during layerwise offloading (e.g. embedders, connectors).
+  Optional — defaults to `[]`.
+
+All attribute names support dotted paths for nested submodules
+(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`).
+
+Both DiT and encoder lists are needed because the offload hooks use
+mutual exclusion: when one group runs, the other moves to CPU.
+
 ### Limitations
 - Cold start latency increases
 - Adds overhead from CPU-GPU transfers between encoder and denoising phases
@@ -116,11 +155,19 @@ class Flux2Transformer2DModel(nn.Module):
 
 **Module Discovery**
 
-The offloader automatically discovers pipeline components:
+The offloader discovers pipeline components in two ways:
 
-- **DiT modules**: `transformer`, `transformer_2`, `dit`
-- **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder`
-- **VAE**: `vae`
+1. **Protocol-based** (preferred): If the pipeline implements
+    `SupportsModuleOffload`, its `_dit_modules`, `_encoder_modules`,
+    `_vae_modules`, and `_resident_modules` class variables are used
+    directly.  All attribute names support dotted paths (e.g.
+    `"pipe.transformer"`, `"bagel.time_embedder"`) for nested submodules.
+
+2. **Fallback attribute scan**: Otherwise, the offloader scans for
+    well-known attribute names:
+    - **DiT modules**: `transformer`, `transformer_2`, `dit`, `sr_dit`, `language_model`, `transformer_blocks`, `model`
+    - **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder`
+    - **VAE**: `vae`, `audio_vae`
 
 **Hook System**
 

tests/diffusion/offloader/test_module_collector.py

@@ -0,0 +1,240 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for ModuleDiscovery and SupportsModuleOffload."""
+
+from typing import ClassVar
+
+import pytest
+from torch import nn
+
+from vllm_omni.diffusion.models.interface import SupportsModuleOffload
+from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+# NOTE: tests for skipped/warned attributes verify the *behavioral*
+# outcome (attribute excluded from results) but do not assert on log
+# output.  vllm's logger sets propagate=False, preventing caplog from
+# capturing records.  See https://github.com/pytest-dev/pytest/issues/3697
+
+
+# ---------------------------------------------------------------------------
+# Test pipelines
+# ---------------------------------------------------------------------------
+
+
+class FallbackPipeline(nn.Module):
+    """Pipeline with standard attribute names (no protocol)."""
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = nn.Linear(10, 10)
+        self.text_encoder = nn.Linear(10, 10)
+        self.text_encoder_2 = nn.Linear(10, 10)
+        self.vae = nn.Linear(10, 10)
+
+
+class NonModuleAttrPipeline(nn.Module):
+    """Pipeline where an attribute is not an nn.Module (fallback path)."""
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = nn.Linear(10, 10)
+        self.text_encoder = "not_a_module"
+        self.vae = nn.Linear(10, 10)
+
+
+class DuplicateAttrPipeline(nn.Module):
+    """Pipeline where two encoder attrs point to the same module."""
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = nn.Linear(10, 10)
+        encoder = nn.Linear(10, 10)
+        self.text_encoder = encoder
+        self.text_encoder_2 = encoder
+        self.vae = nn.Linear(10, 10)
+
+
+class ProtocolPipeline(nn.Module, SupportsModuleOffload):
+    """Pipeline with non-standard names, using the protocol."""
+
+    _dit_modules: ClassVar[list[str]] = ["gen_transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["mllm", "vision_model"]
+    _vae_modules: ClassVar[list[str]] = ["gen_vae"]
+
+    def __init__(self):
+        super().__init__()
+        self.gen_transformer = nn.Linear(10, 10)
+        self.mllm = nn.Linear(10, 10)
+        self.vision_model = nn.Linear(10, 10)
+        self.gen_vae = nn.Linear(10, 10)
+        # Standard name present but NOT declared — should be ignored
+        self.transformer = nn.Linear(10, 10)
+
+
+class MissingAttrPipeline(nn.Module, SupportsModuleOffload):
+    """Pipeline that declares a non-existent attribute."""
+
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["nonexistent_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = nn.Linear(10, 10)
+        self.vae = nn.Linear(10, 10)
+
+
+class MissingIntermediatePipeline(nn.Module, SupportsModuleOffload):
+    """Pipeline with dotted path referencing non-existent intermediate."""
+
+    _dit_modules: ClassVar[list[str]] = ["nonexistent.transformer"]
+    _encoder_modules: ClassVar[list[str]] = []
+    _vae_modules: ClassVar[list[str]] = []
+
+    def __init__(self):
+        super().__init__()
+
+
+class NestedPipeline(nn.Module, SupportsModuleOffload):
+    """Pipeline with nested modules accessed via dotted paths."""
+
+    _dit_modules: ClassVar[list[str]] = ["pipe.transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
+    def __init__(self):
+        super().__init__()
+        self.pipe = nn.Module()
+        self.pipe.transformer = nn.Linear(10, 10)
+        self.pipe.text_encoder = nn.Linear(10, 10)
+        self.vae = nn.Linear(10, 10)
+
+
+class ResidentPipeline(nn.Module, SupportsModuleOffload):
+    """Pipeline with resident modules that must stay on GPU."""
+
+    _dit_modules: ClassVar[list[str]] = ["language_model.model"]
+    _encoder_modules: ClassVar[list[str]] = []
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = [
+        "bagel.time_embedder",
+        "bagel.vae2llm",
+    ]
+
+    def __init__(self):
+        super().__init__()
+        self.language_model = nn.Module()
+        self.language_model.model = nn.Linear(10, 10)
+        self.bagel = nn.Module()
+        self.bagel.time_embedder = nn.Linear(10, 10)
+        self.bagel.vae2llm = nn.Linear(10, 10)
+        self.vae = nn.Linear(10, 10)
+
+
+class MultiVaePipeline(nn.Module, SupportsModuleOffload):
+    """Pipeline with multiple VAEs."""
+
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae", "audio_vae"]
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = nn.Linear(10, 10)
+        self.text_encoder = nn.Linear(10, 10)
+        self.vae = nn.Linear(10, 10)
+        self.audio_vae = nn.Linear(10, 10)
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+class TestFallbackDiscovery:
+    """Test the fallback attribute scan (no SupportsModuleOffload)."""
+
+    def test_discovers_standard_attrs(self):
+        pipeline = FallbackPipeline()
+        result = ModuleDiscovery.discover(pipeline)
+
+        assert not isinstance(pipeline, SupportsModuleOffload)
+        assert result.dit_names == ["transformer"]
+        assert result.dits[0] is pipeline.transformer
+        assert result.encoder_names == ["text_encoder", "text_encoder_2"]
+        assert result.vaes[0] is pipeline.vae
+        assert result.resident_modules == []
+
+    def test_deduplicates_encoders(self):
+        pipeline = DuplicateAttrPipeline()
+        result = ModuleDiscovery.discover(pipeline)
+
+        assert len(result.encoders) == 1
+        assert result.encoder_names == ["text_encoder"]
+
+    def test_skips_non_module_attr(self):
+        pipeline = NonModuleAttrPipeline()
+        result = ModuleDiscovery.discover(pipeline)
+
+        assert len(result.encoders) == 0
+
+
+class TestProtocolDiscovery:
+    """Test discovery via SupportsModuleOffload protocol."""
+
+    def test_discovers_declared_attrs_and_ignores_undeclared(self):
+        pipeline = ProtocolPipeline()
+        result = ModuleDiscovery.discover(pipeline)
+
+        assert isinstance(pipeline, SupportsModuleOffload)
+        assert result.dit_names == ["gen_transformer"]
+        assert result.encoder_names == ["mllm", "vision_model"]
+        assert len(result.vaes) == 1
+        # self.transformer exists but is NOT in _dit_modules
+        assert "transformer" not in result.dit_names
+        # No _resident_modules declared — defaults to empty
+        assert result.resident_modules == []
+
+    def test_skips_missing_attr(self):
+        pipeline = MissingAttrPipeline()
+        result = ModuleDiscovery.discover(pipeline)
+
+        assert len(result.encoders) == 0
+
+    def test_skips_missing_intermediate(self):
+        result = ModuleDiscovery.discover(MissingIntermediatePipeline())
+
+        assert len(result.dits) == 0
+
+    def test_dotted_path_resolves_nested_modules(self):
+        pipeline = NestedPipeline()
+        result = ModuleDiscovery.discover(pipeline)
+
+        assert result.dit_names == ["pipe.transformer"]
+        assert result.dits[0] is pipeline.pipe.transformer
+        assert result.encoder_names == ["pipe.text_encoder"]
+        assert result.encoders[0] is pipeline.pipe.text_encoder
+        assert result.vaes[0] is pipeline.vae
+
+    def test_resident_modules(self):
+        pipeline = ResidentPipeline()
+        result = ModuleDiscovery.discover(pipeline)
+
+        assert result.resident_names == [
+            "bagel.time_embedder",
+            "bagel.vae2llm",
+        ]
+        assert result.resident_modules[0] is pipeline.bagel.time_embedder
+        assert result.resident_modules[1] is pipeline.bagel.vae2llm
+        assert result.dits[0] is pipeline.language_model.model
+
+    def test_multiple_vaes(self):
+        pipeline = MultiVaePipeline()
+        result = ModuleDiscovery.discover(pipeline)
+
+        assert len(result.vaes) == 2
+        assert result.vaes[0] is pipeline.vae
+        assert result.vaes[1] is pipeline.audio_vae

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -12,6 +12,7 @@ from collections.abc import Iterable
 from copy import deepcopy
 from dataclasses import dataclass
 from math import isqrt
+from typing import ClassVar
 
 import numpy as np
 import torch
@@ -26,6 +27,7 @@ from vllm.transformers_utils.configs.bagel import BagelConfig
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.interface import SupportsModuleOffload
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
@@ -148,12 +150,25 @@ class SiglipNaViTWrapper(nn.Module):
         return outputs.last_hidden_state.squeeze(0)
 
 
-class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
+class BagelPipeline(nn.Module, SupportsModuleOffload, DiffusionPipelineProfilerMixin):
     """Bagel generation pipeline (MoT) packaged for vllm-omni diffusion engine.
 
     This pipeline is self-contained and uses the ported Bagel core files.
     """
 
+    _dit_modules: ClassVar[list[str]] = ["language_model.model"]
+    _encoder_modules: ClassVar[list[str]] = []
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = [
+        "bagel.time_embedder",
+        "bagel.vae2llm",
+        "bagel.llm2vae",
+        "bagel.latent_pos_embed",
+        "bagel.vit_model",
+        "bagel.connector",
+        "bagel.vit_pos_embed",
+    ]
+
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
         super().__init__()
         self.od_config = od_config

vllm_omni/diffusion/models/interface.py

@@ -58,6 +58,27 @@ class SupportsStepExecution(Protocol):
         """Decode output after denoise loop."""
 
 
+@runtime_checkable
+class SupportsModuleOffload(Protocol):
+    """Declares which submodules participate in CPU offload.
+
+    All attribute names support dotted paths for nested submodules
+    (e.g. ``"pipe.transformer"``).
+
+    Attributes:
+        _dit_modules: Denoising submodules (on GPU during diffusion).
+        _encoder_modules: Encoder submodules (offloaded during diffusion).
+        _vae_modules: VAE(s) (always on GPU).
+        _resident_modules: Extra modules pinned on GPU during layerwise
+            offloading.  Optional, defaults to ``[]``.
+    """
+
+    _dit_modules: ClassVar[list[str]]
+    _encoder_modules: ClassVar[list[str]]
+    _vae_modules: ClassVar[list[str]]
+    _resident_modules: ClassVar[list[str]] = []
+
+
 def supports_step_execution(pipeline: object) -> bool:
     """Return whether `pipeline` implements :class:`SupportsStepExecution`."""
 

vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py

@@ -9,7 +9,7 @@ import json
 import os
 from collections.abc import Iterable
 from contextlib import nullcontext
-from typing import Any
+from typing import Any, ClassVar
 
 import numpy as np
 import torch
@@ -34,6 +34,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
+from vllm_omni.diffusion.models.interface import SupportsModuleOffload
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.lora.request import LoRARequest
@@ -1152,9 +1153,13 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
         return loader.load_weights(weights)
 
 
-class LTX2TwoStagesPipeline(nn.Module):
+class LTX2TwoStagesPipeline(nn.Module, SupportsModuleOffload):
     """LTX2TwoStagesPipeline is for two stages image to video generation"""
 
+    _dit_modules: ClassVar[list[str]] = ["pipe.transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["pipe.vae", "pipe.audio_vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py

@@ -6,7 +6,7 @@ from __future__ import annotations
 import copy
 import os
 from collections.abc import Iterable
-from typing import Any
+from typing import Any, ClassVar
 
 import numpy as np
 import PIL.Image
@@ -26,6 +26,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
+from vllm_omni.diffusion.models.interface import SupportsModuleOffload
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.lora.request import LoRARequest
 
@@ -733,11 +734,15 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline):
         return DiffusionOutput(output=(video, audio))
 
 
-class LTX2ImageToVideoTwoStagesPipeline(nn.Module):
+class LTX2ImageToVideoTwoStagesPipeline(nn.Module, SupportsModuleOffload):
     """LTXImageToVideoTwoStagesPipeline is for two stages image to video generation"""
 
     support_image_input = True
 
+    _dit_modules: ClassVar[list[str]] = ["pipe.transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["pipe.vae", "pipe.audio_vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py

@@ -8,7 +8,7 @@ import os
 import warnings
 from collections.abc import Iterable
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, ClassVar
 
 import numpy as np
 import PIL.Image
@@ -32,6 +32,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.interface import SupportsModuleOffload
 from vllm_omni.diffusion.models.omnigen2.omnigen2_transformer import (
     OmniGen2RotaryPosEmbed,
     OmniGen2Transformer2DModel,
@@ -620,7 +621,7 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-class OmniGen2Pipeline(CFGParallelMixin, nn.Module):
+class OmniGen2Pipeline(CFGParallelMixin, nn.Module, SupportsModuleOffload):
     """
     Pipeline for text-to-image generation using OmniGen2.
 
@@ -634,6 +635,10 @@ class OmniGen2Pipeline(CFGParallelMixin, nn.Module):
         od_config (OmniDiffusionConfig): The OmniDiffusion configuration.
     """
 
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["mllm"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/offloader/layerwise_backend.py

@@ -298,13 +298,20 @@ class LayerWiseOffloadBackend(OffloadBackend):
         for enc in modules.encoders:
             enc.to(self.device)
 
-        # Move VAE to GPU if available
-        if modules.vae is not None:
+        # Move VAE(s) to GPU if available
+        for vae in modules.vaes:
             try:
-                modules.vae.to(self.device, non_blocking=True)
+                vae.to(self.device, non_blocking=True)
             except Exception as exc:
                 logger.debug("Failed to move VAE to GPU: %s", exc)
 
+        # Move resident modules to GPU (small modules needed every forward)
+        for name, module in zip(modules.resident_names, modules.resident_modules):
+            try:
+                module.to(self.device)
+            except Exception as exc:
+                logger.debug("Failed to move resident module %s to GPU: %s", name, exc)
+
         logger.info("Applying layer-wise offloading on %s", modules.dit_names)
 
         # Apply block-wise offloading hook for each of the blocks in DiT model(s)

vllm_omni/diffusion/offloader/module_collector.py

@@ -1,11 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from operator import attrgetter
 
 from torch import nn
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.models.interface import SupportsModuleOffload
+
 logger = init_logger(__name__)
 
 
@@ -15,15 +18,81 @@ class PipelineModules:
     dit_names: list[str]
     encoders: list[nn.Module]
     encoder_names: list[str]
-    vae: nn.Module | None = None
+    vaes: list[nn.Module]
+    resident_modules: list[nn.Module] = field(default_factory=list)
+    resident_names: list[str] = field(default_factory=list)
 
 
 class ModuleDiscovery:
-    """Discovers pipeline components for offloading"""
+    """Discovers pipeline components for offloading.
 
-    DIT_ATTRS = ["transformer", "transformer_2", "dit", "sr_dit", "language_model", "transformer_blocks", "model"]
-    ENCODER_ATTRS = ["text_encoder", "text_encoder_2", "text_encoder_3", "image_encoder"]
-    VAE_ATTRS = ["vae", "audio_vae"]
+    If the pipeline implements :class:`SupportsModuleOffload`,
+    its ``_dit_modules``, ``_encoder_modules``, and ``_vae_modules``
+    class variables are used directly.  Otherwise, falls back to
+    scanning well-known attribute names.
+    """
+
+    # Fallback attribute names for pipelines that do not implement
+    # SupportsModuleOffload.
+    _FALLBACK_DIT_ATTRS = [
+        "transformer",
+        "transformer_2",
+        "dit",
+        "sr_dit",
+        "language_model",
+        "transformer_blocks",
+        "model",
+    ]
+    _FALLBACK_ENCODER_ATTRS = [
+        "text_encoder",
+        "text_encoder_2",
+        "text_encoder_3",
+        "image_encoder",
+    ]
+    _FALLBACK_VAE_ATTRS = [
+        "vae",
+        "audio_vae",
+    ]
+
+    @staticmethod
+    def _collect_modules(
+        pipeline: nn.Module,
+        attr_names: list[str],
+        *,
+        warn_missing: bool = False,
+    ) -> tuple[list[nn.Module], list[str]]:
+        """Resolve attribute names to (module, name) pairs, skipping missing.
+
+        Supports dotted paths via :func:`operator.attrgetter`.
+        Warns on missing attributes when *warn_missing* is True.
+        """
+        modules: list[nn.Module] = []
+        names: list[str] = []
+        seen: set[int] = set()
+        for attr in attr_names:
+            try:
+                module = attrgetter(attr)(pipeline)
+            except AttributeError:
+                module = None
+            if module is None:
+                if warn_missing:
+                    logger.warning(
+                        "Pipeline declares '%s' as offloadable but the attribute does not exist or is None",
+                        attr,
+                    )
+                continue
+            if not isinstance(module, nn.Module):
+                logger.warning(
+                    "Expected '%s' to be nn.Module, got %r",
+                    attr,
+                    type(module),
+                )
+                continue
+            if id(module) not in seen:
+                seen.add(id(module))
+                modules.append(module)
+                names.append(attr)
+        return modules, names
 
     @staticmethod
     def discover(pipeline: nn.Module) -> PipelineModules:
@@ -35,46 +104,29 @@ class ModuleDiscovery:
         Returns:
             PipelineModules with lists of discovered modules and names
         """
-        # Collect DiT/transformer modules
-        dit_modules: list[nn.Module] = []
-        dit_names: list[str] = []
-        for attr in ModuleDiscovery.DIT_ATTRS:
-            if not hasattr(pipeline, attr):
-                continue
-            module_obj = getattr(pipeline, attr)
-            if module_obj is None:
-                continue
+        declared = isinstance(pipeline, SupportsModuleOffload)
+        if declared:
+            dit_attrs = pipeline._dit_modules
+            enc_attrs = pipeline._encoder_modules
+            vae_attrs = pipeline._vae_modules
+            res_attrs = pipeline._resident_modules
+        else:
+            dit_attrs = ModuleDiscovery._FALLBACK_DIT_ATTRS
+            enc_attrs = ModuleDiscovery._FALLBACK_ENCODER_ATTRS
+            vae_attrs = ModuleDiscovery._FALLBACK_VAE_ATTRS
+            res_attrs = []
 
-            if not isinstance(module_obj, nn.Module):
-                logger.warning(f"Expected {attr} to be nn.Module, got {type(module_obj)!r}")
-                continue
-
-            if module_obj in dit_modules:
-                continue
-
-            dit_modules.append(module_obj)
-            dit_names.append(attr)
-
-        # Collect all encoders
-        encoders: list[nn.Module] = []
-        encoder_names: list[str] = []
-        for attr in ModuleDiscovery.ENCODER_ATTRS:
-            if hasattr(pipeline, attr) and getattr(pipeline, attr) is not None:
-                encoders.append(getattr(pipeline, attr))
-                encoder_names.append(attr)
-
-        # Collect VAE
-        vae = None
-        for attr in ModuleDiscovery.VAE_ATTRS:
-            module = getattr(pipeline, attr, None)
-            if module is not None:
-                vae = module
-                break
+        dit_modules, dit_names = ModuleDiscovery._collect_modules(pipeline, dit_attrs, warn_missing=declared)
+        encoders, encoder_names = ModuleDiscovery._collect_modules(pipeline, enc_attrs, warn_missing=declared)
+        vaes, _ = ModuleDiscovery._collect_modules(pipeline, vae_attrs, warn_missing=declared)
+        residents, resident_names = ModuleDiscovery._collect_modules(pipeline, res_attrs, warn_missing=declared)
 
         return PipelineModules(
             dits=dit_modules,
             dit_names=dit_names,
             encoders=encoders,
             encoder_names=encoder_names,
-            vae=vae,
+            vaes=vaes,
+            resident_modules=residents,
+            resident_names=resident_names,
         )

vllm_omni/diffusion/offloader/sequential_backend.py

@@ -210,10 +210,10 @@ class ModelLevelOffloadBackend(OffloadBackend):
         for enc in modules.encoders:
             enc.to(self.device)
 
-        # Move VAE to GPU if available
-        if modules.vae is not None:
+        # Move VAE(s) to GPU if available
+        for vae in modules.vaes:
             try:
-                modules.vae.to(self.device, non_blocking=True)
+                vae.to(self.device, non_blocking=True)
             except Exception as exc:
                 logger.debug("Failed to move VAE to GPU: %s", exc)