From caa26e1d045ad95f869bde4239f9489bb0e269a4 Mon Sep 17 00:00:00 2001 From: fan2956 Date: Tue, 21 Apr 2026 20:36:08 +0800 Subject: [PATCH] [Cherry-pick][v0.18.0.post1][Bugfix] Fix VAE parallelism dist.gather performance bottleneck on NPU (#2970) Signed-off-by: fan2956 --- .../distributed/autoencoders/distributed_vae_executor.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py index ad60d164a..209f65625 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py +++ b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py @@ -54,9 +54,9 @@ class DistributedVaeExecutor: 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 + gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)] + dist.all_gather(gather_list, tensor, group=self.group) + return gather_list if self.rank == 0 else None def broadcast_tensor(self, tensor: torch.Tensor): dist.broadcast(tensor, src=0, group=self.group)