feat: add CPU staging buffer to access GPU memory via unified buffer

mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h

@@ -195,15 +195,18 @@ class UbContext {
 
     // Polls one JFC and processes the slices internally:
     //   * Aggregates each completion's jetty depth into jetty_depth_set.
-    //   * Successful slices have markSuccess() called in place and are NOT
-    //     returned (they may be recycled by the submitting thread the
+    //   * Successful normal slices have markSuccess() called in place and are
+    //     NOT returned (they may be recycled by the submitting thread the
     //     moment markSuccess() runs).
+    //   * Successful slices that need post-URMA staging work are returned in
+    //     deferred_success_slices and are NOT marked successful yet.
     //   * Failed slices are returned in failed_slices[0..num_failed-1] for
     //     the caller to apply retry / markFailed.
     // Returns the total number of completions polled (>= 0), or a negative
     // error code.
     virtual int poll(int num_entries, Transport::Slice** failed_slices,
                      int& num_failed,
+                     std::vector<Transport::Slice*>& deferred_success_slices,
                      std::unordered_map<volatile int*, int>& jetty_depth_set,
                      int jfc_index = 0) = 0;
 

mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h

@@ -14,8 +14,15 @@
 
 #ifndef UB_TRANSPORT_H
 #define UB_TRANSPORT_H
+#include <atomic>
+#include <cstddef>
+#include <cstdint>
+#include <deque>
+#include <mutex>
 #include <memory>
 #include <string>
+#include <unordered_map>
+#include <utility>
 #include <vector>
 #include "topology.h"
 #include "transfer_metadata.h"
@@ -24,6 +31,7 @@
 namespace mooncake {
 class UbContext;
 class UbEndPoint;
+class UrmaContext;
 class TransferMetadata;
 class UbWorkerpool;
 
@@ -36,6 +44,7 @@ enum UB_ENDPOINT_TYPE { URMA_ENDPOINT = 0, OBMM_ENDPOINT = 1 };
 class UbTransport : public Transport {
     friend class UbContext;
     friend class UbEndPoint;
+    friend class UrmaContext;
     friend class UbWorkerPool;
 
    public:
@@ -82,6 +91,55 @@ class UbTransport : public Transport {
    private:
     int allocateLocalSegmentID();
 
+    struct StagingLease {
+        void* host_ptr = nullptr;
+        size_t size = 0;
+    };
+
+    struct DeviceRegion {
+        uint64_t addr = 0;
+        size_t length = 0;
+        std::string location;
+        bool remote_accessible = true;
+    };
+
+    struct StagingState {
+        void* original_device_ptr = nullptr;
+        void* staging_ptr = nullptr;
+        size_t size = 0;
+        size_t lease_size = 0;
+        TransferRequest::OpCode opcode = TransferRequest::WRITE;
+        std::atomic<uint64_t> completed_slices{0};
+        uint64_t total_slices = 0;
+        std::atomic<bool> failed{false};
+        std::mutex deferred_mutex;
+        std::vector<Slice*> deferred_success_slices;
+    };
+
+    bool stagingEnabled() const;
+    bool isDevicePointer(const void* ptr) const;
+    bool isLogicalDeviceRange(const void* ptr, size_t length) const;
+    int registerLogicalDeviceRegion(void* addr, size_t length,
+                                    const std::string& location,
+                                    bool remote_accessible);
+    int unregisterLogicalDeviceRegion(void* addr);
+    bool copyDeviceToHost(void* dst, const void* src, size_t size) const;
+    bool copyHostToDevice(void* dst, const void* src, size_t size) const;
+    Status acquireStaging(size_t size, StagingLease& lease);
+    void releaseStaging(const StagingLease& lease);
+    bool isStagedSlice(Slice* slice);
+    bool shouldDeferSuccess(Slice* slice);
+    void attachStaging(TransferTask* task,
+                       std::shared_ptr<StagingState> state);
+    void attachStagingSlice(Slice* slice,
+                            const std::shared_ptr<StagingState>& state);
+    void detachStagingSlice(Slice* slice);
+    std::shared_ptr<StagingState> stagingStateForSlice(Slice* slice);
+    void cleanupStagingForTask(TransferTask* task,
+                               bool detach_all_slices = false);
+    void onStagedSliceSuccess(Slice* slice);
+    void onStagedSliceFinalFailure(Slice* slice);
+
    public:
     int onSetupConnections(const HandShakeDesc& peer_desc,
                            HandShakeDesc& local_desc);
@@ -118,6 +176,21 @@ class UbTransport : public Transport {
     std::shared_ptr<Topology> local_topology_;
     UB_ENDPOINT_TYPE endpoint_type_;
     bool runtime_initialized_ = false;
+
+    mutable std::once_flag staging_config_once_;
+    mutable bool staging_enabled_ = false;
+    mutable std::mutex device_region_mutex_;
+    std::vector<DeviceRegion> device_regions_;
+    std::mutex staging_pool_mutex_;
+    void* staging_pool_base_ = nullptr;
+    size_t staging_pool_size_ = 0;
+    size_t staging_pool_offset_ = 0;
+    std::deque<std::pair<void*, size_t>> staging_free_list_;
+    std::mutex staging_state_mutex_;
+    std::unordered_map<TransferTask*, std::shared_ptr<StagingState>>
+        task_staging_map_;
+    std::unordered_map<Slice*, std::shared_ptr<StagingState>>
+        slice_staging_map_;
 };
 }  // namespace mooncake
 

mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h

@@ -62,6 +62,7 @@ class UrmaContext : public UbContext {
     int doProcessContextEvents() override;
     void* retrieveRemoteSeg(const std::string& value) override;
     int poll(int num_entries, Transport::Slice** failed_slices, int& num_failed,
+             std::vector<Transport::Slice*>& deferred_success_slices,
              std::unordered_map<volatile int*, int>& jetty_depth_set,
              int jfc_index) override;
     volatile int* outstandingCount(int jfc_index) override;

mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp

@@ -442,18 +442,19 @@ void UbWorkerPool::performPostSend(int thread_id) {
 void UbWorkerPool::performPoll(int thread_id) {
     int processed_slice_count = 0;
     const static size_t kPollCount = 64;
-    // context_.poll() aggregates each completion's jetty_depth here and
-    // calls markSuccess() on successful slices in place. Successful slices
-    // are NOT returned from poll(), so this worker never dereferences them
-    // after they may have been recycled by the submitting thread.
+    // context_.poll() aggregates each completion's jetty_depth here. Normal
+    // successful slices are published in place; staged READ successes are
+    // returned in deferred_success_slices so H2D can finish before publishing.
     std::unordered_map<volatile int*, int> jetty_depth_set;
     std::vector<UbTransport::Slice*> failed_slices;
+    std::vector<UbTransport::Slice*> deferred_success_slices;
     for (int jfc_index = thread_id; jfc_index < context_.jfcCount();
          jfc_index += kTransferWorkerCount) {
         UbTransport::Slice* failed[kPollCount];
         int num_failed = 0;
         int nr_poll = context_.poll(kPollCount, failed, num_failed,
-                                    jetty_depth_set, jfc_index);
+                                    deferred_success_slices, jetty_depth_set,
+                                    jfc_index);
         if (nr_poll < 0) {
             LOG(ERROR) << "Worker: Failed to poll jetty for complete";
             continue;
@@ -488,6 +489,10 @@ void UbWorkerPool::performPoll(int thread_id) {
     for (auto& entry : jetty_depth_set)
         __sync_fetch_and_sub(entry.first, entry.second);
 
+    for (auto& slice : deferred_success_slices) {
+        context_.engine().onStagedSliceSuccess(slice);
+    }
+
     // Slices that hit max_retry: final markFailed() after all reads (and the
     // jetty depth returns above) are done. Failed slices were never published
     // by poll(), so they remained safe to deref up to this point.
@@ -496,7 +501,11 @@ void UbWorkerPool::performPoll(int thread_id) {
             auto ptr = static_cast<UbEndPoint*>(slice->ub.endpoint);
             context_.deleteEndpointByPtr(ptr);
         }
-        slice->markFailed();
+        if (context_.engine().isStagedSlice(slice)) {
+            context_.engine().onStagedSliceFinalFailure(slice);
+        } else {
+            slice->markFailed();
+        }
         processed_slice_count_++;
     }
 
@@ -633,4 +642,4 @@ void UbWorkerPool::monitorWorker() {
 int UbWorkerPool::doProcessContextEvents() {
     return context_.doProcessContextEvents();
 }
-}  // namespace mooncake
+}  // namespace mooncake