chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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

17
python/sglang/multimodal_gen/runtime/models/vaes/parallel/wan_dist_utils.py
View File

@@ -134,26 +134,29 @@ def halo_exchange(
recv_top_buf = _ensure_recv_buf(recv_top_buf, top_row)
recv_bottom_buf = _ensure_recv_buf(recv_bottom_buf, bottom_row)
reqs = []
# use batched P2P operations
p2p_ops = []
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]
reqs.append(dist.irecv(recv_top_buf, src=prev_rank, group=group))
reqs.append(dist.isend(top_row, dst=prev_rank, group=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))
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]
reqs.append(dist.isend(bottom_row, dst=next_rank, group=group))
reqs.append(dist.irecv(recv_bottom_buf, src=next_rank, group=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))
if rank == 0:
recv_top_buf.zero_()
if rank == world_size - 1:
recv_bottom_buf.zero_()
for req in reqs:
req.wait()
if p2p_ops:
reqs = dist.batch_isend_irecv(p2p_ops)
for req in reqs:
req.wait()
return (
torch.concat([recv_top_buf, x, recv_bottom_buf], dim=-2),