[XPU] xpu support PD disaggregation

Author: ddchenhao66

custom_ops/xpu_ops/src/ops/init_signal_layerwise.cc

@@ -0,0 +1,54 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "ops/remote_cache_kv_ipc.h"
+#include "paddle/extension.h"
+
+#ifndef PD_BUILD_STATIC_OP
+#define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
+#endif
+
+using cache_write_complete_signal_type =
+    RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_meta_data;
+
+paddle::Tensor InitSignalLayerwiseFunc(const paddle::Tensor& kv_signal_metadata,
+                                       const int layer_id) {
+  auto kv_signal_metadata_out =
+      kv_signal_metadata.copy_to(paddle::CPUPlace(), false);
+  kv_signal_metadata_out.data<int64_t>()[0] = static_cast<int64_t>(layer_id);
+  return kv_signal_metadata_out;
+}
+
+std::vector<paddle::Tensor> InitSignalLayerwise(
+    const paddle::Tensor& kv_signal_metadata, const int layer_id) {
+  return {InitSignalLayerwiseFunc(kv_signal_metadata, layer_id)};
+}
+
+std::vector<std::vector<int64_t>> InitSignalLayerwiseShape(
+    const std::vector<int64_t>& kv_signal_metadata_shape, const int layer_id) {
+  return {kv_signal_metadata_shape};
+}
+
+std::vector<paddle::DataType> InitSignalLayerwiseDtype(
+    const paddle::DataType& kv_signal_metadata_dtype, const int layer_id) {
+  return {paddle::DataType::INT64};
+}
+
+PD_BUILD_STATIC_OP(init_signal_layerwise)
+    .Inputs({"kv_signal_metadata"})
+    .Outputs({"kv_signal_metadata_out"})
+    .Attrs({"layer_id: int"})
+    .SetKernelFn(PD_KERNEL(InitSignalLayerwise))
+    .SetInferShapeFn(PD_INFER_SHAPE(InitSignalLayerwiseShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(InitSignalLayerwiseDtype));

custom_ops/xpu_ops/src/ops/open_shm_and_get_meta_signal.cc

@@ -17,22 +17,27 @@
 #include "ops/utility/env.h"
 #include "paddle/extension.h"
 
+#ifndef PD_BUILD_STATIC_OP
+#define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name)
+#endif
+
 XPU_DECLARE_BOOL(fmt_write_cache_completed_signal, false);
 
 using cache_write_complete_signal_type =
     RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_meta_data;
 
 paddle::Tensor OpenShmAndGetMetaSignalFunc(const int rank,
+                                           const int device_id,
                                            const bool keep_pd_step_flag) {
   cache_write_complete_signal_type kv_signal_metadata;
-  const char *fmt_write_cache_completed_signal_str =
+  const char* fmt_write_cache_completed_signal_str =
       std::getenv("FLAGS_fmt_write_cache_completed_signal");
   if (fmt_write_cache_completed_signal_str &&
       (std::strcmp(fmt_write_cache_completed_signal_str, "true") == 0 ||
        std::strcmp(fmt_write_cache_completed_signal_str, "1") == 0)) {
     kv_signal_metadata =
         RemoteCacheKvIpc::open_shm_and_get_complete_signal_meta_data(
-            rank, keep_pd_step_flag);
+            rank, device_id, keep_pd_step_flag);
   }
 
   auto kv_signal_metadata_out =
@@ -46,9 +51,9 @@ paddle::Tensor OpenShmAndGetMetaSignalFunc(const int rank,
   return kv_signal_metadata_out;
 }
 
-void InitKVSignalPerQuery(const paddle::Tensor &seq_lens_encoder_tensor,
-                          const paddle::Tensor &seq_lens_this_time_tensor,
-                          const paddle::Tensor &seq_lens_decoder_tensor,
+void InitKVSignalPerQuery(const paddle::Tensor& seq_lens_encoder_tensor,
+                          const paddle::Tensor& seq_lens_this_time_tensor,
+                          const paddle::Tensor& seq_lens_decoder_tensor,
                           const int rank,
                           const int num_layers) {
   if (FLAGS_fmt_write_cache_completed_signal) {
@@ -68,24 +73,24 @@ void InitKVSignalPerQuery(const paddle::Tensor &seq_lens_encoder_tensor,
 }
 
 std::vector<paddle::Tensor> OpenShmAndGetMetaSignal(
-    const int rank, const bool keep_pd_step_flag) {
-  return {OpenShmAndGetMetaSignalFunc(rank, keep_pd_step_flag)};
+    const int rank, const int device_id, const bool keep_pd_step_flag) {
+  return {OpenShmAndGetMetaSignalFunc(rank, device_id, keep_pd_step_flag)};
 }
 
 std::vector<std::vector<int64_t>> OpenShmAndGetMetaSignalShape(
-    const int rank, const bool keep_pd_step_flag) {
+    const int rank, const int device_id, const bool keep_pd_step_flag) {
   return {{3}};
 }
 
 std::vector<paddle::DataType> OpenShmAndGetMetaSignalDtype(
-    const int rank, const bool keep_pd_step_flag) {
+    const int rank, const int device_id, const bool keep_pd_step_flag) {
   return {paddle::DataType::INT64};
 }
 
-PD_BUILD_OP(open_shm_and_get_meta_signal)
+PD_BUILD_STATIC_OP(open_shm_and_get_meta_signal)
     .Inputs({})
     .Outputs({"kv_signal_metadata"})
-    .Attrs({"rank: int", "keep_pd_step_flag: bool"})
+    .Attrs({"rank: int", "device_id: int", "keep_pd_step_flag: bool"})
     .SetKernelFn(PD_KERNEL(OpenShmAndGetMetaSignal))
     .SetInferShapeFn(PD_INFER_SHAPE(OpenShmAndGetMetaSignalShape))
     .SetInferDtypeFn(PD_INFER_DTYPE(OpenShmAndGetMetaSignalDtype));

custom_ops/xpu_ops/src/ops/remote_cache_kv_ipc.cc

@@ -26,7 +26,7 @@ bool RemoteCacheKvIpc::kv_complete_signal_shmem_opened = false;
 
 RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_meta_data
 RemoteCacheKvIpc::open_shm_and_get_complete_signal_meta_data(
-    const int rank_id, const bool keep_pd_step_flag) {
+    const int rank_id, const int device_id, const bool keep_pd_step_flag) {
   if (RemoteCacheKvIpc::kv_complete_signal_shmem_opened) {
     if (keep_pd_step_flag) {
       return RemoteCacheKvIpc::kv_complete_signal_meta_data;
@@ -47,12 +47,13 @@ RemoteCacheKvIpc::open_shm_and_get_complete_signal_meta_data(
     std::string iflags_server_uuid_env_str(iflags_server_uuid_env_p);
     flags_server_uuid = iflags_server_uuid_env_str;
   }
+
   std::string step_shm_name =
-      ("splitwise_complete_prefilled_step_" + std::to_string(rank_id) + "_" +
-       flags_server_uuid);
+      ("splitwise_complete_prefilled_step_" + std::to_string(rank_id) + "." +
+       std::to_string(device_id));
   std::string layer_shm_name =
-      ("splitwise_complete_prefilled_layer_" + std::to_string(rank_id) + "_" +
-       flags_server_uuid);
+      ("splitwise_complete_prefilled_layer_" + std::to_string(rank_id) + "." +
+       std::to_string(device_id));
   if (const char* use_ep = std::getenv("ENABLE_EP_DP")) {
     if (std::strcmp(use_ep, "1") == 0) {
       step_shm_name = "splitwise_complete_prefilled_step_tprank0_dprank" +

custom_ops/xpu_ops/src/ops/remote_cache_kv_ipc.h

@@ -91,6 +91,7 @@ struct RemoteCacheKvIpc {
 
   static RemoteCacheKvIpc::save_cache_kv_complete_signal_layerwise_meta_data
   open_shm_and_get_complete_signal_meta_data(const int rank_id,
+                                             const int device_id,
                                              const bool keep_pd_step_flag);
   static void save_cache_kv_complete_signal_layerwise(void* meta_data);
   static void save_cache_kv_complete_signal_layerwise_per_query(

custom_ops/xpu_ops/src/ops/share_external_data.cc

@@ -19,26 +19,26 @@
 #include "xpu/plugin.h"
 #include "xpu_multiprocess.h"  // NOLINT(build/include_subdir)
 
-std::vector<paddle::Tensor> ShareExternalData(const paddle::Tensor &input,
+std::vector<paddle::Tensor> ShareExternalData(const paddle::Tensor& input,
                                               const std::string shm_name,
-                                              const std::vector<int> &shape,
+                                              const std::vector<int>& shape,
                                               bool use_ipc) {
   sharedMemoryInfo info;
   int ret = sharedMemoryOpen(shm_name.c_str(), sizeof(shmStruct), &info);
   PD_CHECK(ret == 0, "sharedMemoryOpen failed");
-  volatile shmStruct *shm = static_cast<volatile shmStruct *>(info.addr);
-  void *data_ptr_addr = nullptr;
+  volatile shmStruct* shm = static_cast<volatile shmStruct*>(info.addr);
+  void* data_ptr_addr = nullptr;
   if (use_ipc) {
 #if XPURT_VERSION_MAJOR == 5
     int ret = xpu_ipc_open_memhandle(&data_ptr_addr,
-                                     *(XPUIpcMemHandle *)&shm->memHandle,
+                                     *(XPUIpcMemHandle*)&shm->memHandle,
                                      0x01);  // NOLINT
-    PD_CHECK(ret == XPU_SUCCESS, "xpu_ipc_open_memhandle failed");
+    PD_CHECK(ret == XPU_SUCCESS, "%s xpu_ipc_open_memhandle failed", shm_name);
 #elif XPURT_VERSION_MAJOR == 4
     PD_THROW("kl2 not support prefix cache");
 #endif
   } else {
-    data_ptr_addr = reinterpret_cast<void *>(shm->data_ptr_addr);
+    data_ptr_addr = reinterpret_cast<void*>(shm->data_ptr_addr);
   }
 
   phi::XPUPlace place(phi::backends::xpu::GetXPUCurrentDeviceId());

fastdeploy/cache_manager/cache_messager.py

@@ -25,14 +25,20 @@
 import numpy as np
 import paddle
 
+from fastdeploy.cache_manager.ops import (
+    get_output_kv_signal,
+    get_peer_mem_addr,
+    memory_allocated,
+    set_data_ipc,
+    set_device,
+)
 from fastdeploy.cache_manager.transfer_factory import IPCCommManager, RDMACommManager
 from fastdeploy.config import SpeculativeConfig
 from fastdeploy.inter_communicator import (
     EngineWorkerQueue,
     IPCSignal,
     shared_memory_exists,
 )
-from fastdeploy.model_executor.ops.gpu import get_output_kv_signal, set_data_ipc
 from fastdeploy.utils import envs, get_logger
 
 logger = get_logger("cache_messager", "cache_messager.log")
@@ -155,18 +161,25 @@ 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)
-            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)
 
         # 2. initialize the block_bytes
         cache_shape = key_cache.shape
         max_block_num = cache_shape[0]
         block_bytes = math.prod(cache_shape[1:])
-        if key_cache.dtype == paddle.bfloat16:
+        if key_cache.dtype == paddle.bfloat16 or key_cache.dtype == paddle.float16:
             block_bytes *= 2
         logger.info(
             f"layers {num_layers} cache_shape: {cache_shape}, max_block_num: {max_block_num}, "
@@ -758,7 +771,7 @@ def _handle_connect_task(self):
 def main():
     device = args.device_id
     rank = args.rank
-    paddle.set_device(f"gpu:{device}")
+    set_device(args.rank)
     cache_type = args.cache_dtype
     speculative_config = SpeculativeConfig(args.speculative_config)
     num_extra_layers = speculative_config.num_extra_cache_layer
@@ -818,7 +831,7 @@ def main():
     cache_kv_size_byte = sum([tmp.numel() * 1 for key, tmp in gpu_cache_kvs.items()])
     logger.info(f"device :{device}")
     logger.info(f"cache_kv_size_byte : {cache_kv_size_byte}")
-    logger.info(f"done init cache (full) gmem alloc : {paddle.device.cuda.memory_allocated()}")
+    logger.info(f"done init cache (full) gmem alloc : {memory_allocated}")
 
     if envs.ENABLE_V1_KVCACHE_SCHEDULER:
         cache_messager = CacheMessagerV1(

fastdeploy/cache_manager/ops.py

@@ -6,17 +6,27 @@
     from fastdeploy.model_executor.ops.gpu import (
         cuda_host_alloc,
         cuda_host_free,
+        get_data_ptr_ipc,
+        get_output_kv_signal,
+        ipc_sent_key_value_cache_by_remote_ptr,
+        ipc_sent_key_value_cache_by_remote_ptr_block_sync,
         set_data_ipc,
         share_external_data,
         swap_cache_all_layers,
         unset_data_ipc,
     )
 
     memory_allocated = paddle.device.cuda.memory_allocated
+
+    def get_peer_mem_addr(*args, **kwargs):
+        raise RuntimeError("CUDA no need of get_peer_mem_addr!")
+
 elif current_platform.is_xpu():
     from fastdeploy.model_executor.ops.xpu import (
         cuda_host_alloc,
         cuda_host_free,
+        get_output_kv_signal,
+        get_peer_mem_addr,
         set_data_ipc,
         share_external_data,
         swap_cache_all_layers,
@@ -25,6 +35,15 @@
     unset_data_ipc = None
     memory_allocated = paddle.device.xpu.memory_allocated
 
+    def get_data_ptr_ipc(*args, **kwargs):
+        raise RuntimeError("XPU get_data_ptr_ipc UNIMPLENENTED!")
+
+    def ipc_sent_key_value_cache_by_remote_ptr(*args, **kwargs):
+        raise RuntimeError("XPU ipc_sent_key_value_cache_by_remote_ptr UNIMPLENENTED")
+
+    def ipc_sent_key_value_cache_by_remote_ptr_block_sync(*args, **kwargs):
+        raise RuntimeError("XPU No ipc_sent_key_value_cache_by_remote_ptr UNIMPLENENTED")
+
 else:
     raise RuntimeError("Prefix cache ops only supported CUDA nor XPU platform ")
 
@@ -57,4 +76,9 @@ def share_external_data_(cache, cache_name, cache_shape, use_ipc):
     "unset_data_ipc",  # XPU是 None
     "set_device",
     "memory_allocated",
+    "get_output_kv_signal",
+    "get_data_ptr_ipc",
+    "ipc_sent_key_value_cache_by_remote_ptr",
+    "ipc_sent_key_value_cache_by_remote_ptr_block_sync",
+    "get_peer_mem_addr",
 ]

fastdeploy/cache_manager/transfer_factory/ipc_cache_transfer.py

@@ -16,7 +16,7 @@
 
 import paddle
 
-from fastdeploy.model_executor.ops.gpu import (
+from fastdeploy.cache_manager.ops import (
     get_data_ptr_ipc,
     ipc_sent_key_value_cache_by_remote_ptr,
     ipc_sent_key_value_cache_by_remote_ptr_block_sync,

fastdeploy/model_executor/forward_meta.py

@@ -244,6 +244,8 @@ class XPUForwardMeta(ForwardMeta):
     total_enc_len: Optional[paddle.Tensor] = None
     # position embedding type in rope, supports 'NORMAL' or 'HALF_HEAD_DIM'
     pos_emb_type: Optional[str] = "NORMAL"
+    # for pd_disaggregation
+    kv_signal_sender: Optional[paddle.Tensor] = None
 
 
 @dataclass

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

@@ -29,6 +29,11 @@ def init_signal_layerwise(
     if current_platform.is_cuda():
         from fastdeploy.model_executor.ops.gpu import init_signal_layerwise
 
+        out = init_signal_layerwise(kv_signal_metadata, layer_id)
+        return out
+    elif current_platform.is_xpu():
+        from fastdeploy.model_executor.ops.xpu import init_signal_layerwise
+
         out = init_signal_layerwise(kv_signal_metadata, layer_id)
         return out
     else: