[Bugfix][Diffusion] Fix MiniMax-H3 model-level CPU offload residency (#6072)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -1038,6 +1038,7 @@ def test_model_offload_uses_hooked_text_encoder_call():
     )
     expected = torch.ones(2, 3)
     pipeline.text_encoder = Mock(return_value=expected)
+    pipeline._model_cpu_offload_modules = [pipeline.text_encoder]
     input_ids = torch.tensor([1, 2])
     vision_kwargs = {"pixel_values": torch.ones(1, 4)}
 
@@ -1401,7 +1402,7 @@ def test_distributed_video_vae_encodes_references_sequentially(monkeypatch):
         lambda path: f"frames:{path}",
     )
 
-    rows, shapes = pipeline._encode_video_conditions(None, count=2)
+    rows, shapes = pipeline._encode_video_conditions_resident(None, count=2)
 
     assert pipeline.video_vae.calls == [
         "frames:video-1.mp4",
@@ -1964,7 +1965,7 @@ def test_ref2va_audio_duration_validation_precedes_rank_branch(monkeypatch):
 
     waveform = torch.zeros(1, 10)
     with pytest.raises(ValueError, match="max_duration_seconds must be positive"):
-        pipeline._encode_audio_conditions(
+        pipeline._encode_audio_conditions_resident(
             [(waveform, 10)],
             max_duration_seconds=0,
         )
@@ -1985,7 +1986,7 @@ def test_ref2va_standalone_audio_condition_is_bounded_to_output_duration():
     pipeline.audio_vae = FakeAudioVAE()
     waveform = torch.arange(40, dtype=torch.float32).reshape(1, 40)
 
-    encoded, lengths = pipeline._encode_audio_conditions(
+    encoded, lengths = pipeline._encode_audio_conditions_resident(
         [(waveform, 10)],
         max_duration_seconds=2.5,
     )

tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py

@@ -0,0 +1,308 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import contextmanager
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from PIL import Image
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_h3_profiler_targets_cover_aggregate_reference_encoders():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    targets = MiniMaxH3Pipeline._PROFILER_TARGETS
+    assert "_encode_visual_conditions" in targets
+    assert "_encode_reference_audio_conditions" in targets
+    assert "_encode_video_conditions" not in targets
+    assert "_encode_video_audio_conditions" not in targets
+    assert "_encode_audio_conditions" not in targets
+
+
+@pytest.mark.parametrize("distributed", [False, True])
+def test_manual_component_offload_wins_over_requested_model_offload(distributed):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=True,
+        enable_layerwise_offload=not distributed,
+        enable_distributed_layerwise_offload=distributed,
+    )
+    pipeline._model_cpu_offload_modules = []
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+    component = Mock()
+
+    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
+    with pipeline._component_on_device(component):
+        component.load_to_device.assert_called_once_with()
+        component.offload_to_cpu.assert_not_called()
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
+    component.offload_to_cpu.assert_called_once_with()
+
+
+def test_requested_model_offload_without_backend_uses_resident_components():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=True,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+    )
+    pipeline._model_cpu_offload_modules = []
+    pipeline.text_encoder = Mock()
+    expected = torch.ones(2, 3)
+    pipeline.text_encoder.encode_ids.return_value = expected
+
+    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
+
+    assert actual is expected
+    pipeline.text_encoder.load_to_device.assert_called_once_with()
+    pipeline.text_encoder.encode_ids.assert_called_once()
+
+
+def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.transformer = torch.nn.Linear(2, 2)
+    pipeline.transformers_ref = torch.nn.Linear(2, 2)
+    pipeline.text_encoder = torch.nn.Linear(2, 2)
+    pipeline.video_vae = torch.nn.Linear(2, 2)
+    pipeline.audio_vae = torch.nn.Linear(2, 2)
+    apply_offload = Mock()
+    remove_offload = Mock()
+    monkeypatch.setattr(module, "apply_sequential_offload", apply_offload)
+    monkeypatch.setattr(module, "remove_sequential_offload", remove_offload)
+
+    pipeline.enable_omni_model_cpu_offload(
+        device=torch.device("cpu"),
+        pin_memory=False,
+        use_hsdp=False,
+    )
+
+    dits = [pipeline.transformer, pipeline.transformers_ref]
+    stages = [pipeline.text_encoder, pipeline.video_vae, pipeline.audio_vae]
+    apply_offload.assert_called_once_with(
+        dit_modules=dits,
+        encoder_modules=stages,
+        device=torch.device("cpu"),
+        pin_memory=False,
+        use_hsdp=False,
+        offload_initial_dits=True,
+    )
+
+    pipeline.disable_omni_model_cpu_offload()
+
+    remove_offload.assert_called_once_with([*dits, *stages])
+
+
+@pytest.mark.parametrize("decode_fails", [False, True])
+def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    component = torch.nn.Linear(2, 2)
+    events = []
+
+    @contextmanager
+    def record_component(value):
+        events.append(("activate", value))
+        try:
+            yield
+        finally:
+            events.append(("offload", value))
+
+    monkeypatch.setattr(module, "sequential_offload_component", record_component)
+    pipeline._model_cpu_offload_modules = [component]
+
+    def decode():
+        with pipeline._component_on_device(component):
+            events.append(("decode", component))
+            if decode_fails:
+                raise RuntimeError("decode failed")
+
+    if decode_fails:
+        with pytest.raises(RuntimeError, match="decode failed"):
+            decode()
+    else:
+        decode()
+
+    assert events == [
+        ("activate", component),
+        ("decode", component),
+        ("offload", component),
+    ]
+
+
+@pytest.mark.parametrize("encode_fails", [False, True])
+def test_h3_model_cpu_offload_batches_visual_reference_scope(monkeypatch, encode_fails):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.video_vae = Mock()
+    pipeline._model_cpu_offload_modules = [pipeline.video_vae]
+    events = []
+
+    @contextmanager
+    def record_component(value):
+        events.append(("activate", value))
+        try:
+            yield
+        finally:
+            events.append(("offload", value))
+
+    def encode_image(image):
+        events.append(("encode", image.size))
+        if encode_fails and image.width == 32:
+            raise RuntimeError("encode failed")
+        return torch.ones(image.width // 16, 4)
+
+    pipeline.video_vae.encode_image.side_effect = encode_image
+    monkeypatch.setattr(module, "sequential_offload_component", record_component)
+    monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1))
+    monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value)
+    images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))]
+
+    if encode_fails:
+        with pytest.raises(RuntimeError, match="encode failed"):
+            pipeline._encode_visual_conditions(images, None, video_count=0)
+    else:
+        rows, shapes = pipeline._encode_visual_conditions(images, None, video_count=0)
+        assert rows.shape == (3, 4)
+        assert shapes == [(1, 1, 1), (1, 1, 2)]
+
+    assert events == [
+        ("activate", pipeline.video_vae),
+        ("encode", (16, 16)),
+        ("encode", (32, 16)),
+        ("offload", pipeline.video_vae),
+    ]
+
+
+def test_h3_model_cpu_offload_shares_image_and_video_scope(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.video_vae = Mock()
+    pipeline.video_vae.is_distributed_enabled.return_value = False
+    pipeline._model_cpu_offload_modules = [pipeline.video_vae]
+    events = []
+
+    @contextmanager
+    def record_component(value):
+        events.append(("activate", value))
+        try:
+            yield
+        finally:
+            events.append(("offload", value))
+
+    def encode_image(image):
+        events.append(("image", image.size))
+        return torch.ones(1, 4)
+
+    def encode_video(frames):
+        events.append(("video", frames))
+        return torch.ones(3, 4), (3, 2, 2)
+
+    pipeline.video_vae.encode_image.side_effect = encode_image
+    pipeline.video_vae.encode_video.side_effect = encode_video
+    monkeypatch.setattr(module, "sequential_offload_component", record_component)
+    monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1))
+    monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value)
+    monkeypatch.setattr(module, "load_video_frames", lambda path: f"frames:{path}")
+    images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))]
+    prepared_videos = [{"prepared_path": "reference.mp4"}]
+
+    rows, shapes = pipeline._encode_visual_conditions(
+        images,
+        prepared_videos,
+        video_count=1,
+    )
+
+    assert rows.shape == (5, 4)
+    assert shapes == [(1, 1, 1), (1, 1, 2), (3, 2, 2)]
+    assert events == [
+        ("activate", pipeline.video_vae),
+        ("image", (16, 16)),
+        ("image", (32, 16)),
+        ("video", "frames:reference.mp4"),
+        ("offload", pipeline.video_vae),
+    ]
+
+
+@pytest.mark.parametrize("encode_fails", [False, True])
+def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeypatch, encode_fails):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.audio_vae = Mock()
+    pipeline._model_cpu_offload_modules = [pipeline.audio_vae]
+    events = []
+
+    @contextmanager
+    def record_component(value):
+        events.append(("activate", value))
+        try:
+            yield
+        finally:
+            events.append(("offload", value))
+
+    def encode_embedded(*args, **kwargs):
+        events.append(("embedded",))
+        if encode_fails:
+            raise RuntimeError("audio encode failed")
+        return torch.ones(1, 2), [1]
+
+    def encode_standalone(*args, **kwargs):
+        events.append(("standalone",))
+        return torch.ones(1, 2), [1]
+
+    pipeline._encode_video_audio_conditions_resident = encode_embedded
+    pipeline._encode_audio_conditions_resident = encode_standalone
+    monkeypatch.setattr(module, "sequential_offload_component", record_component)
+
+    def call():
+        return pipeline._encode_reference_audio_conditions(
+            [{"input_has_audio": True}],
+            has_audio=[True],
+            standalone_audios=[(torch.ones(1), 16_000)],
+            max_duration_seconds=1.0,
+        )
+
+    if encode_fails:
+        with pytest.raises(RuntimeError, match="audio encode failed"):
+            call()
+    else:
+        result = call()
+        assert result[0].shape == (1, 2)
+        assert result[2].shape == (1, 2)
+
+    assert events[0] == ("activate", pipeline.audio_vae)
+    assert events[-1] == ("offload", pipeline.audio_vae)
+    assert events.count(("activate", pipeline.audio_vae)) == 1
+    assert events.count(("offload", pipeline.audio_vae)) == 1

tests/diffusion/offloader/test_sequential_backend.py

@@ -8,7 +8,13 @@ import torch
 from torch import nn
 
 from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
-from vllm_omni.diffusion.offloader.sequential_backend import ModelLevelOffloadBackend, SequentialOffloadHook
+from vllm_omni.diffusion.offloader.sequential_backend import (
+    ModelLevelOffloadBackend,
+    SequentialOffloadHook,
+    apply_sequential_offload,
+    remove_sequential_offload,
+    sequential_offload_component,
+)
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -76,6 +82,52 @@ def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None:
     assert pipeline.disable_called is True
 
 
+def test_sequential_offload_can_begin_with_dit_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
+    dit = _create_simple_module()
+    encoder = _create_simple_module()
+    offloaded: list[nn.Module] = []
+    monkeypatch.setattr(
+        SequentialOffloadHook,
+        "_to_cpu",
+        lambda self, module: offloaded.append(module),
+    )
+
+    apply_sequential_offload(
+        dit_modules=[dit],
+        encoder_modules=[encoder],
+        device=torch.device("cpu"),
+        offload_initial_dits=True,
+    )
+
+    assert offloaded == [dit]
+    remove_sequential_offload([dit, encoder])
+
+
+def test_direct_component_activation_failure_still_offloads(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    component = _create_simple_module()
+    apply_sequential_offload(
+        dit_modules=[_create_simple_module()],
+        encoder_modules=[component],
+        device=torch.device("cpu"),
+    )
+    hook = component._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    offloaded: list[nn.Module] = []
+
+    def fail_activation(module: nn.Module) -> None:
+        raise RuntimeError("activation failed")
+
+    monkeypatch.setattr(hook, "pre_forward", fail_activation)
+    monkeypatch.setattr(hook, "_to_cpu", lambda module: offloaded.append(module))
+
+    with pytest.raises(RuntimeError, match="activation failed"):
+        with sequential_offload_component(component):
+            pass
+
+    assert offloaded == [component]
+
+
 class TestMoveParamsPinMemory:
     def test_dtensor_skips_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
         """DTensor should skip pin_memory to avoid RuntimeError."""
@@ -85,7 +137,12 @@ class TestMoveParamsPinMemory:
         original_isinstance = isinstance
 
         def fake_isinstance(obj, cls):
-            if cls.__name__ == "DTensor":
+            # torchada's Tensor.to() checks tuples such as
+            # ``(str, torch.device)``.  Keep the DTensor probe compatible with
+            # the normal isinstance contract instead of breaking those calls.
+            if type(cls) is tuple:
+                return any(fake_isinstance(obj, candidate) for candidate in cls)
+            if getattr(cls, "__name__", None) == "DTensor":
                 return True
             return original_isinstance(obj, cls)
 

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -8,7 +8,7 @@ import math
 import os
 import tempfile
 from collections.abc import Iterable, Mapping, Sequence
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from itertools import groupby
 from pathlib import Path
 from typing import Any, ClassVar
@@ -44,7 +44,13 @@ from vllm_omni.diffusion.models.interface import (
     SupportsComponentDiscovery,
 )
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
-from vllm_omni.diffusion.offloader import OffloadPlan
+from vllm_omni.diffusion.offloader import (
+    OffloadPlan,
+    apply_sequential_offload,
+    remove_sequential_offload,
+    sequential_offload_component,
+)
+from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
     DiffusionPipelineProfilerMixin,
 )
@@ -543,9 +549,8 @@ class MiniMaxH3Pipeline(
     _PROFILER_TARGETS: ClassVar[list[str]] = [
         "_prepare_reference_videos",
         "encode_prompt",
-        "_encode_video_conditions",
-        "_encode_video_audio_conditions",
-        "_encode_audio_conditions",
+        "_encode_visual_conditions",
+        "_encode_reference_audio_conditions",
         "diffuse",
         "decode",
     ]
@@ -1129,9 +1134,7 @@ class MiniMaxH3Pipeline(
         input_ids: torch.Tensor,
         vision_kwargs: dict[str, torch.Tensor],
     ) -> torch.Tensor:
-        if self.od_config.enable_cpu_offload and not getattr(
-            self.od_config, "enable_distributed_layerwise_offload", False
-        ):
+        if getattr(self, "_model_cpu_offload_modules", None):
             # Invoke nn.Module.__call__ so the generic model-level offloader
             # swaps the resident DiT and encoder.
             return self.text_encoder(input_ids, **vision_kwargs)
@@ -1160,8 +1163,48 @@ class MiniMaxH3Pipeline(
             or getattr(od_config, "enable_distributed_layerwise_offload", False)
         )
 
+    def enable_omni_model_cpu_offload(
+        self,
+        *,
+        device: torch.device,
+        pin_memory: bool,
+        use_hsdp: bool,
+    ) -> None:
+        if getattr(self, "_model_cpu_offload_modules", None):
+            return
+
+        components = ModuleDiscovery.discover(self)
+        dits = components.dits
+        stages = [*components.encoders, *components.vaes]
+        modules = [*dits, *stages]
+        apply_sequential_offload(
+            dit_modules=dits,
+            encoder_modules=stages,
+            device=device,
+            pin_memory=pin_memory,
+            use_hsdp=use_hsdp,
+            offload_initial_dits=True,
+        )
+
+        self._model_cpu_offload_modules = modules
+        logger.info(
+            "MiniMax-H3 model-level CPU offload enabled for %d DiT(s), text encoder, video VAE, and audio VAE",
+            len(dits),
+        )
+
+    def disable_omni_model_cpu_offload(self) -> None:
+        modules = getattr(self, "_model_cpu_offload_modules", None)
+        if not modules:
+            return
+        remove_sequential_offload(modules)
+        self._model_cpu_offload_modules = []
+
     @contextmanager
     def _component_on_device(self, component: nn.Module):
+        if getattr(self, "_model_cpu_offload_modules", None):
+            with sequential_offload_component(component):
+                yield
+            return
         staged = self._uses_manual_component_offload()
         if staged:
             component.load_to_device()
@@ -1171,21 +1214,6 @@ class MiniMaxH3Pipeline(
             if staged:
                 component.offload_to_cpu()
 
-    def _encode_visual_condition(
-        self,
-        image: Image.Image,
-    ) -> torch.Tensor:
-        _, rank, _ = _dit_rank_world()
-        rows = None
-        if rank == 0:
-            with self._component_on_device(self.video_vae):
-                rows = self.video_vae.encode_image(image)
-        return _broadcast_tensor(
-            rows,
-            dtype=torch.float32,
-            device=self.device,
-        )
-
     def _encode_visual_conditions(
         self,
         images: list[Image.Image],
@@ -1195,44 +1223,34 @@ class MiniMaxH3Pipeline(
     ) -> tuple[torch.Tensor | None, list[tuple[int, int, int]]]:
         rows: list[torch.Tensor] = []
         shapes: list[tuple[int, int, int]] = []
-        for image in images:
-            rows.append(self._encode_visual_condition(image))
-            shapes.append((1, image.height // 16, image.width // 16))
-        if video_count:
-            video_rows, video_shapes = self._encode_video_conditions(
-                prepared_videos,
-                count=video_count,
-            )
-            rows.append(video_rows)
-            shapes.extend(video_shapes)
+        _, rank, _ = _dit_rank_world()
+        # Keep image and video references in one residency window when both
+        # appear in a request; otherwise the video branch would reload the VAE.
+        needs_video_vae = video_count > 0 or (rank == 0 and bool(images))
+        video_vae_context = self._component_on_device(self.video_vae) if needs_video_vae else nullcontext()
+        with video_vae_context:
+            if images:
+                image_rows = None
+                if rank == 0:
+                    image_rows = torch.cat([self.video_vae.encode_image(image) for image in images])
+                rows.append(
+                    _broadcast_tensor(
+                        image_rows,
+                        dtype=torch.float32,
+                        device=self.device,
+                    )
+                )
+                shapes.extend((1, image.height // 16, image.width // 16) for image in images)
+            if video_count:
+                video_rows, video_shapes = self._encode_video_conditions_resident(
+                    prepared_videos,
+                    count=video_count,
+                )
+                rows.append(video_rows)
+                shapes.extend(video_shapes)
         return (torch.cat(rows) if rows else None), shapes
 
-    def _encode_audio_condition(
-        self,
-        audio: tuple[torch.Tensor, int],
-    ) -> tuple[torch.Tensor, int]:
-        _, rank, _ = _dit_rank_world()
-        rows = None
-        audio_t = 0
-        if rank == 0:
-            with self._component_on_device(self.audio_vae):
-                rows, audio_t = self.audio_vae.encode_waveform(*audio)
-        audio_t_tensor = torch.tensor(
-            [audio_t],
-            dtype=torch.long,
-            device=self.device,
-        )
-        group, _, world_size = _dit_rank_world()
-        if world_size > 1:
-            dist.broadcast(audio_t_tensor, src=0, group=group)
-        rows = _broadcast_tensor(
-            rows,
-            dtype=torch.float32,
-            device=self.device,
-        )
-        return rows, int(audio_t_tensor.item())
-
-    def _encode_audio_conditions(
+    def _encode_audio_conditions_resident(
         self,
         audios: list[tuple[torch.Tensor, int]],
         *,
@@ -1269,15 +1287,6 @@ class MiniMaxH3Pipeline(
             [int(value) for value in lengths.tolist()],
         )
 
-    def _encode_video_conditions(
-        self,
-        prepared_videos: list[dict[str, Any]] | None,
-        *,
-        count: int,
-    ) -> tuple[torch.Tensor, list[tuple[int, int, int]]]:
-        with self._component_on_device(self.video_vae):
-            return self._encode_video_conditions_resident(prepared_videos, count=count)
-
     def _encode_video_conditions_resident(
         self,
         prepared_videos: list[dict[str, Any]] | None,
@@ -1325,15 +1334,6 @@ class MiniMaxH3Pipeline(
             [tuple(int(value) for value in item) for item in shapes.tolist()],
         )
 
-    def _encode_video_audio_conditions(
-        self,
-        prepared_videos: list[dict[str, Any]] | None,
-        *,
-        has_audio: list[bool],
-    ) -> tuple[torch.Tensor | None, list[int]]:
-        with self._component_on_device(self.audio_vae):
-            return self._encode_video_audio_conditions_resident(prepared_videos, has_audio=has_audio)
-
     def _encode_video_audio_conditions_resident(
         self,
         prepared_videos: list[dict[str, Any]] | None,
@@ -1377,6 +1377,34 @@ class MiniMaxH3Pipeline(
             [int(value) for value in lengths.tolist()],
         )
 
+    def _encode_reference_audio_conditions(
+        self,
+        prepared_videos: list[dict[str, Any]] | None,
+        *,
+        has_audio: list[bool],
+        standalone_audios: list[tuple[torch.Tensor, int]],
+        max_duration_seconds: float,
+    ) -> tuple[torch.Tensor | None, list[int], torch.Tensor | None, list[int]]:
+        # Embedded and standalone audio are consecutive direct Audio-VAE
+        # calls. Keep the component resident across both paths.
+        needs_audio_vae = any(has_audio) or bool(standalone_audios)
+        audio_vae_context = self._component_on_device(self.audio_vae) if needs_audio_vae else nullcontext()
+        with audio_vae_context:
+            embedded_condition, embedded_lengths = self._encode_video_audio_conditions_resident(
+                prepared_videos,
+                has_audio=has_audio,
+            )
+            external_condition, external_lengths = self._encode_audio_conditions_resident(
+                standalone_audios,
+                max_duration_seconds=max_duration_seconds,
+            )
+        return (
+            embedded_condition,
+            embedded_lengths,
+            external_condition,
+            external_lengths,
+        )
+
     def _initial_noise(
         self,
         *,
@@ -1755,12 +1783,15 @@ class MiniMaxH3Pipeline(
                     prepared_videos,
                     video_count=video_count,
                 )
-                embedded_audio_condition, embedded_audio_lengths = self._encode_video_audio_conditions(
+                (
+                    embedded_audio_condition,
+                    embedded_audio_lengths,
+                    external_audio_condition,
+                    external_audio_lengths,
+                ) = self._encode_reference_audio_conditions(
                     prepared_videos,
                     has_audio=has_audio,
-                )
-                external_audio_condition, external_audio_lengths = self._encode_audio_conditions(
-                    standalone_audios,
+                    standalone_audios=standalone_audios,
                     max_duration_seconds=float(num_frames) / float(sampling.fps or MINIMAX_H3_FPS),
                 )
                 audio_parts = [

vllm_omni/diffusion/offloader/__init__.py

@@ -8,7 +8,7 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan
 from vllm_omni.platforms import current_omni_platform
 
-from .base import OffloadBackend, OffloadConfig, OffloadStrategy
+from .base import OffloadBackend, OffloadConfig, OffloadStrategy, SupportsModelCpuOffload
 from .block_discovery import get_blocks_attr_names, get_blocks_from_dit, set_blocks_attr_names
 from .distributed_layerwise_backend import (
     DistributedLayerwiseOffloadBackend,
@@ -23,6 +23,7 @@ from .sequential_backend import (
     ModelLevelOffloadBackend,
     apply_sequential_offload,
     remove_sequential_offload,
+    sequential_offload_component,
 )
 from .tensor_utils import (
     dtype_size,
@@ -39,6 +40,7 @@ __all__ = [
     "OffloadConfig",
     "OffloadPlan",
     "OffloadStrategy",
+    "SupportsModelCpuOffload",
     "LayerWiseOffloadBackend",
     "DistributedLayerwiseOffloadBackend",
     "DistributedLayerwiseOffloadHook",
@@ -46,6 +48,7 @@ __all__ = [
     "PinnedModuleStager",
     "apply_sequential_offload",
     "remove_sequential_offload",
+    "sequential_offload_component",
     "apply_distributed_block_hook",
     "remove_distributed_block_hook",
     "get_offload_backend",

vllm_omni/diffusion/offloader/base.py

@@ -4,6 +4,7 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from enum import Enum
+from typing import Protocol, runtime_checkable
 
 import torch
 from torch import nn
@@ -14,6 +15,26 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
 logger = init_logger(__name__)
 
 
+@runtime_checkable
+class SupportsModelCpuOffload(Protocol):
+    """Pipeline-owned lifecycle for model-level CPU offload.
+
+    Pipelines with non-forward component entry points (for example VAE
+    ``decode_latent`` methods) need to activate those stages explicitly, so
+    generic forward-hook discovery cannot manage their full lifecycle.
+    """
+
+    def enable_omni_model_cpu_offload(
+        self,
+        *,
+        device: torch.device,
+        pin_memory: bool,
+        use_hsdp: bool,
+    ) -> None: ...
+
+    def disable_omni_model_cpu_offload(self) -> None: ...
+
+
 class OffloadStrategy(Enum):
     NONE = "none"
     MODEL_LEVEL = "model_level"  # Sequential offloading between DiT and encoders

vllm_omni/diffusion/offloader/sequential_backend.py

@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from collections.abc import Iterator
+from contextlib import contextmanager
+
 import torch
 from torch import nn
 from torch.distributed._tensor import DTensor  # type: ignore[attr-defined]
@@ -9,7 +12,7 @@ from vllm.logger import init_logger
 from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
 from vllm_omni.platforms import current_omni_platform
 
-from .base import OffloadBackend, OffloadConfig
+from .base import OffloadBackend, OffloadConfig, SupportsModelCpuOffload
 from .module_collector import ModuleDiscovery
 
 logger = init_logger(__name__)
@@ -123,6 +126,7 @@ def apply_sequential_offload(
     device: torch.device,
     pin_memory: bool = True,
     use_hsdp: bool = False,
+    offload_initial_dits: bool = False,
 ) -> None:
     """Apply sequential offloading hooks to DiT and encoder modules.
 
@@ -136,6 +140,7 @@ def apply_sequential_offload(
         device: Target GPU device for loading
         pin_memory: Whether to pin CPU memory for faster transfers
         use_hsdp: Whether HSDP is enabled (affects non_blocking behavior)
+        offload_initial_dits: Whether to begin with all DiT modules on CPU.
 
     Example:
         >>> apply_sequential_offload(
@@ -170,6 +175,14 @@ def apply_sequential_offload(
         registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook)
         logger.debug("Registered offload hook for %s", enc.__class__.__name__)
 
+    if offload_initial_dits:
+        try:
+            for dit_mod in dit_modules:
+                _get_sequential_offload_hook(dit_mod)._to_cpu(dit_mod)
+        except Exception:
+            remove_sequential_offload([*dit_modules, *encoder_modules])
+            raise
+
 
 def remove_sequential_offload(modules: list[nn.Module]) -> None:
     """Remove sequential offloading hooks from modules.
@@ -188,6 +201,31 @@ def remove_sequential_offload(modules: list[nn.Module]) -> None:
             logger.debug("Removed offload hook from %s", module.__class__.__name__)
 
 
+def _get_sequential_offload_hook(module: nn.Module) -> SequentialOffloadHook:
+    registry: HookRegistry | None = getattr(module, "_hook_registry", None)
+    hook = registry.get_hook(SequentialOffloadHook._HOOK_NAME) if registry is not None else None
+    if not isinstance(hook, SequentialOffloadHook):
+        raise RuntimeError(f"{module.__class__.__name__} has no sequential offload hook")
+    return hook
+
+
+@contextmanager
+def sequential_offload_component(module: nn.Module) -> Iterator[None]:
+    """Activate and release a hooked component called outside ``forward``."""
+    hook = _get_sequential_offload_hook(module)
+    try:
+        hook.pre_forward(module)
+        yield
+    except BaseException:
+        try:
+            hook._to_cpu(module)
+        except Exception:
+            logger.exception("Failed to release %s after component failure", module.__class__.__name__)
+        raise
+    else:
+        hook._to_cpu(module)
+
+
 class ModelLevelOffloadBackend(OffloadBackend):
     """Model-level (sequential) offloading backend.
 
@@ -197,23 +235,17 @@ class ModelLevelOffloadBackend(OffloadBackend):
     def __init__(self, config: OffloadConfig, device: torch.device):
         super().__init__(config, device)
         self._offload_modules: list[nn.Module] = []  # Track modules with hooks
-        self._custom_pipeline: nn.Module | None = None
+        self._custom_pipeline: SupportsModelCpuOffload | None = None
 
     def enable(self, pipeline: nn.Module) -> None:
         if self.enabled:
             logger.warning("ModelLevelOffloadBackend already enabled")
             return
 
-        # Pipelines with a nested transformer (e.g. Cosmos3's reasoner/generator
-        # pathways) own their own mutual-exclusion swaps and cannot be offloaded
-        # with the generic encoder<->DiT hooks. Delegate to the custom hook.
-        # TODO: Cosmos3 is the only implementer today, so we duck-type the
-        # enable/disable_omni_model_cpu_offload pair via getattr. Once a second
-        # pipeline needs it, formalize this as a Protocol (like
-        # SupportsComponentDiscovery) instead of getattr detection.
-        custom_enable = getattr(pipeline, "enable_omni_model_cpu_offload", None)
-        if callable(custom_enable):
-            custom_enable(
+        # Pipelines with non-forward component entry points own their complete
+        # mutual-exclusion lifecycle. Delegate through the explicit protocol.
+        if isinstance(pipeline, SupportsModelCpuOffload):
+            pipeline.enable_omni_model_cpu_offload(
                 device=self.device,
                 pin_memory=self.config.pin_cpu_memory,
                 use_hsdp=self.config.use_hsdp,
@@ -283,9 +315,7 @@ class ModelLevelOffloadBackend(OffloadBackend):
             return
 
         if self._custom_pipeline is not None:
-            custom_disable = getattr(self._custom_pipeline, "disable_omni_model_cpu_offload", None)
-            if callable(custom_disable):
-                custom_disable()
+            self._custom_pipeline.disable_omni_model_cpu_offload()
             self._custom_pipeline = None
             self.enabled = False
             logger.info("Model-level offloading disabled")