From 00b2eb9399d587204502f248b415498551cac059 Mon Sep 17 00:00:00 2001
From: david6666666 <530634352@qq.com>
Date: Thu, 17 Sep 2026 09:20:18 +0000
Subject: [PATCH] Revert "[Bugfix][MiniMax-H3] Isolate keyframe VAE encoding
 from tile collectives"

This reverts commit 2aac8bd15aaa03b0b10926490380e2f3fe2b02e5.

Signed-off-by: david6666666 <530634352@qq.com>
---
 .../minimax_h3/test_minimax_h3_contract.py    | 19 ++++---------------
 vllm_omni/diffusion/models/minimax_h3/vae.py  | 14 ++++++--------
 2 files changed, 10 insertions(+), 23 deletions(-)

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 1ddc89f6e..0ac14da15 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
@@ -3,7 +3,7 @@
 
 import json
 import sys
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from multiprocessing.reduction import ForkingPickler
 from types import SimpleNamespace
 from typing import Any, TypeVar
@@ -2811,17 +2811,6 @@ 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__()
@@ -2829,7 +2818,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 == ["rank-local-enter", "enter"]
+            assert events == ["enter"]
             assert image is input_image
             assert use_fp16_latent
             assert not self.parallel_tiling
@@ -2847,7 +2836,7 @@ def test_keyframe_encode_enters_backend_context(monkeypatch, fail_encode):
         "latents_std": [1.0],
     }
     video_vae.parallel_size = 4
-    monkeypatch.setattr(video_vae, "_rank_local_tiling", track_rank_local)
+    monkeypatch.setattr(video_vae, "_encoder_tiling_context", lambda _height, _width: nullcontext())
     monkeypatch.setattr(vae_module, "_minimax_h3_keyframe_encode_context", track_context)
 
     if fail_encode:
@@ -2857,7 +2846,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 == ["rank-local-enter", "enter", "encode", "exit", "rank-local-exit"]
+    assert events == ["enter", "encode", "exit"]
     assert video_vae.model.parallel_tiling
 
 
diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py
index f20def08f..68c340c00 100644
--- a/vllm_omni/diffusion/models/minimax_h3/vae.py
+++ b/vllm_omni/diffusion/models/minimax_h3/vae.py
@@ -602,13 +602,11 @@ 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
-        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()
+        # 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
         parameter = next(self.parameters())
         previous_dtype = parameter.dtype
         if previous_dtype != torch.float32:
@@ -616,7 +614,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
         devices = [parameter.device] if parameter.device.type != "cpu" else []
         try:
             with (
-                tiling_context,
+                self._encoder_tiling_context(image.height, image.width),
                 torch.random.fork_rng(
                     devices=devices,
                     device_type=parameter.device.type,