[Bugfix][MiniMax-H3] Stabilize keyframe VAE encoding (#7191)

Signed-off-by: david6666666 <530634352@qq.com>
Signed-off-by: David Chen <530634352@qq.com>

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
 

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: