# 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 for decode parallel, and **Wan2.2** for encode parallel. --- ## Table of Contents - [Overview](#overview) - [Step-by-Step Implementation (Decode)](#step-by-step-implementation-decode) - [Encode Parallel Implementation](#encode-parallel-implementation) - [Testing](#testing) - [Reference Implementations](#reference-implementations) - [Summary](#summary) --- ## Overview ### What is Vae Patch parallel? **VAE Patch Parallelism** is an acceleration technique for both **encoding** and **decoding**. Instead of processing the entire tensor at once, the tensor is: - Split into multiple spatial tiles - Distributed across multiple ranks - Encoded/Decoded in parallel - Merged to reconstruct the final output This approach: - Distributes computation across multiple devices - Reduces peak memory usage per device - Accelerates encoding/decoding latency ### When to Use Encode vs Decode Parallel | Operation | Use Case | Example | | ------------------- | ---------------------------- | --------------------------------- | | **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video | | **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) | ### Architecture We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE encoding/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 (Decode) ### Step 1: Implement DistributedAutoencoderKLQwenImage `QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version: ```python 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`: ```python 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 ``` ```python 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: ```python 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_executor.execute to decode ```python 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_executor.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 ```diff 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) ``` ## Encode Parallel Implementation For models that require VAE encoding (e.g., Image-to-Video), you can also parallelize the encode operation. We use **Wan2.2** as the reference implementation. ### Step 1: Implement encode_tile_split Similar to decode, split the input tensor into tiles. Key considerations: - **Patchify handling**: If the model uses `patch_size`, scale tile parameters accordingly - **Temporal chunking**: Video VAEs may have temporal compression (e.g., 4x) ```python def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]: _, _, num_frames, height, width = x.shape encode_spatial_compression_ratio = self.spatial_compression_ratio # Scale tile parameters for patchified coordinate system tile_sample_min_height = self.tile_sample_min_height tile_sample_min_width = self.tile_sample_min_width 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: # When input is patchified, scale tile parameters accordingly encode_spatial_compression_ratio = self.spatial_compression_ratio // self.config.patch_size tile_sample_min_height = tile_sample_min_height // self.config.patch_size tile_sample_min_width = tile_sample_min_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 latent_height = height // encode_spatial_compression_ratio latent_width = width // encode_spatial_compression_ratio tile_latent_min_height = tile_sample_min_height // encode_spatial_compression_ratio tile_latent_min_width = tile_sample_min_width // encode_spatial_compression_ratio tile_latent_stride_height = tile_sample_stride_height // encode_spatial_compression_ratio tile_latent_stride_width = tile_sample_stride_width // encode_spatial_compression_ratio blend_height = tile_latent_min_height - tile_latent_stride_height blend_width = tile_latent_min_width - tile_latent_stride_width tiletask_list = [] # Use temporal compression ratio from config instead of hardcoding temporal_compression = self.config.scale_factor_temporal for i in range(0, height, tile_sample_stride_height): for j in range(0, width, tile_sample_stride_width): time_list = [] frame_range = 1 + (num_frames - 1) // temporal_compression for k in range(frame_range): if k == 0: tile = x[:, :, :1, i : i + tile_sample_min_height, j : j + tile_sample_min_width] else: tile = x[ :, :, 1 + temporal_compression * (k - 1) : 1 + temporal_compression * k, i : i + tile_sample_min_height, j : j + tile_sample_min_width, ] time_list.append(tile) tiletask_list.append( TileTask(len(tiletask_list), (i // tile_sample_stride_height, j // tile_sample_stride_width), time_list, workload=time_list[0].shape[3] * time_list[0].shape[4]) ) 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={ "latent_height": latent_height, "latent_width": latent_width, "blend_height": blend_height, "blend_width": blend_width, "tile_latent_stride_height": tile_latent_stride_height, "tile_latent_stride_width": tile_latent_stride_width, }, output_dtype=self.dtype, ) return tiletask_list, grid_spec ``` ### Step 2: Implement encode_tile_exec ```python def encode_tile_exec(self, task: TileTask) -> torch.Tensor: """Encode a single sample tile into latent space.""" self.clear_cache() time = [] for k, tile in enumerate(task.tensor): self._enc_conv_idx = [0] encoded = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) encoded = self.quant_conv(encoded) time.append(encoded) result = torch.cat(time, dim=2) self.clear_cache() return result ``` ### Step 3: Implement encode_tile_merge ```python def encode_tile_merge( self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec ) -> torch.Tensor: """Merge encoded tiles into a full latent 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): 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_latent_stride_height"], : grid_spec.tile_spec["tile_latent_stride_width"]]) result_rows.append(torch.cat(result_row, dim=-1)) enc = torch.cat(result_rows, dim=3)[ :, :, :, : grid_spec.tile_spec["latent_height"], : grid_spec.tile_spec["latent_width"] ] return enc ``` ### Step 4: Override tiled_encode method Override `tiled_encode` instead of `encode`. The parent's `_encode()` handles patchify before calling `tiled_encode()`, so input `x` is already patchified. ```python def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: """ Encode using distributed VAE executor. Note: x is already patchified by parent's _encode() before calling this method. """ if not self.is_distributed_enabled(): return super().tiled_encode(x) self.clear_cache() result = self.distributed_executor.execute( x, DistributedOperator( split=self.encode_tile_split, exec=self.encode_tile_exec, merge=self.encode_tile_merge, ), broadcast_result=True, # Latents needed by all ranks for diffusion ) self.clear_cache() return result ``` **Key differences from decode parallel:** | Aspect | Decode Parallel | Encode Parallel | | -------- | ----------------- | ----------------- | | `broadcast_result` | Often `False` (only rank 0 needs output) | `True` (all ranks need latents for diffusion) | | Patchify | Applied in merge (unpatchify) | Handled by parent `_encode()` before `tiled_encode()` | | Temporal chunking | Frame-by-frame | Chunk-based (e.g., 1 + 4n frames) | ## Testing Verify numerical consistency between: - vae_patch_parallel_size = 1 - vae_patch_parallel_size = N Example: ```python 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 | Decode Parallel | Encode Parallel | | ------- | ------ | ----------------- | ----------------- | | **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | ✅ | ❌ | | **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | ✅ | ✅ | | **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | ✅ | ❌ | | **FLUX.2-dev** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_flux2.py` | ✅ | ✅ | --- ## Spatially-Sharded Decode (Wan) The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE additionally supports a **spatially-sharded decode** backend, selected through `DiffusionParallelConfig.vae_parallel_mode` (`"spatial_shard_height"` or `"spatial_shard_width"`; default is `"tile"`). ### How it differs from tile parallel | Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) | | -------- | -------------------------- | ----------------------------------------------- | | Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W | | Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) | | Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) | | Scope | Decode + encode | Decode only | Spatial-shard decode swaps the decoder's spatial convolutions/padding for halo-exchanging variants (`WanDistConv2d`, `WanDistCausalConv3d`, `WanDistZeroPad2d`) so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Implementation lives in `vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py`. ### Wiring `vae_parallel_mode` flows through the same path as `vae_patch_parallel_size`: ```text serve.py (--vae-parallel-mode) / OmniEngineArgs -> DiffusionParallelConfig.vae_parallel_mode -> registry.py: model.vae.set_parallel_size(vae_pp_size, mode=...) -> DistributedVaeExecutor.parallel_mode -> DistributedAutoencoderKLWan.tiled_decode dispatch ``` `DistributedAutoencoderKLWan._spatial_shard_decode_enabled()` gates the path: it requires distributed decode to be enabled, a 5D latent, and `vae_patch_parallel_size == DiT group size`. Otherwise it logs a warning and falls back to tile-parallel decode. ### Notes - The decoder is patched **in place** the first time spatial-shard decode runs and is bound to a single split dimension for the lifetime of the VAE instance. - Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `cards_2`). --- ## Summary Adding VAE Patch Parallel support to diffusion model: 1. **Implement Distributed VAE** - Inherit from base VAE class and `DistributedVaeMixin` 2. **Decode Parallel** - Refactor `tiled_decode` into `tile_split`/`tile_exec`/`tile_merge` 3. **Encode Parallel** (optional) - Implement `encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models 4. **Change VAE model in pipeline** - Use the distributed version 5. **Test** - Verify numerical consistency with `vae_patch_parallel_size=1` vs `N`