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

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

docs/design/feature/vae_parallel.md

@@ -0,0 +1,289 @@
+# VAE Patch Parallelism
+
+This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
+We use **Qwen-Image** as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Testing](#testing)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+### What is Vae Patch parallel?
+
+**VAE Patch Parallelism** is a decoding acceleration technique. Instead of decoding the entire latent tensor at once, the latent tensor is:
+
++ Split into multiple spatial tiles
+
++ Distributed across multiple ranks
+
++ Decoded in parallel
+
++ Merged to reconstruct the final output
+
+This approach:
+
++ Distributes computation across multiple devices
+
++ Reduces peak memory usage per device
+
++ Accelerates decoding latency
+
+### Architecture
+We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE decoding.
+
+The executor is model-agnostic and accepts three function parameters:
+
++ split – Partition the latent into tiles
+
++ exec – Decode a single tile
+
++ merge – Combine decoded tiles into the final output
+
+#### Execution Flow
+
++ Call split(z) to generate a list of TileTask and a GridSpec
+
++ Dispatch tasks across ranks using workload-based balancing
+
++ Each rank executes exec(task) on its assigned tiles
+
++ Gather decoded tile results to rank 0
+
++ Rank 0 performs merge(...)
+
++ (Optional) Broadcast final result to all ranks
+
+This design separates:
+
++ Distributed execution logic
+
++ Model-specific tiling and merging logic
+
+#### Why split / exec / merge is necessary?
+
+The latent tensor cannot be arbitrarily partitioned.
+
+During decoding:
+
++ Each output pixel may depend on neighboring pixels
+
++ The receptive field is model-dependent
+
+Therefore:
+
++ Tiles must include overlap
+
++ Merge must perform blending to avoid seams
+
+## Step-by-Step Implementation
+
+### Step 1: Implement DistributedAutoencoderKLQwenImage
+`QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
+
+
+```
+class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+```
+**Key points**:
++ Inherit both AutoencoderKLQwenImage and DistributedVaeMixin
++ Call init_distributed() after loading weights
+
+### Step 2: Implement split/exec/merge
+Reuse `AutoencoderKLQwenImage.tiled_decode` logic and divide it into three stages. And we need return tiles with `GridSpec` and `TileTask`:
+```
+class GridSpec:
+    split_dims: tuple[int, ...]  # Tensor dimensions being split (e.g., (2, 3) for (B, C, H, W))
+    grid_shape: tuple[int, ...]  # Tile grid layout (num_rows, num_cols)
+    tile_spec: dict = field(default_factory=dict) # Metadata required for merging
+    output_dtype: torch.dtype | None = None # Final output dtype
+```
+```
+class TileTask:
+    tile_id: int # task id
+    grid_coord: tuple[int, ...]  # Tile position in grid
+    tensor: torch.Tensor | list[torch.Tensor]  # The tile tensor
+    workload: int | float = 1 # Used for load balancing (e.g., tile area)
+```
+And tiled base split/exec/merge as follow:
+```
+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
+    result_rows = []
+    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
+```
+
+### Step 3: Override tiled_decode
+We need to override tiled_decode, the main logic is:
++ check distributed is enabled
++ select split/exec/merge
++ Invoke self.distributed_decoder.execute to decode
+```
+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)
+```
+`broadcast_result` is set to True or False depending on the model; when enabled, the result will be used even on ranks other than 0.
+
+### Step 4: Modify Pipeline
+Change vae model from AutoencoderKLQwenImage to DistributedAutoencoderKLQwenImage
+```
+class YourModelPipeline(nn.Module):
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        ...
+-       self.vae = AutoencoderKL.from_pretrained(
+-           model, subfolder="vae", local_files_only=local_files_only).to(self.device)
++       self.vae = DistributedAutoencoderKL.from_pretrained(
++           model, subfolder="vae", local_files_only=local_files_only
++       ).to(self.device)
+```
+
+## Testing
+Verify numerical consistency between:
++ vae_patch_parallel_size = 1
+
++ vae_patch_parallel_size = N
+
+Example:
+torch.allclose(output_1, output_n, atol=1e-5)
+
+Testing requirements:
++ Fix random seed
++ Use identical tiling strategy
+
+```python
+m = Omni(
+        model=model_name,
+        vae_use_tiling=True,
+        parallel_config=DiffusionParallelConfig(
+            tensor_parallel_size=2,
+            vae_patch_parallel_size=1, # or 2
+        ),
+    )
+```
+When vae_patch_parallel_size is larger than the DiT world size, it will automatically fall back to using the DiT world size instead.
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Notes |
+|-------|------|-------|
+| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL |
+| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan |
+| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | Distributed AutoencoderKLQwenImage |
+
+---
+
+## Summary
+
+Adding Vae Patch Parallel support to diffusion model:
+
+1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge
+2. **Change vae model in pipeline to Distributed Vae**
+3. **Test** - Verify with `tensor_parallel_size=N` quality