[Bugfix][TransferEngine] Stage TCP CUDA copies off the legacy stream

mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp

@@ -36,6 +36,7 @@
 #include <mutex>
 #include <optional>
 #include <type_traits>
+#include <unordered_set>
 #include <vector>
 
 #include "common.h"
@@ -552,6 +553,16 @@ Status TcpTransport::submitTransfer(
             std::to_string(batch_id));
     }
 
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX) || defined(USE_SUPA)
+    TcpCallerStreamSync caller_sync;
+    for (const auto& request : entries) {
+        auto status = caller_sync.sync(request);
+        if (!status.ok()) return status;
+    }
+#endif
+
     size_t task_id = batch_desc.task_list.size();
     batch_desc.task_list.resize(task_id + entries.size());
 
@@ -566,6 +577,19 @@ Status TcpTransport::submitTransfer(
 
 Status TcpTransport::submitTransferTask(
     const std::vector<TransferTask*>& task_list) {
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX) || defined(USE_SUPA)
+    TcpCallerStreamSync caller_sync;
+    for (auto* task : task_list) {
+        assert(task && task->request);
+        for (size_t j = 0; j < task->request_count; ++j) {
+            auto status = caller_sync.sync(task->request[j]);
+            if (!status.ok()) return status;
+        }
+    }
+#endif
+
     for (size_t i = 0; i < task_list.size();) {
         auto* task = task_list[i];
         assert(task && task->request);
@@ -594,6 +618,19 @@ Status TcpTransport::submitTransferTask(
 
 Status TcpTransport::submitTransferTaskGroup(
     const std::vector<TransferTask*>& task_list) {
+#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) ||  \
+    defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \
+    defined(USE_COREX) || defined(USE_SUPA)
+    TcpCallerStreamSync caller_sync;
+    for (auto* task : task_list) {
+        assert(task && task->request);
+        for (size_t j = 0; j < task->request_count; ++j) {
+            auto status = caller_sync.sync(task->request[j]);
+            if (!status.ok()) return status;
+        }
+    }
+#endif
+
     std::vector<Slice*> slices;
     slices.reserve(task_list.size());
     for (auto* task : task_list) {

mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport_session_impl.h

@@ -65,6 +65,68 @@ static int getCudaDeviceId(void* addr) {
     return -1;
 }
 
+#ifndef USE_MLU
+static cudaError_t syncTcpCallerStream(int device) {
+    int previous_device = 0;
+    cudaError_t status = cudaGetDevice(&previous_device);
+    if (status != cudaSuccess) return status;
+    if (previous_device != device) {
+        status = cudaSetDevice(device);
+        if (status != cudaSuccess) return status;
+    }
+
+    cudaEvent_t event = nullptr;
+    status = cudaEventCreateWithFlags(&event, cudaEventDisableTiming);
+    if (status == cudaSuccess) {
+        status = cudaEventRecord(event, cudaStreamPerThread);
+    }
+    if (status == cudaSuccess) {
+        status = cudaEventSynchronize(event);
+    }
+    if (event) {
+        const cudaError_t destroy_status = cudaEventDestroy(event);
+        if (status == cudaSuccess) status = destroy_status;
+    }
+    if (previous_device != device) {
+        const cudaError_t restore_status = cudaSetDevice(previous_device);
+        if (status == cudaSuccess) status = restore_status;
+    }
+    return status;
+}
+#endif
+
+// A private staging stream no longer inherits ordering from producer work on
+// the caller's stream. Establish that dependency on the submitting thread,
+// once per device and submission, before any network I/O starts.
+class TcpCallerStreamSync {
+   public:
+    Status sync(const TransferRequest& request) {
+#ifdef USE_MLU
+        // The MLU copy below deliberately stays synchronous because its shim
+        // has no event or per-thread-stream API.
+        (void)request;
+        return Status::OK();
+#else
+        const int device = getCudaDeviceId(request.source);
+        if (device < 0 || !synced_devices_.insert(device).second) {
+            return Status::OK();
+        }
+        const cudaError_t status = syncTcpCallerStream(device);
+        if (status != cudaSuccess) {
+            return Status::Context(
+                "TCP caller-stream synchronization failed: " +
+                std::string(cudaGetErrorString(status)));
+        }
+        return Status::OK();
+#endif
+    }
+
+   private:
+#ifndef USE_MLU
+    std::unordered_set<int> synced_devices_;
+#endif
+};
+
 class TcpStagingBuffer {
    public:
     char* ensure(size_t size) {
@@ -81,8 +143,17 @@ class TcpStagingBuffer {
     std::vector<char> buffer_;
 };
 
-#ifdef USE_MACA
+// cudaMemcpy on the legacy default stream waits behind every kernel the
+// process already queued, and the calling thread sits in it for the whole
+// wait. A private non-blocking stream keeps the copy off that lane; the sync
+// at the end preserves the copy's completion contract for the caller.
 static cudaError_t copyTcpCudaMemory(void* dst, const void* src, size_t size) {
+#ifdef USE_MLU
+    // CNRT exposes neither the event/per-thread dependency API used above nor
+    // CUDA stream flags. Keep the previous synchronous copy so this backend
+    // does not trade producer ordering for a private queue.
+    return cudaMemcpy(dst, src, size, cudaMemcpyDefault);
+#else
     cudaStream_t stream;
     cudaError_t status =
         cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking);
@@ -95,8 +166,8 @@ static cudaError_t copyTcpCudaMemory(void* dst, const void* src, size_t size) {
 
     cudaError_t destroy_status = cudaStreamDestroy(stream);
     return status == cudaSuccess ? destroy_status : status;
-}
 #endif
+}
 #endif
 
 // Forward declaration
@@ -253,14 +324,8 @@ struct ServerSession : public std::enable_shared_from_this<ServerSession> {
         if (cuda_device_ >= 0) {
             dram_buffer = staging_buffer_.ensure(buffer_size);
             cudaSetDevice(cuda_device_);
-#ifdef USE_MACA
             cudaError_t cuda_status = copyTcpCudaMemory(
                 dram_buffer, addr + total_transferred_bytes_, buffer_size);
-#else
-            cudaError_t cuda_status =
-                cudaMemcpy(dram_buffer, addr + total_transferred_bytes_,
-                           buffer_size, cudaMemcpyDefault);
-#endif
             if (cuda_status != cudaSuccess) {
                 LOG(ERROR) << "ServerSession::writeBody failed to copy from "
                               "CUDA memory. "
@@ -342,15 +407,9 @@ struct ServerSession : public std::enable_shared_from_this<ServerSession> {
     defined(USE_COREX) || defined(USE_SUPA)
                 if (cuda_device_ >= 0) {
                     cudaSetDevice(cuda_device_);
-#ifdef USE_MACA
                     cudaError_t cuda_status =
                         copyTcpCudaMemory(addr + total_transferred_bytes_,
                                           dram_buffer, transferred_bytes);
-#else
-                    cudaError_t cuda_status =
-                        cudaMemcpy(addr + total_transferred_bytes_, dram_buffer,
-                                   transferred_bytes, cudaMemcpyDefault);
-#endif
                     if (cuda_status != cudaSuccess) {
                         LOG(ERROR)
                             << "ServerSession::readBody failed to copy to CUDA "
@@ -803,15 +862,9 @@ struct ClientSession : public std::enable_shared_from_this<ClientSession> {
     defined(USE_COREX) || defined(USE_SUPA)
                 if (cuda_device_ >= 0) {
                     cudaSetDevice(cuda_device_);
-#ifdef USE_MACA
                     cudaError_t cuda_status =
                         copyTcpCudaMemory(addr + total_transferred_bytes_,
                                           dram_buffer, transferred_bytes);
-#else
-                    cudaError_t cuda_status =
-                        cudaMemcpy(addr + total_transferred_bytes_, dram_buffer,
-                                   transferred_bytes, cudaMemcpyDefault);
-#endif
                     if (cuda_status != cudaSuccess) {
                         LOG(ERROR)
                             << "ClientSession::readBody failed to copy to CUDA "
@@ -868,14 +921,8 @@ struct ClientSession : public std::enable_shared_from_this<ClientSession> {
         if (cuda_device_ >= 0) {
             dram_buffer = staging_buffer_.ensure(buffer_size);
             cudaSetDevice(cuda_device_);
-#ifdef USE_MACA
             cudaError_t cuda_status = copyTcpCudaMemory(
                 dram_buffer, addr + total_transferred_bytes_, buffer_size);
-#else
-            cudaError_t cuda_status =
-                cudaMemcpy(dram_buffer, addr + total_transferred_bytes_,
-                           buffer_size, cudaMemcpyDefault);
-#endif
             if (cuda_status != cudaSuccess) {
                 LOG(ERROR) << "ClientSession::writeBody failed to copy from "
                               "CUDA memory. "