[XPU] xpu support PD disaggregation in v1 scheduler

Author: ddchenhao66

custom_ops/xpu_ops/src/ops/remote_cache_kv_ipc.h

@@ -72,12 +72,14 @@ struct RemoteCacheKvIpc {
     }
 
     void send_signal() {
-      msg_sed.mtext[1] = layer_id_;
-      if ((msgsnd(msgid, &msg_sed, (MAX_BSZ * 3 + 2) * 4, 0)) == -1) {
-        printf("kv signal full msg buffer\n");
+      if (inited) {
+        msg_sed.mtext[1] = layer_id_;
+        if ((msgsnd(msgid, &msg_sed, (MAX_BSZ * 3 + 2) * 4, 0)) == -1) {
+          printf("kv signal full msg buffer\n");
+        }
+        layer_id_ = (layer_id_ + 1);
+        assert(layer_id_ <= num_layers_);
       }
-      layer_id_ = (layer_id_ + 1);
-      assert(layer_id_ <= num_layers_);
     }
   };
 

fastdeploy/cache_manager/cache_messager.py

@@ -161,9 +161,6 @@ def __init__(
         for layer_idx in range(self.num_layers):
             key_cache = self.gpu_cache_kvs[f"key_caches_{layer_idx}_rank{self.rank}_device{gpu_id}"]
             val_cache = self.gpu_cache_kvs[f"value_caches_{layer_idx}_rank{self.rank}_device{gpu_id}"]
-            logger.info(
-                f"[key_cache: {hex(key_cache.data_ptr())}],[key_cache_mem: {hex(get_peer_mem_addr(key_cache.data_ptr()))}]"
-            )
             cache_k.append(key_cache)
             cache_v.append(val_cache)
             if paddle.is_compiled_with_xpu():
@@ -465,8 +462,12 @@ def __init__(
             val_cache = self.gpu_cache_kvs[f"value_caches_{layer_idx}_rank{self.rank}_device{gpu_id}"]
             cache_k.append(key_cache)
             cache_v.append(val_cache)
-            cache_k_ptr_list.append(key_cache.data_ptr())
-            cache_v_ptr_list.append(val_cache.data_ptr())
+            if paddle.is_compiled_with_xpu():
+                cache_k_ptr_list.append(get_peer_mem_addr(key_cache.data_ptr()))
+                cache_v_ptr_list.append(get_peer_mem_addr(val_cache.data_ptr()))
+            else:
+                cache_k_ptr_list.append(key_cache.data_ptr())
+                cache_v_ptr_list.append(val_cache.data_ptr())
         cache_k_ptr_list = np.array(cache_k_ptr_list)
         cache_v_ptr_list = np.array(cache_v_ptr_list)
 
@@ -771,7 +772,7 @@ def _handle_connect_task(self):
 def main():
     device = args.device_id
     rank = args.rank
-    set_device(args.rank)
+    set_device(device)
     cache_type = args.cache_dtype
     speculative_config = SpeculativeConfig(args.speculative_config)
     num_extra_layers = speculative_config.num_extra_cache_layer
@@ -883,7 +884,6 @@ def main():
     args = parse_args()
     rank_id = args.rank + args.local_data_parallel_id * args.mp_num
     logger = get_logger("cache_messager", f"cache_messager_rank{rank_id}.log")
-
     logger.info("create cache messager...")
     logger.info(f"{args}")
     main()

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.rank)
+        set_device(self.device)
         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.rank)
+                    set_device(self.device)
                     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(rank_id)
+    set_device(args.device_id)
     main()

fastdeploy/cache_manager/ops.py

@@ -67,6 +67,13 @@ def share_external_data_(cache, cache_name, cache_shape, use_ipc):
     return cache
 
 
+def get_all_visible_devices():
+    if current_platform.is_xpu():
+        return "XPU_VISIBLE_DEVICES=0,1,2,3,4,5,6,7"
+    else:
+        return "CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7"
+
+
 __all__ = [
     "cuda_host_alloc",
     "cuda_host_free",
@@ -81,4 +88,5 @@ def share_external_data_(cache, cache_name, cache_shape, use_ipc):
     "ipc_sent_key_value_cache_by_remote_ptr",
     "ipc_sent_key_value_cache_by_remote_ptr_block_sync",
     "get_peer_mem_addr",
+    "get_all_visible_devices",
 ]

fastdeploy/cache_manager/prefix_cache_manager.py

@@ -33,6 +33,7 @@
 from fastdeploy import envs
 from fastdeploy.cache_manager.cache_data import BlockNode, CacheStatus
 from fastdeploy.cache_manager.cache_metrics import CacheMetrics
+from fastdeploy.cache_manager.ops import get_all_visible_devices
 from fastdeploy.inter_communicator import EngineCacheQueue, IPCSignal, PrefixTreeStatus
 from fastdeploy.metrics.metrics import main_process_metrics
 from fastdeploy.utils import get_logger
@@ -243,9 +244,11 @@ def launch_cache_manager(
         # Run command to launch cache transfer managers
         log_dir = envs.FD_LOG_DIR
         cache_manager_processes = []
+        visible_devices = get_all_visible_devices()
         for i in range(tensor_parallel_size):
             launch_cmd = (
-                "FLAGS_allocator_strategy=auto_growth CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7"
+                "FLAGS_allocator_strategy=auto_growth "
+                + visible_devices
                 + " NCCL_MAX_NCHANNELS=1 NCCL_BUFFSIZE=0"
                 + f" FD_ENABLE_SWAP_SPACE_CLEARING={envs.FD_ENABLE_SWAP_SPACE_CLEARING}"
                 + f" {sys.executable} {py_path}"
@@ -328,9 +331,11 @@ def launch_cache_messager(
         py_path = os.path.join(current_dir_path, filename)
         log_dir = envs.FD_LOG_DIR
         cache_messager_processes = []
+        visible_devices = get_all_visible_devices()
         for i in range(tensor_parallel_size):
             launch_cmd = (
-                "FLAGS_allocator_strategy=auto_growth CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7"
+                "FLAGS_allocator_strategy=auto_growth "
+                + visible_devices
                 + " NCCL_MAX_NCHANNELS=1 NCCL_BUFFSIZE=0"
                 + f" {sys.executable} {py_path}"
                 + f" --device_id {int(device_ids[i])}"

fastdeploy/model_executor/layers/attention/ops/init_kv_signal_per_query.py

@@ -32,6 +32,17 @@ def init_kv_signal_per_query(
     if current_platform.is_cuda():
         from fastdeploy.model_executor.ops.gpu import init_kv_signal_per_query
 
+        out = init_kv_signal_per_query(
+            seq_lens_encoder,
+            seq_lens_this_time,
+            seq_lens_decoder,
+            rank,
+            num_layers,
+        )
+        return out
+    elif current_platform.is_xpu():
+        from fastdeploy.model_executor.ops.xpu import init_kv_signal_per_query
+
         out = init_kv_signal_per_query(
             seq_lens_encoder,
             seq_lens_this_time,

fastdeploy/model_executor/layers/attention/xpu_attn_backend.py

@@ -91,6 +91,7 @@ def __init__(
         )
         self.causal: bool = getattr(fd_config.model_config, "causal", True)
         self.keep_pd_step_flag: bool = fd_config.speculative_config.model_type == "mtp"
+        self.num_layers_draft_model: int = int(fd_config.speculative_config.method in ["mtp"])
 
         self.kv_num_heads: int = kv_num_heads
         self.num_heads: int = num_heads
@@ -122,7 +123,7 @@ def init_attention_metadata(self, forward_meta: ForwardMeta):
 
         # pd_disaggregation
         metadata.kv_signal_data_list = [None] * self.num_layers
-        if self.pd_disaggregation_mode == "per_chunk":
+        if self.pd_disaggregation_mode == "per_chunk" and not forward_meta.is_profiling:
             if not self.keep_pd_step_flag:
                 init_kv_signal_per_query(
                     forward_meta.seq_lens_encoder,

fastdeploy/worker/xpu_model_runner.py

@@ -71,6 +71,7 @@ def xpu_pre_process(
     draft_tokens: Optional[paddle.Tensor] = None,
     seq_lens_encoder: Optional[paddle.Tensor] = None,
     seq_lens_decoder: Optional[paddle.Tensor] = None,
+    is_profiling: bool = False,
 ) -> XPUForwardMeta:
     """ """
     max_len = input_ids.shape[1]
@@ -155,6 +156,8 @@ def xpu_pre_process(
 
     share_inputs["ids_remove_padding"] = adjusted_input
     xpu_forward_meta.ids_remove_padding = adjusted_input
+    # Set forward_meta.is_profiling to True to skip init_kv_signal_per_query for attention backends
+    xpu_forward_meta.is_profiling = is_profiling
     return xpu_forward_meta
 
 
@@ -924,6 +927,7 @@ def _prepare_inputs(self, is_dummy_run=False) -> None:
             draft_tokens=None,
             seq_lens_encoder=self.share_inputs["seq_lens_encoder"],
             seq_lens_decoder=self.share_inputs["seq_lens_decoder"],
+            is_profiling=is_dummy_run,
         )
         # Update bad tokens len
         max_bad_tokens_len = paddle.max(self.share_inputs["bad_tokens_len"])
@@ -1188,7 +1192,6 @@ class at the server level, which is too granular for ModelRunner.
             self.kv_signal_sender = create_kv_signal_sender()
         # 1. Prepare inputs of model and decoder.
         self._prepare_inputs(is_dummy_run=is_dummy_run)
-
         # NOTE(wufeisheng): If `not_need_stop`` is False, it means the current worker is in an idle state.
         # This logic is not used in TP (Tensor Parallelism) mode. However, in EP (Expert Parallelism) mode,
         # when there is data on other runner, the current runner is required to execute part of the model.