[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/config/test_environment_variables.py

@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PUBLIC_OMNI: 28,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 65,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 67,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
         EnvironmentVariableCategory.INTERNAL: 2,
     }
@@ -182,7 +182,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
         ModelEnvironmentVariableDisposition.INTERNALIZE: 16,
-        ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 5,
+        ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 7,
     }
 
 

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())

vllm_omni/config/environment_variable_inventory.py

@@ -214,6 +214,8 @@ _MODEL_INTERNALIZE = (
 _MODEL_DEPRECATE_REMOVE = (
     "HIGGS_AUDIO_V2_TOKENIZER_PATH",
     "VLLM_OMNI_LINGBOT_ACTION_ROOT",
+    "VLLM_OMNI_VAE_ENCODE_LEGACY_PREP",
+    "VLLM_OMNI_VAE_LEGACY_TEMPORAL",
     "VLLM_OMNI_VOXCPM_CODE_PATH",
     "XCODEC1_PATH",
     "model_stage",

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -8,7 +8,7 @@ import json
 import math
 import os
 from collections.abc import Iterable, Mapping, Sequence
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from dataclasses import fields, replace
 from itertools import groupby
 from pathlib import Path
@@ -138,9 +138,11 @@ from .time_request import (
     MINIMAX_H3_SHAPE_PLANNER,
     minimax_h3_time_shift_sigmas,
 )
-from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE
+from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE, _VideoVAEPartProxy
 
 if TYPE_CHECKING:
+    from PIL import Image
+
     from vllm_omni.diffusion.worker.input_batch import InputBatch
     from vllm_omni.diffusion.worker.utils import StepRequestState
 
@@ -376,12 +378,20 @@ def _minimax_h3_post_process(output, output_type: str = "np"):
 
 def _prepare_minimax_h3_video_output(video: torch.Tensor) -> torch.Tensor:
     """Quantize decoded frames in place before worker-to-engine transfer."""
-    video = video.detach().float()
+    video = video.detach()
+    if video.dtype == torch.uint8:
+        # Streaming decode already quantized and clamped; only the transfer
+        # layout remains.
+        return video.permute(0, 2, 3, 4, 1).contiguous()
+    video = video.float()
     video.clamp_(0, 1).mul_(255).round_()
-    return video.permute(0, 2, 3, 4, 1).to(
-        dtype=torch.uint8,
-        memory_format=torch.contiguous_format,
-    )
+    permuted = video.permute(0, 2, 3, 4, 1)
+    out = torch.empty(permuted.shape, dtype=torch.uint8, device=video.device)
+    # copy_ fuses the layout change and the cast into one kernel;
+    # ``.to(dtype=uint8, memory_format=contiguous_format)`` materializes a
+    # contiguous FP32 intermediate first (~4.2GB for a 15s clip).
+    out.copy_(permuted)
+    return out
 
 
 def _register_dlo_component_cache(cache: BoundedAllocatorCache, *components: Any) -> None:
@@ -1342,6 +1352,15 @@ class MiniMaxH3Pipeline(
     @contextmanager
     def _component_on_device(self, component: nn.Module):
         if getattr(self, "_model_cpu_offload_modules", None):
+            # Sequential offload hooks whole modules (enable_omni_model_cpu_offload
+            # registers them on the discovered components, e.g. the real
+            # video_vae). Split-residency proxies carry no hook, so unwrap to the
+            # hooked module before entering the context — whole-module movement
+            # has no half-residency benefit anyway. The check is a type check, not
+            # getattr: Mock components auto-create any attribute, which would
+            # unwrap them to a child mock and break scope tracking.
+            if isinstance(component, _VideoVAEPartProxy):
+                component = component.sequential_offload_target
             with sequential_offload_component(component):
                 yield
             return
@@ -1381,6 +1400,50 @@ class MiniMaxH3Pipeline(
         """Whether this rank's output is returned by the diffusion executor."""
         return not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0
 
+    def _encode_visual_conditions(
+        self,
+        images: list[Image.Image],
+        prepared_videos: list[dict[str, Any]] | None,
+        *,
+        video_count: int,
+    ) -> tuple[torch.Tensor | None, list[tuple[int, int, int]]]:
+        rows: list[torch.Tensor] = []
+        shapes: list[tuple[int, int, int]] = []
+        _, 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.
+        # Encoding touches only the CNN encoder half, so the 9GB ViT decoder
+        # stays off the device for the whole window.
+        needs_video_vae = video_count > 0 or (rank == 0 and bool(images))
+        video_vae_context = (
+            self._component_on_device(self.video_vae.encoder_component) 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)
+        # The latents are extracted; the encode's input/staging pages are idle
+        # and must not stay mapped through the denoise and decode peaks.
+        if needs_video_vae:
+            self._release_stage_cache()
+        return (torch.cat(rows) if rows else None), shapes
+
     def _offload_model_cpu_stage_output(self, tensor: torch.Tensor) -> torch.Tensor:
         """Release a decoded output's storage before a later seed reloads the DiT.
 
@@ -1768,6 +1831,26 @@ class MiniMaxH3Pipeline(
             raise ValueError("MiniMax H3 chunked MP4 encoding currently expects one output per decoder")
         return videos[0]
 
+    def _release_stage_cache(self) -> None:
+        """Return idle allocator pages to the device at stage boundaries.
+
+        The bounded component cache releases only past its idle-cache bound
+        (>25% of device capacity), which on large devices lets a finished
+        stage's freed activations -- the DiT's denoise buffers, an encode
+        input, a decoded frame tensor -- stay physically mapped across the
+        next stage's peak. Forcing the release at a boundary is exact (only
+        free pages are returned; live tensors are untouched) and costs one
+        remap per later allocation.
+        """
+        cache = getattr(self, "_dlo_component_cache", None)
+        if cache is not None:
+            try:
+                cache.release_if_needed(force=True)
+            except BaseException:
+                logger.exception("Failed to release retained allocator cache at stage boundary")
+            return
+        current_omni_platform.empty_cache()
+
     def decode(
         self,
         video_latent: torch.Tensor,
@@ -1776,7 +1859,11 @@ class MiniMaxH3Pipeline(
         height: int,
         width: int,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        with self._component_on_device(self.video_vae):
+        # Denoise just ended: its freed activation pages must not stay mapped
+        # through the VAE decode peak. Decoding needs only the ViT decoder
+        # half of the VAE, so the CNN encoder stays off the device.
+        self._release_stage_cache()
+        with self._component_on_device(self.video_vae.decoder_component):
             with current_omni_platform.create_autocast_context(
                 device_type=self.device.type,
                 dtype=torch.float16,
@@ -2199,6 +2286,11 @@ class MiniMaxH3Pipeline(
                 # loop rebinds it. post_decode() rebinds for the same reason.
                 video = self._offload_model_cpu_stage_output(_prepare_minimax_h3_video_output(video))
                 videos.append(video)
+                # The FP32 decoded frames are quantized into the appended uint8
+                # tensor; drop the reference and return the idle pages instead of
+                # holding them through the next output's denoise/decode.
+                del video
+                self._release_stage_cache()
                 audios.append(audio)
         if videos and isinstance(videos[0], bytes):
             video = videos[0] if len(videos) == 1 else videos
@@ -2518,6 +2610,7 @@ class MiniMaxH3Pipeline(
                 width=shape["width"],
             )
             video = self._offload_model_cpu_stage_output(_prepare_minimax_h3_video_output(video))
+            self._release_stage_cache()
         return DiffusionOutput(
             output=(video, audio),
             post_process_func=get_minimax_h3_post_process_func(self.od_config),

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -6,11 +6,13 @@ from __future__ import annotations
 
 import importlib
 import json
+import os
 from collections.abc import Iterator
 from contextlib import AbstractContextManager, contextmanager, nullcontext
 from pathlib import Path
 from typing import Any, ClassVar
 
+import numpy as np
 import torch
 import torch.distributed as dist
 import torch.nn as nn
@@ -33,11 +35,15 @@ from vllm_omni.diffusion.offloader.module_residency import (
 from .chunked_decode import decode_h3_chunks
 from .ops import install_h3_vae_optimizations
 from .packed_tokens import minimax_h3_patchify_video_latent
+from .vae_temporal import install_temporal_stream_patches
 
 MINIMAX_H3_KEYFRAME_ENCODE_SEED = 42
 MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000
 MINIMAX_H3_AUDIO_CHANNELS = 2
 
+# Escape hatch back to the checkpoint's whole-video numpy preparation path.
+_VAE_ENCODE_LEGACY_PREP_ENV = "VLLM_OMNI_VAE_ENCODE_LEGACY_PREP"
+
 
 logger = init_logger(__name__)
 
@@ -63,6 +69,11 @@ def _minimax_h3_keyframe_encode_context(
         yield
 
 
+def _legacy_encode_prep_enabled() -> bool:
+    value = os.environ.get(_VAE_ENCODE_LEGACY_PREP_ENV, "0")
+    return value.strip().lower() not in ("", "0", "false", "off")
+
+
 def _load_component_config(component_path: str) -> dict[str, Any]:
     config_path = Path(component_path) / "config.json"
     config = json.loads(config_path.read_text(encoding="utf-8"))
@@ -264,6 +275,49 @@ class _AudioVAEDeterminismContext(AbstractContextManager):
         return False
 
 
+class _VideoVAEPartProxy(nn.Module):
+    """Residency proxy for one half of the split video VAE.
+
+    The pipeline's ``_component_on_device`` drives whatever object it is
+    given through ``load_to_device``/``offload_to_cpu``. The checkpoint's
+    ViT decoder holds ~9GB of FP32 weights while the CNN encoder holds
+    ~0.7GB, and each is used by exactly one direction, so routing those
+    calls per half keeps the decoder off the device while reference videos
+    encode and the encoder off while latents decode. ``object.__setattr__``
+    keeps the back-reference out of the module registry (registering the
+    adapter as a submodule would recurse through ``parameters()``).
+    """
+
+    def __init__(self, vae: MiniMaxH3VideoVAE, part: str) -> None:
+        super().__init__()
+        object.__setattr__(self, "_vae", vae)
+        object.__setattr__(self, "_part", part)
+
+    def load_to_device(self) -> None:
+        self._vae._load_part_to_device(self._part)
+
+    def offload_to_cpu(self) -> None:
+        self._vae._offload_part_to_cpu(self._part)
+
+    def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None:
+        self._vae.set_omni_component_cache(cache)
+
+    @property
+    def sequential_offload_target(self) -> MiniMaxH3VideoVAE:
+        """The module model-level CPU offload actually hooks.
+
+        ``enable_omni_model_cpu_offload`` registers the sequential hook on the
+        real ``video_vae``, and the hook moves the whole module through
+        ``parameters()`` — entering the sequential context through this proxy
+        would find no ``_hook_registry`` and raise, and whole-module movement
+        has no half-residency benefit anyway. ``_component_on_device`` unwraps
+        through this property before entering its sequential-offload branch;
+        the manual per-half staging above stays proxy-driven, where split
+        residency is the whole point.
+        """
+        return self._vae
+
+
 class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
     """Adapter around the checkpoint's native parallel-tiled video VAE."""
 
@@ -303,14 +357,15 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
                 decoder,
                 device=device,
             )
-        self._stager = None
-        if initial_device.type == "cpu" and device.type not in ("cpu", "meta"):
-            self._stager = PinnedModuleStager(
-                self.remote,
-                device,
-                pin_memory=True,
-            )
+        install_temporal_stream_patches(self.remote.model)
         self.model = self.remote.model
+        self._stager = None
+        self._encoder_stager = None
+        self._decoder_stager = None
+        if initial_device.type == "cpu" and device.type not in ("cpu", "meta"):
+            self._build_residency_stagers(device)
+        self.encoder_component = _VideoVAEPartProxy(self, "encoder")
+        self.decoder_component = _VideoVAEPartProxy(self, "decoder")
         self.use_tiling = True
         self.use_slicing = False
         self.parallel_size = 1
@@ -321,27 +376,97 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         if self._tile_gather_reuse_enabled():
             self._install_persistent_tile_gather()
 
-    def load_to_device(self) -> None:
+    def _build_residency_stagers(self, device: torch.device) -> None:
+        """Stage the encode and decode halves of the remote separately.
+
+        The encode path touches only ``encoder`` + ``quant_conv`` and the
+        decode path only ``post_quant_conv`` + ``decoder`` (verified against
+        the checkpoint's ``AutoencoderKLLegacy``: the one cross reference,
+        ``tiled_decode``'s ``getattr(self.encoder, "mask_enabled")``, is
+        short-circuited by ``self.training`` at inference). No storage is
+        shared across the two groups, so independent staging is exact. When
+        the checkpoint's structure is not discoverable, fall back to
+        whole-module staging (the previous single-stager behavior).
+        """
+        part_names = ("encoder", "quant_conv", "post_quant_conv", "decoder")
+        if all(isinstance(getattr(self.model, name, None), nn.Module) for name in part_names):
+            self._encoder_stager = PinnedModuleStager(
+                [self.model.encoder, self.model.quant_conv],
+                device,
+                pin_memory=True,
+            )
+            self._decoder_stager = PinnedModuleStager(
+                [self.model.post_quant_conv, self.model.decoder],
+                device,
+                pin_memory=True,
+            )
+            return
+        self._stager = PinnedModuleStager(
+            self.remote,
+            device,
+            pin_memory=True,
+        )
+
+    def _part_stager(self, part: str) -> PinnedModuleStager | None:
+        if part == "encoder":
+            return self._encoder_stager
+        if part == "decoder":
+            return self._decoder_stager
+        raise ValueError(f"unknown video VAE part {part!r}")
+
+    def _load_part_to_device(self, part: str) -> None:
+        stager = self._part_stager(part)
+        if stager is not None:
+            stager.load()
+        elif self._stager is not None:
+            self._stager.load()
+        else:
+            # No staged residency (the component lives on the device already
+            # or is fully CPU-resident): fall back to whole-module placement.
+            self.remote.to(self._device_target)
+
+    def _offload_part_to_cpu(self, part: str) -> None:
+        stager = self._part_stager(part)
+        if stager is not None:
+            stager.offload()
+            return
         if self._stager is not None:
+            self._stager.offload()
+            return
+        self.remote.to("cpu")
+        self._release_component_cache()
+
+    def _release_component_cache(self) -> None:
+        cache = getattr(self, "_omni_component_cache", None)
+        if cache is None:
+            torch.accelerator.empty_cache()
+        else:
+            cache.release_if_needed()
+
+    def load_to_device(self) -> None:
+        if self._encoder_stager is not None:
+            self._encoder_stager.load()
+            self._decoder_stager.load()
+        elif self._stager is not None:
             self._stager.load()
         else:
             self.remote.to(self._device_target)
 
     def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None:
         self._omni_component_cache = cache
-        if self._stager is not None:
-            self._stager.set_cache_retention(cache)
+        for stager in (self._encoder_stager, self._decoder_stager, self._stager):
+            if stager is not None:
+                stager.set_cache_retention(cache)
 
     def offload_to_cpu(self) -> None:
-        if self._stager is not None:
+        if self._encoder_stager is not None:
+            self._encoder_stager.offload()
+            self._decoder_stager.offload()
+        elif self._stager is not None:
             self._stager.offload()
         else:
             self.remote.to("cpu")
-            cache = getattr(self, "_omni_component_cache", None)
-            if cache is None:
-                torch.accelerator.empty_cache()
-            else:
-                cache.release_if_needed()
+            self._release_component_cache()
 
     def set_parallel_size(
         self,
@@ -650,6 +775,91 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
             patch_size=(1, 2, 2),
         ).float()
 
+    def _stream_prepare_video_tensor(
+        self,
+        frames: np.ndarray,
+        device: torch.device,
+    ) -> torch.Tensor | None:
+        """Upload a uint8 ``(T, H, W, 3)`` video as one normalized FP32 tensor.
+
+        The checkpoint's numpy path uploads the whole video as FP32, then
+        ``transform_tensor`` materializes a second normalized copy, then
+        ``encode_temporal`` pads through a full-video ``torch.cat``: three
+        resident pixel-scale copies on the device. Here the padding is
+        replicated on the host, a single preallocated ``(3, T', H, W)`` tensor
+        receives clip-by-clip uploads, and the ``÷255 → (x-mean)/std`` chain
+        runs in place on each clip's staging buffer in the same op order as
+        ``convert_numpy_to_tensor`` → ``transform_tensor``, so peak device
+        memory is the output tensor plus one clip instead of three copies.
+
+        Returns ``None`` whenever the checkpoint contract this mirrors (uint8
+        frames, ``clip_length`` alignment, processor ``transform`` constants)
+        is not discoverable; callers fall back to the legacy path.
+        """
+        if frames.dtype != np.uint8 or frames.ndim != 4 or frames.shape[-1] != 3:
+            return None
+        if int(frames.shape[0]) == 0:
+            return None
+        model = self.model
+        clip_length = getattr(model, "clip_length", None)
+        transform = getattr(getattr(model, "processor", None), "transform", None)
+        mean = getattr(transform, "mean", None)
+        std = getattr(transform, "std", None)
+        if not isinstance(clip_length, int) or clip_length <= 0:
+            return None
+        if mean is None or std is None or len(mean) != 3 or len(std) != 3:
+            return None
+        # Mirror the checkpoint's temporal alignment so the device-side
+        # ``get_suitable_video_length`` trim is a no-op and the
+        # ``encode_temporal`` padding ``torch.cat`` never triggers. A
+        # trim-stable length is ``k * clip_length + tail`` (tail = frame
+        # overlap plus the isolated last frame); it also satisfies
+        # encode_temporal only when ``tail % clip_length == offset_frame``.
+        # With an asymmetric checkpoint (isolated last frame but no isolated
+        # first frame) no trim-stable length can satisfy encode_temporal, so
+        # the remote re-pads through a whole-video cat regardless and any
+        # host-side pad would be trimmed away before that -- keep the legacy
+        # frame count there instead of uploading dead frames.
+        isolated_first_frame = bool(getattr(model, "isolated_first_frame", False))
+        frame_pre_padding = int(getattr(model, "frame_pre_padding", 0) or 0)
+        offset = 1 if isolated_first_frame and frame_pre_padding == 0 else 0
+        processor = getattr(model, "processor", None)
+        tail = int(getattr(processor, "frame_overlap", 0) or 0)
+        if bool(getattr(processor, "isolated_last_frame", False)):
+            tail += 1
+        num_frames = int(frames.shape[0])
+        pad = 0
+        if tail % clip_length == offset:
+            align = getattr(processor, "align_video_length", None)
+            if callable(align):
+                pad = max(0, int(align(num_frames, mode="pad", granularity="chunk")))
+            else:
+                chunks = -(-(num_frames - tail) // clip_length)
+                pad = max(max(chunks, 1) * clip_length + tail - num_frames, 0)
+        if pad:
+            frames = np.concatenate([frames, np.repeat(frames[-1:], pad, axis=0)])
+            num_frames += pad
+        mean_t = torch.as_tensor(mean, dtype=torch.float32, device=device).view(3, 1, 1, 1)
+        std_t = torch.as_tensor(std, dtype=torch.float32, device=device).view(3, 1, 1, 1)
+        height, width = int(frames.shape[1]), int(frames.shape[2])
+        out = torch.empty(
+            (3, num_frames, height, width),
+            dtype=torch.float32,
+            device=device,
+        )
+        for start in range(0, num_frames, clip_length):
+            end = min(start + clip_length, num_frames)
+            # Frame slices of a C-contiguous (T, H, W, 3) array stay
+            # contiguous, so each clip uploads at uint8 width (4x less host
+            # traffic than the legacy FP32 upload) before the in-place
+            # normalization chain. permute lands on the checkpoint's
+            # (3, T, H, W) layout directly.
+            chunk = torch.from_numpy(frames[start:end]).to(device=device)
+            chunk = chunk.permute(3, 0, 1, 2).to(torch.float32).div_(255.0)
+            chunk = chunk.sub_(mean_t).div_(std_t)
+            out[:, start:end].copy_(chunk)
+        return out
+
     @torch.inference_mode()
     def encode_video(
         self,
@@ -679,6 +889,14 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
                 for device in devices:
                     with self.device_module.device(device):
                         self.device_module.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
+                prepared = None
+                if isinstance(frames, np.ndarray) and not _legacy_encode_prep_enabled():
+                    prepared = self._stream_prepare_video_tensor(frames, parameter.device)
+                if prepared is not None:
+                    # Tensor inputs already carry the checkpoint's expected
+                    # (3, T, H, W) normalized-FP32 contract, so encode_videos
+                    # skips its own convert/transform/pad whole-video copies.
+                    frames = [prepared]
                 latent = self.model.encode_videos(
                     frames,
                     use_fp16_latent=True,
@@ -738,7 +956,12 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
             raise RuntimeError("MiniMax H3 encode-only video VAE cannot decode latents")
         with self._decode_tiling_context(latent):
             decoded = self.model.decode_base(self._denormalize_latent(latent))
-        return self._normalize_decoded_frames(self.model.processor.revert_tensor(decoded))
+        if decoded.dtype == torch.uint8:
+            # The streaming uint8 write-back already ran the revert and the
+            # output quantizer inside write_part; nothing remains but the
+            # shape contract.
+            return self._normalize_decoded_frames(decoded)
+        return self._normalize_decoded_frames(self._revert_decoded_inplace(decoded))
 
     def _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor:
         channels = int(self.config_dict["latent_channels"])
@@ -755,7 +978,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
             frames = frames.unsqueeze(0).transpose(1, 2)
         if frames.ndim != 5:
             raise ValueError(f"unexpected decoded video shape {tuple(frames.shape)}")
-        return frames.float()
+        return frames if frames.dtype == torch.uint8 else frames.float()
 
     # Chunks are reverted through the checkpoint's processor, which
     # denormalizes and clamps into the unit interval.
@@ -791,6 +1014,33 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         with self._decode_tiling_context(z):
             decode_h3_chunks(self, z, on_chunk, group=group)
 
+    def _revert_decoded_inplace(self, decoded: torch.Tensor) -> torch.Tensor:
+        """In-place counterpart of the processor's ``revert_tensor``.
+
+        The checkpoint version materializes a denormalized copy, a clamped
+        copy, and a contiguous copy of the whole decoded video -- three
+        pixel-scale tensors resident at the decode peak. Decoding owns the
+        tensor here, so the same op order (torchvision ``Normalize`` is
+        ``(x - mean) / std``, then ``clamp(0, 1)``) runs in place on the
+        original ``(B, C, T, H, W)`` layout, which is bit-identical
+        elementwise to normalizing the ``(b t) c h w`` rearrangement and
+        returns with zero whole-video copies.
+
+        Falls back to ``processor.revert_tensor`` when the checkpoint's
+        denormalization constants are not discoverable.
+        """
+        processor = getattr(self.model, "processor", None)
+        transform_rev = getattr(processor, "transform_rev", None)
+        mean = getattr(transform_rev, "mean", None)
+        std = getattr(transform_rev, "std", None)
+        if mean is None or std is None or len(mean) != 3 or len(std) != 3:
+            return processor.revert_tensor(decoded)
+        if bool(getattr(processor, "use_3d_conv", True)) and decoded.ndim == 4:
+            decoded = decoded.unsqueeze(2)
+        mean_t = torch.as_tensor(mean, dtype=decoded.dtype, device=decoded.device).view(1, 3, 1, 1, 1)
+        std_t = torch.as_tensor(std, dtype=decoded.dtype, device=decoded.device).view(1, 3, 1, 1, 1)
+        return decoded.sub_(mean_t).div_(std_t).clamp_(0.0, 1.0)
+
 
 class MiniMaxH3AudioVAE(nn.Module):
     def __init__(

vllm_omni/diffusion/models/minimax_h3/vae_temporal.py

@@ -0,0 +1,271 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Memory-optimized forks of the checkpoint's temporal VAE loops.
+
+Two method replacements installed on the checkpoint's ``AutoencoderKLLegacy``
+at adapter construction:
+
+* ``encode_temporal`` pads misaligned frame counts by concatenating repeated
+  last frames onto the *whole* video before the chunk loop (~4GB copy for a
+  15s clip). The fork pads inside the tail chunk instead: the chunks see
+  bit-identical frames while the copy shrinks to a single clip.
+
+* ``_decode_temporal_streaming`` accumulates the whole decoded video in a
+  floating-point buffer that is denormalized, clamped, and quantized to
+  uint8 only after the final chunk lands. The fork fuses the adapter's
+  in-place revert and the output quantizer into ``write_part`` and
+  accumulates directly into a uint8 buffer -- the same per-element op order
+  (``(x - mean) / std -> clamp(0, 1) -> *255 -> round -> uint8``), so the
+  delivered video is bit-identical while the resident decode buffer shrinks
+  4x.
+
+Set ``VLLM_OMNI_VAE_LEGACY_TEMPORAL=1`` to keep the checkpoint's methods.
+"""
+
+from __future__ import annotations
+
+import os
+from types import MethodType
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_LEGACY_TEMPORAL_ENV = "VLLM_OMNI_VAE_LEGACY_TEMPORAL"
+
+
+def _legacy_temporal_enabled() -> bool:
+    value = os.environ.get(_LEGACY_TEMPORAL_ENV, "0")
+    return value.strip().lower() not in ("", "0", "false", "off")
+
+
+def _encode_temporal_tail_pad(self, x):
+    """Drop-in fork of the checkpoint's ``encode_temporal``.
+
+    Frame handling is identical to the original; only the padding site
+    moves from a whole-video ``torch.cat`` into the tail chunk.
+    """
+    offset_frame = 1 if self.isolated_first_frame and self.frame_pre_padding == 0 else 0
+
+    orig_frames = int(x.shape[2])
+    pad_size = (offset_frame - orig_frames) % self.clip_length
+    padded_frames = orig_frames + pad_size
+    num_chunks = (padded_frames - 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 : min(end_idx, orig_frames), :, :]
+        if end_idx > orig_frames:
+            # The original pads the whole video up front by repeating the
+            # last frame; repeating it inside the tail chunk feeds the
+            # identical frames to the identical encodes.
+            clip_x = torch.cat(
+                [clip_x, clip_x[:, :, -1:].repeat(1, 1, end_idx - orig_frames, 1, 1)],
+                dim=2,
+            )
+
+        if self.isolated_key_frame:
+            key_frame = clip_x[:, :, :1, :, :]
+            z_key = self._adaptive_encode(key_frame)
+
+            if clip_x.shape[2] > 1:
+                video_frames = clip_x[:, :, 1:, :, :]
+                z_video = self._adaptive_encode(video_frames)
+                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:
+        input_first_frame = x[:, :, :1, :, :]
+        z_first_frame = self._adaptive_encode(input_first_frame)
+
+        if self.frame_pre_padding == 0:
+            z = torch.cat([z_first_frame, z], dim=2)
+        else:
+            z = torch.cat([z_first_frame, z[:, :, 1:, :, :]], dim=2)
+
+    if self.isolated_last_frame:
+        # The original indexes into the padded video: pad frames replicate
+        # the last real frame, out-of-range indexes produce an empty slice
+        # (a no-op), and negative indexes wrap. Reproduce that mapping
+        # against the unpadded tensor instead of materializing the pad.
+        last_frame_idx = padded_frames - self.frame_drop + offset_frame
+        if last_frame_idx < 0:
+            last_frame_idx += padded_frames
+        if 0 <= last_frame_idx < padded_frames:
+            src_idx = min(last_frame_idx, orig_frames - 1)
+            input_last_frame = x[:, :, src_idx : src_idx + 1, :, :]
+            z_last_frame = self._adaptive_encode(input_last_frame)
+            z = torch.cat([z, z_last_frame], dim=2)
+
+    return z
+
+
+def _decode_temporal_streaming_uint8(self, z, z_head, z_tail, num_chunks, pad_tokens, temporal_cat_dtype):
+    """Drop-in fork of the checkpoint's ``_decode_temporal_streaming``.
+
+    The chunk loop, overlap blending, frame plan, and assertions are
+    identical to the original. ``write_part`` additionally runs the
+    denormalize -> clamp -> quantize chain in place (the same op order the
+    pipeline applies after decode) and accumulates into a uint8 buffer.
+    """
+    total_frames, pad_frames, output_frames = self._decode_temporal_output_frame_plan(
+        z, z_head, z_tail, num_chunks, pad_tokens
+    )
+    if output_frames <= 0:
+        raise ValueError(
+            f"decode_temporal streaming planned non-positive output_frames={output_frames} "
+            f"total_frames={total_frames} pad_frames={pad_frames}"
+        )
+
+    chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
+    split_count = int(self.token_drop > 0) + 1
+    dec = None
+    dec_overlap = None
+    norm = None
+    write_pos = 0
+    logical_frames = 0
+    dropped_frames = 0
+    decoded_count = 0
+
+    def write_part(part):
+        nonlocal dec, dec_overlap, norm, write_pos, logical_frames, dropped_frames
+        part_frames = int(part.shape[2])
+        if part_frames <= 0:
+            return
+        logical_frames += part_frames
+        if dec is None:
+            out_shape = list(part.shape)
+            out_shape[2] = output_frames
+            dec = torch.empty(out_shape, dtype=torch.uint8, device=part.device)
+            transform_rev = self.processor.transform_rev
+            mean = torch.as_tensor(transform_rev.mean, dtype=part.dtype, device=part.device).view(1, 3, 1, 1, 1)
+            std = torch.as_tensor(transform_rev.std, dtype=part.dtype, device=part.device).view(1, 3, 1, 1, 1)
+            norm = (mean, std)
+        # revert_tensor ((x - mean) / std, clamp) followed by the output
+        # quantizer (*255, round, uint8), in place on the part: the blended
+        # raw values are consumed before this point, and every write into
+        # ``dec`` goes through here, so the final buffer matches the legacy
+        # post-decode chain element for element. Denormalization and clamping
+        # run in the decoded dtype exactly like the legacy adapter; the
+        # quantizer then runs in FP32 because the legacy adapter returned
+        # ``frames.float()`` before the pipeline multiplied and rounded —
+        # rounding fp16/bf16 directly flips borderline pixels by one.
+        mean_t, std_t = norm
+        part.sub_(mean_t).div_(std_t).clamp_(0.0, 1.0)
+        part = part.float()
+        part.mul_(255.0).round_()
+
+        remaining = int(dec.shape[2]) - write_pos
+        copy_frames = min(part_frames, max(0, remaining))
+        if copy_frames > 0:
+            dec[:, :, write_pos : write_pos + copy_frames, :, :].copy_(part[:, :, :copy_frames, :, :])
+            write_pos += copy_frames
+        dropped_frames += part_frames - copy_frames
+
+    for i in range(num_chunks):
+        t_start_idx = i * self.tokens_chunk_size
+        t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap
+        clip_z = z[:, :, t_start_idx:t_end_idx, :, :]
+
+        if i == 0 and z_head is not None:
+            clip_z = torch.cat([z_head, clip_z], dim=2)
+
+        if i == num_chunks - 1 and z_tail is not None:
+            clip_z = torch.cat([clip_z, z_tail], dim=2)
+
+        clip_dec = self._adaptive_decode(clip_z)
+        decoded_count += 1
+        if temporal_cat_dtype is not None and clip_dec.dtype != temporal_cat_dtype:
+            clip_dec = clip_dec.to(temporal_cat_dtype)
+        if clip_dec.device != z.device:
+            clip_dec = clip_dec.to(z.device)
+
+        dec_tail = None
+        if i == 0 and z_head is not None:
+            write_part(clip_dec[:, :, self.vae_ratio_t - 1 : self.vae_ratio_t, :, :])
+            clip_dec = clip_dec[:, :, self.vae_ratio_t :, :, :]
+
+        if i == num_chunks - 1 and z_tail is not None:
+            dec_tail = clip_dec[:, :, -1:, :, :]
+            clip_dec = clip_dec[:, :, : -self.vae_ratio_t, :, :]
+
+        for j in range(split_count):
+            f_start_idx = j * chunk_dec
+            f_end_idx = min(f_start_idx + chunk_dec, clip_dec.shape[2])
+            clip_dec_chunk = clip_dec[:, :, f_start_idx:f_end_idx, :, :]
+            clip_dec_chunk = clip_dec_chunk[:, :, self.frame_pre_padding :, :, :]
+
+            if j == 0:
+                if dec_overlap is not None:
+                    clip_dec_chunk = self.blend(dec_overlap, clip_dec_chunk, self.frame_overlap, dim=-3)
+                    dec_overlap = None
+                write_part(clip_dec_chunk)
+            else:
+                # Break the view's reference to the full decoded clip so earlier
+                # temporal chunks can be released before the final output exists.
+                dec_overlap = clip_dec_chunk.contiguous()
+
+        if i == num_chunks - 1:
+            if dec_overlap is not None:
+                write_part(dec_overlap)
+                dec_overlap = None
+            if dec_tail is not None:
+                write_part(dec_tail)
+
+        del clip_dec, clip_z
+
+    if dec is None:
+        raise RuntimeError("decode_temporal streaming produced no output tensor")
+    if logical_frames != total_frames or dropped_frames != pad_frames or write_pos != output_frames:
+        raise RuntimeError(
+            "decode_temporal streaming frame plan mismatch: "
+            f"logical_frames={logical_frames} total_frames={total_frames} "
+            f"dropped_frames={dropped_frames} pad_frames={pad_frames} "
+            f"write_pos={write_pos} output_frames={output_frames}"
+        )
+
+    return dec
+
+
+def _processor_denorm_available(model) -> bool:
+    transform_rev = getattr(getattr(model, "processor", None), "transform_rev", None)
+    mean = getattr(transform_rev, "mean", None)
+    std = getattr(transform_rev, "std", None)
+    return mean is not None and std is not None and len(mean) == 3 and len(std) == 3
+
+
+def install_temporal_stream_patches(model) -> None:
+    """Install the memory-optimized temporal forks on the checkpoint model.
+
+    Each fork is installed only when the checkpoint contract it mirrors is
+    discoverable; otherwise the checkpoint's own method keeps running (the
+    adapter's post-decode revert already covers that path).
+    """
+    if _legacy_temporal_enabled():
+        logger.info("MiniMax-H3 VAE temporal stream patches disabled by %s", _LEGACY_TEMPORAL_ENV)
+        return
+    if callable(getattr(model, "encode_temporal", None)) and hasattr(model, "clip_length"):
+        model.encode_temporal = MethodType(_encode_temporal_tail_pad, model)
+        logger.info("MiniMax-H3 VAE encode_temporal tail-padding patch installed")
+    else:
+        logger.warning("MiniMax-H3 VAE encode_temporal contract not found; keeping checkpoint method")
+    if callable(getattr(model, "_decode_temporal_streaming", None)) and _processor_denorm_available(model):
+        model._decode_temporal_streaming = MethodType(_decode_temporal_streaming_uint8, model)
+        logger.info("MiniMax-H3 VAE streaming uint8 decode patch installed")
+    else:
+        logger.warning("MiniMax-H3 VAE streaming decode contract not found; keeping checkpoint method")
+
+
+__all__ = ["install_temporal_stream_patches"]