diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index 473058d525..409a031c03 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -354,6 +354,7 @@ class NixlEplbCommunicator(EplbCommunicator): # Per-rank expert_id -> physical row; set by set_transfer_context. self._expert_to_src_row: list[dict[int, int]] | None = None self._layer_idx: int | None = None + self.set_stream(None) nixl_agent_config = nixl_utils.nixl_agent_config config = ( @@ -403,7 +404,7 @@ class NixlEplbCommunicator(EplbCommunicator): return f"eplb-{self._rank}{pp_suffix}-{uid}" def set_stream(self, stream: torch.Stream | None) -> None: - pass + self._stream = stream def add_send( self, @@ -421,6 +422,13 @@ class NixlEplbCommunicator(EplbCommunicator): f"pending transfers from layer {self._layer_idx}; " f"execute() was not called after previous add_recv() calls" ) + # NIXL READs bypass device streams; wait for pending copies out of + # the receive buffer before add_recv overwrites it. + if current_platform.is_rocm(): + with gpu_sync_allowed(): + torch.accelerator.current_stream().synchronize() + if self._stream is not None: + self._stream.synchronize() self._layer_idx = layer_idx n = self._num_local_experts rank_experts = old_indices[: self._world_size * n].reshape(self._world_size, n)