[Bugfix][EPLB] Wait for pending expert buffer copies before NIXL READs (#62081)

Signed-off-by: Djordje Ramic <djoramic@amd.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>

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)