From 37273408ebcde2d16890dc93d4e613825fc22f95 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 13 Feb 2026 22:11:20 +0800 Subject: [PATCH] [diffusion] chore: use batched P2P ops in VAE parallel decoding (#18728) --- .../models/vaes/parallel/wan_dist_utils.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 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 f8aabc44f..83fa12ea0 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 @@ -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),