[Bugfix][MiniMax-H3] Isolate keyframe VAE encoding from tile collectives

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

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -3,7 +3,7 @@
 
 import json
 import sys
-from contextlib import contextmanager, nullcontext
+from contextlib import contextmanager
 from multiprocessing.reduction import ForkingPickler
 from types import SimpleNamespace
 from typing import Any, TypeVar
@@ -2811,6 +2811,17 @@ def test_keyframe_encode_enters_backend_context(monkeypatch, fail_encode):
         finally:
             events.append("exit")
 
+    @contextmanager
+    def track_rank_local():
+        assert video_vae.model.parallel_tiling
+        video_vae.model.parallel_tiling = False
+        events.append("rank-local-enter")
+        try:
+            yield
+        finally:
+            video_vae.model.parallel_tiling = True
+            events.append("rank-local-exit")
+
     class FakeModel(torch.nn.Module):
         def __init__(self):
             super().__init__()
@@ -2818,7 +2829,7 @@ def test_keyframe_encode_enters_backend_context(monkeypatch, fail_encode):
             self.parallel_tiling = True
 
         def encode_images(self, image, *, use_fp16_latent):
-            assert events == ["enter"]
+            assert events == ["rank-local-enter", "enter"]
             assert image is input_image
             assert use_fp16_latent
             assert not self.parallel_tiling
@@ -2836,7 +2847,7 @@ def test_keyframe_encode_enters_backend_context(monkeypatch, fail_encode):
         "latents_std": [1.0],
     }
     video_vae.parallel_size = 4
-    monkeypatch.setattr(video_vae, "_encoder_tiling_context", lambda _height, _width: nullcontext())
+    monkeypatch.setattr(video_vae, "_rank_local_tiling", track_rank_local)
     monkeypatch.setattr(vae_module, "_minimax_h3_keyframe_encode_context", track_context)
 
     if fail_encode:
@@ -2846,7 +2857,7 @@ def test_keyframe_encode_enters_backend_context(monkeypatch, fail_encode):
         rows = video_vae.encode_image(input_image)
         torch.testing.assert_close(rows, torch.ones(1, 4))
 
-    assert events == ["enter", "encode", "exit"]
+    assert events == ["rank-local-enter", "enter", "encode", "exit", "rank-local-exit"]
     assert video_vae.model.parallel_tiling
 
 

vllm_omni/diffusion/models/minimax_h3/vae.py

@@ -602,11 +602,13 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         if getattr(self, "decode_only", False):
             raise RuntimeError("MiniMax H3 decode-only video VAE cannot encode images")
         previous_parallel = self.model.parallel_tiling
-        # Keyframe conditioning is part of the official single-image VAE
-        # contract. Keep it rank-local even when video/reference encoding is
-        # configured for tile parallelism; the tiled collective changes the
-        # reduction order and can move H100 I2VA quality across the threshold.
-        self.model.parallel_tiling = False
+        parallel_size = int(getattr(self, "parallel_size", 1))
+        if parallel_size <= 1:
+            self.model.parallel_tiling = False
+        # A keyframe is a single-image conditioning input. Do not let it use
+        # the multi-rank tiled-VAE collective, whose reduction order depends
+        # on the physical GPU topology; video references keep tile parallelism.
+        tiling_context = self._rank_local_tiling() if parallel_size > 1 else nullcontext()
         parameter = next(self.parameters())
         previous_dtype = parameter.dtype
         if previous_dtype != torch.float32:
@@ -614,7 +616,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         devices = [parameter.device] if parameter.device.type != "cpu" else []
         try:
             with (
-                self._encoder_tiling_context(image.height, image.width),
+                tiling_context,
                 torch.random.fork_rng(
                     devices=devices,
                     device_type=parameter.device.type,