[Feature] Spatially-sharded (SP) decode for the Wan VAE (#4620)

Signed-off-by: Rahul Steiger <rsteiger@nvidia.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py

@@ -49,9 +49,11 @@ class DistributedVaeExecutor:
         self.world_size = dist.get_world_size(self.group)
         self.rank = dist.get_rank(self.group)
         self.parallel_size = 1
+        self.parallel_mode = "tile"
 
-    def set_parallel_size(self, parallel_size: int):
+    def set_parallel_size(self, parallel_size: int, mode: str = "tile"):
         self.parallel_size = parallel_size
+        self.parallel_mode = mode
 
     def gather_tensors(self, tensor: torch.Tensor):
         gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)]
@@ -170,8 +172,8 @@ class DistributedVaeMixin:
     def init_distributed(self):
         self.distributed_executor = DistributedVaeExecutor()
 
-    def set_parallel_size(self, parallel_size: int) -> None:
-        self.distributed_executor.set_parallel_size(parallel_size)
+    def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None:
+        self.distributed_executor.set_parallel_size(parallel_size, mode=mode)
 
     def is_distributed_enabled(self) -> bool:
         if (