[2/N] Add a minimal temporal chunk callback for MiniMax-H3 (#7017)

Signed-off-by: specture724 <specture724@gmail.com>
Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py

@@ -0,0 +1,357 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import gc
+import tempfile
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from vllm_omni.diffusion.models.minimax_h3.chunked_decode import decode_h3_chunks
+from vllm_omni.diffusion.models.minimax_h3.temporal_chunks import decode_temporal_chunks
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture(scope="module")
+def cpu_process_group():
+    if dist.is_initialized():
+        yield dist.group.WORLD
+        return
+
+    with tempfile.NamedTemporaryFile(prefix="h3_chunk_dist_") as rendezvous:
+        init_method = f"file://{rendezvous.name}"
+    dist.init_process_group("gloo", rank=0, world_size=1, init_method=init_method)
+    try:
+        yield dist.group.WORLD
+    finally:
+        dist.destroy_process_group()
+        gc.collect()
+
+
+class _FakeTemporalModel:
+    use_3d_conv = True
+    token_drop = 3
+    tokens_chunk_size = 5
+    token_overlap = 2
+    vae_ratio_t = 4
+    frame_pre_padding = 3
+    frame_overlap = 5
+    isolated_first_frame = False
+    isolated_last_frame = False
+
+    def _decode_temporal_output_frame_plan(self, z, z_head, z_tail, num_chunks, pad_tokens):
+        del z, z_head, z_tail, num_chunks, pad_tokens
+        return 35, 0, 35
+
+    def _adaptive_decode(self, clip):
+        value = float(clip[:, :, 0].mean())
+        return torch.full((1, 1, 24, 2, 2), value)
+
+    @staticmethod
+    def blend(overlap, part, frame_overlap, dim):
+        del overlap, frame_overlap, dim
+        return part
+
+
+class _FakeHost:
+    def __init__(self):
+        self.model = _FakeTemporalModel()
+
+    @staticmethod
+    def _denormalize_latent(latent):
+        return latent
+
+    @staticmethod
+    def _normalize_decoded_frames(frames):
+        return frames.float()
+
+
+def test_temporal_chunks_emit_ordered_frames_and_collect_when_unconsumed():
+    model = _FakeTemporalModel()
+    latent = torch.arange(8, dtype=torch.float32).view(1, 1, 8, 1, 1)
+    chunks = []
+    marker = decode_temporal_chunks(model, latent, chunks.append)
+
+    assert marker.shape == (0,)
+    assert [chunk.shape[2] for chunk in chunks] == [17, 17, 1]
+    assert torch.equal(torch.cat(chunks, dim=2), decode_temporal_chunks(model, latent, None))
+
+
+def test_h3_callback_failure_is_deferred_until_temporal_decode_finishes():
+    host = _FakeHost()
+    latent = torch.zeros(1, 1, 8, 1, 1)
+    seen = []
+
+    def fail_once(frames):
+        seen.append(frames.shape[2])
+        raise RuntimeError("sink failed")
+
+    with pytest.raises(RuntimeError, match="sink failed"):
+        decode_h3_chunks(host, latent, fail_once, group=None)
+    assert seen == [17]
+
+
+def test_h3_without_callback_is_a_plain_full_decode_on_a_vae_group(cpu_process_group):
+    host = _FakeHost()
+    latent = torch.arange(8, dtype=torch.float32).view(1, 1, 8, 1, 1)
+
+    full = decode_h3_chunks(host, latent, None, group=cpu_process_group)
+
+    chunks = []
+    decode_h3_chunks(host, latent, chunks.append, group=cpu_process_group)
+    assert torch.equal(torch.cat(chunks, dim=2), full)
+
+
+class _SimulatedGroup:
+    """Stand in for a VAE group so one process can check every rank's verdict.
+
+    ``all_reduce`` is replaced with the sum the real collective would produce
+    for a group of ``world_size`` ranks, ``supplied`` of which handed in a
+    callback. That sum is what makes a rank-local verdict observable from a
+    single process.
+    """
+
+    def __init__(self, world_size: int, supplied: int):
+        self.world_size = world_size
+        self.supplied = supplied
+
+
+def _patch_simulated_ranks(monkeypatch, simulated: _SimulatedGroup, rank: int) -> None:
+    from vllm_omni.diffusion.models.minimax_h3 import chunked_decode as mod
+
+    def fake_all_reduce(tensor, group=None):
+        del group  # the simulated census below stands in for the real reduction
+        tensor.copy_(torch.tensor([simulated.supplied], dtype=tensor.dtype))
+
+    monkeypatch.setattr(mod.dist, "get_rank", lambda _group: rank)
+    monkeypatch.setattr(mod.dist, "get_world_size", lambda _group: simulated.world_size)
+    monkeypatch.setattr(mod.dist, "all_reduce", fake_all_reduce)
+    monkeypatch.setattr(mod.dist, "broadcast", lambda tensor, src=0, group=None: None)
+    monkeypatch.setattr(mod.dist, "get_global_rank", lambda _group, _rank: 0)
+
+
+@pytest.mark.parametrize("rank", [0, 1, 2])
+def test_h3_every_rank_rejects_a_partially_supplied_callback(monkeypatch, rank):
+    """Ranks that disagree must all refuse, not let one decode alone.
+
+    A rank-local verdict let a peer fall through into the decoder collectives
+    while its peers raised, hanging the stage instead of failing it.
+    """
+    host = _FakeHost()
+    latent = torch.zeros(1, 1, 8, 1, 1)
+    # Two ranks of three handed in a callback.
+    group = _SimulatedGroup(world_size=3, supplied=2)
+    _patch_simulated_ranks(monkeypatch, group, rank)
+
+    with pytest.raises(ValueError, match="every rank"):
+        decode_h3_chunks(host, latent, lambda _frames: None, group=object())
+
+
+@pytest.mark.parametrize("rank", [0, 1, 2])
+def test_h3_only_rank_zero_publishes_when_every_rank_supplies(monkeypatch, rank):
+    """Ownership is positional: all ranks stream, rank 0 publishes."""
+    host = _FakeHost()
+    latent = torch.zeros(1, 1, 8, 1, 1)
+    group = _SimulatedGroup(world_size=3, supplied=3)
+    _patch_simulated_ranks(monkeypatch, group, rank)
+
+    seen: list[int] = []
+    decode_h3_chunks(host, latent, lambda f: seen.append(int(f.shape[2])), group=object())
+
+    # Peers run the loop to stay in the collectives but publish nothing.
+    assert (seen != []) is (rank == 0)
+
+
+def _stub_video_vae(*, parallel_size: int, tile_count: int):
+    """A MiniMaxH3VideoVAE with only the pieces the decode entry points read."""
+    from contextlib import contextmanager
+
+    from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE
+
+    vae = object.__new__(MiniMaxH3VideoVAE)
+    vae.parallel_size = parallel_size
+    vae.model = SimpleNamespace(_adaptive_decode=lambda clip: clip)
+    vae.entered_rank_local = False
+
+    @contextmanager
+    def rank_local():
+        vae.entered_rank_local = True
+        yield
+
+    vae._decoder_tile_count = lambda latent: tile_count
+    vae._rank_local_tiling = rank_local
+    vae.is_distributed_enabled = lambda: False
+    return vae
+
+
+def test_chunked_decode_falls_back_when_tiles_are_fewer_than_ranks(monkeypatch):
+    """Without the fallback, tileless ranks raise while their peers block."""
+    from vllm_omni.diffusion.models.minimax_h3 import vae as vae_mod
+
+    vae = _stub_video_vae(parallel_size=4, tile_count=1)
+    monkeypatch.setattr(vae_mod, "decode_h3_chunks", lambda *args, **kwargs: torch.zeros(1))
+
+    vae.decode_with_chunks(torch.zeros(1, 1, 8, 1, 1), on_chunk=lambda _frames: None)
+
+    assert vae.entered_rank_local, "chunked decode must share decode_latent's too-few-tiles fallback"
+
+
+def test_chunked_decode_keeps_shared_tiling_when_every_rank_has_a_tile(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import vae as vae_mod
+
+    vae = _stub_video_vae(parallel_size=4, tile_count=8)
+    monkeypatch.setattr(vae_mod, "decode_h3_chunks", lambda *args, **kwargs: torch.zeros(1))
+
+    vae.decode_with_chunks(torch.zeros(1, 1, 8, 1, 1), on_chunk=lambda _frames: None)
+
+    assert not vae.entered_rank_local
+
+
+def test_h3_vae_declares_the_chunked_decode_capability():
+    """A consumer must be able to detect the capability and the pixel range."""
+    from vllm_omni.diffusion.models.interface import supports_chunked_vae_decode
+    from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE
+
+    vae = object.__new__(MiniMaxH3VideoVAE)
+
+    assert supports_chunked_vae_decode(vae)
+    # H3 reverts through the checkpoint's processor, which lands in [0, 1];
+    # a Wan VAE publishes [-1, 1], so the range cannot be assumed.
+    assert vae.chunk_value_range == (0.0, 1.0)
+
+
+def _native_decode_temporal_streaming(model, z, z_head, z_tail, num_chunks, pad_tokens):
+    """The released ``_decode_temporal_streaming`` loop, as the parity oracle.
+
+    Transcribed from the checkpoint's ``AutoencoderKLLegacy`` so the fork can be
+    compared against the control flow it mirrors, including the isolated
+    head/tail extraction.
+    """
+    total_frames, pad_frames, output_frames = model._decode_temporal_output_frame_plan(
+        z, z_head, z_tail, num_chunks, pad_tokens
+    )
+    chunk_dec = model.tokens_chunk_size * model.vae_ratio_t
+    split_count = int(model.token_drop > 0) + 1
+    dec = None
+    dec_overlap = None
+    write_pos = 0
+
+    def write_part(part):
+        nonlocal dec, write_pos
+        part_frames = int(part.shape[2])
+        if part_frames <= 0:
+            return
+        if dec is None:
+            out_shape = list(part.shape)
+            out_shape[2] = output_frames
+            dec = torch.empty(out_shape, dtype=part.dtype, device=part.device)
+        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
+
+    for i in range(num_chunks):
+        t_start = i * model.tokens_chunk_size
+        clip_z = z[:, :, t_start : t_start + model.tokens_chunk_size + model.token_overlap]
+        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 = model._adaptive_decode(clip_z)
+
+        dec_tail = None
+        if i == 0 and z_head is not None:
+            write_part(clip_dec[:, :, model.vae_ratio_t - 1 : model.vae_ratio_t])
+            clip_dec = clip_dec[:, :, model.vae_ratio_t :]
+        if i == num_chunks - 1 and z_tail is not None:
+            dec_tail = clip_dec[:, :, -1:]
+            clip_dec = clip_dec[:, :, : -model.vae_ratio_t]
+
+        for j in range(split_count):
+            f_start = j * chunk_dec
+            f_end = min(f_start + chunk_dec, clip_dec.shape[2])
+            chunk = clip_dec[:, :, f_start:f_end][:, :, model.frame_pre_padding :]
+            if j == 0:
+                if dec_overlap is not None:
+                    chunk = model.blend(dec_overlap, chunk, model.frame_overlap, dim=-3)
+                    dec_overlap = None
+                write_part(chunk)
+            else:
+                dec_overlap = 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)
+
+    return dec
+
+
+class _IsolatedFrameModel(_FakeTemporalModel):
+    """Fake decoder with the isolated boundaries the released config can set."""
+
+    frame_pre_padding = 0
+
+    def __init__(self, *, first: bool, last: bool):
+        self.isolated_first_frame = first
+        self.isolated_last_frame = last
+
+    def _decode_temporal_output_frame_plan(self, z, z_head, z_tail, num_chunks, pad_tokens):
+        del pad_tokens
+        frames = num_chunks * self.tokens_chunk_size * self.vae_ratio_t
+        frames += int(z_head is not None) + int(z_tail is not None)
+        del z
+        return frames, 0, frames
+
+    def _adaptive_decode(self, clip):
+        # Every latent token decodes to a distinct, ordered block, so a shifted
+        # or reordered emission is visible in the values themselves.
+        base = float(clip[:, :, 0].mean())
+        frames = int(clip.shape[2]) * self.vae_ratio_t
+        return torch.arange(frames, dtype=torch.float32).view(1, 1, frames, 1, 1) + base
+
+
+@pytest.mark.parametrize(
+    ("first", "last"),
+    [(True, False), (False, True), (True, True)],
+)
+def test_temporal_chunks_match_the_native_loop_at_isolated_boundaries(first, last):
+    """The fork must extract the isolated head/tail the way the native loop does."""
+    model = _IsolatedFrameModel(first=first, last=last)
+    latent = torch.arange(12, dtype=torch.float32).view(1, 1, 12, 1, 1)
+
+    chunks: list[torch.Tensor] = []
+    decode_temporal_chunks(model, latent, chunks.append)
+    streamed = torch.cat(chunks, dim=2)
+
+    # Feed the oracle exactly what the released ``decode_temporal`` would:
+    # split the isolated tokens off, then pad the remainder to a whole chunk.
+    isolated = int(first) + int(last)
+    pseudo_tokens = int(latent.shape[2]) - isolated + model.token_drop
+    remainder = pseudo_tokens % model.tokens_chunk_size
+    pad_tokens = (model.tokens_chunk_size - remainder) if remainder else 0
+    num_chunks = (pseudo_tokens + pad_tokens) // model.tokens_chunk_size - int(model.token_drop > 0)
+
+    body = latent
+    z_head = None
+    if first:
+        z_head = body[:, :, :1]
+        body = body[:, :, 1:]
+    z_tail = None
+    if last:
+        z_tail = body[:, :, -1:]
+        body = body[:, :, :-1]
+    if pad_tokens:
+        body = torch.cat([body, body[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)], dim=2)
+
+    expected = _native_decode_temporal_streaming(model, body, z_head, z_tail, num_chunks, pad_tokens)
+
+    assert streamed.shape == expected.shape
+    assert torch.equal(streamed, expected)

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from contextlib import nullcontext
-from typing import Any
+from typing import Any, ClassVar
 
 import torch
 import torch.distributed as dist
@@ -62,6 +62,10 @@ class OmniAutoencoderKLWan(AutoencoderKLWan):
 
         return self._decode_temporal(z, return_dict=return_dict)
 
+    # Chunks are published straight from the decoder, clamped to the
+    # checkpoint's output range.
+    chunk_value_range: ClassVar[tuple[float, float]] = (-1.0, 1.0)
+
     @apply_forward_hook
     def decode_with_chunks(self, z: torch.Tensor, *, on_chunk: DecodedChunkConsumer) -> None:
         """Decode ``z`` while synchronously delivering temporal chunks.

vllm_omni/diffusion/models/interface.py

@@ -58,8 +58,16 @@ class SupportsChunkedVAEDecode(Protocol):
     Implementations must drain decoding before surfacing callback errors. In a
     distributed VAE, every rank invokes the method so collectives stay in
     lockstep, while ``on_chunk`` is called only on the rank that owns output.
+
+    ``chunk_value_range`` is the closed interval the delivered floats occupy,
+    which differs per checkpoint: a Wan VAE emits ``(-1.0, 1.0)`` while a
+    MiniMax-H3 VAE reverts through its processor to ``(0.0, 1.0)``. A consumer
+    cannot quantize chunks to pixels without it, so it is part of the
+    capability rather than knowledge each consumer hard-codes per model.
     """
 
+    chunk_value_range: ClassVar[tuple[float, float]]
+
     def decode_with_chunks(
         self,
         z: torch.Tensor,

vllm_omni/diffusion/models/minimax_h3/chunked_decode.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Frames-only callback coordination for MiniMax-H3 VAE decoding."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+import torch.distributed as dist
+
+from .temporal_chunks import decode_temporal_chunks
+
+MiniMaxH3VideoChunkCallback = Callable[[torch.Tensor], None]
+
+
+def decode_h3_chunks(
+    host,
+    latent: torch.Tensor,
+    callback: MiniMaxH3VideoChunkCallback | None,
+    *,
+    group: dist.ProcessGroup | None,
+) -> torch.Tensor:
+    """Run the local temporal loop, synchronizing callback errors across ranks."""
+    streaming = callback is not None
+    owner = streaming
+    if group is not None:
+        rank = dist.get_rank(group)
+        world_size = dist.get_world_size(group)
+        # Ownership is positional: every rank supplies the callback and rank 0
+        # publishes. Reduce only whether the ranks agree that this decode
+        # streams, so the verdict is identical everywhere -- a rank-local
+        # verdict would let a peer fall through and enter the decoder
+        # collectives alone, hanging the stage instead of failing it.
+        census = torch.tensor([int(streaming)], dtype=torch.int64, device=latent.device)
+        dist.all_reduce(census, group=group)
+        supplied = int(census.item())
+        if supplied not in (0, world_size):
+            raise ValueError(
+                "MiniMax-H3 chunk callback must be supplied on every rank of the VAE "
+                f"group or on none of them; got {supplied} of {world_size}"
+            )
+        # No rank supplied a callback: every rank runs a plain full decode.
+        streaming = supplied == world_size
+        owner = streaming and rank == 0
+
+    error: BaseException | None = None
+
+    def publish(raw: torch.Tensor) -> None:
+        nonlocal error
+        if not owner or error is not None:
+            return
+        try:
+            processor = getattr(host.model, "processor", None)
+            decoded = raw if processor is None else processor.revert_tensor(raw)
+            frames = host._normalize_decoded_frames(decoded).contiguous()
+            assert callback is not None
+            callback(frames)
+        except BaseException as exc:  # noqa: BLE001
+            error = exc
+
+    # Peer ranks run the temporal loop without materializing a second video.
+    sink = publish if owner else ((lambda _frames: None) if streaming else None)
+    result = decode_temporal_chunks(
+        host.model,
+        host._denormalize_latent(latent),
+        sink,
+    )
+
+    if group is not None:
+        failed = torch.tensor([int(rank == 0 and error is not None)], dtype=torch.int32, device=latent.device)
+        dist.broadcast(failed, src=dist.get_global_rank(group, 0), group=group)
+        if int(failed.item()):
+            if rank == 0:
+                assert error is not None
+                raise error
+            raise RuntimeError("MiniMax-H3 video chunk callback failed on rank zero")
+    elif error is not None:
+        raise error
+
+    if not result.numel():
+        return result
+    processor = getattr(host.model, "processor", None)
+    decoded = result if processor is None else processor.revert_tensor(result)
+    return host._normalize_decoded_frames(decoded)

vllm_omni/diffusion/models/minimax_h3/temporal_chunks.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Small compatibility loop for the released MiniMax-H3 temporal decoder."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+import torch.nn as nn
+
+
+def decode_temporal_chunks(
+    model: nn.Module,
+    latent: torch.Tensor,
+    callback: Callable[[torch.Tensor], None] | None,
+) -> torch.Tensor:
+    """Decode H3 temporal clips and synchronously publish committed frames.
+
+    The released VAE exposes the temporal clip primitives but materializes the
+    final tensor.  This is the same loop with the write point replaced by a
+    frames-only callback.  ``callback=None`` retains the complete-output path.
+    """
+    if latent.ndim != 5 or not bool(getattr(model, "use_3d_conv", False)):
+        raise ValueError("MiniMax-H3 temporal chunk decode requires a rank-5 3D latent")
+
+    token_drop = int(model.token_drop)
+    chunk_size = int(model.tokens_chunk_size)
+    overlap_tokens = int(model.token_overlap)
+    ratio_t = int(model.vae_ratio_t)
+    pre_padding = int(model.frame_pre_padding)
+
+    isolated_first = bool(model.isolated_first_frame and pre_padding == 0)
+    isolated_last = bool(model.isolated_last_frame)
+    z_head = latent[:, :, :1] if isolated_first else None
+    z_tail = latent[:, :, -1:] if isolated_last else None
+    start = int(isolated_first)
+    stop = int(latent.shape[2]) - int(isolated_last)
+    latent = latent[:, :, start:stop]
+
+    pseudo_tokens = int(latent.shape[2]) + token_drop
+    pad_tokens = (-pseudo_tokens) % chunk_size
+    if pad_tokens:
+        latent = torch.cat((latent, latent[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)), dim=2)
+    num_chunks = (pseudo_tokens + pad_tokens) // chunk_size - int(token_drop > 0)
+    if num_chunks <= 0:
+        raise ValueError("MiniMax-H3 temporal chunk plan is empty")
+
+    total_frames, pad_frames, output_frames = model._decode_temporal_output_frame_plan(
+        latent, z_head, z_tail, num_chunks, pad_tokens
+    )
+    output_frames = int(output_frames)
+    if output_frames <= 0:
+        raise ValueError("MiniMax-H3 temporal chunk plan has no output frames")
+
+    collected: list[torch.Tensor] = []
+    overlap: torch.Tensor | None = None
+    written = 0
+
+    def emit(part: torch.Tensor) -> None:
+        nonlocal written
+        frames = int(part.shape[2])
+        if frames <= 0 or written >= output_frames:
+            return
+        frames = min(frames, output_frames - written)
+        part = part[:, :, :frames]
+        if callback is None:
+            collected.append(part)
+        else:
+            callback(part)
+        written += frames
+
+    chunk_frames = chunk_size * ratio_t
+    split_count = int(token_drop > 0) + 1
+    for index in range(num_chunks):
+        begin = index * chunk_size
+        end = begin + chunk_size + overlap_tokens
+        clip = latent[:, :, begin:end]
+        if index == 0 and z_head is not None:
+            clip = torch.cat((z_head, clip), dim=2)
+        if index == num_chunks - 1 and z_tail is not None:
+            clip = torch.cat((clip, z_tail), dim=2)
+
+        decoded = model._adaptive_decode(clip)
+
+        # The isolated boundary latents were prepended/appended to the clip, so
+        # the decoder returns their frames inside this clip's block. Publish
+        # them the way the released loop does and drop their temporal block
+        # before splitting, or every later frame shifts by one block.
+        decoded_tail = None
+        if index == 0 and z_head is not None:
+            emit(decoded[:, :, ratio_t - 1 : ratio_t])
+            decoded = decoded[:, :, ratio_t:]
+        if index == num_chunks - 1 and z_tail is not None:
+            decoded_tail = decoded[:, :, -1:]
+            decoded = decoded[:, :, :-ratio_t]
+
+        for split in range(split_count):
+            begin = split * chunk_frames
+            end = min(begin + chunk_frames, int(decoded.shape[2]))
+            part = decoded[:, :, begin:end]
+            part = part[:, :, pre_padding:]
+            if split == 0:
+                if overlap is not None:
+                    part = model.blend(overlap, part, int(model.frame_overlap), dim=-3)
+                    overlap = None
+                emit(part)
+            else:
+                overlap = part.contiguous()
+
+        if index == num_chunks - 1:
+            if overlap is not None:
+                emit(overlap)
+                overlap = None
+            if decoded_tail is not None:
+                emit(decoded_tail)
+
+    if written != output_frames:
+        raise RuntimeError(f"MiniMax-H3 temporal decode emitted {written}/{output_frames} frames")
+    if callback is not None:
+        return latent.new_empty((0,))
+    return torch.cat(collected, dim=2) if collected else latent.new_empty((0,))

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, ClassVar
 
 import torch
 import torch.distributed as dist
@@ -22,11 +22,13 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
     DistributedVaeMixin,
 )
 from vllm_omni.diffusion.distributed.parallel_state import get_world_group
+from vllm_omni.diffusion.models.interface import DecodedChunkConsumer
 from vllm_omni.diffusion.offloader.module_residency import (
     BoundedAllocatorCache,
     PinnedModuleStager,
 )
 
+from .chunked_decode import decode_h3_chunks
 from .ops import install_h3_vae_optimizations
 from .packed_tokens import minimax_h3_patchify_video_latent
 
@@ -377,25 +379,16 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         ).float()
         return rows, shape
 
-    @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.
+    def _decode_tiling_context(self, latent: torch.Tensor) -> AbstractContextManager:
+        """Pick the tiling mode a decode of ``latent`` can safely use.
+
+        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(
@@ -405,19 +398,64 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
                 num_tiles,
                 self.parallel_size,
             )
-            tiling_context: AbstractContextManager = self._rank_local_tiling()
-        else:
-            tiling_context = nullcontext()
+            return self._rank_local_tiling()
+        return nullcontext()
 
-        with tiling_context:
-            decoded = self.model.decode_base(latent * std + mean)
-        frames = self.model.processor.revert_tensor(decoded)
+    @torch.inference_mode()
+    def decode_latent(self, latent: torch.Tensor) -> torch.Tensor:
+        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))
+
+    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)
+        std = torch.tensor(self.config_dict["latents_std"], device=latent.device, dtype=latent.dtype)
+        shape = (1, channels, 1, 1, 1)
+        return latent * std.view(shape) + mean.view(shape)
+
+    @staticmethod
+    def _normalize_decoded_frames(decoded: torch.Tensor) -> torch.Tensor:
+        """Canonicalize remote H3 decoder output to [B,C,T,H,W]."""
+        frames = 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()
 
+    # Chunks are reverted through the checkpoint's processor, which
+    # denormalizes and clamps into the unit interval.
+    chunk_value_range: ClassVar[tuple[float, float]] = (0.0, 1.0)
+
+    @torch.inference_mode()
+    def decode_with_chunks(self, z: torch.Tensor, *, on_chunk: DecodedChunkConsumer) -> None:
+        """Decode temporal clips and synchronously publish frames-only chunks.
+
+        Implements :class:`SupportsChunkedVAEDecode`. Every rank participating
+        in distributed VAE execution must invoke this method with a callback so
+        the temporal collectives stay in lockstep; ``on_chunk`` is called only
+        on the rank that owns output. Chunks arrive as ``[B, C, T, H, W]``
+        float frames, normalized through the checkpoint's processor to match
+        the complete decode path. After a callback failure, the remaining
+        chunks are decoded and discarded before the exception is re-raised.
+        """
+        if not callable(on_chunk):
+            raise TypeError("on_chunk must be callable")
+        if not callable(getattr(self.model, "_adaptive_decode", None)):
+            raise RuntimeError("Loaded MiniMax-H3 VAE does not expose temporal decode primitives")
+        group = None
+        if self.is_distributed_enabled():
+            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")
+        # Native H3 tiling performs its own collectives for every temporal clip,
+        # so this path needs the same too-few-tiles fallback as the complete
+        # decode: without it a shape that leaves some ranks tileless hangs the
+        # gather instead of decoding rank-locally.
+        with self._decode_tiling_context(z):
+            decode_h3_chunks(self, z, on_chunk, group=group)
+
 
 class MiniMaxH3AudioVAE(nn.Module):
     def __init__(