[Elastic EP] Fix non-contiguous weight transfers (#50641)
Signed-off-by: Itay Alroy <ialroy@nvidia.com>
Itay Alroy committed
Aug 2, 2026 at 18:13 UTC
96add737ff022430b17034ce50dc55504c04094f
1 file changed
+11
-7
vllm/distributed/elastic_ep/elastic_execute.py
+11
-7
@@ -85,16 +85,20 @@ def batch_transfer_weights(
85
assert len(all_params) > 0
86
p2p_ops = []
87
for param in all_params:
88
+ # PyNccl transfers flat memory and does not honor tensor strides.
89
+ transfer_param = param.contiguous()
90
op = object.__new__(P2POp)
89
- if is_sender:
90
- op.op = torch.distributed.isend
91
- op.tensor = param
92
- else:
93
- op.op = torch.distributed.irecv
94
- op.tensor = param
91
+ op.op = torch.distributed.isend if is_sender else torch.distributed.irecv
92
+ op.tensor = transfer_param
93
op.group_peer = peer_rank
94
p2p_ops.append(op)
97
- device_comm.batch_isend_irecv(p2p_ops)
95
+ if transfer_param is not param:
96
+ device_comm.batch_isend_irecv(p2p_ops)
97
+ p2p_ops.clear()
98
+ if not is_sender:
99
+ param.copy_(transfer_param)
100
+ if p2p_ops:
101
+ device_comm.batch_isend_irecv(p2p_ops)
102
103
104
def broadcast_expert_mapping(