[FEATURE] Tile/Patch parallelism refactor for easily support other models (#1366)

Signed-off-by: dengyunyang <584797741@qq.com>

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py

@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+from typing import Any
+
+import torch
+from diffusers.models.autoencoders import AutoencoderKL as Diffusers_AutoencoderKL
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedOperator,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+
+# from vllm_omni.diffusion.models.nextstep_1_1.modeling_flux_vae import AutoencoderKL as Next_Step_AutoencoderKL
+
+logger = init_logger(__name__)
+
+
+# We use base class because some model re-implement AutoencoderKL, but split is share.
+class DistributedAutoencoderKL_base(DistributedVaeMixin):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+
+    def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        # mostly copy from AutoencoderKL
+        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_sample_min_size - blend_extent
+
+        # Split z into overlapping 64x64 tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        tiletask_list = []
+        for i in range(0, z.shape[2], overlap_size):
+            for j in range(0, z.shape[3], overlap_size):
+                tile = z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
+                tiletask_list.append(
+                    TileTask(
+                        len(tiletask_list),
+                        (i // overlap_size, j // overlap_size),
+                        tile,
+                        workload=tile.shape[2] * tile.shape[3],
+                    )
+                )
+
+        tile_spec = {
+            "blend_extent": blend_extent,
+            "row_limit": row_limit,
+        }
+        grid_spec = GridSpec(
+            split_dims=(2, 3),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec=tile_spec,
+        )
+        return tiletask_list, grid_spec
+
+    def tile_exec(self, task: TileTask) -> torch.Tensor:
+        """Decode a single latent tile into RGB space."""
+        tile = task.tensor
+        if self.config.use_post_quant_conv:
+            tile = self.post_quant_conv(tile)
+        decoded = self.decoder(tile)
+        return decoded
+
+    def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
+        """Merge decoded tiles into a full image."""
+
+        grid_h, grid_w = grid_spec.grid_shape
+        result_rows = []
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_extent"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_extent"])
+                result_row.append(tile[:, :, : grid_spec.tile_spec["row_limit"], : grid_spec.tile_spec["row_limit"]])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        dec = torch.cat(result_rows, dim=2)
+        return dec
+
+    def patch_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        overlap_latent = int(self.tile_latent_min_size * float(self.tile_overlap_factor))
+        halo_base = max(0, overlap_latent // 2)
+
+        _, _, latent_h, latent_w = z.shape
+        scale = int(2 ** (len(self.config.block_out_channels) - 1))
+        max_parallel_size = self.distributed_decoder.parallel_size
+
+        root = int(math.sqrt(max_parallel_size))
+        for rows in range(root, 0, -1):
+            if max_parallel_size % rows == 0:
+                grid_rows, grid_cols = rows, max_parallel_size // rows
+                break
+        tiletask_list = []
+        halo_size = dict()
+        for i in range(grid_rows):
+            for j in range(grid_cols):
+                h0 = (i * latent_h) // grid_rows
+                h1 = ((i + 1) * latent_h) // grid_rows
+                w0 = (j * latent_w) // grid_cols
+                w1 = ((j + 1) * latent_w) // grid_cols
+
+                core_h = max(0, h1 - h0)
+                core_w = max(0, w1 - w0)
+                halo = max(halo_base, min(core_h, core_w) // 2)
+                ph0 = max(0, h0 - halo)
+                ph1 = min(latent_h, h1 + halo)
+                pw0 = max(0, w0 - halo)
+                pw1 = min(latent_w, w1 + halo)
+                tile = z[:, :, ph0:ph1, pw0:pw1]
+                tiletask_list.append(TileTask(len(tiletask_list), (i, j), tile, tile.shape[2] * tile.shape[3]))
+                halo_size[(i, j)] = {
+                    "up": h0 - ph0,
+                    "down": ph1 - h1,
+                    "left": w0 - pw0,
+                    "right": pw1 - w1,
+                }
+
+        tile_spec = {
+            "halo_size": halo_size,
+            "scale": scale,
+        }
+        grid_spec = GridSpec(
+            split_dims=(2, 3),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec=tile_spec,
+            output_dtype=self.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def patch_exec(self, task: TileTask) -> torch.Tensor:
+        return self.tile_exec(task)
+
+    def patch_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
+        grid_h, grid_w = grid_spec.grid_shape
+        result_rows = []
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                halo = grid_spec.tile_spec["halo_size"][(i, j)]
+                scale = grid_spec.tile_spec["scale"]
+
+                tile = coord_tensor_map[(i, j)]
+                halo_up = halo["up"] * scale
+                halo_down = halo["down"] * scale
+                halo_left = halo["left"] * scale
+                halo_right = halo["right"] * scale
+
+                core_tile = tile[
+                    :,
+                    :,
+                    halo_up : (None if halo_down == 0 else -halo_down),
+                    halo_left : (None if halo_right == 0 else -halo_right),
+                ]
+                result_row.append(core_tile)
+            result_rows.append(torch.cat(result_row, dim=3))
+        dec = torch.cat(result_rows, dim=2)
+        return dec
+
+    def _strategy_select(self, z: torch.Tensor):
+        tile_latent_min_size = getattr(self, "tile_latent_min_size", None)
+        tile_overlap_factor = getattr(self, "tile_overlap_factor", None)
+        if tile_latent_min_size is None or tile_overlap_factor is None:
+            return None, None, None
+        if z.shape[-1] > tile_latent_min_size or z.shape[-2] > tile_latent_min_size:
+            return self.tile_split, self.tile_exec, self.tile_merge
+
+        return self.patch_split, self.patch_exec, self.patch_merge
+
+    # Normally, we should override tiled_decode. However, since we also need to
+    # support the patch split strategy, we override decode instead.
+    def decode(self, z: torch.Tensor, return_dict: bool = True, *args: Any, **kwargs: Any):
+        if not self.is_distributed_enabled():
+            return super().decode(z, return_dict=return_dict, *args, **kwargs)
+
+        split, exec, merge = self._strategy_select(z)
+
+        if split is not None:
+            strategy = "tile" if split == self.tile_split else "patch"
+            logger.info(f"Decode run with distributed executor, split strategy is {strategy}")
+            result = self.distributed_decoder.execute(
+                z, DistributedOperator(split=split, exec=exec, merge=merge), broadcast_result=False
+            )
+            if not return_dict:
+                return (result,)
+
+            from diffusers.models.autoencoders.vae import DecoderOutput
+
+            return DecoderOutput(sample=result)
+        else:
+            return super().decode(z, return_dict=return_dict, *args, **kwargs)
+
+
+class DistributedAutoencoderKL(DistributedAutoencoderKL_base, Diffusers_AutoencoderKL):
+    pass
+
+
+# Next_Step_AutoencoderKL not support tiling now, so we currently disable it.
+
+# class DistributedAutoencoderKL_NextStep(DistributedAutoencoderKL_base, Next_Step_AutoencoderKL):
+#    pass

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+from diffusers.models.autoencoders import AutoencoderKLQwenImage
+from diffusers.models.autoencoders.vae import DecoderOutput
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedOperator,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+
+logger = init_logger(__name__)
+
+
+class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+
+    def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        # mostly copy from AutoencoderKL
+        _, _, num_frames, height, width = z.shape
+        sample_height = height * self.spatial_compression_ratio
+        sample_width = width * self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+        blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+        # Split z into overlapping tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        tiletask_list = []
+        for i in range(0, height, tile_latent_stride_height):
+            for j in range(0, width, tile_latent_stride_width):
+                time_list = []
+                for k in range(num_frames):
+                    self._conv_idx = [0]
+                    tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
+                    time_list.append(tile)
+                tiletask_list.append(
+                    TileTask(
+                        len(tiletask_list),
+                        (i // tile_latent_stride_height, j // tile_latent_stride_width),
+                        time_list,
+                        workload=time_list[0].shape[3] * time_list[0].shape[4],
+                    )
+                )
+        tile_spec = {
+            "sample_height": sample_height,
+            "sample_width": sample_width,
+            "blend_height": blend_height,
+            "blend_width": blend_width,
+        }
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec=tile_spec,
+            output_dtype=self.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def tile_exec(self, task: TileTask) -> torch.Tensor:
+        """Decode a single latent tile into RGB space."""
+        self.clear_cache()
+        time = []
+        for k in range(len(task.tensor)):
+            self._conv_idx = [0]
+            tile = self.post_quant_conv(task.tensor[k])
+            decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx)
+            time.append(decoded)
+        result = torch.cat(time, dim=2)
+        return result
+
+    def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
+        """Merge decoded tiles into a full image."""
+        grid_h, grid_w = grid_spec.grid_shape
+        self.clear_cache()
+
+        result_rows = []
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+                result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width])
+            result_rows.append(torch.cat(result_row, dim=-1))
+        dec = torch.cat(result_rows, dim=3)[
+            :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
+        ]
+        return dec
+
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
+        if not self.is_distributed_enabled():
+            return super().tiled_decode(z, return_dict=return_dict)
+
+        logger.info("Decode run with distributed executor")
+        result = self.distributed_decoder.execute(
+            z,
+            DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
+            broadcast_result=True,
+        )
+        if not return_dict:
+            return (result,)
+
+        return DecoderOutput(sample=result)

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+from diffusers.models.autoencoders import AutoencoderKLWan
+from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify
+from diffusers.models.autoencoders.vae import DecoderOutput
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedOperator,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+
+logger = init_logger(__name__)
+
+
+class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+
+    def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        _, _, num_frames, height, width = z.shape
+        sample_height = height * self.spatial_compression_ratio
+        sample_width = width * self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+        tile_sample_stride_height = self.tile_sample_stride_height
+        tile_sample_stride_width = self.tile_sample_stride_width
+        if self.config.patch_size is not None:
+            sample_height = sample_height // self.config.patch_size
+            sample_width = sample_width // self.config.patch_size
+            tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size
+            tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size
+            blend_height = self.tile_sample_min_height // self.config.patch_size - tile_sample_stride_height
+            blend_width = self.tile_sample_min_width // self.config.patch_size - tile_sample_stride_width
+        else:
+            blend_height = self.tile_sample_min_height - tile_sample_stride_height
+            blend_width = self.tile_sample_min_width - tile_sample_stride_width
+
+        tiletask_list = []
+        for i in range(0, height, tile_latent_stride_height):
+            for j in range(0, width, tile_latent_stride_width):
+                time_list = []
+                for k in range(num_frames):
+                    self._conv_idx = [0]
+                    tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
+                    time_list.append(tile)
+                tiletask_list.append(
+                    TileTask(
+                        len(tiletask_list),
+                        (i // tile_latent_stride_height, j // tile_latent_stride_width),
+                        time_list,
+                        workload=time_list[0].shape[3] * time_list[0].shape[4],
+                    )
+                )
+        tile_spec = {
+            "sample_height": sample_height,
+            "sample_width": sample_width,
+            "blend_height": blend_height,
+            "blend_width": blend_width,
+            "tile_sample_stride_height": tile_sample_stride_height,
+            "tile_sample_stride_width": tile_sample_stride_width,
+        }
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec=tile_spec,
+            output_dtype=self.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def tile_exec(self, task: TileTask) -> torch.Tensor:
+        """Decode a single latent tile into RGB space."""
+        self.clear_cache()
+        time = []
+        for k in range(len(task.tensor)):
+            self._conv_idx = [0]
+            tile = self.post_quant_conv(task.tensor[k])
+            decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0))
+            time.append(decoded)
+        result = torch.cat(time, dim=2)
+        return result
+
+    def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
+        """Merge decoded tiles into a full image."""
+        grid_h, grid_w = grid_spec.grid_shape
+        result_rows = []
+        self.clear_cache()
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+                result_row.append(
+                    tile[
+                        :,
+                        :,
+                        :,
+                        : grid_spec.tile_spec["tile_sample_stride_height"],
+                        : grid_spec.tile_spec["tile_sample_stride_width"],
+                    ]
+                )
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        dec = torch.cat(result_rows, dim=3)[
+            :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
+        ]
+
+        if self.config.patch_size is not None:
+            dec = unpatchify(dec, patch_size=self.config.patch_size)
+
+        dec = torch.clamp(dec, min=-1.0, max=1.0)
+        return dec
+
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
+        if not self.is_distributed_enabled():
+            return super().tiled_decode(z, return_dict=return_dict)
+
+        logger.info("Decode run with distributed executor")
+        result = self.distributed_decoder.execute(
+            z,
+            DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
+            broadcast_result=False,
+        )
+        if not return_dict:
+            return (result,)
+
+        return DecoderOutput(sample=result)

vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py

@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import torch
+import torch.distributed as dist
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.parallel_state import get_dit_group
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class GridSpec:
+    """The Grid shape split"""
+
+    split_dims: tuple[int, ...]  # For example: (3,4) for (B, C, T, H, W), (2,3) for (B, C, H, W)
+    grid_shape: tuple[int, ...]  # For example: (nh, nw) for (B, C, T, H, W), (nh, nw) for (B, C, H, W)
+    tile_spec: dict = field(default_factory=dict)
+    output_dtype: torch.dtype | None = None
+
+
+@dataclass
+class TileTask:
+    tile_id: int
+    grid_coord: tuple[int, ...]  # The coordinate of the tile in GridSpec.grid_shape
+    tensor: torch.Tensor | list[torch.Tensor]  # The tile tensor
+    workload: int | float = 1
+
+
+@dataclass
+class DistributedOperator:
+    split: callable
+    exec: callable
+    merge: callable
+
+
+class DistributedVaeExecutor:
+    """
+    Abstract util class for distributed patch/tile parallel VAE decoding.
+    """
+
+    def __init__(self):
+        self.group = get_dit_group()
+        self.world_size = dist.get_world_size(self.group)
+        self.rank = dist.get_rank(self.group)
+        self.parallel_size = 1
+
+    def set_parallel_size(self, parallel_size: int):
+        self.parallel_size = parallel_size
+
+    def gather_tensors(self, tensor: torch.Tensor):
+        gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)] if self.rank == 0 else None
+        dist.gather(tensor, gather_list=gather_list, dst=0, group=self.group)
+        return gather_list
+
+    def broadcast_tensor(self, tensor: torch.Tensor):
+        dist.broadcast(tensor, src=0, group=self.group)
+        return tensor
+
+    def _compute_global_padding_shape(self, local_results, output_ndim: int, device):
+        local_tile_max_dims = [0] * output_ndim
+        for _, tile_tensor in local_results:
+            for dim_idx, dim_size in enumerate(tile_tensor.shape):
+                local_tile_max_dims[dim_idx] = max(local_tile_max_dims[dim_idx], dim_size)
+        local_shape_stat = torch.tensor([len(local_results), *local_tile_max_dims], device=device)
+        dist.all_reduce(
+            local_shape_stat,
+            op=dist.ReduceOp.MAX,
+            group=self.group,
+        )
+        return local_shape_stat.tolist()
+
+    def _pack_local_tiles(self, local_results, global_padding_shape, grid_spec, device, dtype):
+        tile_tensor = torch.zeros(global_padding_shape, device=device, dtype=dtype)
+        meta_tensor = torch.full(
+            (global_padding_shape[0], len(grid_spec.split_dims) + 1), -1, device=device, dtype=torch.int64
+        )
+        for idx, (tid, t_tensor) in enumerate(local_results):
+            meta_tensor[idx, 0] = tid
+            for i, dim in enumerate(grid_spec.split_dims):
+                meta_tensor[idx, i + 1] = t_tensor.shape[dim]
+            slices = tuple(slice(0, s) for s in t_tensor.shape)
+            tile_tensor[idx][slices] = t_tensor
+        return tile_tensor, meta_tensor
+
+    def _unpack_tiles(self, meta_gather, tile_gather, grid_spec, tid_coord_map):
+        coord_tensor_map = {}
+        for r in range(self.world_size):
+            meta_src = meta_gather[r]
+            tiles_src = tile_gather[r]
+            for idx in range(meta_src.shape[0]):
+                tid = int(meta_src[idx, 0])
+                if tid < 0:
+                    continue
+                slices = [slice(None)] * tiles_src[idx].ndim
+                for i, dim in enumerate(grid_spec.split_dims):
+                    slices[dim] = slice(0, int(meta_src[idx, i + 1]))
+                slices = tuple(slices)
+                coord_tensor_map[tid_coord_map[tid]] = tiles_src[idx][slices]
+
+        return coord_tensor_map
+
+    def _balance_tasks(self, task_list, num_rank):
+        workloads = [0] * num_rank
+        assigned = [[] for _ in range(num_rank)]
+
+        for task in sorted(task_list, key=lambda t: t.workload, reverse=True):
+            r = workloads.index(min(workloads))
+            assigned[r].append(task)
+            workloads[r] += task.workload
+
+        return assigned
+
+    def execute(self, z: torch.Tensor, operator: DistributedOperator, broadcast_result: bool = True):
+        pp_size = min(self.parallel_size, self.world_size)
+
+        # 1. Split into tiles
+        tiletask_list, grid_spec = operator.split(z)
+        tid_coord_map = {task.tile_id: task.grid_coord for task in tiletask_list}
+
+        # 2. local decode
+        assigned = self._balance_tasks(tiletask_list, pp_size)
+        local_tasks = assigned[self.rank] if pp_size <= self.world_size else []
+        local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks]
+
+        # 3. compute shape per rank
+        global_padding_shape = self._compute_global_padding_shape(local_results, z.ndim, z.device)
+
+        # 5. prepare tile tensors
+        output_dtype = grid_spec.output_dtype if grid_spec.output_dtype is not None else z.dtype
+        local_tile_tensor, local_meta_tensor = self._pack_local_tiles(
+            local_results, global_padding_shape, grid_spec, z.device, output_dtype
+        )
+
+        # 6. gather tiles & meta
+        meta_gather = self.gather_tensors(local_meta_tensor)
+        tile_gather = self.gather_tensors(local_tile_tensor)
+
+        if self.rank != 0:
+            result = torch.empty(0, device=z.device)  # Dummy return for non-zero ranks
+        else:
+            # 7. reconstruct full tensor (rank 0)
+            coord_tensor_map = self._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map)
+            result = operator.merge(coord_tensor_map, grid_spec)
+
+        if broadcast_result:
+            result = self._sync_final_result(result, z.ndim, z.device, output_dtype)
+        return result
+
+    def _sync_final_result(self, rank0_result, output_ndim, output_device, output_dtype):
+        shape_tensor = torch.empty((output_ndim,), device=output_device, dtype=torch.int64)
+        if self.rank == 0:
+            shape_tensor.copy_(torch.tensor(tuple(rank0_result.shape), device=output_device, dtype=torch.int64))
+        dist.broadcast(shape_tensor, src=0, group=self.group)
+
+        if self.rank != 0:
+            sync_result = torch.empty(tuple(shape_tensor.tolist()), device=output_device, dtype=output_dtype)
+        else:
+            sync_result = rank0_result
+        dist.broadcast(sync_result, src=0, group=self.group)
+        return sync_result
+
+
+class DistributedVaeMixin:
+    def init_distributed(self):
+        self.distributed_decoder = DistributedVaeExecutor()
+
+    def set_parallel_size(self, parallel_size: int) -> bool:
+        return self.distributed_decoder.set_parallel_size(parallel_size)
+
+    def is_distributed_enabled(self) -> bool:
+        if (
+            self.distributed_decoder.parallel_size <= 1
+            or not dist.is_initialized()
+            or not getattr(self, "use_tiling", False)
+        ):
+            return False
+        world_size = dist.get_world_size(group=self.distributed_decoder.group)
+        pp_size = min(int(self.distributed_decoder.parallel_size), int(world_size))
+        if pp_size <= 1:
+            return False
+        if self.distributed_decoder.parallel_size > pp_size:
+            logger.warning(
+                f"vae_patch_parallel_size={self.distributed_decoder.parallel_size} "
+                f"is greater than dit_group={world_size};"
+                f" using dit_group size={world_size}"
+            )
+        return True