diff --git a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp index cfb576974e..0311488767 100644 --- a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp +++ b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp @@ -36,6 +36,7 @@ #include #include #include +#include #include #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& 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& 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 slices; slices.reserve(task_list.size()); for (auto* task : task_list) { diff --git a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport_session_impl.h b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport_session_impl.h index 2bbdbfbac4..80c4dca0d3 100644 --- a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport_session_impl.h +++ b/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 synced_devices_; +#endif +}; + class TcpStagingBuffer { public: char* ensure(size_t size) { @@ -81,8 +143,17 @@ class TcpStagingBuffer { std::vector 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 { 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 { 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 { 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 { 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. " diff --git a/mooncake-transfer-engine/tests/tcp_cuda_staging_test.cpp b/mooncake-transfer-engine/tests/tcp_cuda_staging_test.cpp index b4e674a736..7723ac2f35 100644 --- a/mooncake-transfer-engine/tests/tcp_cuda_staging_test.cpp +++ b/mooncake-transfer-engine/tests/tcp_cuda_staging_test.cpp @@ -15,8 +15,10 @@ #include #include +#include #include #include +#include #include #include #include @@ -71,6 +73,36 @@ class DeviceBuffer { void* data_ = nullptr; }; +// a host function blocked until release(), not a long kernel: this file builds +// without nvcc, and the selected stream stays busy however long the copy takes +class StreamGate { + public: + explicit StreamGate(cudaStream_t stream = cudaStreamLegacy) + : stream_(stream) {} + + ~StreamGate() { + release(); + cudaStreamSynchronize(stream_); + } + + cudaError_t close() { return cudaLaunchHostFunc(stream_, &wait, this); } + + void release() { + if (!released_) open_.set_value(); + released_ = true; + } + + private: + static void wait(void* gate) { + static_cast(gate)->opened_.wait(); + } + + cudaStream_t stream_; + std::promise open_; + std::future opened_ = open_.get_future(); + bool released_ = false; +}; + TransferStatusEnum runOne(TransferEngine* engine, const TransferRequest& request) { auto batch_id = engine->allocateBatchID(1); @@ -152,7 +184,7 @@ TEST(TcpCudaStagingTest, ReusesStagingAcrossChunksAndRequests) { write.target_id = segment_id; write.target_offset = remote_base + kTransferSize; ASSERT_EQ(runOne(engine.get(), write), TransferStatusEnum::COMPLETED); - EXPECT_EQ(tcpTransportStagingDeviceQueryCountForTest(), 2u); + EXPECT_EQ(tcpTransportStagingDeviceQueryCountForTest(), 3u); EXPECT_EQ(tcpTransportStagingBufferAllocationCountForTest(), 2u); std::vector actual(kTransferSize); @@ -169,7 +201,7 @@ TEST(TcpCudaStagingTest, ReusesStagingAcrossChunksAndRequests) { read.target_id = segment_id; read.target_offset = remote_base + kTransferSize; ASSERT_EQ(runOne(engine.get(), read), TransferStatusEnum::COMPLETED); - EXPECT_EQ(tcpTransportStagingDeviceQueryCountForTest(), 2u); + EXPECT_EQ(tcpTransportStagingDeviceQueryCountForTest(), 3u); // The client session is new; the persistent server session retains the // staging allocation used by the preceding WRITE. EXPECT_EQ(tcpTransportStagingBufferAllocationCountForTest(), 1u); @@ -183,4 +215,134 @@ TEST(TcpCudaStagingTest, ReusesStagingAcrossChunksAndRequests) { engine.reset(); } +TEST(TcpCudaStagingTest, StagingDoesNotWaitForTheLegacyStream) { + int device_count = 0; + cudaError_t cuda_status = cudaGetDeviceCount(&device_count); + if (cuda_status != cudaSuccess || device_count == 0) { + GTEST_SKIP() << "CUDA device unavailable"; + } + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + + constexpr size_t kTransferSize = 64 * 1024; + constexpr size_t kBufferSize = 2 * kTransferSize; + + DeviceBuffer device_buffer; + ASSERT_EQ(cudaMalloc(device_buffer.out(), kBufferSize), cudaSuccess); + + auto engine = std::make_unique(false); + const std::string server_name = "127.0.0.2:17932"; + const auto hostname_port = parseHostNameWithPort(server_name); + ASSERT_EQ(engine->init(P2PHANDSHAKE, server_name, + hostname_port.first.c_str(), hostname_port.second), + 0); + ASSERT_NE(engine->installTransport("tcp", nullptr), nullptr); + ASSERT_EQ( + engine->registerLocalMemory(device_buffer.get(), kBufferSize, "cuda:0"), + 0); + + const auto segment_id = engine->openSegment(engine->getLocalIpAndPort()); + const auto segment_desc = + engine->getMetadata()->getSegmentDescByID(segment_id); + ASSERT_NE(segment_desc, nullptr); + ASSERT_FALSE(segment_desc->buffers.empty()); + const uint64_t remote_base = segment_desc->buffers[0].addr; + + StreamGate gate; + ASSERT_EQ(gate.close(), cudaSuccess); + + TransferRequest write; + write.opcode = TransferRequest::WRITE; + write.length = kTransferSize; + write.source = device_buffer.get(); + write.target_id = segment_id; + write.target_offset = remote_base + kTransferSize; + EXPECT_EQ(runOne(engine.get(), write), TransferStatusEnum::COMPLETED) + << "a WRITE of device memory waited for the legacy default stream"; + + TransferRequest read = write; + read.opcode = TransferRequest::READ; + EXPECT_EQ(runOne(engine.get(), read), TransferStatusEnum::COMPLETED) + << "a READ of device memory waited for the legacy default stream"; + + gate.release(); + EXPECT_EQ(engine->unregisterLocalMemory(device_buffer.get()), 0); + engine.reset(); +} + +TEST(TcpCudaStagingTest, SubmissionWaitsForCallerProducerBeforeStaging) { + int device_count = 0; + cudaError_t cuda_status = cudaGetDeviceCount(&device_count); + if (cuda_status != cudaSuccess || device_count == 0) { + GTEST_SKIP() << "CUDA device unavailable"; + } + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + + constexpr size_t kTransferSize = 64 * 1024; + constexpr size_t kBufferSize = 2 * kTransferSize; + + DeviceBuffer device_buffer; + ASSERT_EQ(cudaMalloc(device_buffer.out(), kBufferSize), cudaSuccess); + ASSERT_EQ(cudaMemset(device_buffer.get(), 0, kBufferSize), cudaSuccess); + + auto engine = std::make_unique(false); + const std::string server_name = "127.0.0.2:17933"; + const auto hostname_port = parseHostNameWithPort(server_name); + ASSERT_EQ(engine->init(P2PHANDSHAKE, server_name, + hostname_port.first.c_str(), hostname_port.second), + 0); + ASSERT_NE(engine->installTransport("tcp", nullptr), nullptr); + ASSERT_EQ( + engine->registerLocalMemory(device_buffer.get(), kBufferSize, "cuda:0"), + 0); + + const auto segment_id = engine->openSegment(engine->getLocalIpAndPort()); + const auto segment_desc = + engine->getMetadata()->getSegmentDescByID(segment_id); + ASSERT_NE(segment_desc, nullptr); + ASSERT_FALSE(segment_desc->buffers.empty()); + const uint64_t remote_base = segment_desc->buffers[0].addr; + + std::vector pattern(kTransferSize); + for (size_t i = 0; i < pattern.size(); ++i) + pattern[i] = static_cast((i * 37 + 11) & 0xFF); + + StreamGate gate(cudaStreamPerThread); + ASSERT_EQ(gate.close(), cudaSuccess); + ASSERT_EQ( + cudaMemcpyAsync(device_buffer.get(), pattern.data(), pattern.size(), + cudaMemcpyHostToDevice, cudaStreamPerThread), + cudaSuccess); + + TransferRequest write; + write.opcode = TransferRequest::WRITE; + write.length = kTransferSize; + write.source = device_buffer.get(); + write.target_id = segment_id; + write.target_offset = remote_base + kTransferSize; + std::atomic submission_returned{false}; + std::atomic completed_before_release{false}; + std::thread releaser([&] { + std::this_thread::sleep_for(std::chrono::milliseconds(250)); + completed_before_release.store(submission_returned.load()); + gate.release(); + }); + const auto result = runOne(engine.get(), write); + submission_returned.store(true); + releaser.join(); + EXPECT_FALSE(completed_before_release.load()) + << "submission completed before the caller's producer stream"; + ASSERT_EQ(result, TransferStatusEnum::COMPLETED); + + std::vector actual(kTransferSize); + auto* device_bytes = static_cast(device_buffer.get()); + ASSERT_EQ(cudaMemcpy(actual.data(), device_bytes + kTransferSize, + actual.size(), cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_EQ(actual, pattern) + << "staging read the source before the caller's producer finished"; + + EXPECT_EQ(engine->unregisterLocalMemory(device_buffer.get()), 0); + engine.reset(); +} + } // namespace