From 11605767e0fe12fe83127ae44b0127d02725497d Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 4 Jun 2026 10:23:01 +0800 Subject: [PATCH] [diffusion] optimize: skip unused wanvae halo send copies (#27151) --- .../models/vaes/parallel/wan_dist_utils.py | 30 +++++++++++++++---- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/parallel/wan_dist_utils.py b/python/sglang/multimodal_gen/runtime/models/vaes/parallel/wan_dist_utils.py index 15e0f7765..096bf3095 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/parallel/wan_dist_utils.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/parallel/wan_dist_utils.py @@ -103,6 +103,15 @@ def _maybe_contiguous_for_sp_gather(x: torch.Tensor) -> torch.Tensor: return x +def _halo_memory_format(reference: torch.Tensor) -> torch.memory_format: + if reference.dim() > 1 and reference.stride(1) == 1: + if reference.dim() == 5 and hasattr(torch, "channels_last_3d"): + return torch.channels_last_3d + if reference.dim() == 4: + return torch.channels_last + return torch.contiguous_format + + def gather_and_trim_height(x: torch.Tensor, expected_height: int | None): if expected_height is None: return x @@ -115,13 +124,20 @@ def gather_and_trim_height(x: torch.Tensor, expected_height: int | None): def _ensure_recv_buf( recv_buf: torch.Tensor | None, reference: torch.Tensor ) -> torch.Tensor: + memory_format = _halo_memory_format(reference) if ( recv_buf is None or recv_buf.shape != reference.shape or recv_buf.dtype != reference.dtype or recv_buf.device != reference.device + or not recv_buf.is_contiguous(memory_format=memory_format) ): - return torch.empty_like(reference) + return torch.empty( + reference.shape, + dtype=reference.dtype, + device=reference.device, + memory_format=memory_format, + ) return recv_buf @@ -140,11 +156,11 @@ def halo_exchange( group = sp_group.device_group group_ranks = sp_group.ranks - top_row = x[..., :height_halo_size, :].contiguous() - bottom_row = x[..., -height_halo_size:, :].contiguous() + top_row_ref = x[..., :height_halo_size, :] + bottom_row_ref = x[..., -height_halo_size:, :] - recv_top_buf = _ensure_recv_buf(recv_top_buf, top_row) - recv_bottom_buf = _ensure_recv_buf(recv_bottom_buf, bottom_row) + recv_top_buf = _ensure_recv_buf(recv_top_buf, top_row_ref) + recv_bottom_buf = _ensure_recv_buf(recv_bottom_buf, bottom_row_ref) # use batched P2P operations p2p_ops = [] @@ -152,11 +168,15 @@ def halo_exchange( if rank > 0: # has previous neighbor, recv previous rank's data to recv_top_buf and send top_row to it. prev_rank = group_ranks[rank - 1] + top_row = top_row_ref.contiguous(memory_format=_halo_memory_format(top_row_ref)) p2p_ops.append(dist.P2POp(dist.irecv, recv_top_buf, prev_rank, group)) p2p_ops.append(dist.P2POp(dist.isend, top_row, prev_rank, group)) if rank < world_size - 1: # has next neighbor, send bottom_row to next rank and recv next rank's data to recv_bottom_buf. next_rank = group_ranks[rank + 1] + bottom_row = bottom_row_ref.contiguous( + memory_format=_halo_memory_format(bottom_row_ref) + ) p2p_ops.append(dist.P2POp(dist.isend, bottom_row, next_rank, group)) p2p_ops.append(dist.P2POp(dist.irecv, recv_bottom_buf, next_rank, group))