From 53df9fe62b3e05b388916351ecd758cb770e5cb4 Mon Sep 17 00:00:00 2001
From: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
Date: Thu, 20 Aug 2026 16:13:34 +0800
Subject: [PATCH] [Bugfix] Use dedicated WORLD group for distributed VAE
 communication (#6401)

Signed-off-by: david6666666 <530634352@qq.com>
---
 .../test_distributed_vae_executor.py          | 21 +++++++++++++++++++
 .../autoencoders/distributed_vae_executor.py  | 10 +++++----
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/tests/diffusion/distributed/test_distributed_vae_executor.py b/tests/diffusion/distributed/test_distributed_vae_executor.py
index c122e9105..dcaa8e01a 100644
--- a/tests/diffusion/distributed/test_distributed_vae_executor.py
+++ b/tests/diffusion/distributed/test_distributed_vae_executor.py
@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 from types import SimpleNamespace
 
 import pytest
@@ -58,6 +59,11 @@ class E2EOperator:
         return torch.cat(tiles, dim=0)
 
 
+@dataclass
+class FakeWorldGroup:
+    device_group: object
+
+
 class DummyMixin(DistributedVaeMixin):
     def __init__(self):
         self.use_tiling = True
@@ -74,6 +80,16 @@ def mock_dist(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(dist, "broadcast", lambda *args, **kwargs: None)
 
 
+@pytest.fixture(autouse=True)
+def mock_world_group(monkeypatch: pytest.MonkeyPatch):
+    group = object()
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_world_group",
+        lambda: FakeWorldGroup(device_group=group),
+    )
+    return group
+
+
 @pytest.fixture(autouse=True)
 def mock_dist_vae_executor(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(DistributedVaeExecutor, "gather_tensors", lambda self, x: [x])
@@ -85,6 +101,11 @@ def mock_dist_vae_executor(monkeypatch: pytest.MonkeyPatch):
 # ============================
 
 
+def test_uses_dedicated_world_device_group(mock_world_group):
+    executor = DistributedVaeExecutor()
+    assert executor.group is mock_world_group
+
+
 def test_balance_tasks():
     executor = DistributedVaeExecutor()
     workloads = [2, 5, 13, 8, 2, 4]  # each is 17
diff --git a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py
index 9d84cdf53..e8ae6bf32 100644
--- a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py
+++ b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py
@@ -9,6 +9,8 @@ import torch
 import torch.distributed as dist
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.distributed.parallel_state import get_world_group
+
 logger = init_logger(__name__)
 
 
@@ -43,10 +45,10 @@ class DistributedVaeExecutor:
     """
 
     def __init__(self):
-        # VAE tiles are scheduled across the complete worker WORLD.
-        self.group = None
-        self.world_size = dist.get_world_size()
-        self.rank = dist.get_rank()
+        # Use a dedicated process group spanning the complete worker WORLD.
+        self.group = get_world_group().device_group
+        self.world_size = dist.get_world_size(self.group)
+        self.rank = dist.get_rank(self.group)
         self.parallel_size = 1
         self.parallel_mode = "tile"