replace torch.cuda.synchronize() with current_platform.synchronize()

Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>

Author: jikunshang

vllm/distributed/eplb/eplb_state.py

@@ -37,6 +37,7 @@
 from vllm.distributed.parallel_state import get_ep_group, get_node_count
 from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import MixtureOfExperts
+from vllm.platforms import current_platform
 
 from .rebalance_algo import rebalance_experts
 from .rebalance_execute import rearrange_expert_weights_inplace
@@ -348,7 +349,7 @@ def rearrange(self,
         time_start = None
         is_main_rank = ep_rank == 0
         if is_main_rank:
-            torch.cuda.synchronize()
+            current_platform.synchronize()
             time_start = time.perf_counter()
             logger.info("Rearranging experts %s...",
                         "(profile)" if is_profile else "")
@@ -423,7 +424,7 @@ def rearrange(self,
 
         if is_main_rank:
             assert time_start is not None
-            torch.cuda.synchronize()
+            current_platform.synchronize()
             time_end = time.perf_counter()
             logger.info(
                 "Rearranged experts%sin %.2f seconds.",

vllm/distributed/eplb/rebalance_execute.py

@@ -13,6 +13,8 @@
 from torch.distributed import (P2POp, ProcessGroup, all_gather,
                                batch_isend_irecv, get_global_rank)
 
+from vllm.platforms import current_platform
+
 
 def idx_local_to_global(
     local_idx: int,
@@ -292,7 +294,7 @@ def rearrange_expert_weights_inplace(
     for layer in range(num_moe_layers):
         # NOTE(bowen): We need this synchronize to run, but I don't know why.
         # If you figure out the reason, please let me know -- thank you!
-        torch.cuda.synchronize()
+        current_platform.synchronize()
         shuffle_layer(
             num_local_physical_experts,
             ep_rank,