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

[diffusion] chore: use batched P2P ops in VAE parallel decoding (#18728)

Browse Source
This commit is contained in:
Mick
2026-02-13 22:11:20 +08:00
committed by GitHub
parent acc940d302
commit 37273408eb
1 changed files with 10 additions and 7 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
+8 -5
View File
@@ -134,24 +134,27 @@ def halo_exchange(
recv_top_buf = _ensure_recv_buf(recv_top_buf, top_row) recv_top_buf = _ensure_recv_buf(recv_top_buf, top_row)
recv_bottom_buf = _ensure_recv_buf(recv_bottom_buf, bottom_row) recv_bottom_buf = _ensure_recv_buf(recv_bottom_buf, bottom_row)
reqs = [] # use batched P2P operations
p2p_ops = []
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]
reqs.append(dist.irecv(recv_top_buf, src=prev_rank, group=group)) p2p_ops.append(dist.P2POp(dist.irecv, recv_top_buf, prev_rank, group))
reqs.append(dist.isend(top_row, dst=prev_rank, group=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]
reqs.append(dist.isend(bottom_row, dst=next_rank, group=group)) p2p_ops.append(dist.P2POp(dist.isend, bottom_row, next_rank, group))
reqs.append(dist.irecv(recv_bottom_buf, src=next_rank, group=group)) p2p_ops.append(dist.P2POp(dist.irecv, recv_bottom_buf, next_rank, group))
if rank == 0: if rank == 0:
recv_top_buf.zero_() recv_top_buf.zero_()
if rank == world_size - 1: if rank == world_size - 1:
recv_bottom_buf.zero_() recv_bottom_buf.zero_()
if p2p_ops:
reqs = dist.batch_isend_irecv(p2p_ops)
for req in reqs: for req in reqs:
req.wait() req.wait()