[test]: add H3 VAE parity and memory benchmark

tests/local_tests/minimax_h3/README.md

@@ -59,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,11 +102,27 @@ 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 _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(
@@ -169,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), stream_output=True)
+    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
@@ -185,6 +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)