From 34577581148ffa5daf3469521354dd29609b1fe0 Mon Sep 17 00:00:00 2001 From: Huanyu Yang <20242081160@mail.dlut.edu.cn> Date: Fri, 29 May 2026 10:17:34 +0000 Subject: [PATCH] [Bugfix] Fix IndexError in DistributedVaeExecutor when vae_patch_parallel_size < world_size (#3928) Signed-off-by: QingZhou-YangHY <3868850350@qq.com> --- .../distributed/autoencoders/distributed_vae_executor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py index 209f65625..8fb4ebeb3 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py +++ b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py @@ -125,7 +125,7 @@ class DistributedVaeExecutor: # 2. local decode assigned = self._balance_tasks(tiletask_list, pp_size) - local_tasks = assigned[self.rank] if pp_size <= self.world_size else [] + local_tasks = assigned[self.rank] if self.rank < pp_size else [] local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks] # 3. compute shape per rank