[XPU] fix the issue of cache KV transfer process startup failure on non-zero XPU cards

Author: ddchenhao66

custom_ops/xpu_ops/src/ops/share_external_data.cc

@@ -33,7 +33,7 @@ std::vector<paddle::Tensor> ShareExternalData(const paddle::Tensor& input,
     int ret = xpu_ipc_open_memhandle(&data_ptr_addr,
                                      *(XPUIpcMemHandle*)&shm->memHandle,
                                      0x01);  // NOLINT
-    PD_CHECK(ret == XPU_SUCCESS, "%s xpu_ipc_open_memhandle failed", shm_name);
+    PD_CHECK(ret == XPU_SUCCESS, shm_name, " xpu_ipc_open_memhandle failed");
 #elif XPURT_VERSION_MAJOR == 4
     PD_THROW("kl2 not support prefix cache");
 #endif

fastdeploy/cache_manager/cache_transfer_manager.py

@@ -204,7 +204,7 @@ def _init_gpu_cache(self, args):
             logger.info(f"[rank {self.rank}/{self.n_ranks}] OK! Stop waiting.")
 
         logger.info(f"[rank {self.rank}/{self.n_ranks}] Initializing kv cache for all layers.")
-        set_device(self.device)
+        set_device(self.rank)
         for i in range(args.num_layers + self.num_extra_layers):
             num_gpu_blocks = self.num_gpu_blocks if i < args.num_layers else self.num_extra_layer_gpu_blocks
             key_name = f"key_caches_{i}_rank{self.rank}.device{self.device}"
@@ -569,7 +569,7 @@ def clear_or_update_caches(self, args):
                             time.sleep(0.1)
 
                     # clear gpu caches
-                    set_device(self.device)
+                    set_device(self.rank)
                     for name, tensor in self.gpu_cache_kvs.items():
                         unset_data_ipc(tensor, name, True, False)
                     self.gpu_cache_kvs.clear()
@@ -640,5 +640,5 @@ def main():
     args = parse_args()
     rank_id = args.rank + args.local_data_parallel_id * args.mp_num
     logger = get_logger("cache_transfer_manager", f"cache_transfer_manager_rank{rank_id}.log")
-    set_device(args.device_id)
+    set_device(rank_id)
     main()

fastdeploy/model_executor/layers/attention/utils.py

@@ -17,6 +17,7 @@
 import os
 
 from fastdeploy.config import FDConfig
+from fastdeploy.platforms import current_platform
 
 
 def init_rank_and_device_id(fd_config: FDConfig):
@@ -26,7 +27,10 @@ def init_rank_and_device_id(fd_config: FDConfig):
         + fd_config.parallel_config.tensor_parallel_rank
     )
 
-    cuda_visible_devices = os.getenv("CUDA_VISIBLE_DEVICES", None)
+    if current_platform.is_xpu():
+        cuda_visible_devices = os.getenv("XPU_VISIBLE_DEVICES", None)
+    else:  # default cuda
+        cuda_visible_devices = os.getenv("CUDA_VISIBLE_DEVICES", None)
 
     if cuda_visible_devices is None:
         device_id = rank

fastdeploy/worker/xpu_model_runner.py

@@ -1124,7 +1124,6 @@ def _dummy_run(
         self._dummy_prefill_inputs(num_tokens, batch_size)
 
         while True:
-            return None
             self.execute_model(is_dummy_run=True)
 
             if int((self.share_inputs["seq_lens_this_time"] > 0).sum()) == 0: