[Perf][MiniMax-H3] Cut video VAE encode/decode memory peaks for low-memory serving (#7241)

Signed-off-by: brandneway <gyuan4892@gmail.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>

tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py

@@ -86,6 +86,7 @@ def test_decode_releases_only_audio_before_return(monkeypatch):
         audio_vae=Mock(decode_latent=Mock(return_value=audio)),
         _component_on_device=lambda component: nullcontext(),
         _offload_model_cpu_stage_output=Mock(side_effect=lambda value: value),
+        _release_stage_cache=Mock(),
     )
     monkeypatch.setattr(
         module.current_omni_platform,
@@ -152,6 +153,7 @@ def test_forward_releases_video_after_quantization_and_before_next_seed(monkeypa
         diffuse=diffuse,
         decode=decode,
         _offload_model_cpu_stage_output=_recording_release(events, releases),
+        _release_stage_cache=lambda: None,
     )
     request = SimpleNamespace(prompts=["a prompt"], sampling_params=SimpleNamespace())
 
@@ -201,6 +203,7 @@ def test_forward_drops_the_decoded_video_before_the_next_seed_diffuses(monkeypat
         diffuse=diffuse,
         decode=decode,
         _offload_model_cpu_stage_output=lambda value: value.clone(),
+        _release_stage_cache=lambda: None,
     )
     request = SimpleNamespace(prompts=["a prompt"], sampling_params=SimpleNamespace())
 
@@ -231,6 +234,7 @@ def test_post_decode_releases_video_after_quantization(monkeypatch):
         _unpack_denoised_rows=lambda *args, **kwargs: (torch.zeros(1), torch.zeros(1)),
         decode=decode,
         _offload_model_cpu_stage_output=_recording_release(events, releases),
+        _release_stage_cache=lambda: None,
     )
     state = SimpleNamespace(
         latents=torch.zeros(1),
@@ -529,7 +533,7 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp
     torch.nn.Module.__init__(pipeline)
     pipeline.audio_vae = Mock()
     pipeline._model_cpu_offload_modules = [pipeline.audio_vae]
-    events = []
+    events: list[tuple] = []
 
     @contextmanager
     def record_component(value):

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_decode_revert.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU tests for in-place decode revert and stage-boundary cache release.
+
+``_revert_decoded_inplace`` must be numerically identical to the checkpoint's
+``revert_tensor`` (denormalize → clamp → rearrange) with zero whole-video
+copies, and ``_release_stage_cache`` must force the allocator-cache release
+(or fall back to a plain empty_cache when no component cache exists).
+"""
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 as pipeline_module
+from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+DENORM_MEAN = (-2.11628412, -2.03571429, -1.80444444)  # -imagenet_mean/std
+DENORM_STD = (4.36681223, 4.46428571, 4.44444444)  # 1/imagenet_std
+
+
+class _FakeDenormalize:
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+
+class _FakeProcessor:
+    def __init__(self):
+        self.transform_rev = _FakeDenormalize(DENORM_MEAN, DENORM_STD)
+        self.use_3d_conv = True
+        self.revert_calls = 0
+
+    def revert_tensor(self, tensor):
+        self.revert_calls += 1
+        return _legacy_revert(tensor)
+
+
+class _FakeCheckpointModel:
+    def __init__(self):
+        self.processor = _FakeProcessor()
+
+
+def _vae(model=None):
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    vae.model = model if model is not None else _FakeCheckpointModel()
+    return vae
+
+
+def _legacy_revert(tensor: torch.Tensor) -> torch.Tensor:
+    """The checkpoint's revert_tensor, op for op, on the CPU (B == 1)."""
+    if tensor.ndim == 4:
+        tensor = tensor.unsqueeze(2)
+    assert tensor.ndim == 5
+    # Normalize runs on the (b t) c h w rearrangement; at B == 1 that is a
+    # view of (B, C, T, H, W), so the per-channel broadcast is equivalent.
+    mean = torch.tensor(DENORM_MEAN).view(1, 3, 1, 1, 1)
+    std = torch.tensor(DENORM_STD).view(1, 3, 1, 1, 1)
+    tensor = (tensor - mean) / std
+    return tensor.clamp(0, 1)
+
+
+def test_revert_inplace_matches_legacy_bitwise():
+    decoded = torch.randn(1, 3, 7, 6, 8).clamp_(-3, 3)
+    expected = _legacy_revert(decoded.clone())
+    out = _vae()._revert_decoded_inplace(decoded.clone())
+    assert out.shape == expected.shape == (1, 3, 7, 6, 8)
+    assert torch.equal(out, expected)
+
+
+def test_revert_inplace_mutates_without_extra_copies():
+    decoded = torch.randn(1, 3, 5, 4, 4)
+    out = _vae()._revert_decoded_inplace(decoded)
+    assert out is decoded  # in-place: no whole-video copy materializes
+    assert out.is_contiguous()
+
+
+def test_revert_inplace_unsqueezes_4d_input():
+    # 4-D decoded tensors follow the (B, C, H, W) image contract: unsqueeze(2)
+    # lands T=1, matching the checkpoint's revert_tensor branch.
+    decoded = torch.randn(1, 3, 6, 8)
+    expected = _legacy_revert(decoded.clone())
+    out = _vae()._revert_decoded_inplace(decoded.clone())
+    assert torch.equal(out, expected)
+
+
+def test_revert_inplace_falls_back_without_contract():
+    vae = _vae()
+    vae.model.processor.transform_rev = type("Empty", (), {})()
+    decoded = torch.randn(1, 3, 5, 4, 4)
+    out = vae._revert_decoded_inplace(decoded)
+    assert vae.model.processor.revert_calls == 1
+    assert torch.equal(out, _legacy_revert(decoded))
+
+
+def test_revert_inplace_falls_back_on_bad_arity():
+    vae = _vae()
+    vae.model.processor.transform_rev = _FakeDenormalize((0.5, 0.5), (1.0, 1.0, 1.0))
+    vae._revert_decoded_inplace(torch.randn(1, 3, 5, 4, 4))
+    assert vae.model.processor.revert_calls == 1
+
+
+class _RecordingCache:
+    def __init__(self):
+        self.calls = []
+
+    def release_if_needed(self, *, force=False):
+        self.calls.append(force)
+        return True
+
+
+def _pipeline():
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    pipeline._dlo_component_cache = None
+    return pipeline
+
+
+def test_stage_cache_release_forces_component_cache():
+    pipeline = _pipeline()
+    cache = _RecordingCache()
+    pipeline._dlo_component_cache = cache
+    pipeline._release_stage_cache()
+    assert cache.calls == [True]
+
+
+def test_stage_cache_release_falls_back_to_empty_cache(monkeypatch):
+    released = []
+    monkeypatch.setattr(
+        pipeline_module.current_omni_platform,
+        "empty_cache",
+        lambda: released.append(True),
+        raising=True,
+    )
+    pipeline = _pipeline()
+    pipeline._release_stage_cache()
+    assert released == [True]

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_encode_prep.py

@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU tests for the streaming uint8 preparation in the MiniMax-H3 video VAE.
+
+``_stream_prepare_video_tensor`` must be numerically identical to the
+checkpoint's whole-video numpy path (``convert_numpy_to_tensor`` →
+``transform_tensor`` → ``transpose``) while padding on the host, so the
+device-side copies never materialize.
+"""
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+CLIP_LENGTH = 4
+
+
+class _FakeNormalize:
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+
+class _FakeProcessor:
+    # Read by the stream prep via getattr(..., False); annotation only so
+    # tests can attach it per case.
+    isolated_last_frame: bool
+
+    def __init__(self):
+        self.transform = _FakeNormalize(IMAGENET_MEAN, IMAGENET_STD)
+
+
+class _FakeCheckpointModel:
+    def __init__(self, **overrides):
+        self.clip_length = CLIP_LENGTH
+        self.isolated_first_frame = False
+        self.frame_pre_padding = 0
+        self.processor = _FakeProcessor()
+        for key, value in overrides.items():
+            setattr(self, key, value)
+
+
+def _vae(model=None):
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    vae.model = model if model is not None else _FakeCheckpointModel()
+    return vae
+
+
+def _legacy_reference(frames: np.ndarray) -> torch.Tensor:
+    """The checkpoint's numpy path, op for op, on the CPU.
+
+    convert_numpy_to_tensor: astype(fp32) → permute → ÷255
+    transform_tensor: Normalize = (x - mean) / std
+    encode_videos: transpose(0, 1)
+    """
+    tensor = torch.from_numpy(frames.astype(np.float32)).permute(0, 3, 1, 2)
+    tensor = tensor / 255.0
+    mean = torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1)
+    std = torch.tensor(IMAGENET_STD).view(1, 3, 1, 1)
+    tensor = (tensor - mean) / std
+    return tensor.transpose(0, 1).contiguous()
+
+
+def _video(num_frames, height=6, width=8, seed=0):
+    rng = np.random.default_rng(seed)
+    return rng.integers(0, 256, size=(num_frames, height, width, 3), dtype=np.uint8)
+
+
+def _pad_frames(frames: np.ndarray, pad: int) -> np.ndarray:
+    if not pad:
+        return frames
+    return np.concatenate([frames, np.repeat(frames[-1:], pad, axis=0)])
+
+
+def test_stream_prep_matches_legacy_path_bitwise():
+    frames = _video(10)
+    out = _vae()._stream_prepare_video_tensor(frames, torch.device("cpu"))
+    assert out is not None
+    assert out.shape == (3, 12, 6, 8)  # pad = (0 - 10) % 4 == 2
+    assert out.dtype == torch.float32
+    assert torch.equal(out, _legacy_reference(_pad_frames(frames, 2)))
+
+
+def test_stream_prep_isolated_first_frame_without_last_frame_skips_pad():
+    # Asymmetric checkpoint (isolated first frame, no isolated last frame):
+    # trim-stable lengths are k*clip + 0, which never satisfy the
+    # encode_temporal offset of 1, so the remote re-pads through its own
+    # whole-video cat regardless. Padding here would only upload frames the
+    # remote trim discards, so the legacy frame count is kept.
+    model = _FakeCheckpointModel(isolated_first_frame=True)
+    frames = _video(10)
+    out = _vae(model)._stream_prepare_video_tensor(frames, torch.device("cpu"))
+    assert out is not None
+    assert out.shape[1] == 10
+    assert torch.equal(out, _legacy_reference(frames))
+
+
+def test_stream_prep_symmetric_isolated_frames_pad_to_trim_stable():
+    # isolated first + last frames: offset == tail % clip == 1, so padding
+    # to the next trim-stable length makes the remote trim a no-op and keeps
+    # encode_temporal from re-padding: T=10, clip=4, tail=1 → 3*4+1 == 13.
+    model = _FakeCheckpointModel(isolated_first_frame=True)
+    model.processor.isolated_last_frame = True
+    frames = _video(10)
+    out = _vae(model)._stream_prepare_video_tensor(frames, torch.device("cpu"))
+    assert out is not None
+    assert out.shape[1] == 13
+    assert out.shape[1] % 4 == 1
+    assert torch.equal(out, _legacy_reference(_pad_frames(frames, 3)))
+
+
+def test_stream_prep_uses_processor_align_when_available():
+    calls = []
+
+    class _AlignedProcessor(_FakeProcessor):
+        def align_video_length(self, video_length, mode="pad", granularity="chunk"):
+            calls.append((video_length, mode, granularity))
+            return 5
+
+    model = _FakeCheckpointModel()
+    model.processor = _AlignedProcessor()
+    frames = _video(10)
+    out = _vae(model)._stream_prepare_video_tensor(frames, torch.device("cpu"))
+    assert out is not None
+    assert calls == [(10, "pad", "chunk")]
+    assert out.shape[1] == 15
+    assert torch.equal(out, _legacy_reference(_pad_frames(frames, 5)))
+
+
+def test_stream_prep_pads_with_last_frame():
+    frames = _video(2)  # pad = (0 - 2) % 4 == 2
+    out = _vae()._stream_prepare_video_tensor(frames, torch.device("cpu"))
+    padded = np.concatenate([frames, np.repeat(frames[-1:], 2, axis=0)])
+    assert torch.equal(out, _legacy_reference(padded))
+
+
+def test_stream_prep_aligned_input_needs_no_pad():
+    frames = _video(8)  # 8 % 4 == 0
+    out = _vae()._stream_prepare_video_tensor(frames, torch.device("cpu"))
+    assert out.shape[1] == 8
+    assert torch.equal(out, _legacy_reference(frames))
+
+
+@pytest.mark.parametrize(
+    "frames",
+    [
+        _video(5).astype(np.float32),  # not uint8
+        _video(5)[:, :, :, :1],  # single channel
+        np.zeros((5, 6, 8, 4), dtype=np.uint8),  # four channels
+        np.zeros((5, 6, 8), dtype=np.uint8),  # missing channel axis
+        np.zeros((0, 6, 8, 3), dtype=np.uint8),  # empty
+    ],
+)
+def test_stream_prep_rejects_unsupported_inputs(frames):
+    assert _vae()._stream_prepare_video_tensor(frames, torch.device("cpu")) is None
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        _FakeCheckpointModel(clip_length=None),
+        _FakeCheckpointModel(clip_length=0),
+        type("NoProcessor", (), {"clip_length": CLIP_LENGTH})(),
+        type(
+            "NoTransform",
+            (),
+            {"clip_length": CLIP_LENGTH, "processor": type("P", (), {})()},
+        )(),
+    ],
+)
+def test_stream_prep_rejects_missing_checkpoint_contract(model):
+    frames = _video(5)
+    assert _vae(model)._stream_prepare_video_tensor(frames, torch.device("cpu")) is None
+
+
+def test_stream_prep_rejects_bad_norm_arity():
+    model = _FakeCheckpointModel()
+    model.processor.transform = _FakeNormalize((0.5, 0.5), IMAGENET_STD)
+    assert _vae(model)._stream_prepare_video_tensor(_video(5), torch.device("cpu")) is None
+
+
+class _RecordingModel(_FakeCheckpointModel):
+    """Feeds encode_video a well-shaped latent and records its input."""
+
+    def __init__(self):
+        super().__init__()
+        self.received = None
+        self.latent = torch.randn(1, 4, 3, 4, 4)  # (B, C, T', H', W'), even H/W
+
+    def encode_videos(self, videos, use_fp16_latent=False, **kwargs):
+        self.received = videos
+        assert use_fp16_latent is True
+        return [self.latent]
+
+
+def _vae_for_encode_video():
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    vae.model = _RecordingModel()
+    vae.device_module = torch
+    vae.config_dict = {
+        "latent_channels": 4,
+        "latents_mean": [0.0] * 4,
+        "latents_std": [1.0] * 4,
+    }
+    # A plain tensor (not nn.Parameter): nn.Module.__setattr__ refuses
+    # Parameters on instances built via object.__new__.
+    vae._dummy_parameter = torch.zeros(1, dtype=torch.float32)
+    vae.parameters = lambda: iter([vae._dummy_parameter])
+    return vae
+
+
+def test_encode_video_streams_uint8_numpy(monkeypatch):
+    monkeypatch.delenv("VLLM_OMNI_VAE_ENCODE_LEGACY_PREP", raising=False)
+    vae = _vae_for_encode_video()
+    rows, shape = vae.encode_video(_video(10))
+    assert isinstance(vae.model.received, list)
+    assert isinstance(vae.model.received[0], torch.Tensor)
+    assert vae.model.received[0].shape == (3, 12, 6, 8)
+    assert rows.shape[0] == shape[0] * shape[1] * shape[2] // 4  # patchify rows
+
+
+def test_encode_video_legacy_env_escapes_to_checkpoint_path(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_VAE_ENCODE_LEGACY_PREP", "1")
+    vae = _vae_for_encode_video()
+    frames = _video(10)
+    vae.encode_video(frames)
+    assert vae.model.received is frames  # untouched numpy, checkpoint converts
+
+
+def test_encode_video_falls_back_when_preparation_unsupported(monkeypatch):
+    monkeypatch.delenv("VLLM_OMNI_VAE_ENCODE_LEGACY_PREP", raising=False)
+    vae = _vae_for_encode_video()
+    vae.model.clip_length = None  # contract not discoverable
+    frames = _video(10)
+    vae.encode_video(frames)
+    assert vae.model.received is frames

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_split_residency.py

@@ -0,0 +1,218 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU tests for the split encoder/decoder residency of the MiniMax-H3 VAE.
+
+The encode half (CNN encoder + quant_conv) and the decode half
+(post_quant_conv + ViT decoder) stage independently: loading one half must
+rebind only its own parameters while the other half stays on its pinned
+CPU master.
+"""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.models.minimax_h3.vae import (
+    MiniMaxH3VideoVAE,
+    _VideoVAEPartProxy,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _FakeRemoteModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.encoder = nn.Linear(8, 8)
+        self.quant_conv = nn.Linear(4, 4)
+        self.post_quant_conv = nn.Linear(4, 4)
+        self.decoder = nn.Linear(16, 16)
+
+
+class _FakeRemote(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = _FakeRemoteModel()
+
+
+def _vae() -> MiniMaxH3VideoVAE:
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    nn.Module.__init__(vae)  # bare instance: enable Module attribute assignment
+    vae._device_target = torch.device("cpu")
+    vae.remote = _FakeRemote()
+    vae.model = vae.remote.model
+    vae._stager = None
+    vae._encoder_stager = None
+    vae._decoder_stager = None
+    vae._build_residency_stagers(torch.device("cpu"))
+    vae.encoder_component = _VideoVAEPartProxy(vae, "encoder")
+    vae.decoder_component = _VideoVAEPartProxy(vae, "decoder")
+    return vae
+
+
+def _ptr(param: torch.nn.Parameter) -> int:
+    return param.untyped_storage().data_ptr()
+
+
+def test_split_stagers_snapshot_disjoint_halves():
+    vae = _vae()
+    enc_master = _ptr(vae.model.encoder.weight)
+    dec_master = _ptr(vae.model.decoder.weight)
+    pqc_master = _ptr(vae.model.post_quant_conv.weight)
+    assert enc_master != dec_master
+
+    vae._load_part_to_device("encoder")
+    # The encoder half is rebound to freshly allocated staging storage...
+    assert _ptr(vae.model.encoder.weight) != enc_master
+    assert _ptr(vae.model.quant_conv.weight) != enc_master
+    # ...while every decode-half tensor stays on its own CPU master.
+    assert _ptr(vae.model.decoder.weight) == dec_master
+    assert _ptr(vae.model.post_quant_conv.weight) == pqc_master
+
+    vae._offload_part_to_cpu("encoder")
+    assert _ptr(vae.model.encoder.weight) == enc_master
+
+
+def test_part_load_is_idempotent_and_independent():
+    vae = _vae()
+    vae._load_part_to_device("decoder")
+    staged_ptr = _ptr(vae.model.decoder.weight)
+    vae._load_part_to_device("decoder")
+    assert _ptr(vae.model.decoder.weight) == staged_ptr
+    vae._offload_part_to_cpu("decoder")
+
+
+def test_whole_component_load_and_offload_cover_both_halves():
+    vae = _vae()
+    enc_master = _ptr(vae.model.encoder.weight)
+    dec_master = _ptr(vae.model.decoder.weight)
+    vae.load_to_device()
+    assert _ptr(vae.model.encoder.weight) != enc_master
+    assert _ptr(vae.model.decoder.weight) != dec_master
+    vae.offload_to_cpu()
+    assert _ptr(vae.model.encoder.weight) == enc_master
+    assert _ptr(vae.model.decoder.weight) == dec_master
+
+
+def test_part_values_survive_a_load_offload_cycle():
+    vae = _vae()
+    expected = vae.model.encoder.weight.detach().clone()
+    vae._load_part_to_device("encoder")
+    assert torch.equal(vae.model.encoder.weight.detach(), expected)
+    vae._offload_part_to_cpu("encoder")
+    assert torch.equal(vae.model.encoder.weight.detach(), expected)
+
+
+def test_proxy_routes_to_the_matching_half():
+    vae = _vae()
+    enc_master = _ptr(vae.model.encoder.weight)
+    dec_master = _ptr(vae.model.decoder.weight)
+    vae.encoder_component.load_to_device()
+    assert _ptr(vae.model.encoder.weight) != enc_master
+    assert _ptr(vae.model.decoder.weight) == dec_master
+    vae.encoder_component.offload_to_cpu()
+    assert _ptr(vae.model.encoder.weight) == enc_master
+
+
+def test_proxy_does_not_register_the_adapter_as_submodule():
+    vae = _vae()
+    proxy = vae.encoder_component
+    assert proxy._vae is vae
+    # The back-reference must not recurse through module traversal.
+    assert all(child is not vae for child in proxy.children())
+    assert len(list(proxy.parameters())) == 0
+
+
+def test_component_on_device_unwraps_proxy_for_sequential_offload(monkeypatch):
+    """Model-level CPU offload hooks the real video_vae, never the part
+    proxies: _component_on_device must hand the hooked module to the
+    sequential context, or _get_sequential_offload_hook raises on the
+    hook-less proxy on the first encode/decode of every request."""
+    from contextlib import contextmanager
+
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+
+    vae = _vae()
+    seen: list[nn.Module] = []
+
+    @contextmanager
+    def fake_sequential_offload_component(component):
+        seen.append(component)
+        yield
+
+    monkeypatch.setattr(pipeline_module, "sequential_offload_component", fake_sequential_offload_component)
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    nn.Module.__init__(pipeline)
+    pipeline._model_cpu_offload_modules = [vae]
+
+    with pipeline._component_on_device(vae.encoder_component):
+        pass
+    with pipeline._component_on_device(vae.decoder_component):
+        pass
+    # A plain component passes through unchanged.
+    with pipeline._component_on_device(vae):
+        pass
+    # Mocks auto-create any attribute, so the unwrap must be a type check:
+    # a Mock component must reach the sequential context as itself.
+    mock_vae = Mock()
+    pipeline._model_cpu_offload_modules = [mock_vae]
+    with pipeline._component_on_device(mock_vae):
+        pass
+
+    assert seen == [vae, vae, vae, mock_vae]
+
+
+def test_sequential_offload_target_is_the_hooked_module():
+    from vllm_omni.diffusion.offloader.sequential_backend import _get_sequential_offload_hook
+
+    vae = _vae()
+    proxy = vae.encoder_component
+    # The pre-fix failure mode: the proxy carries no hook registry.
+    with pytest.raises(RuntimeError, match="no sequential offload hook"):
+        _get_sequential_offload_hook(proxy)
+    assert proxy.sequential_offload_target is vae
+
+
+def test_cache_retention_forwarded_to_both_stagers():
+    vae = _vae()
+    cache = object()
+    vae.set_omni_component_cache(cache)
+    assert vae._encoder_stager.cache_retention is cache
+    assert vae._decoder_stager.cache_retention is cache
+
+
+def test_unknown_part_raises():
+    vae = _vae()
+    with pytest.raises(ValueError, match="unknown video VAE part"):
+        vae._part_stager("middle")
+
+
+def test_missing_checkpoint_structure_falls_back_to_whole_stager():
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    nn.Module.__init__(vae)
+    vae._device_target = torch.device("cpu")
+
+    class BareRemote(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.model = nn.Linear(4, 4)  # no encoder/decoder submodules
+
+    vae.remote = BareRemote()
+    vae.model = vae.remote.model
+    vae._stager = None
+    vae._encoder_stager = None
+    vae._decoder_stager = None
+    vae._build_residency_stagers(torch.device("cpu"))
+    assert vae._stager is not None
+    assert vae._encoder_stager is None
+    assert vae._decoder_stager is None
+    # Part loads fall back to the whole-module stager.
+    vae._load_part_to_device("decoder")
+    assert vae._stager.loaded
+    vae._offload_part_to_cpu("decoder")
+    assert not vae._stager.loaded

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_temporal_patches.py

@@ -0,0 +1,335 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU tests for the temporal VAE forks and the uint8 output transfer.
+
+The encode fork must feed bit-identical frames to the checkpoint's chunk
+loop (padding moves from a whole-video cat into the tail chunk), the decode
+fork must produce the same video as the legacy fp32-buffer chain
+(revert -> clamp -> *255 -> round -> uint8), and the output transfer must
+match the previous ``.to(uint8, contiguous)`` conversion.
+"""
+
+from typing import Any, cast
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+    _prepare_minimax_h3_video_output,
+)
+from vllm_omni.diffusion.models.minimax_h3.vae_temporal import (
+    _decode_temporal_streaming_uint8,
+    _encode_temporal_tail_pad,
+    install_temporal_stream_patches,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+CLIP = 4
+
+
+# ---------------------------------------------------------------------------
+# encode_temporal fork
+# ---------------------------------------------------------------------------
+
+
+class _EncodeModel:
+    """Minimal stand-in implementing the checkpoint's encode_temporal."""
+
+    # Attached per test case as capability probes for the install checks;
+    # annotations only, the fork must keep working when they are absent.
+    encode_temporal: Any
+    _decode_temporal_streaming: Any
+    processor: Any
+
+    def __init__(self, *, isolated_first=False, isolated_last=False, key_frame=False, frame_drop=0):
+        self.clip_length = CLIP
+        self.isolated_first_frame = isolated_first
+        self.frame_pre_padding = 0
+        self.isolated_key_frame = key_frame
+        self.token_drop = 0
+        self.isolated_last_frame = isolated_last
+        self.frame_drop = frame_drop
+        self.encoded = []
+
+    def _adaptive_encode(self, clip):
+        self.encoded.append(clip.clone())
+        return clip  # identity: z_list entries are the clips themselves
+
+    def reference_encode_temporal(self, x):
+        """The checkpoint's original: whole-video pad, then chunk."""
+        offset_frame = 1 if self.isolated_first_frame and self.frame_pre_padding == 0 else 0
+        if x.shape[2] % self.clip_length != offset_frame:
+            pad_size = (offset_frame - x.shape[2]) % self.clip_length
+            pad_frames = x[:, :, -1:].repeat(1, 1, pad_size, 1, 1)
+            x = torch.cat([x, pad_frames], dim=2)
+        num_chunks = (x.shape[2] - offset_frame) // self.clip_length
+        z_list = []
+        for i in range(num_chunks):
+            start_idx = i * self.clip_length + offset_frame
+            end_idx = (i + 1) * self.clip_length + offset_frame
+            clip_x = x[:, :, start_idx:end_idx, :, :]
+            if self.isolated_key_frame:
+                key_frame = clip_x[:, :, :1, :, :]
+                z_key = self._adaptive_encode(key_frame)
+                if clip_x.shape[2] > 1:
+                    z_video = self._adaptive_encode(clip_x[:, :, 1:, :, :])
+                    z = torch.cat([z_key, z_video], dim=2)
+                else:
+                    z = z_key
+            else:
+                z = self._adaptive_encode(clip_x)
+            z_list.append(z)
+        z = torch.cat(z_list, dim=2)
+        if self.token_drop > 0:
+            z = z[:, :, : -self.token_drop]
+        if self.isolated_first_frame:
+            z_first = self._adaptive_encode(x[:, :, :1, :, :])
+            z = torch.cat([z_first, z], dim=2)
+        if self.isolated_last_frame:
+            frame_num = x.shape[2]
+            last_idx = frame_num - self.frame_drop + offset_frame
+            z_last = self._adaptive_encode(x[:, :, last_idx : last_idx + 1, :, :])
+            z = torch.cat([z, z_last], dim=2)
+        return z
+
+
+def _video_x(t, c=3, h=4, w=4):
+    return torch.arange(t * c * h * w, dtype=torch.float32).reshape(1, c, t, h, w) / 255.0
+
+
+@pytest.mark.parametrize("num_frames", [8, 10, 11, 13])
+@pytest.mark.parametrize(
+    "flags",
+    [
+        {},
+        {"isolated_first": True},
+        {"isolated_last": True},
+        {"isolated_first": True, "isolated_last": True},
+        {"key_frame": True},
+    ],
+)
+def test_encode_fork_matches_whole_video_pad_bitwise(num_frames, flags):
+    x = _video_x(num_frames)
+    ref_model = _EncodeModel(**flags)
+    expected = ref_model.reference_encode_temporal(x.clone())
+    # Empty-slice encodes (out-of-range isolated-last indexes are no-ops in
+    # the checkpoint) carry no content; compare only real calls.
+    ref_calls = [c.clone() for c in ref_model.encoded if c.shape[2] > 0]
+
+    fork_model = _EncodeModel(**flags)
+    out = _encode_temporal_tail_pad(fork_model, x.clone())
+
+    assert out.shape == expected.shape
+    assert torch.equal(out, expected)
+    assert len(fork_model.encoded) == len(ref_calls)
+    for got, want in zip(fork_model.encoded, ref_calls):
+        assert torch.equal(got, want)
+
+
+def test_encode_fork_avoids_whole_video_pad_for_misaligned_input():
+    # 375 frames with an asymmetric config (isolated last frame only):
+    # the checkpoint always pads (375 % 17 != 0); the fork must encode the
+    # tail chunk with an in-chunk pad instead of concatenating the video.
+    model = _EncodeModel(isolated_last=True, frame_drop=16)
+    model.clip_length = 17
+    x = _video_x(375, c=1, h=2, w=2)
+    _encode_temporal_tail_pad(model, x)
+    # 375 % 17 == 1 → pad 16 inside the tail chunk; the isolated last frame
+    # resolves to the final real frame (pad frames replicate it).
+    assert len(model.encoded) == (375 // 17) + 1 + 1
+    assert torch.equal(model.encoded[-1].flatten(), x[:, :, -1:, :, :].flatten())
+
+
+# ---------------------------------------------------------------------------
+# _decode_temporal_streaming uint8 fork
+# ---------------------------------------------------------------------------
+
+DENORM_MEAN = (-2.11628412, -2.03571429, -1.80444444)
+DENORM_STD = (4.36681223, 4.46428571, 4.44444444)
+
+
+class _FakeDenorm:
+    def __init__(self):
+        self.mean = DENORM_MEAN
+        self.std = DENORM_STD
+
+
+class _FakeProcessor:
+    def __init__(self):
+        self.transform_rev = _FakeDenorm()
+
+
+class _DecodeModel:
+    """Minimal stand-in for the checkpoint's streaming decode contract."""
+
+    def __init__(self, *, token_drop=0, frame_overlap=0):
+        self.tokens_chunk_size = 2
+        self.vae_ratio_t = 1
+        self.token_overlap = 1 if token_drop else 0
+        self.token_drop = token_drop
+        self.frame_pre_padding = 0
+        self.frame_overlap = frame_overlap
+        self.processor = _FakeProcessor()
+
+    def _adaptive_decode(self, clip_z):
+        # Deterministic pseudo-decode: map latent tokens to frames 1:1.
+        return torch.sin(clip_z * 3.0)
+
+    def blend(self, a, b, overlap, dim=-3):
+        return (a + b) / 2.0
+
+    def _decode_temporal_pad_frames(self, z, pad_tokens):
+        return int(pad_tokens)
+
+    def _decode_temporal_output_frame_plan(self, z, z_head, z_tail, num_chunks, pad_tokens):
+        total = num_chunks * self.tokens_chunk_size * self.vae_ratio_t
+        if z_head is not None:
+            total += 1
+        if z_tail is not None:
+            total += 1
+        if self.token_drop > 0:
+            total += self.frame_overlap
+        return int(total), int(pad_tokens), int(total - pad_tokens)
+
+    def legacy_streaming(self, z, z_head, z_tail, num_chunks, pad_tokens, temporal_cat_dtype):
+        """The checkpoint's original: accumulate in the cat dtype, quantize later."""
+        total_frames, _, output_frames = self._decode_temporal_output_frame_plan(
+            z, z_head, z_tail, num_chunks, pad_tokens
+        )
+        chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
+        split_count = int(self.token_drop > 0) + 1
+        dec = None
+        dec_overlap = None
+        write_pos = 0
+
+        def write_part(part):
+            nonlocal dec, write_pos
+            if dec is None:
+                out_shape = list(part.shape)
+                out_shape[2] = output_frames
+                dec = torch.empty(out_shape, dtype=part.dtype)
+            n = min(part.shape[2], dec.shape[2] - write_pos)
+            if n > 0:
+                dec[:, :, write_pos : write_pos + n].copy_(part[:, :, :n])
+                write_pos += n
+
+        for i in range(num_chunks):
+            t_start = i * self.tokens_chunk_size
+            t_end = t_start + self.tokens_chunk_size + self.token_overlap
+            clip_z = z[:, :, t_start:t_end]
+            clip_dec = self._adaptive_decode(clip_z)
+            if temporal_cat_dtype is not None and clip_dec.dtype != temporal_cat_dtype:
+                clip_dec = clip_dec.to(temporal_cat_dtype)
+            for j in range(split_count):
+                chunk = clip_dec[:, :, j * chunk_dec : min(j * chunk_dec + chunk_dec, clip_dec.shape[2])]
+                if j == 0:
+                    if dec_overlap is not None:
+                        chunk = self.blend(dec_overlap, chunk, self.frame_overlap, dim=-3)
+                        dec_overlap = None
+                    write_part(chunk)
+                else:
+                    dec_overlap = chunk.contiguous()
+            if i == num_chunks - 1 and dec_overlap is not None:
+                write_part(dec_overlap)
+                dec_overlap = None
+        # Legacy post-decode chain: revert in place in the decoded dtype, then
+        # the adapter's ``frames.float()`` before the pipeline's quantizer.
+        assert dec is not None
+        mean = torch.tensor(DENORM_MEAN, dtype=dec.dtype).view(1, 3, 1, 1, 1)
+        std = torch.tensor(DENORM_STD, dtype=dec.dtype).view(1, 3, 1, 1, 1)
+        dec.sub_(mean).div_(std).clamp_(0.0, 1.0)
+        dec = dec.float()
+        dec.mul_(255.0).round_()
+        return dec.to(torch.uint8)
+
+
+@pytest.mark.parametrize("token_drop,frame_overlap", [(0, 0), (1, 1)])
+@pytest.mark.parametrize("temporal_cat_dtype", [None, torch.float16, torch.bfloat16])
+def test_decode_fork_matches_legacy_chain_bitwise(token_drop, frame_overlap, temporal_cat_dtype):
+    # The quantizer must see FP32 no matter the cat dtype: the legacy adapter
+    # returned frames.float() before the pipeline multiplied and rounded, so
+    # an fp16/bf16 cat dtype only narrows denormalization, never the rounding.
+    torch.manual_seed(0)
+    tokens = 2 * 2 + (1 if token_drop else 0)  # num_chunks * tokens_chunk_size
+    z = torch.randn(1, 3, tokens, 4, 4)
+    legacy = _DecodeModel(token_drop=token_drop, frame_overlap=frame_overlap)
+    expected = legacy.legacy_streaming(z.clone(), None, None, 2, 0, temporal_cat_dtype)
+
+    fork = _DecodeModel(token_drop=token_drop, frame_overlap=frame_overlap)
+    out = _decode_temporal_streaming_uint8(fork, z.clone(), None, None, 2, 0, temporal_cat_dtype)
+    assert out.dtype == torch.uint8
+    assert out.shape == expected.shape
+    assert torch.equal(out, expected)
+
+
+def test_decode_fork_rejects_empty_plan():
+    model = _DecodeModel()
+    model.tokens_chunk_size = 0
+    with pytest.raises(ValueError, match="non-positive output_frames"):
+        _decode_temporal_streaming_uint8(model, torch.randn(1, 3, 0, 4, 4), None, None, 0, 0, None)
+
+
+# ---------------------------------------------------------------------------
+# install + escape hatch
+# ---------------------------------------------------------------------------
+
+
+def test_install_binds_both_forks():
+    model = _EncodeModel()
+    # cast(Any, ...) keeps mypy from narrowing the probes to the lambda type,
+    # which has no __func__ for the bound-method assertions below.
+    model.encode_temporal = cast(Any, lambda x: x)  # placeholder for the capability check
+    model._decode_temporal_streaming = cast(Any, lambda *a: None)
+    model.processor = _FakeProcessor()
+    install_temporal_stream_patches(model)
+    assert model.encode_temporal.__func__ is _encode_temporal_tail_pad
+    assert model._decode_temporal_streaming.__func__ is _decode_temporal_streaming_uint8
+
+
+def test_install_skips_decode_fork_without_denorm_contract():
+    model = _EncodeModel()
+    model.encode_temporal = cast(Any, lambda x: x)
+    model._decode_temporal_streaming = cast(Any, lambda *a: None)
+    install_temporal_stream_patches(model)  # no processor -> keep original
+    assert model.encode_temporal.__func__ is _encode_temporal_tail_pad
+    assert not hasattr(model._decode_temporal_streaming, "__func__")
+
+
+def test_install_respects_legacy_env(monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_VAE_LEGACY_TEMPORAL", "1")
+    model = _EncodeModel()
+    install_temporal_stream_patches(model)
+    assert not hasattr(getattr(model, "encode_temporal", None), "__func__")
+
+
+# ---------------------------------------------------------------------------
+# _prepare_minimax_h3_video_output
+# ---------------------------------------------------------------------------
+
+
+def test_prepare_output_matches_previous_conversion():
+    video = torch.randn(1, 3, 4, 6, 8).clamp_(-1, 2)
+    legacy = (
+        video.detach()
+        .float()
+        .clamp(0, 1)
+        .mul(255)
+        .round()
+        .permute(0, 2, 3, 4, 1)
+        .to(dtype=torch.uint8, memory_format=torch.contiguous_format)
+    )
+    out = _prepare_minimax_h3_video_output(video.clone())
+    assert out.dtype == torch.uint8
+    assert out.shape == (1, 4, 6, 8, 3)
+    assert out.is_contiguous()
+    assert torch.equal(out, legacy)
+
+
+def test_prepare_output_uint8_passthrough():
+    video = torch.randint(0, 256, (1, 3, 4, 6, 8), dtype=torch.uint8)
+    out = _prepare_minimax_h3_video_output(video)
+    assert out.dtype == torch.uint8
+    assert out.is_contiguous()
+    assert torch.equal(out, video.permute(0, 2, 3, 4, 1).contiguous())