[Model][MiniMax-H3] Support released VAE chunks directly

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import hashlib
 import multiprocessing as mp
 import time
 from pathlib import Path
@@ -14,6 +15,11 @@ import torch
 import torch.distributed as dist
 import torch.nn as nn
 
+import vllm_omni.diffusion.models.minimax_h3.temporal_chunks as temporal_chunks_module
+import vllm_omni.diffusion.models.minimax_h3.vae as vae_module
+from vllm_omni.diffusion.models.minimax_h3.temporal_chunks import (
+    MINIMAX_H3_RELEASED_BUNDLE_SHA256,
+)
 from vllm_omni.diffusion.models.minimax_h3.vae import (
     MiniMaxH3ChunkCallbackPeerError,
     MiniMaxH3ChunkedDecodeUnsupportedError,
@@ -29,81 +35,145 @@ class _IdentityProcessor:
         return tensor
 
 
-class _FakeVendorTemporalModel(nn.Module):
-    temporal_chunk_callback_api_version = 1
+class _FakeCompatTemporalModel(nn.Module):
+    use_3d_conv = True
+    clip_length = 17
+    tokens_chunk_size = 5
+    token_overlap = 2
+    token_drop = 3
+    vae_ratio_t = 4
+    frame_pre_padding = 3
+    frame_overlap = 5
+    isolated_first_frame = False
+    isolated_last_frame = False
     vae_ratio = 1
 
     def __init__(
         self,
+        latent_t: int,
         *,
         collective: bool = False,
-        metadata_mode: str = "valid",
+        tile_count: int = 2,
     ) -> None:
         super().__init__()
-        values = torch.linspace(0, 1, 124).view(1, 1, 124, 1, 1)
-        self.expected = values.expand(1, 3, -1, 1, 1).contiguous()
-        self.processor = _IdentityProcessor()
         self.parallel_tiling = collective
         self.collective = collective
-        self.metadata_mode = metadata_mode
+        self.tile_count = tile_count
+        pseudo_tokens = latent_t + self.token_drop
+        remainder = pseudo_tokens % self.tokens_chunk_size
+        self.pad_tokens = 0 if remainder == 0 else self.tokens_chunk_size - remainder
+        self.num_chunks = (pseudo_tokens + self.pad_tokens) // self.tokens_chunk_size - 1
+        ideal_frames = self.num_chunks * 17 + 5
+        output_frames = {
+            33: 111,
+            34: 115,
+            35: 119,
+            36: 120,
+        }.get(latent_t, ideal_frames)
+        values = torch.linspace(0, 1, ideal_frames).view(1, 1, ideal_frames, 1, 1)
+        self.ideal = values.expand(1, 3, -1, 1, 1).contiguous()
+        self.expected = self.ideal[:, :, :output_frames].clone()
+        self.processor = _IdentityProcessor()
         self.decode_base_calls = 0
-        self.vendor_callback_calls = 0
+        self.adaptive_decode_calls = 0
+        self.blend_calls = 0
+        self.eval()
 
-    @staticmethod
-    def split_tiles(size: int, tiled: bool):
+    def split_tiles(self, size: int, tiled: bool):
         del size, tiled
-        return [0, 1], None, None
+        return list(range(self.tile_count)), None, None
 
-    def decode_base(
-        self,
-        latent: torch.Tensor,
-        *,
-        temporal_chunk_callback=None,
-    ) -> torch.Tensor:
+    def decode_base(self, latent: torch.Tensor) -> torch.Tensor:
         del latent
         self.decode_base_calls += 1
-        output = self.expected.clone()
-        if temporal_chunk_callback is None:
-            return output
+        return self.expected.clone()
 
-        frame_start = 0
-        for chunk_index, chunk_size in enumerate([17] * 7 + [5]):
-            if self.collective:
-                collective_probe = torch.tensor([chunk_index], dtype=torch.int32)
-                dist.all_reduce(collective_probe)
-            self.vendor_callback_calls += 1
-            metadata_index = chunk_index
-            metadata_total = 8
-            metadata_start = frame_start
-            is_final = chunk_index == 7
-            if self.metadata_mode == "gap" and chunk_index == 1:
-                metadata_start += 1
-            elif self.metadata_mode == "bad_total" and chunk_index == 1:
-                metadata_total += 1
-            elif self.metadata_mode == "no_final":
-                is_final = False
-            temporal_chunk_callback(
-                self.expected[:, :, frame_start : frame_start + chunk_size].clone(),
-                chunk_index=metadata_index,
-                total_chunks=metadata_total,
-                frame_start=metadata_start,
-                is_final=is_final,
-            )
-            frame_start += chunk_size
-        return output
+    def _adaptive_decode(self, latent: torch.Tensor) -> torch.Tensor:
+        del latent
+        call_index = self.adaptive_decode_calls % self.num_chunks
+        self.adaptive_decode_calls += 1
+        if self.collective and self.parallel_tiling:
+            collective_probe = torch.tensor([call_index], dtype=torch.int32)
+            dist.all_reduce(collective_probe)
+        decoded = torch.full((1, 3, 28, 1, 1), -1.0)
+        body_start = call_index * 17
+        decoded[:, :, 3:20].copy_(self.ideal[:, :, body_start : body_start + 17])
+        if call_index == self.num_chunks - 1:
+            decoded[:, :, 23:28].copy_(self.ideal[:, :, -5:])
+        return decoded
+
+    def _decode_temporal_output_frame_plan(
+        self,
+        latent: torch.Tensor,
+        head: torch.Tensor | None,
+        tail: torch.Tensor | None,
+        num_chunks: int,
+        pad_tokens: int,
+    ) -> tuple[int, int, int]:
+        del latent
+        assert head is None
+        assert tail is None
+        assert num_chunks == self.num_chunks
+        assert pad_tokens == self.pad_tokens
+        total_frames = num_chunks * 17 + 5
+        pad_frames = total_frames - int(self.expected.shape[2])
+        return total_frames, pad_frames, int(self.expected.shape[2])
+
+    def _decode_temporal_pad_frames(
+        self,
+        latent: torch.Tensor,
+        pad_tokens: int,
+    ) -> int:
+        del latent
+        assert pad_tokens == self.pad_tokens
+        return int(self.ideal.shape[2] - self.expected.shape[2])
+
+    def blend(
+        self,
+        previous: torch.Tensor,
+        current: torch.Tensor,
+        extent: int,
+        *,
+        dim: int,
+    ) -> torch.Tensor:
+        del previous
+        assert extent == self.frame_overlap
+        assert dim == -3
+        self.blend_calls += 1
+        return current
 
 
-def _vae(
+def _install_released_source_fingerprint() -> None:
+    temporal_chunks_module._source_bundle_sha256 = lambda model: (
+        "fake/video_vae",
+        MINIMAX_H3_RELEASED_BUNDLE_SHA256,
+    )
+
+
+def _allow_released_source(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        temporal_chunks_module,
+        "_source_bundle_sha256",
+        lambda model: (
+            "fake/video_vae",
+            MINIMAX_H3_RELEASED_BUNDLE_SHA256,
+        ),
+    )
+
+
+def _compat_vae(
+    latent_t: int,
     *,
     parallel_size: int = 1,
     collective: bool = False,
-    metadata_mode: str = "valid",
-) -> tuple[MiniMaxH3VideoVAE, _FakeVendorTemporalModel]:
+    tile_count: int = 2,
+) -> tuple[MiniMaxH3VideoVAE, _FakeCompatTemporalModel]:
     vae = object.__new__(MiniMaxH3VideoVAE)
     nn.Module.__init__(vae)
-    model = _FakeVendorTemporalModel(
+    model = _FakeCompatTemporalModel(
+        latent_t,
         collective=collective,
-        metadata_mode=metadata_mode,
+        tile_count=tile_count,
     )
     vae.model = model
     vae.config_dict = {
@@ -115,22 +185,57 @@ def _vae(
     return vae, model
 
 
-def _latent() -> torch.Tensor:
-    return torch.zeros(1, 3, 37, 1, 1)
+def _latent(latent_t: int = 37) -> torch.Tensor:
+    return torch.zeros(1, 3, latent_t, 1, 1)
 
 
-def test_default_decode_path_does_not_request_vendor_chunks() -> None:
-    vae, model = _vae()
+def test_default_decode_path_does_not_request_chunks() -> None:
+    vae, model = _compat_vae(37)
 
     output = vae.decode_latent(_latent())
 
     torch.testing.assert_close(output, model.expected, rtol=0, atol=0)
     assert model.decode_base_calls == 1
-    assert model.vendor_callback_calls == 0
+    assert model.adaptive_decode_calls == 0
 
 
-def test_callback_chunks_are_owned_and_reconstruct_the_full_decode() -> None:
-    vae, model = _vae()
+def test_default_decode_bypasses_compatibility_source_checks(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    vae, model = _compat_vae(37)
+    monkeypatch.setattr(
+        temporal_chunks_module,
+        "_source_bundle_sha256",
+        lambda model: pytest.fail("default decode fingerprinted remote code"),
+    )
+
+    output = vae.decode_latent(_latent())
+
+    torch.testing.assert_close(output, model.expected, rtol=0, atol=0)
+    assert model.decode_base_calls == 1
+    assert model.adaptive_decode_calls == 0
+
+
+@pytest.mark.parametrize(
+    ("latent_t", "expected_chunk_sizes"),
+    [
+        (37, [17] * 7 + [5]),
+        (62, [17] * 12 + [5]),
+        (72, [17] * 14 + [5]),
+        (107, [17] * 21 + [5]),
+        (36, [17] * 7 + [1]),
+        (35, [17] * 7),
+        (34, [17] * 6 + [13]),
+        (33, [17] * 6 + [9]),
+    ],
+)
+def test_released_compatibility_chunks_reconstruct_full_decode(
+    monkeypatch: pytest.MonkeyPatch,
+    latent_t: int,
+    expected_chunk_sizes: list[int],
+) -> None:
+    _allow_released_source(monkeypatch)
+    vae, model = _compat_vae(latent_t)
     chunks: list[torch.Tensor] = []
     snapshots: list[torch.Tensor] = []
     metadata: list[tuple[int, int, int, bool]] = []
@@ -150,18 +255,35 @@ def test_callback_chunks_are_owned_and_reconstruct_the_full_decode() -> None:
         snapshots.append(frames.clone())
         metadata.append((chunk_index, total_chunks, frame_start, is_final))
 
-    output = vae.decode_latent_with_chunks(_latent(), consume)
+    output = vae.decode_latent_with_chunks(_latent(latent_t), consume)
 
     torch.testing.assert_close(output, model.expected, rtol=0, atol=0)
     torch.testing.assert_close(torch.cat(chunks, dim=2), output, rtol=0, atol=0)
+    assert [int(chunk.shape[2]) for chunk in chunks] == expected_chunk_sizes
+    expected_metadata = []
+    frame_start = 0
+    for chunk_index, chunk_size in enumerate(expected_chunk_sizes):
+        expected_metadata.append(
+            (
+                chunk_index,
+                len(expected_chunk_sizes),
+                frame_start,
+                chunk_index == len(expected_chunk_sizes) - 1,
+            )
+        )
+        frame_start += chunk_size
+    assert metadata == expected_metadata
     for chunk, snapshot in zip(chunks, snapshots, strict=True):
         torch.testing.assert_close(chunk, snapshot, rtol=0, atol=0)
     assert len({chunk.untyped_storage().data_ptr() for chunk in chunks}) == len(chunks)
-    assert metadata == [(index, 8, index * 17, index == 7) for index in range(8)]
+    assert model.adaptive_decode_calls == model.num_chunks
 
 
-def test_callback_mutation_does_not_change_returned_full_output() -> None:
-    vae, model = _vae()
+def test_compat_callback_mutation_does_not_change_returned_full_output(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _allow_released_source(monkeypatch)
+    vae, model = _compat_vae(37)
 
     def mutate(frames: torch.Tensor, **metadata: Any) -> None:
         del metadata
@@ -172,29 +294,77 @@ def test_callback_mutation_does_not_change_returned_full_output() -> None:
     torch.testing.assert_close(output, model.expected, rtol=0, atol=0)
 
 
-def test_callback_failure_is_deferred_until_vendor_decode_finishes() -> None:
-    vae, model = _vae()
-    callback_calls = 0
+def test_compat_source_fingerprint_is_cached_after_success(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    calls = 0
 
-    def fail(*args, **kwargs) -> None:
-        nonlocal callback_calls
-        del args, kwargs
-        callback_calls += 1
-        raise LookupError("sink failed")
+    def fingerprint(model):
+        nonlocal calls
+        del model
+        calls += 1
+        return "fake/video_vae", MINIMAX_H3_RELEASED_BUNDLE_SHA256
 
-    with pytest.raises(LookupError, match="sink failed"):
-        vae.decode_latent_with_chunks(_latent(), fail)
+    monkeypatch.setattr(
+        temporal_chunks_module,
+        "_source_bundle_sha256",
+        fingerprint,
+    )
+    vae, model = _compat_vae(37)
 
-    assert callback_calls == 1
-    assert model.vendor_callback_calls == 8
-    torch.testing.assert_close(vae.decode_latent(_latent()), model.expected, rtol=0, atol=0)
+    for _ in range(2):
+        output = vae.decode_latent_with_chunks(
+            _latent(),
+            lambda *args, **kwargs: None,
+        )
+        torch.testing.assert_close(output, model.expected, rtol=0, atol=0)
+
+    assert calls == 1
 
 
-def test_terminal_callback_failure_preserves_the_original_exception() -> None:
-    vae, model = _vae()
-    seen: list[tuple[int, int, int, bool]] = []
+def test_compat_callback_publishes_after_temporal_blending(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _allow_released_source(monkeypatch)
+    vae, model = _compat_vae(37)
+    chunks: list[torch.Tensor] = []
 
-    def fail_final(
+    def mark_overlap(
+        previous: torch.Tensor,
+        current: torch.Tensor,
+        extent: int,
+        *,
+        dim: int,
+    ) -> torch.Tensor:
+        del previous
+        assert extent == 5
+        assert dim == -3
+        blended = current.clone()
+        blended[:, :, :extent].fill_(0.25)
+        return blended
+
+    monkeypatch.setattr(model, "blend", mark_overlap)
+    output = vae.decode_latent_with_chunks(
+        _latent(),
+        lambda frames, **metadata: chunks.append(frames.clone()),
+    )
+
+    expected_overlap = torch.full_like(chunks[1][:, :, :5], 0.25)
+    torch.testing.assert_close(chunks[1][:, :, :5], expected_overlap, rtol=0, atol=0)
+    torch.testing.assert_close(output[:, :, 17:22], expected_overlap, rtol=0, atol=0)
+    torch.testing.assert_close(torch.cat(chunks, dim=2), output, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("fail_final", [False, True])
+def test_compat_callback_failure_finishes_decode_and_recovers(
+    monkeypatch: pytest.MonkeyPatch,
+    fail_final: bool,
+) -> None:
+    _allow_released_source(monkeypatch)
+    vae, model = _compat_vae(37)
+    seen: list[tuple[int, bool]] = []
+
+    def fail(
         frames: torch.Tensor,
         *,
         chunk_index: int,
@@ -202,48 +372,165 @@ def test_terminal_callback_failure_preserves_the_original_exception() -> None:
         frame_start: int,
         is_final: bool,
     ) -> None:
-        del frames
-        seen.append((chunk_index, total_chunks, frame_start, is_final))
-        if is_final:
-            raise LookupError("final sink failed")
+        del frames, total_chunks, frame_start
+        seen.append((chunk_index, is_final))
+        if is_final == fail_final:
+            raise LookupError("compat sink failed")
 
-    with pytest.raises(LookupError, match="final sink failed"):
-        vae.decode_latent_with_chunks(_latent(), fail_final)
+    with pytest.raises(LookupError, match="compat sink failed"):
+        vae.decode_latent_with_chunks(_latent(), fail)
 
-    assert seen == [(index, 8, index * 17, index == 7) for index in range(8)]
-    assert model.vendor_callback_calls == 8
+    assert model.adaptive_decode_calls == model.num_chunks
+    if fail_final:
+        assert seen == [(index, index == 7) for index in range(8)]
+    else:
+        assert seen == [(0, False)]
+
+    recovered_chunks: list[torch.Tensor] = []
+    recovered = vae.decode_latent_with_chunks(
+        _latent(),
+        lambda frames, **metadata: recovered_chunks.append(frames),
+    )
+    torch.testing.assert_close(recovered, model.expected, rtol=0, atol=0)
+    torch.testing.assert_close(
+        torch.cat(recovered_chunks, dim=2),
+        model.expected,
+        rtol=0,
+        atol=0,
+    )
+    assert model.adaptive_decode_calls == model.num_chunks * 2
+
+
+@pytest.mark.parametrize("drift", ["changed", "missing"])
+def test_compat_source_drift_fails_closed_before_decode(
+    monkeypatch: pytest.MonkeyPatch,
+    drift: str,
+) -> None:
+    _allow_released_source(monkeypatch)
+    if drift == "changed":
+        monkeypatch.setattr(
+            temporal_chunks_module,
+            "_source_bundle_sha256",
+            lambda model: ("fake/video_vae", "0" * 64),
+        )
+    else:
+        monkeypatch.setattr(
+            temporal_chunks_module,
+            "_source_bundle_sha256",
+            lambda model: (None, None),
+        )
+    vae, model = _compat_vae(37)
+
+    with pytest.raises(
+        MiniMaxH3ChunkedDecodeUnsupportedError,
+        match="42ed227e",
+    ):
+        vae.decode_latent_with_chunks(_latent(), lambda *args, **kwargs: None)
+
+    assert model.adaptive_decode_calls == 0
+    torch.testing.assert_close(vae.decode_latent(_latent()), model.expected, rtol=0, atol=0)
+
+
+def test_uninspectable_compat_source_maps_to_unsupported(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(
+        temporal_chunks_module.inspect,
+        "getsourcefile",
+        lambda object_type: (_ for _ in ()).throw(TypeError("no source")),
+    )
+    vae, model = _compat_vae(37)
+
+    with pytest.raises(MiniMaxH3ChunkedDecodeUnsupportedError, match="bundle=None"):
+        vae.decode_latent_with_chunks(_latent(), lambda *args, **kwargs: None)
+
+    assert model.adaptive_decode_calls == 0
+
+
+def test_source_bundle_manifest_is_deterministic_and_excludes_init(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    package = tmp_path / "video_vae"
+    package.mkdir()
+    (package / "z.py").write_text("z = 1\n", encoding="utf-8")
+    (package / "klvae.py").write_text("k = 2\n", encoding="utf-8")
+    (package / "a.py").write_text("a = 3\n", encoding="utf-8")
+    (package / "__init__.py").write_text("ignored = 1\n", encoding="utf-8")
+    monkeypatch.setattr(
+        temporal_chunks_module.inspect,
+        "getsourcefile",
+        lambda object_type: str(package / "klvae.py"),
+    )
+
+    expected = hashlib.sha256()
+    for filename in ("a.py", "klvae.py", "z.py"):
+        expected.update(filename.encode())
+        expected.update(b"\0")
+        expected.update((package / filename).read_bytes())
+        expected.update(b"\0")
+    _, first = temporal_chunks_module._source_bundle_sha256(nn.Linear(1, 1))
+    assert first == expected.hexdigest()
+
+    (package / "__init__.py").write_text("ignored = 2\n", encoding="utf-8")
+    _, after_init = temporal_chunks_module._source_bundle_sha256(nn.Linear(1, 1))
+    assert after_init == first
+    (package / "a.py").write_text("a = 4\n", encoding="utf-8")
+    _, after_change = temporal_chunks_module._source_bundle_sha256(nn.Linear(1, 1))
+    assert after_change != first
 
 
 @pytest.mark.parametrize(
-    ("metadata_mode", "message"),
+    ("attribute", "value"),
     [
-        ("gap", "discontinuous chunk metadata"),
-        ("bad_total", "discontinuous chunk metadata"),
-        ("no_final", "discontinuous chunk metadata"),
+        ("tokens_chunk_size", 6),
+        ("token_overlap", 1),
+        ("isolated_first_frame", True),
+        ("isolated_last_frame", True),
     ],
 )
-def test_adapter_rejects_a_broken_vendor_callback_contract(
-    metadata_mode: str,
-    message: str,
+def test_compat_config_drift_fails_closed(
+    monkeypatch: pytest.MonkeyPatch,
+    attribute: str,
+    value: object,
 ) -> None:
-    vae, model = _vae(metadata_mode=metadata_mode)
+    _allow_released_source(monkeypatch)
+    vae, model = _compat_vae(37)
+    setattr(model, attribute, value)
 
-    with pytest.raises(RuntimeError, match=message):
-        vae.decode_latent_with_chunks(_latent(), lambda *args, **kwargs: None)
-
-    assert model.vendor_callback_calls == 8
-
-
-def test_callback_request_fails_closed_without_the_vendor_api() -> None:
-    vae, model = _vae()
-    model.temporal_chunk_callback_api_version = -1
-
-    with pytest.raises(MiniMaxH3ChunkedDecodeUnsupportedError, match="official MiniMax"):
+    with pytest.raises(
+        MiniMaxH3ChunkedDecodeUnsupportedError,
+        match="released temporal configuration",
+    ):
         vae.decode_latent_with_chunks(_latent(), lambda *args, **kwargs: None)
 
+    assert model.adaptive_decode_calls == 0
     torch.testing.assert_close(vae.decode_latent(_latent()), model.expected, rtol=0, atol=0)
 
 
+@pytest.mark.parametrize("invalid", ["rank", "empty_time", "training", "non_3d"])
+def test_compat_structural_invalids_fail_before_decode(
+    monkeypatch: pytest.MonkeyPatch,
+    invalid: str,
+) -> None:
+    _allow_released_source(monkeypatch)
+    vae, model = _compat_vae(37)
+    latent = _latent()
+    if invalid == "rank":
+        latent = latent[:, :, 0]
+    elif invalid == "empty_time":
+        latent = latent[:, :, :0]
+    elif invalid == "training":
+        model.train()
+    else:
+        model.use_3d_conv = False
+
+    with pytest.raises(MiniMaxH3ChunkedDecodeUnsupportedError):
+        vae.decode_latent_with_chunks(latent, lambda *args, **kwargs: None)
+
+    assert model.adaptive_decode_calls == 0
+
+
 def _distributed_worker(
     rank: int,
     init_file: str,
@@ -256,11 +543,105 @@ def _distributed_worker(
         world_size=2,
     )
     try:
-        vae, model = _vae(parallel_size=2, collective=True)
-        vae._native_parallel_state = MethodType(
-            lambda self: {"sp_process_group": dist.group.WORLD},
-            vae,
+        vae_module.get_data_parallel_world_size = lambda: 1
+        vae_module.model_parallel_is_initialized = lambda: True
+
+        def attach_group(vae: MiniMaxH3VideoVAE) -> None:
+            vae._native_parallel_state = MethodType(
+                lambda self: {"sp_process_group": dist.group.WORLD},
+                vae,
+            )
+
+        _install_released_source_fingerprint()
+        source_vae, source_model = _compat_vae(
+            37,
+            parallel_size=2,
+            collective=True,
         )
+        attach_group(source_vae)
+        if rank == 1:
+            temporal_chunks_module._source_bundle_sha256 = lambda model: (
+                "fake/video_vae",
+                "0" * 64,
+            )
+        try:
+            source_vae.decode_latent_with_chunks(
+                _latent(),
+                (lambda *args, **kwargs: None) if rank == 0 else None,
+            )
+        except BaseException as exc:  # noqa: BLE001
+            source_mismatch_type = type(exc).__name__
+        else:
+            source_mismatch_type = "none"
+        source_decode_calls = source_model.adaptive_decode_calls
+
+        _install_released_source_fingerprint()
+        planner_vae, planner_model = _compat_vae(
+            37,
+            parallel_size=2,
+            collective=True,
+        )
+        attach_group(planner_vae)
+        if rank == 1:
+
+            def fail_plan(self, *args, **kwargs):
+                del self, args, kwargs
+                raise RuntimeError("rank-local planner failed")
+
+            setattr(
+                planner_model,
+                "_decode_temporal_output_frame_plan",
+                MethodType(fail_plan, planner_model),
+            )
+        try:
+            planner_vae.decode_latent_with_chunks(
+                _latent(),
+                (lambda *args, **kwargs: None) if rank == 0 else None,
+            )
+        except BaseException as exc:  # noqa: BLE001
+            planner_failure_type = type(exc).__name__
+        else:
+            planner_failure_type = "none"
+        planner_decode_calls = planner_model.adaptive_decode_calls
+
+        dtype_vae, dtype_model = _compat_vae(
+            37,
+            parallel_size=2,
+            collective=True,
+        )
+        attach_group(dtype_vae)
+        temporal_chunks_module.resolve_minimax_h3_temporal_cat_dtype = (
+            (lambda model: torch.float16) if rank == 0 else (lambda model: None)
+        )
+        try:
+            dtype_vae.decode_latent_with_chunks(
+                _latent(),
+                (lambda *args, **kwargs: None) if rank == 0 else None,
+            )
+        except BaseException as exc:  # noqa: BLE001
+            dtype_mismatch_type = type(exc).__name__
+        else:
+            dtype_mismatch_type = "none"
+        dtype_decode_calls = dtype_model.adaptive_decode_calls
+        temporal_chunks_module.resolve_minimax_h3_temporal_cat_dtype = lambda model: None
+
+        vae_module.get_data_parallel_world_size = lambda: 2
+        dp_vae, dp_model = _compat_vae(37, parallel_size=2, collective=True)
+        attach_group(dp_vae)
+        try:
+            dp_vae.decode_latent_with_chunks(
+                _latent(),
+                (lambda *args, **kwargs: None) if rank == 0 else None,
+            )
+        except BaseException as exc:  # noqa: BLE001
+            dp_vae_type = type(exc).__name__
+        else:
+            dp_vae_type = "none"
+        dp_decode_calls = dp_model.adaptive_decode_calls
+        vae_module.get_data_parallel_world_size = lambda: 1
+
+        vae, model = _compat_vae(37, parallel_size=2, collective=True)
+        attach_group(vae)
 
         try:
             vae.decode_latent_with_chunks(
@@ -298,7 +679,26 @@ def _distributed_worker(
             succeed if rank == 0 else None,
         )
 
-        local_vae, local_model = _vae(parallel_size=1)
+        fallback_vae, fallback_model = _compat_vae(
+            37,
+            parallel_size=2,
+            collective=True,
+            tile_count=1,
+        )
+        attach_group(fallback_vae)
+        fallback_calls = 0
+
+        def consume_fallback(*args, **kwargs) -> None:
+            nonlocal fallback_calls
+            del args, kwargs
+            fallback_calls += 1
+
+        fallback_output = fallback_vae.decode_latent_with_chunks(
+            _latent(),
+            consume_fallback if rank == 0 else None,
+        )
+
+        local_vae, local_model = _compat_vae(37, parallel_size=1)
         local_calls = 0
 
         def consume_local(*args, **kwargs) -> None:
@@ -316,11 +716,26 @@ def _distributed_worker(
         result_queue.put(
             {
                 "rank": rank,
+                "source_mismatch_type": source_mismatch_type,
+                "source_decode_calls": source_decode_calls,
+                "planner_failure_type": planner_failure_type,
+                "planner_decode_calls": planner_decode_calls,
+                "dtype_mismatch_type": dtype_mismatch_type,
+                "dtype_decode_calls": dtype_decode_calls,
+                "dp_vae_type": dp_vae_type,
+                "dp_decode_calls": dp_decode_calls,
                 "invalid_owner_type": invalid_owner_type,
                 "failure_type": failure_type,
-                "vendor_calls": model.vendor_callback_calls,
+                "adaptive_calls": model.adaptive_decode_calls,
                 "success_calls": success_calls,
                 "output_matches": torch.equal(output, model.expected),
+                "fallback_calls": fallback_calls,
+                "fallback_decode_calls": fallback_model.adaptive_decode_calls,
+                "fallback_output_matches": torch.equal(
+                    fallback_output,
+                    fallback_model.expected,
+                ),
+                "fallback_tiling_restored": fallback_model.parallel_tiling,
                 "local_calls": local_calls,
                 "local_output_matches": torch.equal(local_output, local_model.expected),
                 "recovery": int(recovery.item()),
@@ -367,13 +782,37 @@ def test_two_rank_callback_failure_is_symmetric_and_group_recovers(
         "LookupError",
         MiniMaxH3ChunkCallbackPeerError.__name__,
     ]
+    assert [result["source_mismatch_type"] for result in results] == [
+        MiniMaxH3ChunkedDecodeUnsupportedError.__name__,
+        MiniMaxH3ChunkedDecodeUnsupportedError.__name__,
+    ]
+    assert [result["source_decode_calls"] for result in results] == [0, 0]
+    assert [result["planner_failure_type"] for result in results] == [
+        MiniMaxH3ChunkedDecodeUnsupportedError.__name__,
+        "RuntimeError",
+    ]
+    assert [result["planner_decode_calls"] for result in results] == [0, 0]
+    assert [result["dtype_mismatch_type"] for result in results] == [
+        MiniMaxH3ChunkedDecodeUnsupportedError.__name__,
+        MiniMaxH3ChunkedDecodeUnsupportedError.__name__,
+    ]
+    assert [result["dtype_decode_calls"] for result in results] == [0, 0]
+    assert [result["dp_vae_type"] for result in results] == [
+        MiniMaxH3ChunkedDecodeUnsupportedError.__name__,
+        MiniMaxH3ChunkedDecodeUnsupportedError.__name__,
+    ]
+    assert [result["dp_decode_calls"] for result in results] == [0, 0]
     assert [result["invalid_owner_type"] for result in results] == [
         "ValueError",
         "ValueError",
     ]
-    assert [result["vendor_calls"] for result in results] == [16, 16]
+    assert [result["adaptive_calls"] for result in results] == [14, 14]
     assert [result["success_calls"] for result in results] == [8, 0]
     assert all(result["output_matches"] for result in results)
+    assert [result["fallback_calls"] for result in results] == [8, 0]
+    assert [result["fallback_decode_calls"] for result in results] == [7, 7]
+    assert all(result["fallback_output_matches"] for result in results)
+    assert all(result["fallback_tiling_restored"] for result in results)
     # With VAE parallelism disabled, the caller still designates one request owner.
     assert [result["local_calls"] for result in results] == [8, 0]
     assert all(result["local_output_matches"] for result in results)

vllm_omni/diffusion/models/minimax_h3/temporal_chunks.py

@@ -0,0 +1,408 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Source-gated temporal chunk compatibility for the released MiniMax H3 VAE."""
+
+from __future__ import annotations
+
+import hashlib
+import importlib
+import inspect
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Protocol
+
+import torch
+import torch.nn as nn
+
+MINIMAX_H3_RELEASED_REVISION = "42ed227ee7df40d41602854ae760620d6eb651fe"
+# Compatibility pins detect upstream code drift; they are not a trust or
+# security boundary for remote code execution.
+MINIMAX_H3_RELEASED_BUNDLE_SHA256 = "30e64afbaa940696bf8bfe2c86003a4b1666c84b22fa3aa12403a3e6a03f705d"
+
+_RELEASED_TEMPORAL_CONFIG = {
+    "clip_length": 17,
+    "tokens_chunk_size": 5,
+    "token_overlap": 2,
+    "token_drop": 3,
+    "vae_ratio_t": 4,
+    "frame_pre_padding": 3,
+    "frame_overlap": 5,
+    "isolated_first_frame": False,
+    "isolated_last_frame": False,
+}
+
+
+class MiniMaxH3TemporalChunkCompatibilityError(RuntimeError):
+    """Raised when remote H3 code cannot use the local compatibility path."""
+
+
+class _RawTemporalChunkCallback(Protocol):
+    def __call__(
+        self,
+        frames: torch.Tensor,
+        *,
+        chunk_index: int,
+        total_chunks: int,
+        frame_start: int,
+        is_final: bool,
+    ) -> None: ...
+
+
+@dataclass(slots=True)
+class MiniMaxH3TemporalDecodePlan:
+    latent: torch.Tensor
+    num_chunks: int
+    total_frames: int
+    pad_frames: int
+    output_frames: int
+    chunk_frames: tuple[int, ...]
+    temporal_dtype: torch.dtype | None
+
+
+class _RawChunkAssembly:
+    def __init__(
+        self,
+        *,
+        callback: _RawTemporalChunkCallback,
+        plan: MiniMaxH3TemporalDecodePlan,
+    ) -> None:
+        self.callback = callback
+        self.plan = plan
+        self.frames: torch.Tensor | None = None
+        self.write_pos = 0
+        self.logical_frames = 0
+        self.dropped_frames = 0
+        self.chunk_index = 0
+        self.pending_final: tuple[torch.Tensor, int, int] | None = None
+        self.contract_error: RuntimeError | None = None
+
+    def write(self, part: torch.Tensor) -> None:
+        part_frames = int(part.shape[2])
+        if part_frames <= 0:
+            return
+        self.logical_frames += part_frames
+        remaining = self.plan.output_frames - self.write_pos
+        copy_frames = min(part_frames, max(0, remaining))
+        if copy_frames <= 0:
+            self.dropped_frames += part_frames
+            return
+        if self.chunk_index >= len(self.plan.chunk_frames):
+            self.contract_error = RuntimeError("MiniMax-H3 compatibility decode emitted an extra chunk")
+        else:
+            expected_frames = self.plan.chunk_frames[self.chunk_index]
+            if copy_frames != expected_frames and self.contract_error is None:
+                self.contract_error = RuntimeError(
+                    "MiniMax-H3 compatibility chunk size mismatch: "
+                    f"chunk={self.chunk_index}, frames={copy_frames}/{expected_frames}"
+                )
+
+        emitted = part[:, :, :copy_frames]
+        if self.frames is None:
+            output_shape = list(emitted.shape)
+            output_shape[2] = self.plan.output_frames
+            self.frames = torch.empty(
+                output_shape,
+                dtype=emitted.dtype,
+                device=emitted.device,
+            )
+        frame_start = self.write_pos
+        frame_stop = frame_start + copy_frames
+        self.frames[:, :, frame_start:frame_stop].copy_(emitted)
+        if frame_stop == self.plan.output_frames:
+            self.pending_final = (emitted, self.chunk_index, frame_start)
+        elif self.contract_error is None:
+            self.callback(
+                emitted,
+                chunk_index=self.chunk_index,
+                total_chunks=len(self.plan.chunk_frames),
+                frame_start=frame_start,
+                is_final=False,
+            )
+        self.chunk_index += 1
+        self.write_pos = frame_stop
+        self.dropped_frames += part_frames - copy_frames
+
+    def finish(self) -> torch.Tensor:
+        if self.contract_error is not None:
+            raise self.contract_error
+        if self.frames is None or self.pending_final is None:
+            raise RuntimeError("MiniMax-H3 compatibility decode produced no terminal chunk")
+        if (
+            self.logical_frames != self.plan.total_frames
+            or self.dropped_frames != self.plan.pad_frames
+            or self.write_pos != self.plan.output_frames
+            or self.chunk_index != len(self.plan.chunk_frames)
+        ):
+            raise RuntimeError(
+                "MiniMax-H3 compatibility frame plan mismatch: "
+                f"logical={self.logical_frames}/{self.plan.total_frames}, "
+                f"dropped={self.dropped_frames}/{self.plan.pad_frames}, "
+                f"written={self.write_pos}/{self.plan.output_frames}, "
+                f"chunks={self.chunk_index}/{len(self.plan.chunk_frames)}"
+            )
+
+        final_frames, final_index, final_start = self.pending_final
+        self.callback(
+            final_frames,
+            chunk_index=final_index,
+            total_chunks=len(self.plan.chunk_frames),
+            frame_start=final_start,
+            is_final=True,
+        )
+        return self.frames
+
+
+def _source_bundle_sha256(model: nn.Module) -> tuple[str | None, str | None]:
+    try:
+        source = inspect.getsourcefile(type(model))
+    except TypeError:
+        return None, None
+    if source is None:
+        return None, None
+    package_dir = Path(source).parent
+    files = sorted(path for path in package_dir.glob("*.py") if path.name != "__init__.py")
+    digest = hashlib.sha256()
+    try:
+        for path in files:
+            digest.update(path.name.encode())
+            digest.update(b"\0")
+            digest.update(path.read_bytes())
+            digest.update(b"\0")
+    except OSError:
+        return str(package_dir), None
+    return str(package_dir), digest.hexdigest()
+
+
+def _validate_released_sources(model: nn.Module) -> None:
+    source_dir, actual = _source_bundle_sha256(model)
+    if actual != MINIMAX_H3_RELEASED_BUNDLE_SHA256:
+        raise MiniMaxH3TemporalChunkCompatibilityError(
+            "The vLLM-owned MiniMax-H3 chunk adapter supports only remote code "
+            f"from MiniMaxAI/MiniMax-H3@{MINIMAX_H3_RELEASED_REVISION}: "
+            f"bundle={actual!r} ({source_dir}). "
+            "The normal complete-output decode remains available."
+        )
+
+
+def _validate_released_config(model: nn.Module) -> None:
+    actual = {name: getattr(model, name, None) for name in _RELEASED_TEMPORAL_CONFIG}
+    if actual != _RELEASED_TEMPORAL_CONFIG:
+        raise MiniMaxH3TemporalChunkCompatibilityError(
+            "The vLLM-owned MiniMax-H3 chunk adapter requires the released temporal "
+            f"configuration {_RELEASED_TEMPORAL_CONFIG}, found {actual}. "
+            "The normal complete-output decode remains available."
+        )
+
+
+def _validate_compat_contract(model: nn.Module, latent: torch.Tensor) -> None:
+    if latent.ndim != 5 or int(latent.shape[2]) <= 0:
+        raise MiniMaxH3TemporalChunkCompatibilityError(
+            f"MiniMax-H3 video latent must have positive rank-5 time, got {tuple(latent.shape)}"
+        )
+    if not bool(getattr(model, "use_3d_conv", False)):
+        raise MiniMaxH3TemporalChunkCompatibilityError("MiniMax-H3 temporal chunk decode requires the 3D decoder")
+    if bool(getattr(model, "training", False)):
+        raise MiniMaxH3TemporalChunkCompatibilityError("MiniMax-H3 temporal chunk decode is inference-only")
+    required_methods = (
+        "_adaptive_decode",
+        "_decode_temporal_output_frame_plan",
+        "_decode_temporal_pad_frames",
+        "blend",
+    )
+    missing = [name for name in required_methods if not callable(getattr(model, name, None))]
+    if missing:
+        raise MiniMaxH3TemporalChunkCompatibilityError(
+            "Loaded MiniMax-H3 VAE lacks required released-code methods: " + ", ".join(missing)
+        )
+
+
+def _truncate_chunk_frames(
+    logical_chunk_frames: list[int],
+    output_frames: int,
+) -> tuple[int, ...]:
+    remaining = output_frames
+    result: list[int] = []
+    for logical_frames in logical_chunk_frames:
+        emitted_frames = min(logical_frames, max(0, remaining))
+        if emitted_frames > 0:
+            result.append(emitted_frames)
+            remaining -= emitted_frames
+    if remaining != 0:
+        raise RuntimeError(f"MiniMax-H3 compatibility chunk plan cannot cover output: remaining={remaining}")
+    return tuple(result)
+
+
+def _build_decode_plan(
+    model: nn.Module,
+    latent: torch.Tensor,
+    temporal_dtype: torch.dtype | None,
+) -> MiniMaxH3TemporalDecodePlan:
+    pseudo_total_tokens = int(latent.shape[2]) + int(model.token_drop)  # type: ignore[attr-defined]
+    tokens_chunk_size = int(model.tokens_chunk_size)  # type: ignore[attr-defined]
+    remainder = pseudo_total_tokens % tokens_chunk_size
+    pad_tokens = 0 if remainder == 0 else tokens_chunk_size - remainder
+    pseudo_total_tokens += pad_tokens
+    num_chunks = pseudo_total_tokens // tokens_chunk_size - int(model.token_drop > 0)  # type: ignore[attr-defined]
+    if num_chunks <= 0:
+        raise MiniMaxH3TemporalChunkCompatibilityError(
+            f"MiniMax-H3 temporal chunk plan is empty for latent T={latent.shape[2]}"
+        )
+
+    temporal_latent = latent
+    if pad_tokens:
+        pad = temporal_latent[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)
+        temporal_latent = torch.cat((temporal_latent, pad), dim=2)
+    total_frames, pad_frames, output_frames = model._decode_temporal_output_frame_plan(  # type: ignore[attr-defined]
+        temporal_latent,
+        None,
+        None,
+        num_chunks,
+        pad_tokens,
+    )
+    if output_frames <= 0:
+        raise MiniMaxH3TemporalChunkCompatibilityError(
+            f"MiniMax-H3 temporal chunk plan has {output_frames} output frames"
+        )
+
+    chunk_decoded_frames = tokens_chunk_size * int(model.vae_ratio_t)  # type: ignore[attr-defined]
+    split_count = int(model.token_drop > 0) + 1  # type: ignore[attr-defined]
+    logical_chunk_frames: list[int] = []
+    final_overlap_frames = 0
+    for temporal_index in range(num_chunks):
+        token_start = temporal_index * tokens_chunk_size
+        token_stop = token_start + tokens_chunk_size + int(model.token_overlap)  # type: ignore[attr-defined]
+        clip_tokens = max(
+            0,
+            min(token_stop, int(temporal_latent.shape[2])) - min(token_start, int(temporal_latent.shape[2])),
+        )
+        clip_frames = clip_tokens * int(model.vae_ratio_t)  # type: ignore[attr-defined]
+        for split_index in range(split_count):
+            frame_start = split_index * chunk_decoded_frames
+            frame_stop = min(frame_start + chunk_decoded_frames, clip_frames)
+            part_frames = max(
+                0,
+                frame_stop - frame_start - int(model.frame_pre_padding),  # type: ignore[attr-defined]
+            )
+            if split_index == 0:
+                if part_frames > 0:
+                    logical_chunk_frames.append(part_frames)
+            else:
+                final_overlap_frames = part_frames
+    if final_overlap_frames > 0:
+        logical_chunk_frames.append(final_overlap_frames)
+    if sum(logical_chunk_frames) != int(total_frames):
+        raise RuntimeError(
+            "MiniMax-H3 compatibility planner disagrees with released remote code: "
+            f"frames={sum(logical_chunk_frames)}/{total_frames}"
+        )
+    chunk_frames = _truncate_chunk_frames(logical_chunk_frames, int(output_frames))
+    return MiniMaxH3TemporalDecodePlan(
+        latent=temporal_latent,
+        num_chunks=num_chunks,
+        total_frames=int(total_frames),
+        pad_frames=int(pad_frames),
+        output_frames=int(output_frames),
+        chunk_frames=chunk_frames,
+        temporal_dtype=temporal_dtype,
+    )
+
+
+def resolve_minimax_h3_temporal_cat_dtype(model: nn.Module) -> torch.dtype | None:
+    module = importlib.import_module(type(model).__module__)
+    resolver = getattr(module, "_resolve_temporal_cat_dtype", None)
+    if resolver is None:
+        return None
+    if not callable(resolver):
+        raise MiniMaxH3TemporalChunkCompatibilityError(
+            "Loaded MiniMax-H3 VAE has a non-callable temporal dtype resolver"
+        )
+    try:
+        dtype = resolver()
+    except ValueError as exc:
+        raise MiniMaxH3TemporalChunkCompatibilityError(str(exc)) from None
+    if dtype is not None and not isinstance(dtype, torch.dtype):
+        raise MiniMaxH3TemporalChunkCompatibilityError(
+            f"Loaded MiniMax-H3 VAE returned invalid temporal dtype {dtype!r}"
+        )
+    return dtype
+
+
+def decode_minimax_h3_temporal_chunks_compat(
+    model: nn.Module,
+    plan: MiniMaxH3TemporalDecodePlan,
+    callback: _RawTemporalChunkCallback,
+) -> torch.Tensor:
+    """Mirror the released Apache-2.0 H3 temporal assembler with callbacks.
+
+    Keep this implementation and the source/config fingerprints above in sync
+    with ``MiniMaxAI/MiniMax-H3@42ed227e``. Unknown remote code fails closed
+    before this function is called.
+    """
+
+    chunk_decoded_frames = int(model.tokens_chunk_size) * int(model.vae_ratio_t)  # type: ignore[attr-defined]
+    split_count = int(model.token_drop > 0) + 1  # type: ignore[attr-defined]
+    assembly = _RawChunkAssembly(callback=callback, plan=plan)
+    overlap: torch.Tensor | None = None
+
+    for temporal_index in range(plan.num_chunks):
+        token_start = temporal_index * int(model.tokens_chunk_size)  # type: ignore[attr-defined]
+        token_stop = token_start + int(model.tokens_chunk_size) + int(model.token_overlap)  # type: ignore[attr-defined]
+        chunk_latent = plan.latent[:, :, token_start:token_stop]
+        decoded = model._adaptive_decode(chunk_latent)  # type: ignore[attr-defined]
+        if plan.temporal_dtype is not None and decoded.dtype != plan.temporal_dtype:
+            decoded = decoded.to(plan.temporal_dtype)
+        if decoded.device != plan.latent.device:
+            decoded = decoded.to(plan.latent.device)
+
+        for split_index in range(split_count):
+            frame_start = split_index * chunk_decoded_frames
+            frame_stop = min(frame_start + chunk_decoded_frames, int(decoded.shape[2]))
+            part = decoded[:, :, frame_start:frame_stop]
+            part = part[:, :, int(model.frame_pre_padding) :]  # type: ignore[attr-defined]
+            if split_index == 0:
+                if overlap is not None:
+                    part = model.blend(  # type: ignore[attr-defined]
+                        overlap,
+                        part,
+                        int(model.frame_overlap),  # type: ignore[attr-defined]
+                        dim=-3,
+                    )
+                    overlap = None
+                assembly.write(part)
+            else:
+                overlap = part.contiguous()
+            del part
+
+        if temporal_index == plan.num_chunks - 1 and overlap is not None:
+            assembly.write(overlap)
+            overlap = None
+        del decoded, chunk_latent
+
+    return assembly.finish()
+
+
+def prepare_minimax_h3_temporal_chunks_compat(
+    model: nn.Module,
+    latent: torch.Tensor,
+    *,
+    validate_sources: bool = True,
+) -> MiniMaxH3TemporalDecodePlan:
+    """Build a validated plan without entering decoder collectives."""
+
+    if validate_sources:
+        _validate_released_sources(model)
+    _validate_released_config(model)
+    _validate_compat_contract(model, latent)
+    temporal_dtype = resolve_minimax_h3_temporal_cat_dtype(model)
+    return _build_decode_plan(model, latent, temporal_dtype)
+
+
+__all__ = [
+    "MINIMAX_H3_RELEASED_REVISION",
+    "MINIMAX_H3_RELEASED_BUNDLE_SHA256",
+    "MiniMaxH3TemporalDecodePlan",
+    "MiniMaxH3TemporalChunkCompatibilityError",
+    "decode_minimax_h3_temporal_chunks_compat",
+    "prepare_minimax_h3_temporal_chunks_compat",
+    "resolve_minimax_h3_temporal_cat_dtype",
+]

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -21,7 +21,11 @@ from vllm.logger import init_logger
 from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
     DistributedVaeMixin,
 )
-from vllm_omni.diffusion.distributed.parallel_state import get_world_group
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_data_parallel_world_size,
+    get_world_group,
+    model_parallel_is_initialized,
+)
 from vllm_omni.diffusion.offloader.module_residency import (
     BoundedAllocatorCache,
     PinnedModuleStager,
@@ -29,6 +33,12 @@ from vllm_omni.diffusion.offloader.module_residency import (
 
 from .ops import install_h3_vae_optimizations
 from .packed_tokens import minimax_h3_patchify_video_latent
+from .temporal_chunks import (
+    MiniMaxH3TemporalChunkCompatibilityError,
+    MiniMaxH3TemporalDecodePlan,
+    decode_minimax_h3_temporal_chunks_compat,
+    prepare_minimax_h3_temporal_chunks_compat,
+)
 
 MINIMAX_H3_KEYFRAME_ENCODE_SEED = 42
 MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000
@@ -60,7 +70,7 @@ class MiniMaxH3VideoChunkCallback(Protocol):
 
 
 class MiniMaxH3ChunkedDecodeUnsupportedError(RuntimeError):
-    """Raised when the loaded MiniMax H3 VAE lacks the vendor chunk API."""
+    """Raised when loaded MiniMax H3 code cannot honor the chunk contract."""
 
 
 class MiniMaxH3ChunkCallbackPeerError(RuntimeError):
@@ -338,6 +348,10 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
     def _chunk_process_group(self) -> dist.ProcessGroup | None:
         if not self.is_distributed_enabled():
             return None
+        if model_parallel_is_initialized() and get_data_parallel_world_size() > 1:
+            raise MiniMaxH3ChunkedDecodeUnsupportedError(
+                "MiniMax-H3 temporal chunks do not yet support combining data parallelism with VAE patch parallelism"
+            )
         group = self._native_parallel_state().get("sp_process_group")
         if group is None or dist.get_world_size(group) != self.parallel_size:
             raise RuntimeError("MiniMax-H3 VAE chunk decode has an invalid spatial-parallel group")
@@ -395,14 +409,48 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
             raise callback_error
         raise MiniMaxH3ChunkCallbackPeerError("MiniMax-H3 video chunk callback failed on this VAE group's output rank")
 
-    def _validate_vendor_chunk_api(self) -> None:
-        version = getattr(self.model, "temporal_chunk_callback_api_version", None)
-        if version != 1:
-            raise MiniMaxH3ChunkedDecodeUnsupportedError(
-                "Loaded MiniMax-H3 VAE does not provide temporal_chunk_callback API v1 "
-                f"(found {version!r}); install a model snapshot containing the official "
-                "MiniMax temporal chunk callback"
+    @staticmethod
+    def _synchronize_chunk_preflight_error(
+        preflight_error: Exception | None,
+        temporal_dtype: torch.dtype | None,
+        *,
+        group: dist.ProcessGroup | None,
+        device: torch.device,
+    ) -> None:
+        if group is None:
+            if preflight_error is not None:
+                raise preflight_error
+            return
+        failed = torch.tensor(
+            [int(preflight_error is not None)],
+            dtype=torch.int32,
+            device=device,
+        )
+        dist.all_reduce(failed, op=dist.ReduceOp.MAX, group=group)
+        if int(failed.item()) == 0:
+            dtype_codes = {
+                None: 0,
+                torch.float16: 1,
+                torch.bfloat16: 2,
+                torch.float32: 3,
+            }
+            code = torch.tensor(
+                [dtype_codes.get(temporal_dtype, -1)],
+                dtype=torch.int32,
+                device=device,
             )
+            lowest = code.clone()
+            highest = code.clone()
+            dist.all_reduce(lowest, op=dist.ReduceOp.MIN, group=group)
+            dist.all_reduce(highest, op=dist.ReduceOp.MAX, group=group)
+            if int(lowest.item()) == int(highest.item()) and int(lowest.item()) >= 0:
+                return
+            raise MiniMaxH3ChunkedDecodeUnsupportedError(
+                "MiniMax-H3 VAE ranks selected incompatible temporal concat dtypes"
+            )
+        if preflight_error is not None:
+            raise preflight_error
+        raise MiniMaxH3ChunkedDecodeUnsupportedError("A peer VAE rank rejected the MiniMax-H3 temporal chunk preflight")
 
     @torch.inference_mode()
     def encode_image(self, image: Image.Image) -> torch.Tensor:
@@ -506,11 +554,11 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         latent: torch.Tensor,
         chunk_callback: MiniMaxH3VideoChunkCallback | None,
     ) -> torch.Tensor:
-        """Decode a video while publishing vendor-committed temporal chunks.
+        """Decode a video while publishing committed temporal chunks.
 
         Every rank collaborating in the configured H3 VAE tile group must call
         this method, with ``chunk_callback`` supplied only on group rank 0. A
-        callback failure is deferred until the vendor decoder finishes its
+        callback failure is deferred until the compatibility decoder finishes its
         remaining collectives, then propagated to every group rank. With
         ``parallel_size=1`` the runtime designates each request's output owner
         by supplying a callback there and ``None`` on replica peers.
@@ -522,8 +570,41 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         requested-size crop as the complete H3 pipeline output.
         """
 
-        self._validate_vendor_chunk_api()
         group = self._chunk_process_group()
+        plan: MiniMaxH3TemporalDecodePlan | None = None
+        preflight_error: Exception | None = None
+        if latent.ndim != 5:
+            preflight_error = MiniMaxH3ChunkedDecodeUnsupportedError(
+                f"MiniMax-H3 video latent must be rank 5, got {tuple(latent.shape)}"
+            )
+        else:
+            try:
+                denormalized = self._denormalize_latent(latent)
+                plan = prepare_minimax_h3_temporal_chunks_compat(
+                    self.model,
+                    denormalized,
+                    validate_sources=not getattr(
+                        self,
+                        "_minimax_h3_chunk_sources_validated",
+                        False,
+                    ),
+                )
+                self._minimax_h3_chunk_sources_validated = True
+                del denormalized
+            except MiniMaxH3TemporalChunkCompatibilityError as exc:
+                preflight_error = MiniMaxH3ChunkedDecodeUnsupportedError(str(exc))
+            except Exception as exc:  # noqa: BLE001
+                # Every VAE rank must reach the preflight reduction even when
+                # one node's private planner or environment fails locally.
+                exc.__traceback__ = None
+                preflight_error = exc
+        self._synchronize_chunk_preflight_error(
+            preflight_error,
+            plan.temporal_dtype if plan is not None else None,
+            group=group,
+            device=latent.device,
+        )
+        assert plan is not None
         is_output_rank = self._is_chunk_output_rank(
             chunk_callback,
             group=group,
@@ -560,7 +641,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
                     or is_final != (chunk_index + 1 == total_chunks)
                 ):
                     raise RuntimeError(
-                        "MiniMax-H3 vendor callback emitted discontinuous chunk metadata: "
+                        "MiniMax-H3 temporal backend emitted discontinuous chunk metadata: "
                         f"index={chunk_index}/{next_chunk_index}, "
                         f"total={total_chunks}/{expected_total_chunks}, "
                         f"frame_start={frame_start}/{next_frame_start}, "
@@ -570,7 +651,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
                 frames = frames.clone(memory_format=torch.contiguous_format)
                 frame_count = int(frames.shape[2])
                 if frame_count <= 0:
-                    raise RuntimeError("MiniMax-H3 vendor callback emitted an empty chunk")
+                    raise RuntimeError("MiniMax-H3 temporal backend emitted an empty chunk")
                 next_chunk_index += 1
                 next_frame_start += frame_count
                 saw_final = is_final
@@ -589,9 +670,10 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
                 callback_error = exc
 
         with self._decode_tiling_context(latent):
-            decoded = self.model.decode_base(
-                self._denormalize_latent(latent),
-                temporal_chunk_callback=publish,
+            decoded = decode_minimax_h3_temporal_chunks_compat(
+                self.model,
+                plan,
+                publish,
             )
         frames = self._normalize_decoded_frames(decoded)
         if (
@@ -600,7 +682,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
             and (not saw_final or next_chunk_index != expected_total_chunks or next_frame_start != int(frames.shape[2]))
         ):
             callback_error = RuntimeError(
-                "MiniMax-H3 vendor callback did not reconstruct the full decode: "
+                "MiniMax-H3 temporal backend did not reconstruct the full decode: "
                 f"final={saw_final}, chunks={next_chunk_index}/{expected_total_chunks}, "
                 f"frames={next_frame_start}/{frames.shape[2]}"
             )