[Model][MiniMax-H3] Delegate temporal VAE chunks

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

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py

@@ -0,0 +1,380 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import multiprocessing as mp
+import time
+from pathlib import Path
+from types import MethodType
+from typing import Any
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+
+from vllm_omni.diffusion.models.minimax_h3.vae import (
+    MiniMaxH3ChunkCallbackPeerError,
+    MiniMaxH3ChunkedDecodeUnsupportedError,
+    MiniMaxH3VideoVAE,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _IdentityProcessor:
+    @staticmethod
+    def revert_tensor(tensor: torch.Tensor) -> torch.Tensor:
+        return tensor
+
+
+class _FakeVendorTemporalModel(nn.Module):
+    temporal_chunk_callback_api_version = 1
+    vae_ratio = 1
+
+    def __init__(
+        self,
+        *,
+        collective: bool = False,
+        metadata_mode: str = "valid",
+    ) -> 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.decode_base_calls = 0
+        self.vendor_callback_calls = 0
+
+    @staticmethod
+    def split_tiles(size: int, tiled: bool):
+        del size, tiled
+        return [0, 1], None, None
+
+    def decode_base(
+        self,
+        latent: torch.Tensor,
+        *,
+        temporal_chunk_callback=None,
+    ) -> torch.Tensor:
+        del latent
+        self.decode_base_calls += 1
+        output = self.expected.clone()
+        if temporal_chunk_callback is None:
+            return output
+
+        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 _vae(
+    *,
+    parallel_size: int = 1,
+    collective: bool = False,
+    metadata_mode: str = "valid",
+) -> tuple[MiniMaxH3VideoVAE, _FakeVendorTemporalModel]:
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    nn.Module.__init__(vae)
+    model = _FakeVendorTemporalModel(
+        collective=collective,
+        metadata_mode=metadata_mode,
+    )
+    vae.model = model
+    vae.config_dict = {
+        "latent_channels": 3,
+        "latents_mean": [0.0, 0.0, 0.0],
+        "latents_std": [1.0, 1.0, 1.0],
+    }
+    vae.parallel_size = parallel_size
+    return vae, model
+
+
+def _latent() -> torch.Tensor:
+    return torch.zeros(1, 3, 37, 1, 1)
+
+
+def test_default_decode_path_does_not_request_vendor_chunks() -> None:
+    vae, model = _vae()
+
+    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
+
+
+def test_callback_chunks_are_owned_and_reconstruct_the_full_decode() -> None:
+    vae, model = _vae()
+    chunks: list[torch.Tensor] = []
+    snapshots: list[torch.Tensor] = []
+    metadata: list[tuple[int, int, int, bool]] = []
+
+    def consume(
+        frames: torch.Tensor,
+        *,
+        chunk_index: int,
+        total_chunks: int,
+        frame_start: int,
+        is_final: bool,
+    ) -> None:
+        assert frames.dtype is torch.float32
+        assert frames.is_contiguous()
+        assert torch.all((0 <= frames) & (frames <= 1))
+        chunks.append(frames)
+        snapshots.append(frames.clone())
+        metadata.append((chunk_index, total_chunks, frame_start, is_final))
+
+    output = vae.decode_latent_with_chunks(_latent(), 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)
+    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)]
+
+
+def test_callback_mutation_does_not_change_returned_full_output() -> None:
+    vae, model = _vae()
+
+    def mutate(frames: torch.Tensor, **metadata: Any) -> None:
+        del metadata
+        frames.fill_(1234)
+
+    output = vae.decode_latent_with_chunks(_latent(), mutate)
+
+    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 fail(*args, **kwargs) -> None:
+        nonlocal callback_calls
+        del args, kwargs
+        callback_calls += 1
+        raise LookupError("sink failed")
+
+    with pytest.raises(LookupError, match="sink failed"):
+        vae.decode_latent_with_chunks(_latent(), fail)
+
+    assert callback_calls == 1
+    assert model.vendor_callback_calls == 8
+    torch.testing.assert_close(vae.decode_latent(_latent()), model.expected, rtol=0, atol=0)
+
+
+def test_terminal_callback_failure_preserves_the_original_exception() -> None:
+    vae, model = _vae()
+    seen: list[tuple[int, int, int, bool]] = []
+
+    def fail_final(
+        frames: torch.Tensor,
+        *,
+        chunk_index: int,
+        total_chunks: int,
+        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")
+
+    with pytest.raises(LookupError, match="final sink failed"):
+        vae.decode_latent_with_chunks(_latent(), fail_final)
+
+    assert seen == [(index, 8, index * 17, index == 7) for index in range(8)]
+    assert model.vendor_callback_calls == 8
+
+
+@pytest.mark.parametrize(
+    ("metadata_mode", "message"),
+    [
+        ("gap", "discontinuous chunk metadata"),
+        ("bad_total", "discontinuous chunk metadata"),
+        ("no_final", "discontinuous chunk metadata"),
+    ],
+)
+def test_adapter_rejects_a_broken_vendor_callback_contract(
+    metadata_mode: str,
+    message: str,
+) -> None:
+    vae, model = _vae(metadata_mode=metadata_mode)
+
+    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"):
+        vae.decode_latent_with_chunks(_latent(), lambda *args, **kwargs: None)
+
+    torch.testing.assert_close(vae.decode_latent(_latent()), model.expected, rtol=0, atol=0)
+
+
+def _distributed_worker(
+    rank: int,
+    init_file: str,
+    result_queue: Any,
+) -> None:
+    dist.init_process_group(
+        "gloo",
+        init_method=f"file://{init_file}",
+        rank=rank,
+        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,
+        )
+
+        try:
+            vae.decode_latent_with_chunks(
+                _latent(),
+                lambda *args, **kwargs: None,
+            )
+        except BaseException as exc:  # noqa: BLE001
+            invalid_owner_type = type(exc).__name__
+        else:
+            invalid_owner_type = "none"
+
+        def fail(*args, **kwargs) -> None:
+            del args, kwargs
+            raise LookupError("distributed sink failed")
+
+        try:
+            vae.decode_latent_with_chunks(
+                _latent(),
+                fail if rank == 0 else None,
+            )
+        except BaseException as exc:  # noqa: BLE001
+            failure_type = type(exc).__name__
+        else:
+            failure_type = "none"
+
+        success_calls = 0
+
+        def succeed(*args, **kwargs) -> None:
+            nonlocal success_calls
+            del args, kwargs
+            success_calls += 1
+
+        output = vae.decode_latent_with_chunks(
+            _latent(),
+            succeed if rank == 0 else None,
+        )
+
+        local_vae, local_model = _vae(parallel_size=1)
+        local_calls = 0
+
+        def consume_local(*args, **kwargs) -> None:
+            nonlocal local_calls
+            del args, kwargs
+            local_calls += 1
+
+        local_output = local_vae.decode_latent_with_chunks(
+            _latent(),
+            consume_local if rank == 0 else None,
+        )
+
+        recovery = torch.tensor([rank + 1], dtype=torch.int32)
+        dist.all_reduce(recovery)
+        result_queue.put(
+            {
+                "rank": rank,
+                "invalid_owner_type": invalid_owner_type,
+                "failure_type": failure_type,
+                "vendor_calls": model.vendor_callback_calls,
+                "success_calls": success_calls,
+                "output_matches": torch.equal(output, model.expected),
+                "local_calls": local_calls,
+                "local_output_matches": torch.equal(local_output, local_model.expected),
+                "recovery": int(recovery.item()),
+            }
+        )
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.parallel
+def test_two_rank_callback_failure_is_symmetric_and_group_recovers(
+    tmp_path: Path,
+) -> None:
+    context = mp.get_context("spawn")
+    result_queue = context.Queue()
+    init_file = str(tmp_path / "gloo-init")
+    processes = [
+        context.Process(
+            target=_distributed_worker,
+            args=(rank, init_file, result_queue),
+        )
+        for rank in range(2)
+    ]
+    for process in processes:
+        process.start()
+    deadline = time.monotonic() + 90
+    for process in processes:
+        process.join(timeout=max(0, deadline - time.monotonic()))
+    for process in processes:
+        if process.is_alive():
+            process.join(timeout=1)
+    hung = [process for process in processes if process.is_alive()]
+    for process in hung:
+        process.terminate()
+        process.join(timeout=5)
+    assert not hung, "two-rank chunk callback test deadlocked"
+    assert [process.exitcode for process in processes] == [0, 0]
+
+    results = sorted(
+        [result_queue.get(timeout=5) for _ in range(2)],
+        key=lambda result: result["rank"],
+    )
+    assert [result["failure_type"] for result in results] == [
+        "LookupError",
+        MiniMaxH3ChunkCallbackPeerError.__name__,
+    ]
+    assert [result["invalid_owner_type"] for result in results] == [
+        "ValueError",
+        "ValueError",
+    ]
+    assert [result["vendor_calls"] for result in results] == [16, 16]
+    assert [result["success_calls"] for result in results] == [8, 0]
+    assert all(result["output_matches"] 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)
+    assert [result["recovery"] for result in results] == [3, 3]

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -9,7 +9,7 @@ import json
 from collections.abc import Iterator
 from contextlib import AbstractContextManager, contextmanager, nullcontext
 from pathlib import Path
-from typing import Any
+from typing import Any, Protocol
 
 import torch
 import torch.distributed as dist
@@ -38,6 +38,35 @@ MINIMAX_H3_AUDIO_CHANNELS = 2
 logger = init_logger(__name__)
 
 
+class MiniMaxH3VideoChunkCallback(Protocol):
+    """Consume one ordered, temporally committed H3 video chunk.
+
+    ``frames`` is an owned, contiguous float32 RGB tensor in ``BCTHW`` layout
+    and the ``[0, 1]`` value range. The callback runs synchronously on the
+    producer's current device stream. Before reading it on another stream, the
+    consumer must establish ordering from the producer stream (for example via
+    an event) and retain the allocation for that stream's lifetime.
+    """
+
+    def __call__(
+        self,
+        frames: torch.Tensor,
+        *,
+        chunk_index: int,
+        total_chunks: int,
+        frame_start: int,
+        is_final: bool,
+    ) -> None: ...
+
+
+class MiniMaxH3ChunkedDecodeUnsupportedError(RuntimeError):
+    """Raised when the loaded MiniMax H3 VAE lacks the vendor chunk API."""
+
+
+class MiniMaxH3ChunkCallbackPeerError(RuntimeError):
+    """Raised on peer VAE ranks when the output-rank callback fails."""
+
+
 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"))
@@ -265,6 +294,116 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
     def is_distributed_enabled(self) -> bool:
         return self.parallel_size > 1 and dist.is_initialized()
 
+    def _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor:
+        channels = int(self.config_dict["latent_channels"])
+        mean = torch.tensor(
+            self.config_dict["latents_mean"],
+            device=latent.device,
+            dtype=latent.dtype,
+        ).view(1, channels, 1, 1, 1)
+        std = torch.tensor(
+            self.config_dict["latents_std"],
+            device=latent.device,
+            dtype=latent.dtype,
+        ).view(1, channels, 1, 1, 1)
+        return latent * std + mean
+
+    def _decode_tiling_context(self, latent: torch.Tensor) -> AbstractContextManager:
+        # The checkpoint hands rank r the tiles ``range(r, num_tiles, sp_size)``
+        # and then rejects an empty share inside the gather. A rank with no
+        # tiles raises and leaves the collective while the others block in it
+        # forever, so too few tiles hangs the whole stage rather than failing
+        # it. Tile count depends only on the latent shape, so every rank takes
+        # this branch together.
+        num_tiles = self._decoder_tile_count(latent)
+        if self.parallel_size > 1 and num_tiles < self.parallel_size:
+            logger.warning_once(
+                "MiniMax-H3 VAE decode splits into %d tile(s) but the tile group has "
+                "%d ranks; decoding rank-locally for this shape instead, which is "
+                "slower but avoids ranks without tiles hanging the collective.",
+                num_tiles,
+                self.parallel_size,
+            )
+            return self._rank_local_tiling()
+        return nullcontext()
+
+    def _normalize_decoded_frames(self, decoded: torch.Tensor) -> torch.Tensor:
+        frames = self.model.processor.revert_tensor(decoded)
+        if frames.ndim == 4:
+            frames = frames.unsqueeze(0).transpose(1, 2)
+        if frames.ndim != 5:
+            raise ValueError(f"unexpected decoded video shape {tuple(frames.shape)}")
+        return frames.float()
+
+    def _chunk_process_group(self) -> dist.ProcessGroup | None:
+        if not self.is_distributed_enabled():
+            return None
+        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")
+        return group
+
+    @staticmethod
+    def _is_chunk_output_rank(
+        chunk_callback: MiniMaxH3VideoChunkCallback | None,
+        *,
+        group: dist.ProcessGroup | None,
+        device: torch.device,
+    ) -> bool:
+        if group is None:
+            return chunk_callback is not None
+
+        group_rank = dist.get_rank(group)
+        callback_owners = torch.tensor(
+            [int(chunk_callback is not None)],
+            dtype=torch.int32,
+            device=device,
+        )
+        dist.all_reduce(callback_owners, group=group)
+        valid_owner = int(callback_owners.item()) == 1 and ((group_rank == 0) == (chunk_callback is not None))
+        if not valid_owner:
+            raise ValueError("MiniMax-H3 chunk decode requires a callback on exactly VAE group rank 0")
+        return group_rank == 0
+
+    @staticmethod
+    def _synchronize_chunk_callback_error(
+        callback_error: BaseException | None,
+        *,
+        group: dist.ProcessGroup | None,
+        device: torch.device,
+    ) -> None:
+        if group is None:
+            if callback_error is not None:
+                raise callback_error
+            return
+
+        is_output_rank = dist.get_rank(group) == 0
+        failed = torch.tensor(
+            [int(is_output_rank and callback_error is not None)],
+            dtype=torch.int32,
+            device=device,
+        )
+        dist.broadcast(
+            failed,
+            src=dist.get_global_rank(group, 0),
+            group=group,
+        )
+        if int(failed.item()) == 0:
+            return
+        if is_output_rank:
+            assert callback_error is not None
+            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"
+            )
+
     @torch.inference_mode()
     def encode_image(self, image: Image.Image) -> torch.Tensor:
         previous_parallel = self.model.parallel_tiling
@@ -357,44 +496,120 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
 
     @torch.inference_mode()
     def decode_latent(self, latent: torch.Tensor) -> torch.Tensor:
-        channels = int(self.config_dict["latent_channels"])
-        mean = torch.tensor(
-            self.config_dict["latents_mean"],
-            device=latent.device,
-            dtype=latent.dtype,
-        ).view(1, channels, 1, 1, 1)
-        std = torch.tensor(
-            self.config_dict["latents_std"],
-            device=latent.device,
-            dtype=latent.dtype,
-        ).view(1, channels, 1, 1, 1)
-        # The checkpoint hands rank r the tiles ``range(r, num_tiles, sp_size)``
-        # and then rejects an empty share inside the gather. A rank with no
-        # tiles raises and leaves the collective while the others block in it
-        # forever, so too few tiles hangs the whole stage rather than failing
-        # it. Tile count depends only on the latent shape, so every rank takes
-        # this branch together.
-        num_tiles = self._decoder_tile_count(latent)
-        if self.parallel_size > 1 and num_tiles < self.parallel_size:
-            logger.warning_once(
-                "MiniMax-H3 VAE decode splits into %d tile(s) but the tile group has "
-                "%d ranks; decoding rank-locally for this shape instead, which is "
-                "slower but avoids ranks without tiles hanging the collective.",
-                num_tiles,
-                self.parallel_size,
-            )
-            tiling_context: AbstractContextManager = self._rank_local_tiling()
-        else:
-            tiling_context = nullcontext()
+        with self._decode_tiling_context(latent):
+            decoded = self.model.decode_base(self._denormalize_latent(latent))
+        return self._normalize_decoded_frames(decoded)
 
-        with tiling_context:
-            decoded = self.model.decode_base(latent * std + mean)
-        frames = self.model.processor.revert_tensor(decoded)
-        if frames.ndim == 4:
-            frames = frames.unsqueeze(0).transpose(1, 2)
-        if frames.ndim != 5:
-            raise ValueError(f"unexpected decoded video shape {tuple(frames.shape)}")
-        return frames.float()
+    @torch.inference_mode()
+    def decode_latent_with_chunks(
+        self,
+        latent: torch.Tensor,
+        chunk_callback: MiniMaxH3VideoChunkCallback | None,
+    ) -> torch.Tensor:
+        """Decode a video while publishing vendor-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
+        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.
+
+        This is a model-local producer seam before the typed media boundary in
+        RFC #6541 / PR #6615. Representation conversion, transport,
+        presentation policy, and encoding remain runtime-owned. Chunks retain
+        the VAE canvas dimensions; a runtime consumer must apply the same
+        requested-size crop as the complete H3 pipeline output.
+        """
+
+        self._validate_vendor_chunk_api()
+        group = self._chunk_process_group()
+        is_output_rank = self._is_chunk_output_rank(
+            chunk_callback,
+            group=group,
+            device=latent.device,
+        )
+        callback_error: BaseException | None = None
+        next_chunk_index = 0
+        expected_total_chunks: int | None = None
+        next_frame_start = 0
+        saw_final = False
+
+        def publish(
+            decoded_chunk: torch.Tensor,
+            *,
+            chunk_index: int,
+            total_chunks: int,
+            frame_start: int,
+            is_final: bool,
+        ) -> None:
+            nonlocal callback_error, expected_total_chunks
+            nonlocal next_chunk_index, next_frame_start, saw_final
+            if not is_output_rank or callback_error is not None:
+                return
+            try:
+                if expected_total_chunks is None:
+                    expected_total_chunks = total_chunks
+                if (
+                    saw_final
+                    or total_chunks <= 0
+                    or total_chunks != expected_total_chunks
+                    or chunk_index != next_chunk_index
+                    or chunk_index >= total_chunks
+                    or frame_start != next_frame_start
+                    or is_final != (chunk_index + 1 == total_chunks)
+                ):
+                    raise RuntimeError(
+                        "MiniMax-H3 vendor callback 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}, "
+                        f"final={is_final}, after_final={saw_final}"
+                    )
+                frames = self._normalize_decoded_frames(decoded_chunk)
+                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")
+                next_chunk_index += 1
+                next_frame_start += frame_count
+                saw_final = is_final
+                assert chunk_callback is not None
+                chunk_callback(
+                    frames,
+                    chunk_index=chunk_index,
+                    total_chunks=total_chunks,
+                    frame_start=frame_start,
+                    is_final=is_final,
+                )
+            except BaseException as exc:  # noqa: BLE001
+                # Keep a failed GPU chunk and its callback locals out of the
+                # traceback while peer ranks finish the remaining collectives.
+                exc.__traceback__ = None
+                callback_error = exc
+
+        with self._decode_tiling_context(latent):
+            decoded = self.model.decode_base(
+                self._denormalize_latent(latent),
+                temporal_chunk_callback=publish,
+            )
+        frames = self._normalize_decoded_frames(decoded)
+        if (
+            is_output_rank
+            and callback_error is None
+            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: "
+                f"final={saw_final}, chunks={next_chunk_index}/{expected_total_chunks}, "
+                f"frames={next_frame_start}/{frames.shape[2]}"
+            )
+        self._synchronize_chunk_callback_error(
+            callback_error,
+            group=group,
+            device=frames.device,
+        )
+        return frames
 
 
 class MiniMaxH3AudioVAE(nn.Module):
@@ -516,4 +731,10 @@ class MiniMaxH3AudioVAE(nn.Module):
         return waveform.permute(1, 0, 2).contiguous().float()
 
 
-__all__ = ["MiniMaxH3AudioVAE", "MiniMaxH3VideoVAE"]
+__all__ = [
+    "MiniMaxH3AudioVAE",
+    "MiniMaxH3ChunkCallbackPeerError",
+    "MiniMaxH3ChunkedDecodeUnsupportedError",
+    "MiniMaxH3VideoChunkCallback",
+    "MiniMaxH3VideoVAE",
+]