mooncake-track
/
Mooncake
0
0
Fork 12
Code Issues Pull Requests 13 Actions Packages Projects Releases Wiki Activity

[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.
This commit is contained in:
Anatolii Rozanov 2025-12-16 12:27:47 +01:00 committed by GitHub
parent 284a3edfec
commit 38bd84f3c3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
19 changed files with 933 additions and 95 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
docs/source/design/transfer-engine/index.md
View File

@ -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

16
mooncake-transfer-engine/example/transfer_engine_bench.cpp
View File

@ -41,7 +41,7 @@
#include <cassert>
#ifdef USE_MNNVL
#include <transport/nvlink_transport/nvlink_transport.h>
#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;

16
mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp
View File

@ -41,7 +41,7 @@
#include <cassert>
#ifdef USE_MNNVL
#include <transport/nvlink_transport/nvlink_transport.h>
#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;

16
mooncake-transfer-engine/example/transfer_engine_validator.cpp
View File

@ -41,7 +41,7 @@
#include <cassert>
#ifdef USE_MNNVL
#include <transport/nvlink_transport/nvlink_transport.h>
#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;

41
mooncake-transfer-engine/include/common/hash_utils.h Normal file
View File

@ -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 <cstddef>
#include <functional>
#include <utility>
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 <typename T1, typename T2>
std::size_t operator()(const std::pair<T1, T2>& p) const {
std::size_t h1 = std::hash<T1>{}(p.first);
std::size_t h2 = std::hash<T2>{}(p.second);
return h1 ^ (h2 << 1);
}
};
} // namespace mooncake
#endif // HASH_UTILS_H_

95
mooncake-transfer-engine/include/common/serialization.h Normal file
View File

@ -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 <cstddef>
#include <stdexcept>
#include <string>
#include <vector>
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<const unsigned char *>(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<unsigned char> &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<unsigned char>((high << 4) | low));
}
}
} // namespace mooncake
#endif // SERIALIZATION_H_

29
mooncake-transfer-engine/include/gpu_vendor/mnnvl.h Normal file
View File

@ -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 <transport/hip_transport/hip_transport.h>
#define allocateFabricMemory(size) \
mooncake::HipTransport::allocatePinnedLocalMemory(size)
#define freeFabricMemory(addr) \
mooncake::HipTransport::freePinnedLocalMemory(addr)
#else
#include <transport/nvlink_transport/nvlink_transport.h>
#define allocateFabricMemory(size) \
mooncake::NvlinkTransport::allocatePinnedLocalMemory(size)
#define freeFabricMemory(addr) \
mooncake::NvlinkTransport::freePinnedLocalMemory(addr)
#endif

2
mooncake-transfer-engine/include/transfer_metadata.h
View File

@ -54,7 +54,7 @@ class TransferMetadata {
uint64_t length;
std::vector<uint32_t> lkey; // for rdma
std::vector<uint32_t> rkey; // for rdma
std::string shm_name; // for nvlink
std::string shm_name; // for nvlink and hip
uint64_t offset; // for cxl
};

82
mooncake-transfer-engine/include/transport/hip_transport/hip_transport.h Normal file
View File

@ -0,0 +1,82 @@
// Copyright(C) 2025 Advanced Micro Devices, Inc. All rights reserved.
#ifndef HIP_TRANSPORT_H_
#define HIP_TRANSPORT_H_
#include <hip/hip_runtime.h>
#include <functional>
#include <iostream>
#include <memory>
#include <mutex>
#include <string>
#include <vector>
#include <utility>
#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<TransferRequest>& entries) override;
Status submitTransferTask(
const std::vector<TransferTask*>& 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<TransferMetadata> meta,
std::shared_ptr<Topology> 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<BufferEntry>& buffer_list,
const std::string& location) override;
int unregisterLocalMemoryBatch(
const std::vector<void*>& 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<std::pair<uint64_t, uint64_t>, OpenedShmEntry, PairHash>
remap_entries_;
RWSpinlock remap_lock_;
bool use_fabric_mem_;
std::mutex register_mutex_;
};
} // namespace mooncake
#endif // HIP_TRANSPORT_H_

10
mooncake-transfer-engine/include/transport/nvlink_transport/nvlink_transport.h
View File

@ -14,6 +14,7 @@
#include <vector>
#include <utility>
#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 <typename T1, typename T2>
std::size_t operator()(const std::pair<T1, T2>& p) const {
std::size_t h1 = std::hash<T1>{}(p.first);
std::size_t h2 = std::hash<T2>{}(p.second);
return h1 ^ (h2 << 1);
}
};
class NvlinkTransport : public Transport {
public:
NvlinkTransport();

12
mooncake-transfer-engine/src/multi_transport.cpp
View File

@ -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();

9
mooncake-transfer-engine/src/transfer_engine.cpp
View File

@ -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 &&

5
mooncake-transfer-engine/src/transfer_metadata.cpp
View File

@ -238,7 +238,8 @@ int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc,
rankInfoJSON["pid"] = static_cast<Json::UInt64>(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();

9
mooncake-transfer-engine/src/transport/CMakeLists.txt
View File

@ -38,6 +38,11 @@ if (USE_ASCEND_HETEROGENEOUS)
endif()
if (USE_MNNVL)
add_subdirectory(nvlink_transport)
target_sources(transport PUBLIC $<TARGET_OBJECTS:nvlink_transport>)
if (USE_HIP)
add_subdirectory(hip_transport)
target_sources(transport PUBLIC $<TARGET_OBJECTS:hip_transport>)
else()
add_subdirectory(nvlink_transport)
target_sources(transport PUBLIC $<TARGET_OBJECTS:nvlink_transport>)
endif()
endif()

10
mooncake-transfer-engine/src/transport/hip_transport/CMakeLists.txt Normal file
View File

@ -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
)

603
mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp Normal file
View File

@ -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 <glog/logging.h>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <memory>
#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<unsigned char> &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<unsigned char> &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<hipMemAccessDesc> 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<TransferMetadata> metadata,
std::shared_ptr<Topology> topology) {
metadata_ = metadata;
local_server_name_ = local_server_name;
auto desc = std::make_shared<SegmentDesc>();
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<TransferRequest> &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<TransferTask *> &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<std::mutex> 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<unsigned char> 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<Transport::BufferEntry> &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<void *> &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(&currentDev, 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<hipMemAccessDesc> 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

8
mooncake-transfer-engine/src/transport/nvlink_transport/CMakeLists.txt
View File

@ -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()

45
mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp
View File

@ -26,6 +26,7 @@
#include <memory>
#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<const unsigned char *>(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<unsigned char> &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<unsigned char>((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 {

15
mooncake-transfer-engine/tests/nvlink_transport_test.cpp
View File

@ -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<TransferEngine>(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<TransferEngine>(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);