[perf] Reduce MiniMax-H3 VAE peak memory (#1703)

fastvideo/models/vaes/minimax_h3_video.py

@@ -7,6 +7,7 @@ This module intentionally uses only PyTorch and FastVideo configuration types.
 """
 
 import math
+from collections.abc import Iterator
 from dataclasses import dataclass
 
 import torch
@@ -747,43 +748,140 @@ class AutoencoderKLMiniMaxH3(nn.Module):
             moments = moments[:, :, :-self.config.token_drop]
         return moments
 
-    def _decode(self, z: torch.Tensor) -> torch.Tensor:
+    def _encode_pixels(self, pixels: torch.Tensor) -> torch.Tensor:
-        tokens_chunk_size = self.tokens_chunk_size
+        """Encode unnormalized pixels while keeping full videos off the accelerator."""
+        clip_length = self.config.clip_length
+        moments = []
+        for frame_start in range(0, pixels.shape[2], clip_length):
+            clip = pixels[:, :, frame_start:frame_start + clip_length].to(
+                device=self.pixel_mean.device,
+                dtype=torch.float32,
+            )
+            if pixels.dtype == torch.uint8:
+                clip = clip / 255.0
+            if clip.shape[2] < clip_length:
+                pad_frames = clip[:, :, -1:].repeat(1, 1, clip_length - clip.shape[2], 1, 1)
+                clip = torch.cat([clip, pad_frames], dim=2)
+            clip = self.normalize_pixels(clip)
+            moments.append(self._encode_clip(clip))
+            del clip
+        encoded = torch.cat(moments, dim=2)
+        if self.config.token_drop > 0:
+            encoded = encoded[:, :, :-self.config.token_drop]
+        return encoded
+
+    def _temporal_decode_plan(self, latent_num_frames: int) -> tuple[int, int, int]:
+        """Return pad tokens, chunk count, and exact decoded frame count."""
+        if latent_num_frames <= 0:
+            raise ValueError(f"MiniMax-H3 latent frame count must be positive, got {latent_num_frames}.")
+
         token_drop = self.config.token_drop
+        tokens_chunk_size = self.tokens_chunk_size
         temporal_ratio = self.temporal_compression_ratio
-        chunk_num_frames = tokens_chunk_size * temporal_ratio
+        num_tokens = latent_num_frames + token_drop
-        num_tokens = z.shape[2] + token_drop
         pad_tokens = (-num_tokens) % tokens_chunk_size
         num_chunks = (num_tokens + pad_tokens) // tokens_chunk_size - int(token_drop > 0)
+        if num_chunks < 1:
+            pad_tokens += tokens_chunk_size
+            num_chunks = 1
+
+        decoded_num_frames = num_chunks * (tokens_chunk_size * temporal_ratio - self.frame_pre_padding)
+        if token_drop > 0:
+            decoded_num_frames += self.frame_overlap
+        if pad_tokens > 0:
+            intra_tail = self.config.clip_length % temporal_ratio
+            pad_frames = sum(intra_tail if intra_tail and (latent_num_frames + offset) %
+                             tokens_chunk_size == 0 else temporal_ratio for offset in range(pad_tokens))
+            decoded_num_frames -= pad_frames
+        if decoded_num_frames <= 0:
+            raise RuntimeError(
+                f"MiniMax-H3 decode plan produced {decoded_num_frames} frames for {latent_num_frames} latent "
+                "frames; the clip_length/token_drop configuration is inconsistent.")
+        return pad_tokens, num_chunks, decoded_num_frames
+
+    def _decode_chunks(self, z: torch.Tensor) -> Iterator[torch.Tensor]:
+        """Yield finalized temporal chunks in decode order."""
+        tokens_chunk_size = self.tokens_chunk_size
+        chunk_num_frames = tokens_chunk_size * self.temporal_compression_ratio
+        pad_tokens, num_chunks, output_num_frames = self._temporal_decode_plan(z.shape[2])
         if pad_tokens > 0:
             z = torch.cat([z, z[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)], dim=2)
 
-        decoded_chunks = []
+        output_frame_start = 0
         overlap = None
         for index in range(num_chunks):
             start = index * tokens_chunk_size
             clip = self._decode_clip(z[:, :, start:start + tokens_chunk_size + self.token_overlap])
-            for overlap_index in range(int(token_drop > 0) + 1):
+            chunk = clip[:, :, self.frame_pre_padding:chunk_num_frames]
-                frame_start = overlap_index * chunk_num_frames
+            next_overlap = None
-                chunk = clip[:, :, frame_start:frame_start + chunk_num_frames]
+            if self.config.token_drop > 0:
-                chunk = chunk[:, :, self.frame_pre_padding:]
+                next_overlap = clip[:, :, chunk_num_frames + self.frame_pre_padding:].clone()
-                if overlap_index == 0:
+            if overlap is not None:
-                    if overlap is not None:
+                chunk = self._blend(overlap, chunk, self.frame_overlap, dim=-3)
-                        chunk = self._blend(overlap, chunk, self.frame_overlap, dim=-3)
-                    decoded_chunks.append(chunk)
-                else:
-                    overlap = chunk
-        if overlap is not None:
-            decoded_chunks.append(overlap)
-        decoded = torch.cat(decoded_chunks, dim=2)
 
-        if pad_tokens > 0:
+            num_frames = min(chunk.shape[2], output_num_frames - output_frame_start)
-            intra_tail = self.config.clip_length % temporal_ratio
+            if num_frames > 0:
-            num_tokens_before_pad = z.shape[2] - pad_tokens
+                yield chunk[:, :, :num_frames]
-            pad_frames = sum(intra_tail if intra_tail and (num_tokens_before_pad + offset) %
+                output_frame_start += num_frames
-                             tokens_chunk_size == 0 else temporal_ratio for offset in range(pad_tokens))
+            overlap = next_overlap
-            decoded = decoded[:, :, :-pad_frames]
+
-        return decoded
+        if overlap is not None and output_frame_start < output_num_frames:
+            yield overlap[:, :, :output_num_frames - output_frame_start]
+
+    def decoded_pixel_shape(self, latent_shape: torch.Size | tuple[int, ...]) -> tuple[int, int, int, int, int]:
+        """Return the exact CPU pixel-buffer shape for a latent tensor shape."""
+        if len(latent_shape) != 5:
+            raise ValueError(f"MiniMax-H3 latents must be five-dimensional, got shape {tuple(latent_shape)}.")
+        batch_size, channels, latent_num_frames, latent_height, latent_width = map(int, latent_shape)
+        if channels != self.latent_channels:
+            raise ValueError(f"MiniMax-H3 latents must have {self.latent_channels} channels, got {channels}.")
+        _, _, decoded_num_frames = self._temporal_decode_plan(latent_num_frames)
+        return (
+            batch_size,
+            int(self.config.out_channels),
+            decoded_num_frames,
+            latent_height * self.spatial_compression_ratio,
+            latent_width * self.spatial_compression_ratio,
+        )
+
+    @staticmethod
+    def _streams_chunk_copies(z: torch.Tensor, output: torch.Tensor) -> bool:
+        """Whether finalized chunks copy to ``output`` asynchronously on the current CUDA stream."""
+        return z.device.type == "cuda" and output.is_pinned()
+
+    def _decode_to_pixels(self, z: torch.Tensor, output: torch.Tensor) -> None:
+        """Decode temporal chunks and immediately copy finalized pixels to CPU.
+
+        Device-to-host copies run per (batch, channel) plane: the temporal
+        slice of ``output`` is strided across channels, but each plane is
+        contiguous on both sides, so every transfer stays a direct memcpy
+        instead of staging through a pageable CPU temporary. With a pinned
+        ``output`` the copies are additionally asynchronous and overlap the
+        next chunk's decode; ``decode_to_pixels`` synchronizes once before
+        returning.
+        """
+        non_blocking = self._streams_chunk_copies(z, output)
+        output_frame_start = 0
+        for chunk in self._decode_chunks(z):
+            num_frames = chunk.shape[2]
+            pixels = self.denormalize_pixels(chunk.float()).clamp_(0, 1)
+            target = output[:, :, output_frame_start:output_frame_start + num_frames]
+            if z.device.type == "cuda":
+                pixels = pixels.contiguous()
+                for batch_index in range(pixels.shape[0]):
+                    for channel_index in range(pixels.shape[1]):
+                        target[batch_index, channel_index].copy_(pixels[batch_index, channel_index],
+                                                                 non_blocking=non_blocking)
+            else:
+                target.copy_(pixels)
+            output_frame_start += num_frames
+        if output_frame_start != output.shape[2]:
+            raise RuntimeError(
+                f"MiniMax-H3 decode wrote {output_frame_start} frames into an output buffer expecting "
+                f"{output.shape[2]}.")
+
+    def _decode(self, z: torch.Tensor) -> torch.Tensor:
+        return torch.cat(list(self._decode_chunks(z)), dim=2)
 
     def encode(
         self,
@@ -799,6 +897,34 @@ class AutoencoderKLMiniMaxH3(nn.Module):
             return (posterior, )
         return AutoencoderKLOutput(latent_dist=posterior)
 
+    def encode_pixels(
+        self,
+        pixels: torch.Tensor,
+        return_dict: bool = True,
+    ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
+        """Encode CPU-resident pixels one VAE clip at a time.
+
+        ``pixels`` stays on CPU as ``uint8`` in ``[0, 255]`` or floating point
+        in ``[0, 1]``; each clip is moved to the VAE device, normalized, and
+        encoded so only one clip of pixels is resident on the accelerator.
+        """
+        if pixels.ndim != 5 or pixels.shape[1] != self.config.in_channels or pixels.shape[2] <= 0:
+            raise ValueError(
+                f"`pixels` must have shape [B, {self.config.in_channels}, T, H, W] with T > 0, "
+                f"got {tuple(pixels.shape)}.")
+        if pixels.device.type != "cpu":
+            raise ValueError(f"`pixels` must remain on CPU, got device={pixels.device}.")
+        if pixels.dtype != torch.uint8 and not pixels.is_floating_point():
+            raise TypeError(f"`pixels` must use uint8 or a floating-point dtype, got {pixels.dtype}.")
+        if self.use_slicing and pixels.shape[0] > 1:
+            moments = torch.cat([self._encode_pixels(pixel_slice) for pixel_slice in pixels.split(1)])
+        else:
+            moments = self._encode_pixels(pixels)
+        posterior = DiagonalGaussianDistribution(moments)
+        if not return_dict:
+            return (posterior, )
+        return AutoencoderKLOutput(latent_dist=posterior)
+
     def encode_keyframe(
         self,
         x: torch.Tensor,
@@ -825,6 +951,26 @@ class AutoencoderKLMiniMaxH3(nn.Module):
             return (decoded, )
         return DecoderOutput(sample=decoded)
 
+    def decode_to_pixels(self, z: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
+        """Stream decoded ``[0, 1]`` FP32 pixels into a caller-owned CPU buffer."""
+        expected_shape = self.decoded_pixel_shape(z.shape)
+        if output.device.type != "cpu" or output.dtype != torch.float32 or tuple(output.shape) != expected_shape:
+            raise ValueError(
+                "`output` must be a CPU float32 tensor with shape "
+                f"{expected_shape}, got device={output.device}, dtype={output.dtype}, shape={tuple(output.shape)}.")
+        try:
+            if self.use_slicing and z.shape[0] > 1:
+                for batch_index, z_slice in enumerate(z.split(1)):
+                    self._decode_to_pixels(z_slice, output[batch_index:batch_index + 1])
+            else:
+                self._decode_to_pixels(z, output)
+        finally:
+            # Drain async chunk copies before the caller (or an exception
+            # handler) can read or release the pinned buffer.
+            if self._streams_chunk_copies(z, output):
+                torch.cuda.current_stream(z.device).synchronize()
+        return output
+
     def forward(
         self,
         sample: torch.Tensor,

fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py

@@ -7,7 +7,7 @@ from typing import Any
 
 import torch
 
-from fastvideo.distributed import get_local_torch_device
+from fastvideo.distributed import get_local_torch_device, get_world_group, model_parallel_is_initialized
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAE
 from fastvideo.models.vaes.minimax_h3_video import AutoencoderKLMiniMaxH3
@@ -21,6 +21,7 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
+from fastvideo.utils import is_pin_memory_available
 
 
 def _layout(batch: ForwardBatch) -> MiniMaxH3PackedLayout:
@@ -54,6 +55,13 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
 
     @torch.no_grad()
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if model_parallel_is_initialized() and not get_world_group().is_first_rank:
+            # Distributed executors consume rank 0's ForwardBatch. Keep a
+            # verifier-compatible placeholder on other ranks and avoid
+            # duplicating the full VAE decode and CPU output buffer.
+            batch.output = torch.empty((0, 3, 0, 0, 0), device="cpu", dtype=torch.float32)
+            return batch
+
         layout = _layout(batch)
         if batch.latents is None or batch.raw_latent_shape is None or len(batch.raw_latent_shape) != 5:
             raise ValueError("MiniMax-H3 video latents or raw geometry are missing at decode.")
@@ -74,10 +82,16 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
                 batch.output = latents.detach().float().cpu()
                 return batch
 
+            output = torch.empty(
+                self.vae.decoded_pixel_shape(latents.shape),
+                device="cpu",
+                dtype=torch.float32,
+                pin_memory=fastvideo_args.pin_cpu_memory and is_pin_memory_available(),
+            )
             # The published decode recipe uses FP16 autocast over FP32 weights.
             with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"):
-                video = self.vae.decode(latents).sample
+                self.vae.decode_to_pixels(latents, output)
-            batch.output = self.vae.denormalize_pixels(video.float()).clamp_(0, 1).cpu()
+            batch.output = output
             return batch
         finally:
             if fastvideo_args.vae_cpu_offload:
@@ -107,6 +121,12 @@ class MiniMaxH3AudioDecodingStage(PipelineStage):
 
     @torch.no_grad()
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        if model_parallel_is_initialized() and not get_world_group().is_first_rank:
+            batch.extra["audio"] = torch.empty((0, 2), device="cpu", dtype=torch.float32)
+            batch.extra["audio_sample_rate"] = self.audio_vae.sampling_rate
+            self._clear_runtime(batch)
+            return batch
+
         layout = _layout(batch)
         if batch.audio_latents is None:
             raise ValueError("MiniMax-H3 audio latents are missing at decode.")

fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py

@@ -119,9 +119,8 @@ class MiniMaxH3LatentPreparationStage(PipelineStage):
                 if reference.frames is None:
                     raise ValueError("MiniMax-H3 reference video frames are missing.")
                 frames = reference.frames[:trim_reference_num_frames(reference.frames.shape[0])]
-                pixels = torch.from_numpy(frames.copy()).permute(3, 0, 1, 2)[None]
+                pixels = torch.from_numpy(np.ascontiguousarray(frames)).permute(3, 0, 1, 2)[None]
-                pixels = pixels.to(device=device, dtype=torch.float32).div_(255.0)
+                posterior = self.vae.encode_pixels(pixels).latent_dist
-                posterior = self.vae.encode(self.vae.normalize_pixels(pixels)).latent_dist
                 latents = self.vae.normalize_latents(_sample_visual_posterior(posterior).to(
                     torch.float16).float()).cpu()
             reference.num_latent_frames = int(latents.shape[2])

fastvideo/tests/stages/test_minimax_h3_vae_streaming.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+
+import numpy as np
+import torch
+
+from fastvideo.pipelines.basic.minimax_h3.packing import (
+    MiniMaxH3PackedLayout,
+    patchify_video_latents,
+)
+from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference
+from fastvideo.pipelines.basic.minimax_h3.stages import minimax_h3_decoding
+from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_decoding import MiniMaxH3VideoDecodingStage
+from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import (
+    MINIMAX_H3_LAYOUT_KEY,
+    MiniMaxH3LatentPreparationStage,
+)
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+
+def _layout(rows: int, latent_shape: tuple[int, ...]) -> MiniMaxH3PackedLayout:
+    empty = torch.empty(0, dtype=torch.long)
+    return MiniMaxH3PackedLayout(
+        sequence_length=rows,
+        position_ids=empty,
+        token_tags=empty,
+        video_indices=empty,
+        audio_indices=empty,
+        text_indices=empty,
+        num_condition_video_rows=0,
+        num_condition_audio_rows=0,
+        num_video_latent_frames=latent_shape[2],
+        latent_height=latent_shape[3],
+        latent_width=latent_shape[4],
+        num_audio_latents=0,
+    )
+
+
+def test_reference_video_encode_keeps_pixels_on_cpu() -> None:
+    observed = {}
+
+    class VAE:
+
+        def encode_pixels(self, pixels):
+            observed["pixels"] = pixels
+            posterior = SimpleNamespace(sample=lambda generator=None: torch.zeros(1, 4, 7, 4, 4))
+            return SimpleNamespace(latent_dist=posterior)
+
+        def normalize_latents(self, latents):
+            return latents
+
+    stage = MiniMaxH3LatentPreparationStage(
+        transformer=SimpleNamespace(patch_size=(1, 1, 1)),
+        vae=VAE(),
+        audio_vae=None,
+        scheduler=None,
+        ref2va=True,
+    )
+    reference = MiniMaxH3PreparedReference(
+        media_type="video",
+        frames=np.zeros((22, 16, 16, 3), dtype=np.uint8),
+    )
+    rows = stage._encode_visual_rows([reference], torch.device("cpu"))
+
+    assert observed["pixels"].dtype == torch.uint8
+    assert observed["pixels"].device.type == "cpu"
+    assert rows[0].shape == (7 * 4 * 4, 4)
+
+
+def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None:
+    latent_shape = (1, 4, 2, 4, 4)
+    latents = torch.randn(latent_shape)
+    rows = patchify_video_latents(latents, (1, 1, 1))
+    batch = ForwardBatch(data_type="video", latents=rows, raw_latent_shape=latent_shape)
+    batch.extra[MINIMAX_H3_LAYOUT_KEY] = _layout(rows.shape[0], latent_shape)
+    observed = {}
+
+    class VAE:
+
+        def to(self, device):
+            return self
+
+        def denormalize_latents(self, decoded_latents):
+            return decoded_latents
+
+        def decoded_pixel_shape(self, shape):
+            assert tuple(shape) == latent_shape
+            return (1, 3, 5, 16, 16)
+
+        def decode_to_pixels(self, decoded_latents, output):
+            observed["latents"] = decoded_latents
+            observed["output"] = output
+            output.fill_(0.25)
+
+    monkeypatch.setattr(minimax_h3_decoding, "get_local_torch_device", lambda: torch.device("cpu"))
+    result = MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace(patch_size=(1, 1, 1))).forward(
+        batch,
+        SimpleNamespace(output_type="pil", pin_cpu_memory=False, vae_cpu_offload=False),
+    )
+
+    torch.testing.assert_close(observed["latents"], latents)
+    assert observed["output"] is result.output
+    assert result.output.device.type == "cpu"
+    assert torch.all(result.output == 0.25)
+
+
+def test_decode_stages_skip_vae_on_non_output_rank(monkeypatch) -> None:
+    class VAE:
+
+        sampling_rate = 32000
+
+        def to(self, device):
+            raise AssertionError("non-output ranks must not execute a VAE")
+
+    monkeypatch.setattr(minimax_h3_decoding, "model_parallel_is_initialized", lambda: True)
+    monkeypatch.setattr(minimax_h3_decoding, "get_world_group", lambda: SimpleNamespace(is_first_rank=False))
+    args = SimpleNamespace(output_type="pil", pin_cpu_memory=False, vae_cpu_offload=True)
+
+    video = MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace()).forward(ForwardBatch(data_type="video"), args)
+    assert video.output.shape == (0, 3, 0, 0, 0)
+
+    audio_batch = ForwardBatch(data_type="audio", latents=torch.zeros(1), audio_latents=torch.zeros(1))
+    audio_batch.extra[MINIMAX_H3_LAYOUT_KEY] = object()
+    audio = minimax_h3_decoding.MiniMaxH3AudioDecodingStage(VAE()).forward(audio_batch, args)
+    assert audio.extra["audio"].shape == (0, 2)
+    assert audio.extra["audio_sample_rate"] == 32000
+    assert audio.latents is None
+    assert audio.audio_latents is None
+    assert MINIMAX_H3_LAYOUT_KEY not in audio.extra

fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py

@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+import pytest
+import torch
+from torch.testing import assert_close
+
+from fastvideo.configs.models.vaes.minimax_h3_video import (
+    MiniMaxH3VideoVAEArchConfig,
+    MiniMaxH3VideoVAEConfig,
+)
+from fastvideo.models.vaes.minimax_h3_video import AutoencoderKLMiniMaxH3
+
+
+def _tiny_vae() -> AutoencoderKLMiniMaxH3:
+    arch = MiniMaxH3VideoVAEArchConfig(
+        latent_channels=4,
+        block_out_channels=(32, 32),
+        layers_per_block=1,
+        spatial_downsample_factors=(2, 2),
+        temporal_downsample_factors=(2, 2),
+        decoder_num_layers=1,
+        decoder_num_attention_heads=1,
+        decoder_attention_head_dim=8,
+        decoder_num_register_tokens=2,
+        decoder_ffn_mult=1,
+        latents_mean=(0.0, ) * 4,
+        latents_std=(1.0, ) * 4,
+    )
+    return AutoencoderKLMiniMaxH3(
+        MiniMaxH3VideoVAEConfig(
+            arch_config=arch,
+            use_tiling=False,
+            use_temporal_tiling=False,
+            use_parallel_tiling=False,
+        )).eval()
+
+
+@torch.inference_mode()
+def test_encode_pixels_matches_encode() -> None:
+    torch.manual_seed(20260810)
+    vae = _tiny_vae()
+    pixels = torch.randint(0, 256, (1, 3, 22, 16, 16), dtype=torch.uint8)
+    expected = vae.encode(vae.normalize_pixels(pixels.float().div(255))).latent_dist.parameters
+    assert_close(vae.encode_pixels(pixels).latent_dist.parameters, expected, atol=0.0, rtol=0.0)
+
+    float_pixels = torch.rand(1, 3, 22, 16, 16)
+    original_pixels = float_pixels.clone()
+    expected = vae.encode(vae.normalize_pixels(float_pixels)).latent_dist.parameters
+    actual = vae.encode_pixels(float_pixels).latent_dist.parameters
+    assert_close(float_pixels, original_pixels, atol=0.0, rtol=0.0)
+    assert_close(actual, expected, atol=0.0, rtol=0.0)
+
+    with pytest.raises(ValueError, match="must remain on CPU"):
+        vae.encode_pixels(torch.empty(1, 3, 1, 16, 16, device="meta"))
+
+
+def _legacy_decode(vae: AutoencoderKLMiniMaxH3, z: torch.Tensor) -> torch.Tensor:
+    """Verbatim pre-streaming ``_decode`` (main @ 8208536cd) as a reference oracle."""
+    tokens_chunk_size = vae.tokens_chunk_size
+    token_drop = vae.config.token_drop
+    temporal_ratio = vae.temporal_compression_ratio
+    chunk_num_frames = tokens_chunk_size * temporal_ratio
+    num_tokens = z.shape[2] + token_drop
+    pad_tokens = (-num_tokens) % tokens_chunk_size
+    num_chunks = (num_tokens + pad_tokens) // tokens_chunk_size - int(token_drop > 0)
+    if pad_tokens > 0:
+        z = torch.cat([z, z[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)], dim=2)
+
+    decoded_chunks = []
+    overlap = None
+    for index in range(num_chunks):
+        start = index * tokens_chunk_size
+        clip = vae._decode_clip(z[:, :, start:start + tokens_chunk_size + vae.token_overlap])
+        for overlap_index in range(int(token_drop > 0) + 1):
+            frame_start = overlap_index * chunk_num_frames
+            chunk = clip[:, :, frame_start:frame_start + chunk_num_frames]
+            chunk = chunk[:, :, vae.frame_pre_padding:]
+            if overlap_index == 0:
+                if overlap is not None:
+                    chunk = vae._blend(overlap, chunk, vae.frame_overlap, dim=-3)
+                decoded_chunks.append(chunk)
+            else:
+                overlap = chunk
+    if overlap is not None:
+        decoded_chunks.append(overlap)
+    decoded = torch.cat(decoded_chunks, dim=2)
+
+    if pad_tokens > 0:
+        intra_tail = vae.config.clip_length % temporal_ratio
+        num_tokens_before_pad = z.shape[2] - pad_tokens
+        pad_frames = sum(intra_tail if intra_tail and (num_tokens_before_pad + offset) %
+                         tokens_chunk_size == 0 else temporal_ratio for offset in range(pad_tokens))
+        decoded = decoded[:, :, :-pad_frames]
+    return decoded
+
+
+@pytest.mark.parametrize("latent_frames", (2, 12))
+@torch.inference_mode()
+def test_decode_to_pixels_matches_decode(latent_frames: int) -> None:
+    torch.manual_seed(20260810)
+    vae = _tiny_vae()
+    latents = torch.randn(1, 4, latent_frames, 4, 4)
+    expected = vae.denormalize_pixels(vae.decode(latents).sample.float()).clamp_(0, 1)
+    actual = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+
+    vae.decode_to_pixels(latents, actual)
+
+    assert_close(actual, expected, atol=0.0, rtol=0.0)
+
+
+# 3: one chunk with pad tokens; 6: pad hitting the intra-clip tail;
+# 12: two blended chunks without padding; 13: three chunks plus pad trim.
+@pytest.mark.parametrize("latent_frames", (3, 6, 12, 13))
+@torch.inference_mode()
+def test_decode_matches_legacy_algorithm(latent_frames: int) -> None:
+    """The chunk iterator must stay bit-exact with the pre-streaming decode."""
+    torch.manual_seed(20260811 + latent_frames)
+    vae = _tiny_vae()
+    latents = torch.randn(1, 4, latent_frames, 4, 4)
+    expected = _legacy_decode(vae, latents)
+
+    decoded = vae.decode(latents).sample
+    assert_close(decoded, expected, atol=0.0, rtol=0.0)
+
+    streamed = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+    vae.decode_to_pixels(latents, streamed)
+    assert_close(streamed, vae.denormalize_pixels(expected.float()).clamp_(0, 1), atol=0.0, rtol=0.0)
+
+
+@torch.inference_mode()
+def test_streaming_slicing_matches_unbatched() -> None:
+    torch.manual_seed(20260812)
+    vae = _tiny_vae()
+    vae.enable_slicing()
+
+    latents = torch.randn(2, 4, 7, 4, 4)
+    expected = vae.denormalize_pixels(vae.decode(latents).sample.float()).clamp_(0, 1)
+    actual = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+    vae.decode_to_pixels(latents, actual)
+    assert_close(actual, expected, atol=0.0, rtol=0.0)
+
+    pixels = torch.randint(0, 256, (2, 3, 22, 16, 16), dtype=torch.uint8)
+    expected_moments = vae.encode(vae.normalize_pixels(pixels.float().div(255))).latent_dist.parameters
+    assert_close(vae.encode_pixels(pixels).latent_dist.parameters, expected_moments, atol=0.0, rtol=0.0)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="pinned-buffer streaming requires CUDA")
+@torch.inference_mode()
+def test_decode_to_pixels_pinned_buffer_matches_dense_on_cuda() -> None:
+    """Async chunk copies into a pinned buffer must equal the dense decode."""
+    torch.manual_seed(20260813)
+    vae = _tiny_vae().to("cuda")
+    latents = torch.randn(1, 4, 12, 4, 4, device="cuda")
+    expected = vae.denormalize_pixels(vae.decode(latents).sample.float()).clamp_(0, 1).cpu()
+
+    pinned = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32, pin_memory=True)
+    vae.decode_to_pixels(latents, pinned)
+
+    assert_close(pinned, expected, atol=0.0, rtol=0.0)
+
+
+def test_decode_to_pixels_rejects_incomplete_output(monkeypatch) -> None:
+    vae = _tiny_vae()
+    latents = torch.randn(1, 4, 2, 4, 4)
+    output = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+    monkeypatch.setattr(vae, "_decode_chunks", lambda _: iter(()))
+
+    with pytest.raises(RuntimeError, match="wrote 0 frames"):
+        vae.decode_to_pixels(latents, output)

tests/local_tests/minimax_h3/PORT_STATUS.md

@@ -14,6 +14,7 @@
 | Qwen3-VL encoder | exact text/image/video hidden states through the production loader | complete |
 | FL2VA and Ref2VA DiTs | exact video/audio heads for both model partitions | complete |
 | Video VAE | exact encode, normalization, and decode through the production loader | complete |
+| Video VAE streaming | exact chunked encode/decode and output-rank-only distributed decode | complete |
 | Audio VAE | exact encode and normalization; decode maximum absolute drift `2.4e-7` | complete |
 | Video/audio schedulers | pinned `12/3` schedule parity | complete |
 | FL2VA packing | pinned row, position, tag, timestep, and RNG parity | complete |
@@ -34,6 +35,7 @@ T2VA, FL2VA, and Ref2VA match the official video/audio latents exactly.
 - Load `transformer/` for T2VA/FL2VA and `transformer_ref/` for Ref2VA.
 - Keep `last_image`, `references`, and `audio_latents` on the typed request path.
 - Treat the published component folders as the loading boundary.
+- Keep reference videos on CPU between VAE clips and decode final pixels only on the executor's output rank.
 
 ## Evidence boundary
 

tests/local_tests/minimax_h3/README.md

@@ -12,6 +12,14 @@ FastVideo-owned unit contracts belong under `fastvideo/tests/`.
 The reference helper verifies the pinned source and import origin. A missing checkout may skip a source-parity module;
 that skip is not parity evidence.
 
+## FastVideo unit contracts
+
+```bash
+pytest \
+  fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py \
+  fastvideo/tests/stages/test_minimax_h3_vae_streaming.py -q
+```
+
 ## Registry smoke
 
 ```bash
@@ -51,5 +59,30 @@ pytest \
 With a gate enabled, missing CUDA, source, or weights is a failure. Recorded component evidence is exact for both DiT
 partitions, the video VAE, and all Qwen3-VL hidden states; audio decode has maximum absolute drift `2.4e-7`.
 
+The video VAE test verifies the reference checkout at commit
+`abc5e9bf71fd38f53cd471bc3acaa84bc5ecbfdc` and compares the production CPU `uint8` `encode_pixels()` path against
+the official posterior element by element.
+
+## Video VAE memory benchmark
+
+The benchmark uses one warmup and three measured runs with `vae_cpu_offload=True`. It reports absolute and
+stage-incremental allocated/reserved CUDA peaks for every rank. For SP runs, the reported aggregate is explicitly the
+sum of rank-local maxima, not a simultaneous node peak.
+
+```bash
+python tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py \
+  --source-root "$PWD" --model-root "$MINIMAX_H3_MODEL_ROOT" \
+  --revision-label candidate --operation encode
+
+python -m torch.distributed.run --nproc_per_node=4 \
+  tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py \
+  --source-root "$PWD" --model-root "$MINIMAX_H3_MODEL_ROOT" \
+  --revision-label candidate-sp4 --operation decode
+```
+
+Run the same script with `--source-root` pointed at the base checkout for a comparable baseline. The default workload
+is deterministic `124 x 768 x 1344` video geometry with seed `20260803`; the JSON record includes source/model
+revisions, software/allocator metadata, exact measurement boundaries, per-repetition values, and output shapes.
+
 FastVideo joint audio/video generation and SP=1/SP=4 latent consistency have been validated. T2VA, FL2VA, and
 Ref2VA video/audio latents match the pinned Diffusers pipeline exactly.

tests/local_tests/minimax_h3/_reference.py

@@ -3,6 +3,7 @@
 import hashlib
 import inspect
 import os
+import subprocess
 from pathlib import Path
 
 import pytest
@@ -13,6 +14,28 @@ REFERENCE_SRC = REFERENCE_ROOT / "src"
 PINNED_COMMIT = "abc5e9bf71fd38f53cd471bc3acaa84bc5ecbfdc"
 
 
+def _run_git(*args: str) -> str:
+    try:
+        result = subprocess.run(
+            ["git", "-C", str(REFERENCE_ROOT), *args],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+    except (FileNotFoundError, subprocess.CalledProcessError) as error:
+        raise RuntimeError(f"Could not verify the MiniMax-H3 reference checkout at {REFERENCE_ROOT}.") from error
+    return result.stdout.strip()
+
+
+def assert_reference_revision() -> None:
+    actual_commit = _run_git("rev-parse", "HEAD")
+    if actual_commit != PINNED_COMMIT:
+        raise RuntimeError(f"MiniMax-H3 parity requires Diffusers commit {PINNED_COMMIT}, got {actual_commit}.")
+    dirty_source = _run_git("status", "--short", "--untracked-files=no", "--", "src/diffusers")
+    if dirty_source:
+        raise RuntimeError(f"MiniMax-H3 reference source has tracked changes:\n{dirty_source}")
+
+
 def assert_pinned_reference(relative_path: str, sha256: str) -> Path:
     path = REFERENCE_ROOT / relative_path
     if not path.is_file():

tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py

@@ -0,0 +1,300 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Measure MiniMax-H3 production VAE stage memory with CPU offload enabled."""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import hashlib
+import json
+import os
+from pathlib import Path
+import statistics
+import subprocess
+import sys
+import time
+from types import SimpleNamespace
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--source-root", type=Path, required=True)
+    parser.add_argument("--model-root", type=Path, required=True)
+    parser.add_argument("--revision-label", required=True)
+    parser.add_argument("--operation", choices=("encode", "decode"), required=True)
+    parser.add_argument("--height", type=int, default=768)
+    parser.add_argument("--width", type=int, default=1344)
+    parser.add_argument("--num-frames", type=int, default=124)
+    parser.add_argument("--seed", type=int, default=20260803)
+    parser.add_argument("--warmups", type=int, default=1)
+    parser.add_argument("--repetitions", type=int, default=3)
+    return parser.parse_args()
+
+
+def _git(source_root: Path, *args: str) -> str:
+    result = subprocess.run(
+        ["git", "-C", str(source_root), *args],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    return result.stdout.strip()
+
+
+def _sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as file:
+        for block in iter(lambda: file.read(1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def _model_snapshot(model_root: Path) -> str | None:
+    parts = model_root.resolve().parts
+    try:
+        return parts[parts.index("snapshots") + 1]
+    except (ValueError, IndexError):
+        return None
+
+
+def _make_layout(rows, latent_shape):
+    import torch
+
+    from fastvideo.pipelines.basic.minimax_h3.packing import MiniMaxH3PackedLayout
+
+    empty = torch.empty(0, dtype=torch.long, device=rows.device)
+    return MiniMaxH3PackedLayout(
+        sequence_length=rows.shape[0],
+        position_ids=empty,
+        token_tags=empty,
+        video_indices=empty,
+        audio_indices=empty,
+        text_indices=empty,
+        num_condition_video_rows=0,
+        num_condition_audio_rows=0,
+        num_video_latent_frames=latent_shape[2],
+        latent_height=latent_shape[3],
+        latent_width=latent_shape[4],
+        num_audio_latents=0,
+    )
+
+
+def _build_operation(args, vae, device):
+    import numpy as np
+    import torch
+
+    from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config
+    from fastvideo.pipelines.basic.minimax_h3.packing import patchify_video_latents, video_latent_num_frames
+    from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference
+    from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_decoding import MiniMaxH3VideoDecodingStage
+    from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import (
+        MINIMAX_H3_LAYOUT_KEY,
+        MiniMaxH3LatentPreparationStage,
+    )
+    from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+    patch_size = MiniMaxH3Config().arch_config.patch_size
+    transformer = SimpleNamespace(patch_size=patch_size)
+    runtime_args = SimpleNamespace(output_type="pil", pin_cpu_memory=False, vae_cpu_offload=True)
+
+    if args.operation == "encode":
+        if int(os.environ.get("WORLD_SIZE", "1")) != 1:
+            raise ValueError("The encode benchmark is single-rank; use one process.")
+        frames = np.random.default_rng(args.seed).integers(
+            0,
+            256,
+            size=(args.num_frames, args.height, args.width, 3),
+            dtype=np.uint8,
+        )
+        stage = MiniMaxH3LatentPreparationStage(
+            transformer=transformer,
+            vae=vae,
+            audio_vae=None,
+            scheduler=None,
+            ref2va=True,
+        )
+
+        def run_once():
+            reference = MiniMaxH3PreparedReference(media_type="video", frames=frames)
+            vae.to(device)
+            try:
+                return stage._encode_visual_rows([reference], device)[0]
+            finally:
+                vae.to("cpu")
+
+        return run_once, {
+            "input": "NumPy PCG64 CPU uint8 RGB pixels",
+            "input_shape": [1, 3, args.num_frames, args.height, args.width],
+            "boundary": "before VAE CPU-to-GPU transfer through post-encode VAE CPU offload",
+        }
+
+    latent_shape = (
+        1,
+        vae.latent_channels,
+        video_latent_num_frames(args.num_frames),
+        args.height // vae.spatial_compression_ratio,
+        args.width // vae.spatial_compression_ratio,
+    )
+    generator = torch.Generator(device=device).manual_seed(args.seed)
+    latents = torch.randn(latent_shape, generator=generator, device=device, dtype=torch.float32)
+    rows = patchify_video_latents(latents, patch_size)
+    layout = _make_layout(rows, latent_shape)
+    stage = MiniMaxH3VideoDecodingStage(vae, transformer)
+
+    def run_once():
+        batch = ForwardBatch(data_type="video", latents=rows, raw_latent_shape=latent_shape)
+        batch.extra[MINIMAX_H3_LAYOUT_KEY] = layout
+        return stage.forward(batch, runtime_args).output
+
+    return run_once, {
+        "input": "PyTorch Philox normalized FP32 latents",
+        "input_shape": list(latent_shape),
+        "boundary": "MiniMaxH3VideoDecodingStage.forward including VAE CPU-to-GPU transfer and CPU offload",
+    }
+
+
+def _measure(run_once, device, warmups: int, repetitions: int) -> dict:
+    import torch
+
+    for _ in range(warmups):
+        with torch.inference_mode():
+            result = run_once()
+        torch.cuda.synchronize(device)
+        del result
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    records = []
+    for repetition in range(repetitions):
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize(device)
+        start_allocated = torch.cuda.memory_allocated(device)
+        start_reserved = torch.cuda.memory_reserved(device)
+        torch.cuda.reset_peak_memory_stats(device)
+        started = time.perf_counter()
+        with torch.inference_mode():
+            result = run_once()
+        torch.cuda.synchronize(device)
+        elapsed = time.perf_counter() - started
+        records.append({
+            "repetition": repetition,
+            "elapsed_seconds": elapsed,
+            "start_allocated_bytes": start_allocated,
+            "peak_allocated_bytes": torch.cuda.max_memory_allocated(device),
+            "incremental_peak_allocated_bytes": torch.cuda.max_memory_allocated(device) - start_allocated,
+            "start_reserved_bytes": start_reserved,
+            "peak_reserved_bytes": torch.cuda.max_memory_reserved(device),
+            "incremental_peak_reserved_bytes": torch.cuda.max_memory_reserved(device) - start_reserved,
+            "output_shape": list(result.shape),
+        })
+        del result
+    return {
+        "repetitions": records,
+        "median_elapsed_seconds": statistics.median(record["elapsed_seconds"] for record in records),
+        "max_peak_allocated_bytes": max(record["peak_allocated_bytes"] for record in records),
+        "max_incremental_peak_allocated_bytes": max(
+            record["incremental_peak_allocated_bytes"] for record in records),
+        "max_peak_reserved_bytes": max(record["peak_reserved_bytes"] for record in records),
+        "max_incremental_peak_reserved_bytes": max(
+            record["incremental_peak_reserved_bytes"] for record in records),
+    }
+
+
+def main() -> None:
+    args = _parse_args()
+    if args.warmups < 0 or args.repetitions < 1:
+        raise ValueError("warmups must be non-negative and repetitions must be positive.")
+    source_root = args.source_root.resolve()
+    sys.path.insert(0, str(source_root))
+
+    import torch
+    import torch.distributed as dist
+
+    import fastvideo
+    from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig
+    from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
+    from fastvideo.models.loader.component_loader import VAELoader
+
+    imported_root = Path(fastvideo.__file__).resolve().parents[1]
+    if imported_root != source_root:
+        raise RuntimeError(f"Imported FastVideo from {imported_root}, expected {source_root}.")
+
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29673")
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("LOCAL_RANK", "0")
+    world_size = int(os.environ["WORLD_SIZE"])
+    maybe_init_distributed_environment_and_model_parallel(1, world_size)
+    rank = dist.get_rank()
+    device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
+
+    component_dir = args.model_root.resolve() / "vae"
+    config_path = component_dir / "config.json"
+    index_path = component_dir / "diffusion_pytorch_model.safetensors.index.json"
+    for path in (config_path, index_path):
+        if not path.is_file():
+            raise FileNotFoundError(path)
+
+    loader_args = SimpleNamespace(
+        pipeline_config=MiniMaxH3PipelineConfig(),
+        model_paths={},
+        vae_cpu_offload=True,
+    )
+    vae = VAELoader().load(str(component_dir), loader_args)
+    parameter_bytes = sum(parameter.numel() * parameter.element_size() for parameter in vae.parameters())
+    run_once, workload = _build_operation(args, vae, device)
+    rank_result = _measure(run_once, device, args.warmups, args.repetitions)
+    rank_result.update({
+        "rank": rank,
+        "device_index": device.index,
+        "device_name": torch.cuda.get_device_name(device),
+        "device_total_memory_bytes": torch.cuda.get_device_properties(device).total_memory,
+    })
+    rank_results = [None] * world_size
+    dist.all_gather_object(rank_results, rank_result)
+
+    if rank == 0:
+        result = {
+            "schema_version": 1,
+            "exit_status": 0,
+            "revision_label": args.revision_label,
+            "source_root": str(source_root),
+            "source_git_head": _git(source_root, "rev-parse", "HEAD"),
+            "source_tracked_status": _git(source_root, "status", "--short", "--untracked-files=no"),
+            "benchmark_script_sha256": _sha256(Path(__file__).resolve()),
+            "model_root": str(args.model_root.resolve()),
+            "model_snapshot": _model_snapshot(args.model_root),
+            "model_config_sha256": _sha256(config_path),
+            "model_index_sha256": _sha256(index_path),
+            "vae_parameter_bytes": parameter_bytes,
+            "operation": args.operation,
+            "vae_cpu_offload": True,
+            "pin_cpu_memory": False,
+            "warmups": args.warmups,
+            "measurement_repetitions": args.repetitions,
+            "seed": args.seed,
+            "workload": workload,
+            "world_size": world_size,
+            "torch_version": torch.__version__,
+            "cuda_version": torch.version.cuda,
+            "allocator_config": os.environ.get("PYTORCH_ALLOC_CONF")
+            or os.environ.get("PYTORCH_CUDA_ALLOC_CONF"),
+            "rank_results": rank_results,
+            "sum_rank_local_max_peak_allocated_bytes": sum(
+                item["max_peak_allocated_bytes"] for item in rank_results),
+            "sum_rank_local_max_incremental_peak_allocated_bytes": sum(
+                item["max_incremental_peak_allocated_bytes"] for item in rank_results),
+            "sum_rank_local_max_peak_reserved_bytes": sum(
+                item["max_peak_reserved_bytes"] for item in rank_results),
+            "sum_rank_local_max_incremental_peak_reserved_bytes": sum(
+                item["max_incremental_peak_reserved_bytes"] for item in rank_results),
+        }
+        print("MINIMAX_H3_VAE_MEMORY=" + json.dumps(result, sort_keys=True), flush=True)
+
+    cleanup_dist_env_and_memory()
+
+
+if __name__ == "__main__":
+    main()

tests/local_tests/vaes/test_minimax_h3_video_vae_parity.py

@@ -62,8 +62,9 @@ def _require_assets() -> tuple[torch.device, Path]:
 
 def _load_official(component_dir: Path, device: torch.device) -> torch.nn.Module:
     from diffusers.models.autoencoders.autoencoder_kl_minimax_h3 import AutoencoderKLMiniMaxH3
-    from tests.local_tests.minimax_h3._reference import assert_reference_source
+    from tests.local_tests.minimax_h3._reference import assert_reference_revision, assert_reference_source
 
+    assert_reference_revision()
     assert_reference_source(
         AutoencoderKLMiniMaxH3,
         "src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py",
@@ -101,14 +102,35 @@ def _load_fastvideo(component_dir: Path) -> torch.nn.Module:
     return model
 
 
-def _make_video() -> torch.Tensor:
+def _make_pixels() -> torch.Tensor:
     generator = torch.Generator(device="cpu").manual_seed(20260803)
     # Two logical 17-frame clips after padding are required for H3's three-token
     # temporal overlap contract; 32 px remains safely above reflect-pad minima.
-    return torch.randn(1, 3, 22, 32, 32, generator=generator, dtype=torch.float32)
+    return torch.randint(0, 256, (1, 3, 22, 32, 32), generator=generator, dtype=torch.uint8)
 
 
-def _run(model: torch.nn.Module, video: torch.Tensor) -> dict[str, torch.Tensor]:
+def _normalize_pixels(pixels: torch.Tensor, device: torch.device) -> torch.Tensor:
+    video = pixels.to(device=device, dtype=torch.float32).div_(255.0)
+    pixel_mean = torch.tensor((0.485, 0.456, 0.406), device=device).view(1, -1, 1, 1, 1)
+    pixel_std = torch.tensor((0.229, 0.224, 0.225), device=device).view(1, -1, 1, 1, 1)
+    return (video - pixel_mean) / pixel_std
+
+
+def _run_encode_pixels(model: torch.nn.Module, pixels: torch.Tensor) -> dict[str, torch.Tensor]:
+    with torch.inference_mode():
+        posterior = model.encode_pixels(pixels, return_dict=False)[0]
+    return {
+        "mean": posterior.mean.detach().cpu(),
+        "logvar": posterior.logvar.detach().cpu(),
+    }
+
+
+def _run(
+    model: torch.nn.Module,
+    video: torch.Tensor,
+    *,
+    stream_output: bool = False,
+) -> dict[str, torch.Tensor]:
     with torch.inference_mode():
         posterior = model.encode(video, return_dict=False)[0]
         latents = posterior.mode()
@@ -121,12 +143,25 @@ def _run(model: torch.nn.Module, video: torch.Tensor) -> dict[str, torch.Tensor]
                                dtype=latents.dtype).view(1, -1, 1, 1, 1)
             normalized = (latents - mean) / std
         decoded = model.decode(latents, return_dict=False)[0]
-    return {
+        pixel_mean = torch.tensor((0.485, 0.456, 0.406), device=decoded.device,
+                                  dtype=torch.float32).view(1, -1, 1, 1, 1)
+        pixel_std = torch.tensor((0.229, 0.224, 0.225), device=decoded.device,
+                                 dtype=torch.float32).view(1, -1, 1, 1, 1)
+        pixels = (decoded.float() * pixel_std + pixel_mean).clamp_(0, 1)
+        streamed = None
+        if stream_output:
+            streamed = torch.empty(model.decoded_pixel_shape(latents.shape), dtype=torch.float32, device="cpu")
+            model.decode_to_pixels(latents, streamed)
+    result = {
         "mean": posterior.mean.detach().cpu(),
         "logvar": posterior.logvar.detach().cpu(),
         "normalized": normalized.detach().cpu(),
         "decoded": decoded.detach().cpu(),
+        "pixels": pixels.detach().cpu(),
     }
+    if streamed is not None:
+        result["streamed"] = streamed
+    return result
 
 
 def _reclaim_vram() -> None:
@@ -151,15 +186,16 @@ def _assert_tensor_parity(name: str, actual: torch.Tensor, expected: torch.Tenso
 def test_minimax_h3_video_vae_parity() -> None:
     """Match posterior, normalization, geometry, and deterministic decode."""
     device, component_dir = _require_assets()
-    video = _make_video()
+    pixels = _make_pixels()
 
     official = _load_official(component_dir, device)
-    expected = _run(official, video.to(device))
+    expected = _run(official, _normalize_pixels(pixels, device))
     del official
     _reclaim_vram()
 
     fastvideo = _load_fastvideo(component_dir)
-    actual = _run(fastvideo, video.to(device))
+    actual = _run(fastvideo, _normalize_pixels(pixels, device), stream_output=True)
+    streaming_encode = _run_encode_pixels(fastvideo, pixels)
     assert fastvideo.temporal_compression_ratio == 4
     assert fastvideo.spatial_compression_ratio == 16
     del fastvideo
@@ -167,5 +203,8 @@ def test_minimax_h3_video_vae_parity() -> None:
 
     _assert_tensor_parity("video_vae.mean", actual["mean"], expected["mean"], 0.0)
     _assert_tensor_parity("video_vae.logvar", actual["logvar"], expected["logvar"], 0.0)
+    _assert_tensor_parity("video_vae.streaming_encode.mean", streaming_encode["mean"], expected["mean"], 0.0)
+    _assert_tensor_parity("video_vae.streaming_encode.logvar", streaming_encode["logvar"], expected["logvar"], 0.0)
     _assert_tensor_parity("video_vae.normalized", actual["normalized"], expected["normalized"], 0.0)
     _assert_tensor_parity("video_vae.decode", actual["decoded"], expected["decoded"], 0.0)
+    _assert_tensor_parity("video_vae.streaming_decode", actual["streamed"], expected["pixels"], 0.0)