minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] optimize: skip unused wanvae halo send copies (#27151)

Browse Source
This commit is contained in:
Mick
2026-06-04 10:23:01 +08:00
committed by GitHub
parent 736263f3dc
commit 11605767e0
1 changed files with 25 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/models/vaes/parallel/wan_dist_utils.py
+25 -5
View File
@@ -103,6 +103,15 @@ def _maybe_contiguous_for_sp_gather(x: torch.Tensor) -> torch.Tensor:
return x 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): def gather_and_trim_height(x: torch.Tensor, expected_height: int | None):
if expected_height is None: if expected_height is None:
return x return x
@@ -115,13 +124,20 @@ def gather_and_trim_height(x: torch.Tensor, expected_height: int | None):
def _ensure_recv_buf( def _ensure_recv_buf(
recv_buf: torch.Tensor | None, reference: torch.Tensor recv_buf: torch.Tensor | None, reference: torch.Tensor
) -> torch.Tensor: ) -> torch.Tensor:
memory_format = _halo_memory_format(reference)
if ( if (
recv_buf is None recv_buf is None
or recv_buf.shape != reference.shape or recv_buf.shape != reference.shape
or recv_buf.dtype != reference.dtype or recv_buf.dtype != reference.dtype
or recv_buf.device != reference.device 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 return recv_buf
@@ -140,11 +156,11 @@ def halo_exchange(
group = sp_group.device_group group = sp_group.device_group
group_ranks = sp_group.ranks group_ranks = sp_group.ranks
top_row = x[..., :height_halo_size, :].contiguous() top_row_ref = x[..., :height_halo_size, :]
bottom_row = x[..., -height_halo_size:, :].contiguous() bottom_row_ref = x[..., -height_halo_size:, :]
recv_top_buf = _ensure_recv_buf(recv_top_buf, top_row) recv_top_buf = _ensure_recv_buf(recv_top_buf, top_row_ref)
recv_bottom_buf = _ensure_recv_buf(recv_bottom_buf, bottom_row) recv_bottom_buf = _ensure_recv_buf(recv_bottom_buf, bottom_row_ref)
# use batched P2P operations # use batched P2P operations
p2p_ops = [] p2p_ops = []
@@ -152,11 +168,15 @@ def halo_exchange(
if rank > 0: if rank > 0:
# has previous neighbor, recv previous rank's data to recv_top_buf and send top_row to it. # has previous neighbor, recv previous rank's data to recv_top_buf and send top_row to it.
prev_rank = group_ranks[rank - 1] 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.irecv, recv_top_buf, prev_rank, group))
p2p_ops.append(dist.P2POp(dist.isend, top_row, prev_rank, group)) p2p_ops.append(dist.P2POp(dist.isend, top_row, prev_rank, group))
if rank < world_size - 1: if rank < world_size - 1:
# has next neighbor, send bottom_row to next rank and recv next rank's data to recv_bottom_buf. # 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] 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.isend, bottom_row, next_rank, group))
p2p_ops.append(dist.P2POp(dist.irecv, recv_bottom_buf, next_rank, group)) p2p_ops.append(dist.P2POp(dist.irecv, recv_bottom_buf, next_rank, group))