[PG] force register local memory for P2P memory regions (kvcache-ai#1690)

Similar to kvcache-ai#1656.

Author: UNIDY2002

mooncake-pg/include/p2p_proxy.h

@@ -66,6 +66,7 @@ class P2PProxy {
         int rank = 0;
         int size = 0;
         int cuda_device_index = -1;
+        std::string location;
     };
 
     struct SendOp {
@@ -231,6 +232,7 @@ class P2PProxy {
     int rank_ = 0;
     int size_ = 0;
     int cuda_device_index_ = -1;
+    std::string location_;
     P2PResources resources_;
 
     std::queue<SendOpContext> send_queue_;

mooncake-pg/src/mooncake_backend.cpp

@@ -185,6 +185,7 @@ MooncakeBackend::MooncakeBackend(
                      .rank = rank_,
                      .size = size_,
                      .cuda_device_index = cuda_device_index,
+                     .location = location,
                  });
     p2p_device_worker_->registerProxy(p2p_proxy_);
 

mooncake-pg/src/p2p_proxy.cpp

@@ -40,7 +40,8 @@ P2PProxy::P2PProxy(TransferEngine* engine, const Options& options)
       is_cpu_(options.is_cpu),
       rank_(options.rank),
       size_(options.size),
-      cuda_device_index_(options.cuda_device_index) {
+      cuda_device_index_(options.cuda_device_index),
+      location_(options.location) {
     if (!is_cpu_ && cuda_device_index_ < 0) {
         int current_device = -1;
         const cudaError_t get_device_error = cudaGetDevice(&current_device);
@@ -71,15 +72,15 @@ void P2PProxy::AllocateResources() {
         resources_.send_buffer_ = std::malloc(kP2PTotalBufferSize);
         TORCH_CHECK(resources_.send_buffer_ != nullptr,
                     "Failed to allocate CPU P2P send buffer");
-        int rc = engine_->registerLocalMemory(
-            resources_.send_buffer_, kP2PTotalBufferSize, kWildcardLocation);
+        int rc = engine_->registerLocalMemory(resources_.send_buffer_,
+                                              kP2PTotalBufferSize, location_);
         TORCH_CHECK(rc == 0, "Failed to register CPU P2P send buffer");
 
         resources_.recv_buffer_ = std::malloc(kP2PTotalBufferSize);
         TORCH_CHECK(resources_.recv_buffer_ != nullptr,
                     "Failed to allocate CPU P2P recv buffer");
-        rc = engine_->registerLocalMemory(
-            resources_.recv_buffer_, kP2PTotalBufferSize, kWildcardLocation);
+        rc = engine_->registerLocalMemory(resources_.recv_buffer_,
+                                          kP2PTotalBufferSize, location_);
         TORCH_CHECK(rc == 0, "Failed to register CPU P2P recv buffer");
     } else {
         SetCudaDeviceIfNeeded(
@@ -89,15 +90,15 @@ void P2PProxy::AllocateResources() {
             cudaMalloc(&resources_.send_buffer_, kP2PTotalBufferSize);
         TORCH_CHECK(err == cudaSuccess,
                     "Failed to allocate CUDA P2P send buffer");
-        int rc = engine_->registerLocalMemory(
-            resources_.send_buffer_, kP2PTotalBufferSize, kWildcardLocation);
+        int rc = engine_->registerLocalMemory(resources_.send_buffer_,
+                                              kP2PTotalBufferSize, location_);
         TORCH_CHECK(rc == 0, "Failed to register CUDA P2P send buffer");
 
         err = cudaMalloc(&resources_.recv_buffer_, kP2PTotalBufferSize);
         TORCH_CHECK(err == cudaSuccess,
                     "Failed to allocate CUDA P2P recv buffer");
-        rc = engine_->registerLocalMemory(
-            resources_.recv_buffer_, kP2PTotalBufferSize, kWildcardLocation);
+        rc = engine_->registerLocalMemory(resources_.recv_buffer_,
+                                          kP2PTotalBufferSize, location_);
         TORCH_CHECK(rc == 0, "Failed to register CUDA P2P recv buffer");
     }
 
@@ -152,12 +153,12 @@ void P2PProxy::AllocateResources() {
 
     int rc = engine_->registerLocalMemory(resources_.ctrl_send_region_,
                                           kMaxNumRanks * sizeof(P2PControlSlot),
-                                          kWildcardLocation);
+                                          location_);
     TORCH_CHECK(rc == 0, "Failed to register P2P ctrl send region");
 
     rc = engine_->registerLocalMemory(resources_.ctrl_recv_region_,
                                       kMaxNumRanks * sizeof(P2PControlSlot),
-                                      kWildcardLocation);
+                                      location_);
     TORCH_CHECK(rc == 0, "Failed to register P2P ctrl recv region");
 }