From 38bd84f3c38905db6cdc66b7de6b08d0d8796d09 Mon Sep 17 00:00:00 2001 From: Anatolii Rozanov Date: Tue, 16 Dec 2025 12:27:47 +0100 Subject: [PATCH] [TE]: Add HIP transport for AMD GPUs support (#1208) * [TE]: Add HIP transport for AMD GPUs support Forked from nvlink_transport and adapted for HIP/AMD GPUs. * [TE/HIP] Addressed review comments * [TE] Move NVLINK and HIP common functions to common files * [TE] Fix incorrect length assignment in relocateSharedMemoryAddress Use entry.length instead of length parameter when storing OpenedShmEntry. The length parameter represents the requested transfer length, while entry.length represents the actual buffer's full length, which is the correct value to store and is consistent with openShareableHandle usage. --- docs/source/design/transfer-engine/index.md | 5 +- .../example/transfer_engine_bench.cpp | 16 +- .../transfer_engine_bench_with_notify.cpp | 16 +- .../example/transfer_engine_validator.cpp | 16 +- .../include/common/hash_utils.h | 41 ++ .../include/common/serialization.h | 95 +++ .../include/gpu_vendor/mnnvl.h | 29 + .../include/transfer_metadata.h | 2 +- .../transport/hip_transport/hip_transport.h | 82 +++ .../nvlink_transport/nvlink_transport.h | 10 +- .../src/multi_transport.cpp | 12 +- .../src/transfer_engine.cpp | 9 + .../src/transfer_metadata.cpp | 5 +- .../src/transport/CMakeLists.txt | 9 +- .../transport/hip_transport/CMakeLists.txt | 10 + .../transport/hip_transport/hip_transport.cpp | 603 ++++++++++++++++++ .../transport/nvlink_transport/CMakeLists.txt | 8 - .../nvlink_transport/nvlink_transport.cpp | 45 +- .../tests/nvlink_transport_test.cpp | 15 +- 19 files changed, 933 insertions(+), 95 deletions(-) create mode 100644 mooncake-transfer-engine/include/common/hash_utils.h create mode 100644 mooncake-transfer-engine/include/common/serialization.h create mode 100644 mooncake-transfer-engine/include/gpu_vendor/mnnvl.h create mode 100644 mooncake-transfer-engine/include/transport/hip_transport/hip_transport.h create mode 100644 mooncake-transfer-engine/src/transport/hip_transport/CMakeLists.txt create mode 100644 mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp diff --git a/docs/source/design/transfer-engine/index.md b/docs/source/design/transfer-engine/index.md index aad56649..82348b4f 100644 --- a/docs/source/design/transfer-engine/index.md +++ b/docs/source/design/transfer-engine/index.md @@ -11,7 +11,7 @@ Mooncake Transfer Engine is a high-performance, zero-copy data transfer library As shown in the diagram, each specific client corresponds to a `TransferEngine`, which not only includes a RAM Segment but also integrates management for high-speed transfers across multiple threads and network cards. The RAM Segment, in principle, corresponds to the entire virtual address space of this `TransferEngine`, but in reality, only parts of it (known as a `Buffer`) are registered for (GPUDirect) RDMA Read/Write. Each Buffer can have separate permissions (corresponding to RDMA `rkey`, etc.) and network card affinity (e.g., preferred NICs for different types of memory). -Mooncake Transfer Engine provides interfaces through the `TransferEngine` class (located in `mooncake-transfer-engine/include/transfer_engine.h`), where the specific data transfer functions for different backends are implemented by the `Transport` class, currently supporting `TcpTransport`, `RdmaTransport`, and `NVMeoFTransport`. +Mooncake Transfer Engine provides interfaces through the `TransferEngine` class (located in `mooncake-transfer-engine/include/transfer_engine.h`), where the specific data transfer functions for different backends are implemented by the `Transport` class, currently supporting `TcpTransport`, `RdmaTransport`, `NVMeoFTransport`, `NvlinkTransport`, and `HipTransport`. ### Segment Segment represents a collection of source address ranges and target address ranges available during the data transfer process in Transfer Engine. That is, all local and remote addresses involved in `BatchTransfer` requests must be within the valid segment range. Transfer Engine supports the following two types of Segments. @@ -41,6 +41,7 @@ With the help of Transfer Engine, Mooncake Store can achieve local DRAM/VRAM rea - Local memcpy: If the target Segment is actually in the local DRAM/VRAM, direct data copy interfaces such as memcpy, cudaMemcpy are used. - TCP: Supports data transfer between local DRAM and remote DRAM. - RDMA: Supports data transfer between local DRAM/VRAM and remote DRAM. It supports multi-network card pooling and retry functions in implementation. +- HIP: Supports intra-node data transfers between GPU VRAM and GPU VRAM, as well as between GPU VRAM and CPU DRAM, using IPC handles or Shareable handles for ROCm. - cuFile (GPUDirect Storage): Implements data transfer between local DRAM/VRAM and Local/Remote NVMeof. The BatchTransfer API uses an array of requests, which specify the operation type (READ or WRITE), data length, and local and remote memory addresses. The transfer operation is applicable to DRAM and GPU VRAM. The completion of these operations can be asynchronously monitored through the `getTransferStatus` API. @@ -154,7 +155,7 @@ The following video shows a normal run as described above, with the Target on th ![transfer-engine-running](../image/transfer-engine-running.gif) ## Transfer Engine C/C++ API -Transfer Engine provides interfaces through the `TransferEngine` class (located in `mooncake-transfer-engine/include/transfer_engine.h`), where the specific data transfer functions for different backends are implemented by the `Transport` class, currently supporting `TcpTransport`, `RdmaTransport` and `NVMeoFTransport`. +Transfer Engine provides interfaces through the `TransferEngine` class (located in `mooncake-transfer-engine/include/transfer_engine.h`), where the specific data transfer functions for different backends are implemented by the `Transport` class, currently supporting `TcpTransport`, `RdmaTransport`, `NVMeoFTransport`, `NvlinkTransport` (for NVIDIA GPUs), and `HipTransport` (for AMD GPUs). ### Data Transfer diff --git a/mooncake-transfer-engine/example/transfer_engine_bench.cpp b/mooncake-transfer-engine/example/transfer_engine_bench.cpp index 3d8e200f..6fe02f35 100644 --- a/mooncake-transfer-engine/example/transfer_engine_bench.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_bench.cpp @@ -41,7 +41,7 @@ #include #ifdef USE_MNNVL -#include +#include "gpu_vendor/mnnvl.h" #endif static void checkCudaError(cudaError_t result, const char *message) { @@ -66,7 +66,7 @@ DEFINE_string(mode, "initiator", "data blocks from target node"); DEFINE_string(operation, "read", "Operation type: read or write"); -DEFINE_string(protocol, "rdma", "Transfer protocol: rdma|barex|tcp"); +DEFINE_string(protocol, "rdma", "Transfer protocol: rdma|barex|tcp|nvlink|hip"); DEFINE_string(device_name, "mlx5_2", "Device name to use, valid if protocol=rdma"); @@ -105,7 +105,7 @@ static void *allocateMemoryPool(size_t size, int buffer_id, LOG(INFO) << "Allocating memory on GPU " << gpu_id; checkCudaError(cudaSetDevice(gpu_id), "Failed to set device"); #ifdef USE_MNNVL - d_buf = mooncake::NvlinkTransport::allocatePinnedLocalMemory(size); + d_buf = allocateFabricMemory(size); #else checkCudaError(cudaMalloc(&d_buf, size), "Failed to allocate device memory"); @@ -128,7 +128,7 @@ static void freeMemoryPool(void *addr, size_t size) { #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) #ifdef USE_MNNVL if (FLAGS_use_vram) { - mooncake::NvlinkTransport::freePinnedLocalMemory(addr); + freeFabricMemory(addr); return; } #endif // USE_MNNVL @@ -319,6 +319,8 @@ int initiator() { xport = engine->installTransport("tcp", nullptr); } else if (FLAGS_protocol == "nvlink") { xport = engine->installTransport("nvlink", nullptr); + } else if (FLAGS_protocol == "hip") { + xport = engine->installTransport("hip", nullptr); } else { LOG(ERROR) << "Unsupported protocol"; } @@ -445,6 +447,8 @@ int target() { engine->installTransport("tcp", nullptr); } else if (FLAGS_protocol == "nvlink") { engine->installTransport("nvlink", nullptr); + } else if (FLAGS_protocol == "hip") { + engine->installTransport("hip", nullptr); } else { LOG(ERROR) << "Unsupported protocol"; } @@ -502,11 +506,7 @@ int target() { while (target_running) sleep(1); for (int i = 0; i < buffer_num; ++i) { engine->unregisterLocalMemory(addr[i]); -#ifdef USE_MNNVL - mooncake::NvlinkTransport::freePinnedLocalMemory(addr[i]); -#else freeMemoryPool(addr[i], FLAGS_buffer_size); -#endif } return 0; diff --git a/mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp b/mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp index e01aed21..db60fc96 100644 --- a/mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp @@ -41,7 +41,7 @@ #include #ifdef USE_MNNVL -#include +#include "gpu_vendor/mnnvl.h" #endif static void checkCudaError(cudaError_t result, const char *message) { @@ -68,7 +68,7 @@ DEFINE_string(mode, "initiator", "data blocks from target node"); DEFINE_string(operation, "read", "Operation type: read or write"); -DEFINE_string(protocol, "rdma", "Transfer protocol: rdma|tcp"); +DEFINE_string(protocol, "rdma", "Transfer protocol: rdma|tcp|nvlink|hip"); DEFINE_string(device_name, "mlx5_2", "Device name to use, valid if protocol=rdma"); @@ -101,7 +101,7 @@ static void *allocateMemoryPool(size_t size, int socket_id, void *d_buf; checkCudaError(cudaSetDevice(gpu_id), "Failed to set device"); #ifdef USE_MNNVL - d_buf = mooncake::NvlinkTransport::allocatePinnedLocalMemory(size); + d_buf = allocateFabricMemory(size); #else checkCudaError(cudaMalloc(&d_buf, size), "Failed to allocate device memory"); @@ -123,7 +123,7 @@ static void freeMemoryPool(void *addr, size_t size) { #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) #ifdef USE_MNNVL if (FLAGS_use_vram) { - mooncake::NvlinkTransport::freePinnedLocalMemory(addr); + freeFabricMemory(addr); return; } #endif // USE_MNNVL @@ -313,6 +313,8 @@ int initiator() { xport = engine->installTransport("tcp", nullptr); } else if (FLAGS_protocol == "nvlink") { xport = engine->installTransport("nvlink", nullptr); + } else if (FLAGS_protocol == "hip") { + xport = engine->installTransport("hip", nullptr); } else { LOG(ERROR) << "Unsupported protocol"; } @@ -408,6 +410,8 @@ int target() { engine->installTransport("tcp", nullptr); } else if (FLAGS_protocol == "nvlink") { engine->installTransport("nvlink", nullptr); + } else if (FLAGS_protocol == "hip") { + engine->installTransport("hip", nullptr); } else { LOG(ERROR) << "Unsupported protocol"; } @@ -447,11 +451,7 @@ int target() { } for (int i = 0; i < buffer_num; ++i) { engine->unregisterLocalMemory(addr[i]); -#ifdef USE_MNNVL - mooncake::NvlinkTransport::freePinnedLocalMemory(addr[i]); -#else freeMemoryPool(addr[i], FLAGS_buffer_size); -#endif } return 0; diff --git a/mooncake-transfer-engine/example/transfer_engine_validator.cpp b/mooncake-transfer-engine/example/transfer_engine_validator.cpp index 7533af75..6f5f98d9 100644 --- a/mooncake-transfer-engine/example/transfer_engine_validator.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_validator.cpp @@ -41,7 +41,7 @@ #include #ifdef USE_MNNVL -#include +#include "gpu_vendor/mnnvl.h" #endif static void checkCudaError(cudaError_t result, const char *message) { @@ -65,7 +65,7 @@ DEFINE_string(mode, "initiator", "Running mode: initiator or target. Initiator node read/write " "data blocks from target node"); -DEFINE_string(protocol, "rdma", "Transfer protocol: rdma|tcp"); +DEFINE_string(protocol, "rdma", "Transfer protocol: rdma|tcp|nvlink|hip"); DEFINE_string(device_name, "mlx5_2", "Device name to use, valid if protocol=rdma"); @@ -103,7 +103,7 @@ static void *allocateMemoryPool(size_t size, int buffer_id, LOG(INFO) << "Allocating memory on GPU " << gpu_id; checkCudaError(cudaSetDevice(gpu_id), "Failed to set device"); #ifdef USE_MNNVL - d_buf = mooncake::NvlinkTransport::allocatePinnedLocalMemory(size); + d_buf = allocateFabricMemory(size); #else checkCudaError(cudaMalloc(&d_buf, size), "Failed to allocate device memory"); @@ -118,7 +118,7 @@ static void freeMemoryPool(void *addr, size_t size) { #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) #ifdef USE_MNNVL if (FLAGS_use_vram) { - mooncake::NvlinkTransport::freePinnedLocalMemory(addr); + freeFabricMemory(addr); return; } #endif // USE_MNNVL @@ -397,6 +397,8 @@ int initiator() { xport = engine->installTransport("tcp", nullptr); } else if (FLAGS_protocol == "nvlink") { xport = engine->installTransport("nvlink", nullptr); + } else if (FLAGS_protocol == "hip") { + xport = engine->installTransport("hip", nullptr); } else { LOG(ERROR) << "Unsupported protocol"; } @@ -517,6 +519,8 @@ int target() { engine->installTransport("tcp", nullptr); } else if (FLAGS_protocol == "nvlink") { engine->installTransport("nvlink", nullptr); + } else if (FLAGS_protocol == "hip") { + engine->installTransport("hip", nullptr); } else { LOG(ERROR) << "Unsupported protocol"; } @@ -574,11 +578,7 @@ int target() { while (target_running) sleep(1); for (int i = 0; i < buffer_num; ++i) { engine->unregisterLocalMemory(addr[i]); -#ifdef USE_MNNVL - mooncake::NvlinkTransport::freePinnedLocalMemory(addr[i]); -#else freeMemoryPool(addr[i], FLAGS_buffer_size); -#endif } return 0; diff --git a/mooncake-transfer-engine/include/common/hash_utils.h b/mooncake-transfer-engine/include/common/hash_utils.h new file mode 100644 index 00000000..b5ffabe3 --- /dev/null +++ b/mooncake-transfer-engine/include/common/hash_utils.h @@ -0,0 +1,41 @@ +// Copyright 2025 KVCache.AI +// +// 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. + +#ifndef HASH_UTILS_H_ +#define HASH_UTILS_H_ + +#include +#include +#include + +namespace mooncake { + +/** + * @brief Hash functor for std::pair to be used with std::unordered_map + * + * This struct provides a hash function for pairs by combining the hashes + * of both elements using XOR and bit shifting. + */ +struct PairHash { + template + std::size_t operator()(const std::pair& p) const { + std::size_t h1 = std::hash{}(p.first); + std::size_t h2 = std::hash{}(p.second); + return h1 ^ (h2 << 1); + } +}; + +} // namespace mooncake + +#endif // HASH_UTILS_H_ diff --git a/mooncake-transfer-engine/include/common/serialization.h b/mooncake-transfer-engine/include/common/serialization.h new file mode 100644 index 00000000..8757f13c --- /dev/null +++ b/mooncake-transfer-engine/include/common/serialization.h @@ -0,0 +1,95 @@ +// Copyright 2025 KVCache.AI +// +// 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. + +#ifndef SERIALIZATION_H_ +#define SERIALIZATION_H_ + +#include +#include +#include +#include + +namespace mooncake { + +/** + * @brief Convert a hexadecimal character to its integer value + * + * @param c Hexadecimal character ('0'-'9', 'A'-'F', 'a'-'f') + * @return Integer value (0-15) + * @throws std::invalid_argument if character is not a valid hex digit + */ +inline int hexCharToValue(char c) { + if (c >= '0' && c <= '9') return c - '0'; + if (c >= 'A' && c <= 'F') return 10 + c - 'A'; + if (c >= 'a' && c <= 'f') return 10 + c - 'a'; + throw std::invalid_argument("Invalid hexadecimal character"); +} + +/** + * @brief Serialize binary data to hexadecimal string representation + * + * This function converts binary data (e.g., IPC memory handles) into a + * hexadecimal string for transmission or storage. + * + * @param data Pointer to binary data + * @param length Length of data in bytes + * @return Hexadecimal string representation (2 characters per byte) + * @throws std::invalid_argument if data pointer is null + */ +inline std::string serializeBinaryData(const void *data, size_t length) { + if (!data) { + throw std::invalid_argument("Data pointer cannot be null"); + } + + std::string hexString; + hexString.reserve(length * 2); + + const auto *byteData = static_cast(data); + for (size_t i = 0; i < length; ++i) { + hexString.push_back("0123456789ABCDEF"[(byteData[i] >> 4) & 0x0F]); + hexString.push_back("0123456789ABCDEF"[byteData[i] & 0x0F]); + } + + return hexString; +} + +/** + * @brief Deserialize hexadecimal string back to binary data + * + * This function converts a hexadecimal string back into binary data, + * typically for reconstructing IPC memory handles. + * + * @param hexString Hexadecimal string to deserialize + * @param buffer Output buffer to store binary data + * @throws std::invalid_argument if input string length is not even + */ +inline void deserializeBinaryData(const std::string &hexString, + std::vector &buffer) { + if (hexString.length() % 2 != 0) { + throw std::invalid_argument("Input string length must be even"); + } + + buffer.clear(); + buffer.reserve(hexString.length() / 2); + + for (size_t i = 0; i < hexString.length(); i += 2) { + int high = hexCharToValue(hexString[i]); + int low = hexCharToValue(hexString[i + 1]); + buffer.push_back(static_cast((high << 4) | low)); + } +} + +} // namespace mooncake + +#endif // SERIALIZATION_H_ diff --git a/mooncake-transfer-engine/include/gpu_vendor/mnnvl.h b/mooncake-transfer-engine/include/gpu_vendor/mnnvl.h new file mode 100644 index 00000000..8e617a6b --- /dev/null +++ b/mooncake-transfer-engine/include/gpu_vendor/mnnvl.h @@ -0,0 +1,29 @@ +// Copyright 2025 KVCache.AI +// +// 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. + +#pragma once + +#ifdef USE_HIP +#include +#define allocateFabricMemory(size) \ + mooncake::HipTransport::allocatePinnedLocalMemory(size) +#define freeFabricMemory(addr) \ + mooncake::HipTransport::freePinnedLocalMemory(addr) +#else +#include +#define allocateFabricMemory(size) \ + mooncake::NvlinkTransport::allocatePinnedLocalMemory(size) +#define freeFabricMemory(addr) \ + mooncake::NvlinkTransport::freePinnedLocalMemory(addr) +#endif diff --git a/mooncake-transfer-engine/include/transfer_metadata.h b/mooncake-transfer-engine/include/transfer_metadata.h index 8eb16d9a..3e21c6ff 100644 --- a/mooncake-transfer-engine/include/transfer_metadata.h +++ b/mooncake-transfer-engine/include/transfer_metadata.h @@ -54,7 +54,7 @@ class TransferMetadata { uint64_t length; std::vector lkey; // for rdma std::vector rkey; // for rdma - std::string shm_name; // for nvlink + std::string shm_name; // for nvlink and hip uint64_t offset; // for cxl }; diff --git a/mooncake-transfer-engine/include/transport/hip_transport/hip_transport.h b/mooncake-transfer-engine/include/transport/hip_transport/hip_transport.h new file mode 100644 index 00000000..5f4cfac2 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/hip_transport/hip_transport.h @@ -0,0 +1,82 @@ +// Copyright(C) 2025 Advanced Micro Devices, Inc. All rights reserved. + +#ifndef HIP_TRANSPORT_H_ +#define HIP_TRANSPORT_H_ + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "common/hash_utils.h" +#include "topology.h" +#include "transfer_metadata.h" +#include "transport/transport.h" + +namespace mooncake { + +class TransferMetadata; + +class HipTransport : public Transport { + public: + HipTransport(); + + ~HipTransport(); + + Status submitTransfer(BatchID batch_id, + const std::vector& entries) override; + + Status submitTransferTask( + const std::vector& task_list) override; + + Status getTransferStatus(BatchID batch_id, size_t task_id, + TransferStatus& status) override; + + static void* allocatePinnedLocalMemory(size_t length); + + static void freePinnedLocalMemory(void* addr); + + protected: + int install(std::string& local_server_name, + std::shared_ptr meta, + std::shared_ptr topo) override; + + int registerLocalMemory(void* addr, size_t length, + const std::string& location, bool remote_accessible, + bool update_metadata = true) override; + + int unregisterLocalMemory(void* addr, bool update_metadata = true) override; + + int registerLocalMemoryBatch(const std::vector& buffer_list, + const std::string& location) override; + + int unregisterLocalMemoryBatch( + const std::vector& addr_list) override; + + int relocateSharedMemoryAddress(uint64_t& dest_addr, uint64_t length, + uint64_t target_id); + + const char* getName() const override { return "hip"; } + + private: + struct OpenedShmEntry { + void* shm_addr; + uint64_t length; + }; + + std::unordered_map, OpenedShmEntry, PairHash> + remap_entries_; + RWSpinlock remap_lock_; + bool use_fabric_mem_; + + std::mutex register_mutex_; +}; + +} // namespace mooncake + +#endif // HIP_TRANSPORT_H_ diff --git a/mooncake-transfer-engine/include/transport/nvlink_transport/nvlink_transport.h b/mooncake-transfer-engine/include/transport/nvlink_transport/nvlink_transport.h index 2f7d80cf..68c6b933 100644 --- a/mooncake-transfer-engine/include/transport/nvlink_transport/nvlink_transport.h +++ b/mooncake-transfer-engine/include/transport/nvlink_transport/nvlink_transport.h @@ -14,6 +14,7 @@ #include #include +#include "common/hash_utils.h" #include "topology.h" #include "transfer_metadata.h" #include "transport/transport.h" @@ -22,15 +23,6 @@ namespace mooncake { class TransferMetadata; -struct PairHash { - template - std::size_t operator()(const std::pair& p) const { - std::size_t h1 = std::hash{}(p.first); - std::size_t h2 = std::hash{}(p.second); - return h1 ^ (h2 << 1); - } -}; - class NvlinkTransport : public Transport { public: NvlinkTransport(); diff --git a/mooncake-transfer-engine/src/multi_transport.cpp b/mooncake-transfer-engine/src/multi_transport.cpp index 3f207894..5d999e1f 100644 --- a/mooncake-transfer-engine/src/multi_transport.cpp +++ b/mooncake-transfer-engine/src/multi_transport.cpp @@ -37,8 +37,12 @@ #include "transport/ascend_transport/heterogeneous_rdma_transport.h" #endif #ifdef USE_MNNVL +#ifdef USE_HIP +#include "transport/hip_transport/hip_transport.h" +#else #include "transport/nvlink_transport/nvlink_transport.h" #endif +#endif #ifdef USE_CXL #include "transport/cxl_transport/cxl_transport.h" #endif @@ -236,10 +240,16 @@ Transport *MultiTransport::installTransport(const std::string &proto, } #endif #ifdef USE_MNNVL +#ifdef USE_HIP + else if (std::string(proto) == "hip") { + transport = new HipTransport(); + } +#else else if (std::string(proto) == "nvlink") { transport = new NvlinkTransport(); } -#endif +#endif // USE_HIP +#endif // USE_MNNVL #ifdef USE_CXL else if (std::string(proto) == "cxl") { transport = new CxlTransport(); diff --git a/mooncake-transfer-engine/src/transfer_engine.cpp b/mooncake-transfer-engine/src/transfer_engine.cpp index b6223f87..cf559c5a 100644 --- a/mooncake-transfer-engine/src/transfer_engine.cpp +++ b/mooncake-transfer-engine/src/transfer_engine.cpp @@ -245,12 +245,21 @@ int TransferEngine::init(const std::string &metadata_conn_string, return -1; } } else { +#ifdef USE_HIP + Transport *hip_transport = + multi_transports_->installTransport("hip", nullptr); + if (!hip_transport) { + LOG(ERROR) << "Failed to install HIP transport"; + return -1; + } +#else Transport *nvlink_transport = multi_transports_->installTransport("nvlink", nullptr); if (!nvlink_transport) { LOG(ERROR) << "Failed to install NVLink transport"; return -1; } +#endif } #else if (local_topology_->getHcaList().size() > 0 && diff --git a/mooncake-transfer-engine/src/transfer_metadata.cpp b/mooncake-transfer-engine/src/transfer_metadata.cpp index 1d132b9d..40b97f3e 100644 --- a/mooncake-transfer-engine/src/transfer_metadata.cpp +++ b/mooncake-transfer-engine/src/transfer_metadata.cpp @@ -238,7 +238,8 @@ int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc, rankInfoJSON["pid"] = static_cast(desc.rank_info.pid); segmentJSON["rank_info"] = rankInfoJSON; - } else if (segmentJSON["protocol"] == "nvlink") { + } else if (segmentJSON["protocol"] == "nvlink" || + segmentJSON["protocol"] == "hip") { Json::Value buffersJSON(Json::arrayValue); for (const auto &buffer : desc.buffers) { Json::Value bufferJSON; @@ -377,7 +378,7 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, } desc->buffers.push_back(buffer); } - } else if (desc->protocol == "nvlink") { + } else if (desc->protocol == "nvlink" || desc->protocol == "hip") { for (const auto &bufferJSON : segmentJSON["buffers"]) { BufferDesc buffer; buffer.name = bufferJSON["name"].asString(); diff --git a/mooncake-transfer-engine/src/transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/CMakeLists.txt index d2a4746c..446add9e 100644 --- a/mooncake-transfer-engine/src/transport/CMakeLists.txt +++ b/mooncake-transfer-engine/src/transport/CMakeLists.txt @@ -38,6 +38,11 @@ if (USE_ASCEND_HETEROGENEOUS) endif() if (USE_MNNVL) - add_subdirectory(nvlink_transport) - target_sources(transport PUBLIC $) + if (USE_HIP) + add_subdirectory(hip_transport) + target_sources(transport PUBLIC $) + else() + add_subdirectory(nvlink_transport) + target_sources(transport PUBLIC $) + endif() endif() diff --git a/mooncake-transfer-engine/src/transport/hip_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/hip_transport/CMakeLists.txt new file mode 100644 index 00000000..1213ba2d --- /dev/null +++ b/mooncake-transfer-engine/src/transport/hip_transport/CMakeLists.txt @@ -0,0 +1,10 @@ +file(GLOB HIP_TRANSPORT_SOURCES "*.cpp") + +add_library(hip_transport OBJECT ${HIP_TRANSPORT_SOURCES}) + +target_include_directories(hip_transport PUBLIC ${HIP_INCLUDE_DIRS}) + +# Treat warnings as errors for hip_transport +target_compile_options(hip_transport PRIVATE + -Werror +) diff --git a/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp b/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp new file mode 100644 index 00000000..26bc0c61 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp @@ -0,0 +1,603 @@ +// Copyright(C) 2025 Advanced Micro Devices, Inc. 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 "transport/hip_transport/hip_transport.h" + +#include + +#include +#include +#include +#include + +#include "common.h" +#include "common/serialization.h" +#include "config.h" +#include "transfer_metadata.h" +#include "transport/transport.h" + +// HIP-specific type aliases +constexpr auto HIPX_MEM_HANDLE_TYPE_FABRIC = + hipMemHandleTypePosixFileDescriptor; +using hipxFabricHandle = int; + +namespace mooncake { +static bool checkHip(hipError_t result, const char *message) { + if (result != hipSuccess) { + LOG(ERROR) << message << " (Error code: " << result << " - " + << hipGetErrorString(result) << ")"; + return false; + } + return true; +} + +static int openIPCHandle(const std::vector &buffer, + void **shm_addr) { + hipIpcMemHandle_t handle; + memcpy(&handle, buffer.data(), sizeof(handle)); + if (!checkHip(hipIpcOpenMemHandle(shm_addr, handle, + hipIpcMemLazyEnablePeerAccess), + "HipTransport: hipIpcOpenMemHandle failed")) { + return -1; + } + return 0; +} + +static int openShareableHandle(const std::vector &buffer, + size_t length, void **shm_addr) { + hipxFabricHandle export_handle; + memcpy(&export_handle, buffer.data(), sizeof(export_handle)); + + hipMemGenericAllocationHandle_t handle; + if (!checkHip(hipMemImportFromShareableHandle(&handle, &export_handle, + HIPX_MEM_HANDLE_TYPE_FABRIC), + "HipTransport: hipMemImportFromShareableHandle failed")) { + return -1; + } + + if (!checkHip(hipMemAddressReserve((hipDeviceptr_t *)shm_addr, length, 0, + nullptr, 0), + "HipTransport: hipMemAddressReserve failed")) { + return -1; + } + + if (!checkHip(hipMemMap((hipDeviceptr_t)*shm_addr, length, 0, handle, 0), + "HipTransport: hipMemMap failed")) { + (void)hipMemAddressFree((hipDeviceptr_t)*shm_addr, length); + return -1; + } + + int device_count = 0; + (void)hipGetDeviceCount(&device_count); + std::vector accessDesc(device_count); + for (int device_id = 0; device_id < device_count; ++device_id) { + accessDesc[device_id].location.type = hipMemLocationTypeDevice; + accessDesc[device_id].location.id = device_id; + accessDesc[device_id].flags = hipMemAccessFlagsProtReadWrite; + } + + if (!checkHip(hipMemSetAccess((hipDeviceptr_t)*shm_addr, length, + accessDesc.data(), device_count), + "HipTransport: hipMemSetAccess failed")) { + (void)hipMemUnmap((hipDeviceptr_t)*shm_addr, length); + (void)hipMemAddressFree((hipDeviceptr_t)*shm_addr, length); + return -1; + } + + return 0; +} + +static bool supportFabricMem() { + // For HIP transport, prefer HIP-specific env var, but also check NVLINK for + // backward compatibility. + if (getenv("MC_USE_HIP_IPC") || getenv("MC_USE_NVLINK_IPC")) return false; + + int num_devices = 0; + if (!checkHip(hipGetDeviceCount(&num_devices), + "HipTransport: hipGetDeviceCount failed")) { + return false; + } + + if (num_devices == 0) { + LOG(ERROR) << "HipTransport: no device found"; + return false; + } + + // Check if all devices support virtual memory management, + // which is required for fabric memory operations + for (int device_id = 0; device_id < num_devices; ++device_id) { + hipDevice_t device; + if (!checkHip(hipDeviceGet(&device, device_id), + "HipTransport: hipDeviceGet failed")) { + return false; + } + + int vmm_supported = 0; + hipError_t result = hipDeviceGetAttribute( + &vmm_supported, hipDeviceAttributeVirtualMemoryManagementSupported, + device); + if (result != hipSuccess || !vmm_supported) { + LOG(WARNING) << "HipTransport: Device " << device_id + << " does not support virtual memory management, " + << "falling back to IPC mode"; + return false; + } + } + + return true; +} + +HipTransport::HipTransport() : use_fabric_mem_(supportFabricMem()) {} + +HipTransport::~HipTransport() { + if (use_fabric_mem_) { + for (auto &entry : remap_entries_) { + freePinnedLocalMemory(entry.second.shm_addr); + } + } else { + for (auto &entry : remap_entries_) { + (void)hipIpcCloseMemHandle(entry.second.shm_addr); + } + } + remap_entries_.clear(); +} + +int HipTransport::install(std::string &local_server_name, + std::shared_ptr metadata, + std::shared_ptr topology) { + metadata_ = metadata; + local_server_name_ = local_server_name; + + auto desc = std::make_shared(); + if (!desc) return ERR_MEMORY; + + desc->name = local_server_name_; + desc->protocol = "hip"; + + metadata_->addLocalSegment(LOCAL_SEGMENT_ID, local_server_name_, + std::move(desc)); + return 0; +} + +Status HipTransport::submitTransfer( + BatchID batch_id, const std::vector &entries) { + auto &batch_desc = *((BatchDesc *)(batch_id)); + if (batch_desc.task_list.size() + entries.size() > batch_desc.batch_size) { + LOG(ERROR) << "HipTransport: Exceed the limitation of current batch's " + "capacity"; + return Status::InvalidArgument( + "HipTransport: Exceed the limitation of capacity, batch id: " + + std::to_string(batch_id)); + } + + size_t task_id = batch_desc.task_list.size(); + batch_desc.task_list.resize(task_id + entries.size()); + + for (auto &request : entries) { + TransferTask &task = batch_desc.task_list[task_id]; + ++task_id; + + uint64_t dest_addr = request.target_offset; + if (request.target_id != LOCAL_SEGMENT_ID) { + int rc = relocateSharedMemoryAddress(dest_addr, request.length, + request.target_id); + if (rc) return Status::Memory("device memory not registered"); + } + + task.total_bytes = request.length; + + // Allocate and configure slice + Slice *slice = getSliceCache().allocate(); + slice->source_addr = (char *)request.source; + slice->local.dest_addr = (char *)dest_addr; + slice->length = request.length; + slice->opcode = request.opcode; + slice->task = &task; + slice->target_id = request.target_id; + slice->status = Slice::PENDING; + __sync_fetch_and_add(&task.slice_count, 1); + + // Execute memory copy + hipError_t err; + if (slice->opcode == TransferRequest::READ) { + err = hipMemcpy(slice->source_addr, (void *)slice->local.dest_addr, + slice->length, hipMemcpyDefault); + } else { + err = hipMemcpy((void *)slice->local.dest_addr, slice->source_addr, + slice->length, hipMemcpyDefault); + } + + // Mark completion status + if (err != hipSuccess) { + slice->markFailed(); + } else { + slice->markSuccess(); + } + } + + return Status::OK(); +} + +Status HipTransport::getTransferStatus(BatchID batch_id, size_t task_id, + TransferStatus &status) { + auto &batch_desc = *((BatchDesc *)(batch_id)); + const size_t task_count = batch_desc.task_list.size(); + + if (task_id >= task_count) { + return Status::InvalidArgument( + "HipTransport::getTransferStatus invalid argument, batch id: " + + std::to_string(batch_id)); + } + + // Get task and status info + auto &task = batch_desc.task_list[task_id]; + status.transferred_bytes = task.transferred_bytes; + uint64_t success_slice_count = task.success_slice_count; + uint64_t failed_slice_count = task.failed_slice_count; + + // Determine completion status + if (success_slice_count + failed_slice_count == task.slice_count) { + if (failed_slice_count) { + status.s = TransferStatusEnum::FAILED; + } else { + status.s = TransferStatusEnum::COMPLETED; + } + task.is_finished = true; + } else { + status.s = TransferStatusEnum::WAITING; + } + + return Status::OK(); +} + +Status HipTransport::submitTransferTask( + const std::vector &task_list) { + for (auto *task_ptr : task_list) { + assert(task_ptr); + auto &task = *task_ptr; + assert(task.request); + auto &request = *task.request; + + uint64_t dest_addr = request.target_offset; + if (request.target_id != LOCAL_SEGMENT_ID) { + int rc = relocateSharedMemoryAddress(dest_addr, request.length, + request.target_id); + if (rc) return Status::Memory("device memory not registered"); + } + + task.total_bytes = request.length; + + // Allocate and configure slice + Slice *slice = getSliceCache().allocate(); + slice->source_addr = (char *)request.source; + slice->local.dest_addr = (char *)dest_addr; + slice->length = request.length; + slice->opcode = request.opcode; + slice->task = &task; + slice->target_id = request.target_id; + slice->status = Slice::PENDING; + task.slice_list.push_back(slice); + __sync_fetch_and_add(&task.slice_count, 1); + + // Execute memory copy + hipError_t err; + if (slice->opcode == TransferRequest::READ) { + err = hipMemcpy(slice->source_addr, (void *)slice->local.dest_addr, + slice->length, hipMemcpyDefault); + } else { + err = hipMemcpy((void *)slice->local.dest_addr, slice->source_addr, + slice->length, hipMemcpyDefault); + } + + // Mark completion status + if (err != hipSuccess) { + slice->markFailed(); + } else { + slice->markSuccess(); + } + } + + return Status::OK(); +} + +int HipTransport::registerLocalMemory(void *addr, size_t length, + const std::string &location, + bool remote_accessible, + bool update_metadata) { + std::lock_guard lock(register_mutex_); + + if (globalConfig().trace) { + LOG(INFO) << "register memory: addr " << addr << ", length " << length; + } + + // IPC-based memory registration + if (!use_fabric_mem_) { + // Validate memory type + hipPointerAttribute_t attr; + if (!checkHip(hipPointerGetAttributes(&attr, addr), + "HipTransport: hipPointerGetAttributes failed")) { + return -1; + } + + if (attr.type != hipMemoryTypeDevice) { + LOG(ERROR) << "Unsupported memory type, " << addr << " " + << attr.type; + return -1; + } + + // Get IPC handle + hipIpcMemHandle_t handle; + if (!checkHip(hipIpcGetMemHandle(&handle, addr), + "HipTransport: hipIpcGetMemHandle failed")) { + return -1; + } + + // Register buffer with metadata + (void)remote_accessible; + BufferDesc desc; + desc.addr = (uint64_t)addr; + desc.length = length; + desc.name = location; + desc.shm_name = serializeBinaryData(&handle, sizeof(hipIpcMemHandle_t)); + return metadata_->addLocalMemoryBuffer(desc, true); + } + + // Fabric memory registration + else { + // Retain allocation handle + hipMemGenericAllocationHandle_t handle; + hipError_t result = hipMemRetainAllocationHandle(&handle, addr); + if (result != hipSuccess) { + LOG(WARNING) << "Memory region " << addr + << " is not allocated by hipMemCreate, " + << "but it can be used as local buffer"; + return 0; + } + + // Find whole physical page for memory registration + void *real_addr; + size_t real_size; + result = hipMemGetAddressRange((hipDeviceptr_t *)&real_addr, &real_size, + (hipDeviceptr_t)addr); + if (result != hipSuccess) { + LOG(WARNING) << "HipTransport: hipMemGetAddressRange failed: " + << result; + const uint64_t granularity = 2ULL * 1024 * 1024; + real_addr = addr; + real_size = (length + granularity - 1) & ~(granularity - 1); + } + + // Export shareable handle + hipxFabricHandle export_handle_raw; + if (!checkHip( + hipMemExportToShareableHandle(&export_handle_raw, handle, + HIPX_MEM_HANDLE_TYPE_FABRIC, 0), + "HipTransport: hipMemExportToShareableHandle failed")) { + return -1; + } + + (void)remote_accessible; + BufferDesc desc; + desc.addr = (uint64_t)real_addr; + desc.length = real_size; + desc.name = location; + desc.shm_name = serializeBinaryData((const void *)&export_handle_raw, + sizeof(hipxFabricHandle)); + return metadata_->addLocalMemoryBuffer(desc, true); + } +} + +int HipTransport::unregisterLocalMemory(void *addr, bool update_metadata) { + return metadata_->removeLocalMemoryBuffer(addr, update_metadata); +} + +int HipTransport::relocateSharedMemoryAddress(uint64_t &dest_addr, + uint64_t length, + uint64_t target_id) { + auto desc = metadata_->getSegmentDescByID(target_id); + + // Search for matching buffer entry + for (auto &entry : desc->buffers) { + if (!entry.shm_name.empty() && entry.addr <= dest_addr && + dest_addr + length <= entry.addr + entry.length) { + // Check if already remapped (shared lock) + remap_lock_.lockShared(); + if (remap_entries_.count(std::make_pair(target_id, entry.addr))) { + auto shm_addr = + remap_entries_[std::make_pair(target_id, entry.addr)] + .shm_addr; + remap_lock_.unlockShared(); + dest_addr = dest_addr - entry.addr + ((uint64_t)shm_addr); + return 0; + } + remap_lock_.unlockShared(); + + RWSpinlock::WriteGuard lock_guard(remap_lock_); + if (!remap_entries_.count(std::make_pair(target_id, entry.addr))) { + std::vector output_buffer; + deserializeBinaryData(entry.shm_name, output_buffer); + + void *shm_addr = nullptr; + int rc = -1; + + if (output_buffer.size() == sizeof(hipIpcMemHandle_t) && + !use_fabric_mem_) { + rc = openIPCHandle(output_buffer, &shm_addr); + } else if (output_buffer.size() == sizeof(hipxFabricHandle) && + use_fabric_mem_) { + rc = openShareableHandle(output_buffer, entry.length, + &shm_addr); + } else { + LOG(ERROR) << "Mismatched HIP data transfer method"; + return -1; + } + + if (rc != 0) { + return -1; + } + + OpenedShmEntry shm_entry; + shm_entry.shm_addr = shm_addr; + shm_entry.length = entry.length; + remap_entries_[std::make_pair(target_id, entry.addr)] = + shm_entry; + } + + // Calculate relocated address + auto shm_addr = + remap_entries_[std::make_pair(target_id, entry.addr)].shm_addr; + dest_addr = dest_addr - entry.addr + ((uint64_t)shm_addr); + return 0; + } + } + LOG(ERROR) << "Requested address " << (void *)dest_addr << " to " + << (void *)(dest_addr + length) << " not found!"; + return ERR_INVALID_ARGUMENT; +} + +int HipTransport::registerLocalMemoryBatch( + const std::vector &buffer_list, + const std::string &location) { + for (auto &buffer : buffer_list) { + int rc = registerLocalMemory(buffer.addr, buffer.length, location, true, + false); + if (rc < 0) return rc; + } + return metadata_->updateLocalSegmentDesc(); +} + +int HipTransport::unregisterLocalMemoryBatch( + const std::vector &addr_list) { + for (auto &addr : addr_list) { + int rc = unregisterLocalMemory(addr, false); + if (rc < 0) return rc; + } + return metadata_->updateLocalSegmentDesc(); +} + +void *HipTransport::allocatePinnedLocalMemory(size_t size) { + if (!supportFabricMem()) { + void *ptr = nullptr; + if (!checkHip(hipMalloc(&ptr, size), + "HipTransport: hipMalloc failed")) { + return nullptr; + } + return ptr; + } + + size_t granularity = 0; + hipDevice_t currentDev; + hipMemAllocationProp prop = {}; + hipMemGenericAllocationHandle_t handle; + void *ptr = nullptr; + int hipDev; + int flag = 0; + + if (!checkHip(hipGetDevice(&hipDev), "HipTransport: hipGetDevice failed")) { + return nullptr; + } + + if (!checkHip(hipDeviceGet(¤tDev, hipDev), + "HipTransport: hipDeviceGet failed")) { + return nullptr; + } + + prop.type = hipMemAllocationTypePinned; + prop.location.type = hipMemLocationTypeDevice; + prop.requestedHandleType = HIPX_MEM_HANDLE_TYPE_FABRIC; + prop.location.id = currentDev; + + hipError_t result = hipDeviceGetAttribute( + &flag, hipDeviceAttributeVirtualMemoryManagementSupported, currentDev); + if (!checkHip(result, "HipTransport: hipDeviceGetAttribute failed")) { + return nullptr; + } + + if (flag) prop.allocFlags.gpuDirectRDMACapable = 1; + + result = hipMemGetAllocationGranularity(&granularity, &prop, + hipMemAllocationGranularityMinimum); + if (!checkHip(result, + "HipTransport: hipMemGetAllocationGranularity failed")) { + return nullptr; + } + + size = (size + granularity - 1) & ~(granularity - 1); + if (size == 0) size = granularity; + + result = hipMemCreate(&handle, size, &prop, 0); + if (!checkHip(result, "HipTransport: hipMemCreate failed")) { + return nullptr; + } + + result = hipMemAddressReserve((hipDeviceptr_t *)&ptr, size, granularity, + nullptr, 0); + if (!checkHip(result, "HipTransport: hipMemAddressReserve failed")) { + (void)hipMemRelease(handle); + return nullptr; + } + + result = hipMemMap((hipDeviceptr_t)ptr, size, 0, handle, 0); + if (!checkHip(result, "HipTransport: hipMemMap failed")) { + (void)hipMemAddressFree((hipDeviceptr_t)ptr, size); + (void)hipMemRelease(handle); + return nullptr; + } + + int device_count = 0; + (void)hipGetDeviceCount(&device_count); + std::vector accessDesc(device_count); + for (int idx = 0; idx < device_count; ++idx) { + accessDesc[idx].location.type = hipMemLocationTypeDevice; + accessDesc[idx].location.id = idx; + accessDesc[idx].flags = hipMemAccessFlagsProtReadWrite; + } + + result = hipMemSetAccess((hipDeviceptr_t)ptr, size, accessDesc.data(), + device_count); + if (!checkHip(result, "HipTransport: hipMemSetAccess failed")) { + (void)hipMemUnmap((hipDeviceptr_t)ptr, size); + (void)hipMemAddressFree((hipDeviceptr_t)ptr, size); + (void)hipMemRelease(handle); + return nullptr; + } + + return ptr; +} + +void HipTransport::freePinnedLocalMemory(void *ptr) { + if (!supportFabricMem()) { + (void)hipFree(ptr); + return; + } + + hipMemGenericAllocationHandle_t handle; + size_t size = 0; + + if (!checkHip(hipMemRetainAllocationHandle(&handle, ptr), + "HipTransport: hipMemRetainAllocationHandle failed")) { + return; + } + + hipDeviceptr_t base = 0; + hipError_t result = + hipMemGetAddressRange(&base, &size, (hipDeviceptr_t)ptr); + if (checkHip(result, "HipTransport: hipMemGetAddressRange")) { + (void)hipMemUnmap((hipDeviceptr_t)ptr, size); + (void)hipMemAddressFree((hipDeviceptr_t)ptr, size); + } + + (void)hipMemRelease(handle); +} +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/nvlink_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/nvlink_transport/CMakeLists.txt index 0b9c03a3..88461699 100644 --- a/mooncake-transfer-engine/src/transport/nvlink_transport/CMakeLists.txt +++ b/mooncake-transfer-engine/src/transport/nvlink_transport/CMakeLists.txt @@ -1,15 +1,7 @@ file(GLOB NVLINK_SOURCES "*.cpp") -if (USE_HIP) - hipify_files(NVLINK_SOURCES) -endif() - add_library(nvlink_transport OBJECT ${NVLINK_SOURCES}) if (USE_CUDA) target_include_directories(nvlink_transport PUBLIC CUDA::cudart "/usr/local/cuda/include") endif() - -if (USE_HIP) -target_include_directories(nvlink_transport PUBLIC ${HIP_INCLUDE_DIRS}) -endif() diff --git a/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp b/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp index b7213c06..7cf80652 100644 --- a/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp +++ b/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp @@ -26,6 +26,7 @@ #include #include "common.h" +#include "common/serialization.h" #include "config.h" #include "transfer_engine.h" #include "transfer_metadata.h" @@ -302,46 +303,6 @@ Status NvlinkTransport::submitTransferTask( return Status::OK(); } -int hexCharToValue(char c) { - if (c >= '0' && c <= '9') return c - '0'; - if (c >= 'A' && c <= 'F') return 10 + c - 'A'; - if (c >= 'a' && c <= 'f') return 10 + c - 'a'; - throw std::invalid_argument("Invalid hexadecimal character"); -} - -std::string serializeBinaryData(const void *data, size_t length) { - if (!data) { - throw std::invalid_argument("Data pointer cannot be null"); - } - - std::string hexString; - hexString.reserve(length * 2); - - const unsigned char *byteData = static_cast(data); - for (size_t i = 0; i < length; ++i) { - hexString.push_back("0123456789ABCDEF"[(byteData[i] >> 4) & 0x0F]); - hexString.push_back("0123456789ABCDEF"[byteData[i] & 0x0F]); - } - - return hexString; -} - -void deserializeBinaryData(const std::string &hexString, - std::vector &buffer) { - if (hexString.length() % 2 != 0) { - throw std::invalid_argument("Input string length must be even"); - } - - buffer.clear(); - buffer.reserve(hexString.length() / 2); - - for (size_t i = 0; i < hexString.length(); i += 2) { - int high = hexCharToValue(hexString[i]); - int low = hexCharToValue(hexString[i + 1]); - buffer.push_back(static_cast((high << 4) | low)); - } -} - int NvlinkTransport::registerLocalMemory(void *addr, size_t length, const std::string &location, bool remote_accessible, @@ -464,7 +425,7 @@ int NvlinkTransport::relocateSharedMemoryAddress(uint64_t &dest_addr, } OpenedShmEntry shm_entry; shm_entry.shm_addr = shm_addr; - shm_entry.length = length; + shm_entry.length = entry.length; remap_entries_[std::make_pair(target_id, entry.addr)] = shm_entry; } else if (output_buffer.size() == sizeof(CUmemFabricHandle) && @@ -518,7 +479,7 @@ int NvlinkTransport::relocateSharedMemoryAddress(uint64_t &dest_addr, } OpenedShmEntry shm_entry; shm_entry.shm_addr = shm_addr; - shm_entry.length = length; + shm_entry.length = entry.length; remap_entries_[std::make_pair(target_id, entry.addr)] = shm_entry; } else { diff --git a/mooncake-transfer-engine/tests/nvlink_transport_test.cpp b/mooncake-transfer-engine/tests/nvlink_transport_test.cpp index bce7029b..f486e48a 100644 --- a/mooncake-transfer-engine/tests/nvlink_transport_test.cpp +++ b/mooncake-transfer-engine/tests/nvlink_transport_test.cpp @@ -11,6 +11,13 @@ using namespace mooncake; +// Select protocol based on build configuration +#ifdef USE_HIP +#define MNNVL_PROTOCOL "hip" +#else +#define MNNVL_PROTOCOL "nvlink" +#endif + DEFINE_string(metadata_server, "127.0.0.1:2379", "etcd server host address"); DEFINE_string(local_server_name, "cuda_server:12345", "Local server name"); DEFINE_string(segment_id, "cuda_server:12345", "Segment ID to access data"); @@ -44,9 +51,9 @@ TEST(NvlinkTransportTest, WriteAndRead) { auto server_engine = std::make_unique(false); server_engine->init(FLAGS_metadata_server, FLAGS_local_server_name); - // Install NvlinkTransport on server + // Install MNNVL transport (nvlink or hip) on server Transport* server_transport = - server_engine->installTransport("nvlink", nullptr); + server_engine->installTransport(MNNVL_PROTOCOL, nullptr); ASSERT_NE(server_transport, nullptr); void* server_buffer = allocateCudaBuffer(kDataLength * 2, gpu_id); @@ -60,9 +67,9 @@ TEST(NvlinkTransportTest, WriteAndRead) { auto client_engine = std::make_unique(false); client_engine->init(FLAGS_metadata_server, "cuda_client:12346"); - // Install NvlinkTransport on client + // Install MNNVL transport (nvlink or hip) on client Transport* client_transport = - client_engine->installTransport("nvlink", nullptr); + client_engine->installTransport(MNNVL_PROTOCOL, nullptr); ASSERT_NE(client_transport, nullptr); void* client_buffer = allocateCudaBuffer(kDataLength * 2, gpu_id);