diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index cb12ce9e5..114f5eab3 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -3,6 +3,7 @@ import json import sys +from contextlib import contextmanager from multiprocessing.reduction import ForkingPickler from types import SimpleNamespace from typing import Any, TypeVar @@ -1550,6 +1551,97 @@ def test_video_vae_keeps_reference_fp32_weights(monkeypatch): assert next(video_vae.parameters()).dtype == torch.float32 +def test_keyframe_encode_pins_and_restores_cudnn_settings(): + from vllm_omni.diffusion.models.minimax_h3.vae import ( + _minimax_h3_keyframe_encode_context, + ) + + cudnn = torch.backends.cudnn + original = ( + cudnn.enabled, + cudnn.benchmark, + cudnn.deterministic, + cudnn.allow_tf32, + ) + try: + cudnn.enabled = False + cudnn.benchmark = True + cudnn.deterministic = False + cudnn.allow_tf32 = False + + with _minimax_h3_keyframe_encode_context(torch.device("cuda")): + assert cudnn.enabled + assert not cudnn.benchmark + assert cudnn.deterministic + assert cudnn.allow_tf32 + + assert not cudnn.enabled + assert cudnn.benchmark + assert not cudnn.deterministic + assert not cudnn.allow_tf32 + finally: + ( + cudnn.enabled, + cudnn.benchmark, + cudnn.deterministic, + cudnn.allow_tf32, + ) = original + + +@pytest.mark.parametrize("fail_encode", [False, True]) +def test_keyframe_encode_enters_backend_context(monkeypatch, fail_encode): + from PIL import Image + + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + events = [] + input_image = Image.new("RGB", (32, 32)) + + @contextmanager + def track_context(device): + assert device == torch.device("cpu") + events.append("enter") + try: + yield + finally: + events.append("exit") + + class FakeModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(1)) + self.parallel_tiling = True + + def encode_images(self, image, *, use_fp16_latent): + assert events == ["enter"] + assert image is input_image + assert use_fp16_latent + events.append("encode") + if fail_encode: + raise RuntimeError("encode failed") + return [torch.ones(1, 1, 2, 2)] + + video_vae = object.__new__(vae_module.MiniMaxH3VideoVAE) + torch.nn.Module.__init__(video_vae) + video_vae.model = FakeModel() + video_vae.config_dict = { + "latent_channels": 1, + "latents_mean": [0.0], + "latents_std": [1.0], + } + monkeypatch.setattr(vae_module, "_minimax_h3_keyframe_encode_context", track_context) + + if fail_encode: + with pytest.raises(RuntimeError, match="encode failed"): + video_vae.encode_image(input_image) + else: + rows = video_vae.encode_image(input_image) + torch.testing.assert_close(rows, torch.ones(1, 4)) + + assert events == ["enter", "encode", "exit"] + assert video_vae.model.parallel_tiling + + def test_video_vae_can_load_on_cpu_for_staged_gpu_residency(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index 2e830fa04..e98659409 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -38,6 +38,27 @@ MINIMAX_H3_AUDIO_CHANNELS = 2 logger = init_logger(__name__) +@contextmanager +def _minimax_h3_keyframe_encode_context( + device: torch.device, +) -> Iterator[None]: + if device.type != "cuda": + yield + return + + # The official keyframe latent uses cuDNN's TF32 convolution path. The + # default non-deterministic algorithm can select numerically different + # reductions on H100s, and the difference is amplified by the denoiser. + # Pin both the algorithm and math mode for this sensitive encode only. + with torch.backends.cudnn.flags( + enabled=True, + benchmark=False, + deterministic=True, + allow_tf32=True, + ): + yield + + def _load_component_config(component_path: str) -> dict[str, Any]: config_path = Path(component_path) / "config.json" config = json.loads(config_path.read_text(encoding="utf-8")) @@ -280,10 +301,11 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): for device in devices: with self.device_module.device(device): self.device_module.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED) - latent = self.model.encode_images( - image, - use_fp16_latent=True, - )[0] + with _minimax_h3_keyframe_encode_context(parameter.device): + latent = self.model.encode_images( + image, + use_fp16_latent=True, + )[0] finally: self.model.parallel_tiling = previous_parallel if previous_dtype != torch.float32: