diff --git a/docs/source/design/transfer-engine/index.md b/docs/source/design/transfer-engine/index.md index 0a9a4827..019ae1a2 100644 --- a/docs/source/design/transfer-engine/index.md +++ b/docs/source/design/transfer-engine/index.md @@ -1,6 +1,6 @@ # Transfer Engine -## Overview +## Overview Mooncake Transfer Engine is a high-performance, zero-copy data transfer library designed around two core abstractions: Segment and BatchTransfer. - [**Segment**](#segment) represents a contiguous address space that can be remotely read and written, which can be either non-persistent storage provided by DRAM or VRAM, known as **RAM Segment**, or persistent storage provided by NVMeof, known as **NVMeof Segment**. @@ -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`, `NVMeoFTransport`, `NvlinkTransport`, and `HipTransport`. +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`, `IntraNodeNvlinkTransport`, 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. @@ -49,9 +49,9 @@ The BatchTransfer API uses an array of requests, which specify the operation typ ### Topology Aware Path Selection Modern inference servers often consist of multiple CPU sockets, DRAM, GPUs, and RDMA NIC devices. Although it's technically possible to transfer data from local DRAM or VRAM to a remote location using any RDMA NIC, these transfers can be limited by the bandwidth constraints of the Ultra Path Interconnect (UPI) or PCIe Switch. To overcome these limitations, Transfer Engine implements a topology-aware path selection algorithm. -Before processing requests, each server generates a topology matrix and broadcasts it across the cluster. -This matrix categorizes network interface cards (NICs) into preferred and secondary lists for various types of memory, which types are specified during memory registration. -Under normal conditions, a NIC from the preferred list is selected for transfers, facilitating RDMA operations within the local NUMA or GPU Direct RDMA through the local PCIe switch only. +Before processing requests, each server generates a topology matrix and broadcasts it across the cluster. +This matrix categorizes network interface cards (NICs) into preferred and secondary lists for various types of memory, which types are specified during memory registration. +Under normal conditions, a NIC from the preferred list is selected for transfers, facilitating RDMA operations within the local NUMA or GPU Direct RDMA through the local PCIe switch only. In case of failures, NICs from both lists may be utilized. The process involves identifying the appropriate local and target NICs based on the memory addresses, establishing a connection, and executing the data transfer. @@ -59,7 +59,7 @@ The process involves identifying the appropriate local and target NICs based on For instance, as illustrated in figure above, to transfer data from buffer 0 (assigned to cpu:0) in the local node to buffer 1 (assigned to cpu:1) in the target node, the engine first identifies the preferred NICs for cpu:0 using the local server's topology matrix and selects one, such as mlx5_1, as the local NIC. Similarly, the target NIC, such as mlx5_3, is selected based on the target memory address. This setup enables establishing an RDMA connection from mlx5_1@local to mlx5_3@target to carry out RDMA read and write operations. -To further maximize bandwidth utilization, if a single request's transfer is internally divided into multiple slices if its length exceeds 64KB. +To further maximize bandwidth utilization, if a single request's transfer is internally divided into multiple slices if its length exceeds 64KB. Each slice might use a different path, enabling collaborative work among all RDMA NICs. ### Endpoint Management @@ -140,7 +140,7 @@ After successfully compiling Transfer Engine, the test program `transfer_engine_ ``` The meanings of the various parameters are as follows (the rest are the same as before): - `--segment_id` is the segment name of target node. It needs to be consistent with the value passed to `--local_server_name` when starting the target node (if any). - + Under normal circumstances, the initiator node will start the transfer operation, wait for 10 seconds, and then display the "Test completed" message, indicating that the test is complete. The initiator node can also configure the following test parameters: `--operation` (can be `"read"` or `"write"`), `batch_size`, `block_size`, `duration`, `threads`, etc. @@ -347,7 +347,7 @@ Key = mooncake/nvmeof/[segment_name] Value = { 'server_name': server_name, 'protocol': nvmeof, - 'buffers':[ + 'buffers':[ { 'length': 1073741824, 'file_path': "/mnt/nvme0" // The file path on this machine @@ -358,7 +358,7 @@ Value = { }， { 'length': 1073741824, - 'file_path': "/mnt/nvme1", + 'file_path': "/mnt/nvme1", 'local_path_map': { "node02": "/mnt/transfer_engine/node02/nvme1", ..... diff --git a/docs/source/getting_started/build.md b/docs/source/getting_started/build.md index f002c1ba..5c486d76 100644 --- a/docs/source/getting_started/build.md +++ b/docs/source/getting_started/build.md @@ -153,6 +153,7 @@ The following options can be used during `cmake ..` to specify whether to compil - `-DUSE_CUDA=[ON|OFF]`: Enable GPU Direct RDMA and NVMe-of support - `-DUSE_MUSA=[ON|OFF]`: Enable Moore Threads GPU support via MUSA - `-DUSE_HIP=[ON|OFF]`: Enable AMD GPU support via HIP/ROCm +- `-DUSE_INTRA_NVLINK=[ON|OFF]`: Enable intranode nvlink transport - `-DUSE_CXL=[ON|OFF]`: Enable CXL support - `-DWITH_STORE=[ON|OFF]`: Build Mooncake Store component - `-DWITH_P2P_STORE=[ON|OFF]`: Enable Golang support and build P2P Store component, require go 1.23+ diff --git a/mooncake-common/common.cmake b/mooncake-common/common.cmake index a31a4158..f5a657d8 100644 --- a/mooncake-common/common.cmake +++ b/mooncake-common/common.cmake @@ -80,6 +80,7 @@ option(USE_EVENT_DRIVEN_COMPLETION "option for using event-driven completion (st option(USE_TENT "option for building Mooncake TENT" OFF) option(USE_LRU_MASTER "option for using LRU in master service" OFF) +option(USE_INTRA_NVLINK "option for using IntraNode nvlink transport" OFF) set(LRU_MAX_CAPACITY 1000) if (USE_LRU_MASTER) diff --git a/mooncake-integration/CMakeLists.txt b/mooncake-integration/CMakeLists.txt index 7e3542b3..fccc7b01 100644 --- a/mooncake-integration/CMakeLists.txt +++ b/mooncake-integration/CMakeLists.txt @@ -33,13 +33,13 @@ set(PYTHON_PACKAGE_NAME "mooncake") if (WITH_TE) pybind11_add_module(engine ${SOURCES} ${CACHE_ALLOCATOR_SOURCES} - transfer_engine/transfer_engine_py.cpp + transfer_engine/transfer_engine_py.cpp ) set_target_properties(engine PROPERTIES - INSTALL_RPATH "$ORIGIN" + INSTALL_RPATH "$ORIGIN" ) - - target_link_libraries(engine PRIVATE + + target_link_libraries(engine PRIVATE $ ) @@ -113,6 +113,10 @@ if (USE_MNNVL) ) endif() +if(USE_INTRA_NVLINK) + message(STATUS "USE_INTRA_NVLINK is enabled, IntraNode nvlink is now activated") +endif() + install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/store/async_store.py" DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}) # Install Python scripts from mooncake-wheel/mooncake/ directory diff --git a/mooncake-integration/transfer_engine/transfer_engine_py.cpp b/mooncake-integration/transfer_engine/transfer_engine_py.cpp index 3c2977cf..c5eccf6b 100644 --- a/mooncake-integration/transfer_engine/transfer_engine_py.cpp +++ b/mooncake-integration/transfer_engine/transfer_engine_py.cpp @@ -23,17 +23,57 @@ #ifdef USE_MNNVL #include "transport/nvlink_transport/nvlink_transport.h" -static void *allocateMemory(size_t size) { - return mooncake::NvlinkTransport::allocatePinnedLocalMemory(size); -} -static void freeMemory(void *ptr) { - mooncake::NvlinkTransport::freePinnedLocalMemory(ptr); -} -#else -static void *allocateMemory(size_t size) { return malloc(size); } -static void freeMemory(void *ptr) { free(ptr); } #endif +#ifdef USE_INTRA_NVLINK +#include "transport/intranode_nvlink_transport/intranode_nvlink_transport.h" +#endif + +static void *(*allocateMemory)(size_t) = nullptr; +static void (*freeMemory)(void *) = nullptr; +static std::string g_protocol; + +// Handle allocateMemory function pointer based on protocol +void initMemoryAllocator(const char *protocol) { + if (allocateMemory != nullptr) { + LOG(WARNING) << "Memory allocator already initialized with: " + << g_protocol; + return; + } + g_protocol = protocol; + if (strcmp(protocol, "nvlink") == 0) { +#ifdef USE_MNNVL + allocateMemory = [](size_t s) -> void * { + return mooncake::NvlinkTransport::allocatePinnedLocalMemory(s); + }; + freeMemory = [](void *p) { + mooncake::NvlinkTransport::freePinnedLocalMemory(p); + }; + LOG(INFO) << "Selected MNNVL (NVLink) memory allocator"; +#else + LOG(ERROR) << "Protocol 'nvlink' requires -DUSE_MNNVL=ON"; +#endif + } else if (strcmp(protocol, "nvlink_intra") == 0) { +#ifdef USE_INTRA_NVLINK + allocateMemory = [](size_t s) -> void * { + return mooncake::IntraNodeNvlinkTransport:: + allocatePinnedLocalMemory(s); + }; + freeMemory = [](void *p) { + mooncake::IntraNodeNvlinkTransport::freePinnedLocalMemory(p); + }; + LOG(INFO) << "Selected Intra-NVLink memory allocator"; +#else + LOG(ERROR) << "Protocol 'nvlink_intra' requires -DUSE_INTRA_NVLINK=ON"; +#endif + } else { + // default fallback + allocateMemory = malloc; + freeMemory = free; + LOG(WARNING) << "Using default malloc/free for protocol: " << protocol; + } +} + TransferEnginePy::TransferEnginePy() { const int64_t kNanosPerSecond = 1000 * 1000 * 1000; if (getenv("MC_TRANSFER_TIMEOUT")) { @@ -102,6 +142,8 @@ int TransferEnginePy::initialize(const char *local_hostname, const char *metadata_server, const char *protocol, const char *device_name) { + initMemoryAllocator(protocol); + auto conn_string = parseConnectionString(metadata_server); return initializeExt(local_hostname, conn_string.second.c_str(), protocol, device_name, conn_string.first.c_str()); diff --git a/mooncake-transfer-engine/example/transfer_engine_bench.cpp b/mooncake-transfer-engine/example/transfer_engine_bench.cpp index 559d9ffb..481d1e85 100644 --- a/mooncake-transfer-engine/example/transfer_engine_bench.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_bench.cpp @@ -49,6 +49,10 @@ #include "gpu_vendor/mnnvl.h" #endif +#ifdef USE_INTRA_NVLINK +#include "gpu_vendor/intra_nvlink.h" +#endif + static void checkCudaError(cudaError_t result, const char *message) { if (result != cudaSuccess) { LOG(ERROR) << message << " (Error code: " << result << " - " @@ -71,7 +75,8 @@ 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|nvlink|hip"); +DEFINE_string(protocol, "rdma", + "Transfer protocol: rdma|barex|tcp|nvlink|nvlink_intra|hip"); DEFINE_string(device_name, "mlx5_2", "Device name to use, valid if protocol=rdma"); @@ -110,12 +115,27 @@ static void *allocateMemoryPool(size_t size, int buffer_id, void *d_buf; LOG(INFO) << "Allocating memory on GPU " << gpu_id; checkCudaError(cudaSetDevice(gpu_id), "Failed to set device"); + if (FLAGS_protocol == "nvlink") { #ifdef USE_MNNVL - d_buf = allocateFabricMemory(size); + d_buf = allocateFabricMemory(size); + LOG(INFO) << "Using MNNVL fabric memory allocation"; #else - checkCudaError(cudaMalloc(&d_buf, size), - "Failed to allocate device memory"); + LOG(ERROR) << "--protocol=nvlink requires USE_MNNVL=ON"; + return nullptr; #endif + } else if (FLAGS_protocol == "nvlink_intra") { +#ifdef USE_INTRA_NVLINK + d_buf = allocateFabricMemory_intra(size); + LOG(INFO) << "Using intra-NVLink memory allocation"; +#else + LOG(ERROR) + << "--protocol=nvlink_intra requires USE_INTRA_NVLINK=ON"; + return nullptr; +#endif + } else { + checkCudaError(cudaMalloc(&d_buf, size), + "Failed to allocate device memory"); + } if (FLAGS_init_mem) { checkCudaError(cudaMemset(d_buf, 0xCC, size), @@ -132,25 +152,35 @@ static void *allocateMemoryPool(size_t size, int buffer_id, static void freeMemoryPool(void *addr, size_t size) { #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) + if (FLAGS_protocol == "nvlink") { #ifdef USE_MNNVL - if (FLAGS_use_vram) { - freeFabricMemory(addr); - return; - } + if (FLAGS_use_vram) { + freeFabricMemory(addr); + return; + } #endif // USE_MNNVL - - // check pointer on GPU - cudaPointerAttributes attributes; - checkCudaError(cudaPointerGetAttributes(&attributes, addr), - "Failed to get pointer attributes"); - - if (attributes.type == cudaMemoryTypeDevice) { - cudaFree(addr); - } else if (attributes.type == cudaMemoryTypeHost || - attributes.type == cudaMemoryTypeUnregistered) { - numa_free(addr, size); + } else if (FLAGS_protocol == "nvlink_intra") { +#ifdef USE_INTRA_NVLINK + if (FLAGS_use_vram) { + freeFabricMemory_intra(addr); + return; + } +#endif } else { - LOG(ERROR) << "Unknown memory type, " << addr << " " << attributes.type; + // check pointer on GPU + cudaPointerAttributes attributes; + checkCudaError(cudaPointerGetAttributes(&attributes, addr), + "Failed to get pointer attributes"); + + if (attributes.type == cudaMemoryTypeDevice) { + cudaFree(addr); + } else if (attributes.type == cudaMemoryTypeHost || + attributes.type == cudaMemoryTypeUnregistered) { + numa_free(addr, size); + } else { + LOG(ERROR) << "Unknown memory type, " << addr << " " + << attributes.type; + } } #else numa_free(addr, size); @@ -373,7 +403,7 @@ static Transport *installTransportFromFlags(TransferEngine *engine) { args.get()[1] = nullptr; xport = engine->installTransport(FLAGS_protocol.c_str(), args.get()); } else if (FLAGS_protocol == "tcp" || FLAGS_protocol == "nvlink" || - FLAGS_protocol == "hip") { + FLAGS_protocol == "hip" || FLAGS_protocol == "nvlink_intra") { xport = engine->installTransport(FLAGS_protocol.c_str(), nullptr); } else { LOG(ERROR) << "Unsupported protocol: " << FLAGS_protocol; 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 db60fc96..e29cdee9 100644 --- a/mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp @@ -44,6 +44,10 @@ #include "gpu_vendor/mnnvl.h" #endif +#ifdef USE_INTRA_NVLINK +#include "gpu_vendor/intra_nvlink.h" +#endif + static void checkCudaError(cudaError_t result, const char *message) { if (result != cudaSuccess) { LOG(ERROR) << message << " (Error code: " << result << " - " @@ -313,6 +317,8 @@ int initiator() { xport = engine->installTransport("tcp", nullptr); } else if (FLAGS_protocol == "nvlink") { xport = engine->installTransport("nvlink", nullptr); + } else if (FLAGS_protocol == "nvlink_intra") { + xport = engine->installTransport("nvlink_intra", nullptr); } else if (FLAGS_protocol == "hip") { xport = engine->installTransport("hip", nullptr); } else { diff --git a/mooncake-transfer-engine/example/transfer_engine_validator.cpp b/mooncake-transfer-engine/example/transfer_engine_validator.cpp index 6f5f98d9..3b6a1c0a 100644 --- a/mooncake-transfer-engine/example/transfer_engine_validator.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_validator.cpp @@ -44,6 +44,10 @@ #include "gpu_vendor/mnnvl.h" #endif +#ifdef USE_INTRA_NVLINK +#include "gpu_vendor/intra_nvlink.h" +#endif + static void checkCudaError(cudaError_t result, const char *message) { if (result != cudaSuccess) { LOG(ERROR) << message << " (Error code: " << result << " - " @@ -397,6 +401,8 @@ int initiator() { xport = engine->installTransport("tcp", nullptr); } else if (FLAGS_protocol == "nvlink") { xport = engine->installTransport("nvlink", nullptr); + } else if (FLAGS_protocol == "nvlink_intra") { + xport = engine->installTransport("nvlink_intra", nullptr); } else if (FLAGS_protocol == "hip") { xport = engine->installTransport("hip", nullptr); } else { diff --git a/mooncake-transfer-engine/include/gpu_vendor/intra_nvlink.h b/mooncake-transfer-engine/include/gpu_vendor/intra_nvlink.h new file mode 100644 index 00000000..11c44f92 --- /dev/null +++ b/mooncake-transfer-engine/include/gpu_vendor/intra_nvlink.h @@ -0,0 +1,21 @@ +// 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 +#include + +#define allocateFabricMemory_intra(size) \ + mooncake::IntraNodeNvlinkTransport::allocatePinnedLocalMemory(size) +#define freeFabricMemory_intra(addr) \ + mooncake::IntraNodeNvlinkTransport::freePinnedLocalMemory(addr) diff --git a/mooncake-transfer-engine/include/transport/intranode_nvlink_transport/intranode_nvlink_transport.h b/mooncake-transfer-engine/include/transport/intranode_nvlink_transport/intranode_nvlink_transport.h new file mode 100644 index 00000000..d76b20c3 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/intranode_nvlink_transport/intranode_nvlink_transport.h @@ -0,0 +1,85 @@ +// Copyright 2024 KVCache.AI + +#ifndef INTRANODE_NVLINK_TRANSPORT_H_ +#define INTRANODE_NVLINK_TRANSPORT_H_ + +#include "cuda_alike.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 IntraNodeNvlinkTransport : public Transport { + public: + IntraNodeNvlinkTransport(); + + ~IntraNodeNvlinkTransport(); + + 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 "nvlink_intraNode"; } + + private: + std::atomic_bool running_; + + 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 // INTRANODE_NVLINK_TRANSPORT_H_ diff --git a/mooncake-transfer-engine/src/CMakeLists.txt b/mooncake-transfer-engine/src/CMakeLists.txt index 2fba37a0..1225449e 100644 --- a/mooncake-transfer-engine/src/CMakeLists.txt +++ b/mooncake-transfer-engine/src/CMakeLists.txt @@ -2,85 +2,103 @@ file(GLOB ENGINE_SOURCES "*.cpp") add_subdirectory(common) add_subdirectory(transport) -SET(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE) +set(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE) -if (USE_HIP) +if(USE_HIP) hipify_files(ENGINE_SOURCES) endif() add_library(transfer_engine ${ENGINE_SOURCES} $) -if (BUILD_SHARED_LIBS) +if(BUILD_SHARED_LIBS) install(TARGETS transfer_engine DESTINATION lib) endif() add_compile_definitions(transfer_engine PUBLIC MOONCAKE_USE_ETCD CMAKE_INCLUDE) -if (USE_ETCD) - if (USE_ETCD_LEGACY) - if (USE_STATIC_ETCD_CPP_API) +if(USE_ETCD) + if(USE_ETCD_LEGACY) + if(USE_STATIC_ETCD_CPP_API) find_package(PkgConfig REQUIRED) pkg_check_modules(GRPCPP REQUIRED grpc++) pkg_check_modules(GRPC REQUIRED grpc) - target_link_libraries(transfer_engine PUBLIC etcd-cpp-api-core protobuf ${GRPC_LDFLAGS} ${GRPCPP_LDFLAGS}) + target_link_libraries( + transfer_engine PUBLIC etcd-cpp-api-core protobuf ${GRPC_LDFLAGS} + ${GRPCPP_LDFLAGS}) else() target_link_libraries(transfer_engine PUBLIC etcd-cpp-api) endif() - else () + else() add_dependencies(transfer_engine build_etcd_wrapper) - set(ETCD_WRAPPER_LIB ${CMAKE_CURRENT_BINARY_DIR}/../../mooncake-common/etcd/libetcd_wrapper.so) + set(ETCD_WRAPPER_LIB + ${CMAKE_CURRENT_BINARY_DIR}/../../mooncake-common/etcd/libetcd_wrapper.so + ) target_link_libraries(transfer_engine PRIVATE ${ETCD_WRAPPER_LIB}) endif() endif() -if (USE_REDIS) +if(USE_REDIS) target_link_libraries(transfer_engine PUBLIC hiredis) endif() -if (USE_HTTP) +if(USE_HTTP) find_package(CURL REQUIRED) target_link_libraries(transfer_engine PUBLIC ${CURL_LIBRARIES}) endif() target_link_libraries( - transfer_engine - PUBLIC - base transport rdma_transport ibverbs glog::glog gflags::gflags pthread JsonCpp::JsonCpp numa yalantinglibs::yalantinglibs - ) + transfer_engine + PUBLIC base + transport + rdma_transport + ibverbs + glog::glog + gflags::gflags + pthread + JsonCpp::JsonCpp + numa + yalantinglibs::yalantinglibs) -if (USE_BAREX) +if(USE_BAREX) target_link_libraries(transfer_engine PUBLIC barex_transport) endif() -if (USE_CUDA) +if(USE_CUDA) target_include_directories(transfer_engine PRIVATE /usr/local/cuda/include) target_link_libraries(transfer_engine PUBLIC cuda cudart rt) - if (USE_NVMEOF) + if(USE_NVMEOF) target_link_libraries(transfer_engine PUBLIC nvmeof_transport cufile) endif() endif() -if (USE_MUSA) +if(USE_MUSA) target_include_directories(transfer_engine PRIVATE /usr/local/musa/include) target_link_libraries(transfer_engine PUBLIC musa musart rt) endif() -if (USE_HIP) +if(USE_HIP) target_include_directories(transfer_engine PRIVATE ${HIP_INCLUDE_DIRS}) target_link_libraries(transfer_engine PUBLIC hip::host rt) endif() -if (USE_ASCEND) - target_link_libraries(transfer_engine PUBLIC ascendcl hccl ascend_transport MPI::MPI) +if(USE_ASCEND) + target_link_libraries(transfer_engine PUBLIC ascendcl hccl ascend_transport + MPI::MPI) endif() -if (USE_ASCEND_DIRECT) +if(USE_ASCEND_DIRECT) target_link_libraries(transfer_engine PUBLIC ascend_transport) endif() -if (USE_ASCEND_HETEROGENEOUS) - file(GLOB ASCEND_TOOLKIT_ROOT "/usr/local/Ascend/ascend-toolkit/latest/*-linux") +if(USE_ASCEND_HETEROGENEOUS) + file(GLOB ASCEND_TOOLKIT_ROOT + "/usr/local/Ascend/ascend-toolkit/latest/*-linux") set(ASCEND_LIB_DIR "${ASCEND_TOOLKIT_ROOT}/lib64") link_directories(${ASCEND_LIB_DIR}) target_link_libraries(transfer_engine PUBLIC ascendcl ascend_transport) endif() -if (USE_TENT) +if(USE_TENT) add_compile_definitions(transfer_engine PUBLIC USE_TENT) target_link_libraries(transfer_engine PUBLIC tent) -endif() \ No newline at end of file +endif() + +if(USE_INTRA_NVLINK) + message(STATUS "Enabled USE_INTRA_NVLINK support") + target_compile_definitions(transfer_engine PUBLIC USE_INTRA_NVLINK) +endif() diff --git a/mooncake-transfer-engine/src/multi_transport.cpp b/mooncake-transfer-engine/src/multi_transport.cpp index a38a00bc..be87db5e 100644 --- a/mooncake-transfer-engine/src/multi_transport.cpp +++ b/mooncake-transfer-engine/src/multi_transport.cpp @@ -36,6 +36,9 @@ #ifdef USE_ASCEND_HETEROGENEOUS #include "transport/ascend_transport/heterogeneous_rdma_transport.h" #endif +#ifdef USE_INTRA_NVLINK +#include "transport/intranode_nvlink_transport/intranode_nvlink_transport.h" +#endif #ifdef USE_MNNVL #ifdef USE_HIP #include "transport/hip_transport/hip_transport.h" @@ -249,6 +252,13 @@ Transport *MultiTransport::installTransport(const std::string &proto, transport = new HeterogeneousRdmaTransport(); } #endif + +#ifdef USE_INTRA_NVLINK + else if (std::string(proto) == "nvlink_intra") { + transport = new IntraNodeNvlinkTransport(); + } +#endif + #ifdef USE_MNNVL #ifdef USE_HIP else if (std::string(proto) == "hip") { diff --git a/mooncake-transfer-engine/src/transfer_engine.cpp b/mooncake-transfer-engine/src/transfer_engine.cpp index 54c977f7..f4bcd03c 100644 --- a/mooncake-transfer-engine/src/transfer_engine.cpp +++ b/mooncake-transfer-engine/src/transfer_engine.cpp @@ -576,4 +576,4 @@ std::shared_ptr TransferEngine::getLocalTopology() { } } // namespace mooncake -#endif \ No newline at end of file +#endif diff --git a/mooncake-transfer-engine/src/transfer_engine_impl.cpp b/mooncake-transfer-engine/src/transfer_engine_impl.cpp index 1353b25e..454bfc5c 100644 --- a/mooncake-transfer-engine/src/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/src/transfer_engine_impl.cpp @@ -239,23 +239,38 @@ int TransferEngineImpl::init(const std::string& metadata_conn_string, LOG(ERROR) << "Failed to install Ascend transport"; return -1; } -#elif defined(USE_MNNVL) - if (local_topology_->getHcaList().size() > 0 && - !getenv("MC_FORCE_MNNVL")) { - Transport* rdma_transport = - multi_transports_->installTransport("rdma", local_topology_); - if (!rdma_transport) { - LOG(ERROR) << "Failed to install RDMA transport"; - return -1; - } - } else { - Transport* nvlink_transport = +#elif defined(USE_MNNVL) || defined(USE_INTRA_NVLINK) + + const char* force_mnnvl = getenv("MC_FORCE_MNNVL"); + const char* intra_env = getenv("MC_INTRANODE_NVLINK"); + if (force_mnnvl || local_topology_->getHcaList().empty()) { + Transport* t = multi_transports_->installTransport("nvlink", nullptr); - if (!nvlink_transport) { + if (!t) { LOG(ERROR) << "Failed to install NVLink transport"; return -1; } + LOG(INFO) << "Using cross-node NVLink transport " + << "(MC_FORCE_MNNVL or no HCA detected)"; + } else if (intra_env) { + Transport* t = + multi_transports_->installTransport("nvlink_intra", nullptr); + if (!t) { + LOG(ERROR) << "Failed to install Intra-Node NVLink transport"; + return -1; + } + LOG(INFO) << "Using Intra-Node NVLink transport " + "(MC_INTRANODE_NVLINK set)"; + } else { + Transport* t = + multi_transports_->installTransport("rdma", local_topology_); + if (!t) { + LOG(ERROR) << "Failed to install RDMA transport"; + return -1; + } + LOG(INFO) << "Using RDMA transport (RoCE/iWARP)"; } + #else if (local_topology_->getHcaList().size() > 0 && !getenv("MC_FORCE_TCP") || diff --git a/mooncake-transfer-engine/src/transfer_metadata.cpp b/mooncake-transfer-engine/src/transfer_metadata.cpp index 40b97f3e..b6a07d76 100644 --- a/mooncake-transfer-engine/src/transfer_metadata.cpp +++ b/mooncake-transfer-engine/src/transfer_metadata.cpp @@ -239,6 +239,7 @@ int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc, segmentJSON["rank_info"] = rankInfoJSON; } else if (segmentJSON["protocol"] == "nvlink" || + segmentJSON["protocol"] == "nvlink_intra" || segmentJSON["protocol"] == "hip") { Json::Value buffersJSON(Json::arrayValue); for (const auto &buffer : desc.buffers) { @@ -378,7 +379,8 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, } desc->buffers.push_back(buffer); } - } else if (desc->protocol == "nvlink" || desc->protocol == "hip") { + } else if (desc->protocol == "nvlink" || desc->protocol == "nvlink_intra" || + 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 446add9e..0fc7ca9a 100644 --- a/mooncake-transfer-engine/src/transport/CMakeLists.txt +++ b/mooncake-transfer-engine/src/transport/CMakeLists.txt @@ -46,3 +46,8 @@ if (USE_MNNVL) target_sources(transport PUBLIC $) endif() endif() + +if (USE_INTRA_NVLINK) + add_subdirectory(intranode_nvlink_transport) + target_sources(transport PUBLIC $) +endif() diff --git a/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/CMakeLists.txt new file mode 100644 index 00000000..524891a6 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/CMakeLists.txt @@ -0,0 +1,7 @@ +file(GLOB INTRANODE_NVLINK_SOURCES "*.cpp") + +add_library(intranode_nvlink_transport OBJECT ${INTRANODE_NVLINK_SOURCES}) + +if (USE_CUDA) +target_include_directories(intranode_nvlink_transport PUBLIC CUDA::cudart "/usr/local/cuda/include") +endif() diff --git a/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp b/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp new file mode 100644 index 00000000..17d3dd4a --- /dev/null +++ b/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp @@ -0,0 +1,413 @@ +// Copyright 2024 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. + +#include "transport/intranode_nvlink_transport/intranode_nvlink_transport.h" + +#include +#include "cuda_alike.h" +#include + +#include +#include +#include +#include +#include +#include + +#include "common.h" +#include "common/serialization.h" +#include "config.h" +#include "transfer_engine.h" +#include "transfer_metadata.h" +#include "transport/transport.h" + +static bool checkCudaErrorReturn(cudaError_t result, const char *message) { + if (result != cudaSuccess) { + LOG(ERROR) << message << " (Error code: " << result << " - " + << cudaGetErrorString(result) << ")" << std::endl; + return false; + } + return true; +} + +namespace mooncake { +static int getNumDevices() { + static int cached_num_devices = -1; + if (cached_num_devices == -1) { + if (!checkCudaErrorReturn( + cudaGetDeviceCount(&cached_num_devices), + "IntraNodeNvlinkTransport: cudaGetDeviceCount failed")) { + return 0; + } + } + return cached_num_devices; +} + +static bool enableP2PAccess(int src_device_id, int dst_device_id) { + int canAccessPeer = 0; + if (!checkCudaErrorReturn( + cudaDeviceCanAccessPeer(&canAccessPeer, src_device_id, + dst_device_id), + "IntraNodeNvlinkTransport: failed to query peer access")) { + return false; + } + + if (!canAccessPeer) { + LOG(ERROR) << "IntraNodeNvlinkTransport: device " << src_device_id + << " cannot p2p access device " << dst_device_id; + return false; + } + + // enable src->dst p2p access + if (!checkCudaErrorReturn( + cudaSetDevice(src_device_id), + "IntraNodeNvlinkTransport: failed to set device")) { + return false; + } + cudaError_t result = cudaDeviceEnablePeerAccess(dst_device_id, 0); + + if (result != cudaSuccess && result != cudaErrorPeerAccessAlreadyEnabled) { + LOG(ERROR) << "IntraNodeNvlinkTransport: failed to enable p2p access " + "(Error code: " + << result << " - " << cudaGetErrorString(result) << ")" + << std::endl; + + return false; + } + + // enable dst->src p2p access + if (!checkCudaErrorReturn( + cudaSetDevice(dst_device_id), + "IntraNodeNvlinkTransport: failed to set device")) { + return false; + } + result = cudaDeviceEnablePeerAccess(src_device_id, 0); + + if (result != cudaSuccess && result != cudaErrorPeerAccessAlreadyEnabled) { + LOG(ERROR) << "IntraNodeNvlinkTransport: failed to enable p2p access " + "(Error code: " + << result << " - " << cudaGetErrorString(result) << ")" + << std::endl; + + return false; + } + + return true; +} +IntraNodeNvlinkTransport::IntraNodeNvlinkTransport() {} + +// IntraNodeNvlinkTransport::IntraNodeNvlinkTransport() : +// use_fabric_mem_(supportFabricMem()) {} +// int num_devices = getNumDevices(); +// if (globalConfig().trace) { +// LOG(INFO) << "IntraNodeNvlinkTransport: use_fabric_mem_:" << +// use_fabric_mem_ +// << ", num_devices: " << num_devices; +// } + +// for (int src_device_id = 0; src_device_id < num_devices; ++src_device_id) +// { +// for (int dst_device_id = src_device_id + 1; dst_device_id < +// num_devices; +// ++dst_device_id) { +// if (enableP2PAccess(src_device_id, dst_device_id)) { +// if (globalConfig().trace) { +// LOG(INFO) +// << "IntraNodeNvlinkTransport: enabled p2p access +// between device " +// << src_device_id << " and " << dst_device_id; +// } +// } else { +// LOG(ERROR) << "IntraNodeNvlinkTransport: failed to enable p2p +// access " +// "between device " +// << src_device_id << " and " << dst_device_id; +// } +// } +// } +// } + +IntraNodeNvlinkTransport::~IntraNodeNvlinkTransport() { + for (auto &entry : remap_entries_) { + cudaIpcCloseMemHandle(entry.second.shm_addr); + } + remap_entries_.clear(); +} + +int IntraNodeNvlinkTransport::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 = "nvlink_intra"; + metadata_->addLocalSegment(LOCAL_SEGMENT_ID, local_server_name_, + std::move(desc)); + return 0; +} + +Status IntraNodeNvlinkTransport::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) << "IntraNodeNvlinkTransport: Exceed the limitation of " + "current batch's " + "capacity"; + return Status::InvalidArgument( + "IntraNodeNvlinkTransport: 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; + 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); + cudaError_t err; + if (slice->opcode == TransferRequest::READ) + err = cudaMemcpy(slice->source_addr, (void *)slice->local.dest_addr, + slice->length, cudaMemcpyDefault); + else + err = cudaMemcpy((void *)slice->local.dest_addr, slice->source_addr, + slice->length, cudaMemcpyDefault); + if (err != cudaSuccess) + slice->markFailed(); + else + slice->markSuccess(); + } + + return Status::OK(); +} + +Status IntraNodeNvlinkTransport::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( + "IntraNodeNvlinkTransport::getTransportStatus invalid argument, " + "batch id: " + + std::to_string(batch_id)); + } + 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; + 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 IntraNodeNvlinkTransport::submitTransferTask( + const std::vector &task_list) { + for (size_t index = 0; index < task_list.size(); ++index) { + assert(task_list[index]); + auto &task = *task_list[index]; + 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; + 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); + cudaError_t err; + if (slice->opcode == TransferRequest::READ) + err = cudaMemcpy(slice->source_addr, (void *)slice->local.dest_addr, + slice->length, cudaMemcpyDefault); + else + err = cudaMemcpy((void *)slice->local.dest_addr, slice->source_addr, + slice->length, cudaMemcpyDefault); + if (err != cudaSuccess) + slice->markFailed(); + else + slice->markSuccess(); + } + return Status::OK(); +} + +int IntraNodeNvlinkTransport::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; + } + cudaPointerAttributes attr; + cudaError_t err = cudaPointerGetAttributes(&attr, addr); + if (err != cudaSuccess) { + LOG(ERROR) + << "IntraNodeNvlinkTransport: cudaPointerGetAttributes failed"; + return -1; + } + + if (attr.type != cudaMemoryTypeDevice) { + LOG(ERROR) << "Unsupported memory type, " << addr << " " << attr.type; + return -1; + } + + cudaIpcMemHandle_t handle; + err = cudaIpcGetMemHandle(&handle, addr); + if (err != cudaSuccess) { + LOG(ERROR) << "IntraNodeNvlinkTransport: cudaIpcGetMemHandle failed"; + return -1; + } + + (void)remote_accessible; + BufferDesc desc; + desc.addr = (uint64_t)addr; + desc.length = length; + desc.name = location; + desc.shm_name = serializeBinaryData(&handle, sizeof(cudaIpcMemHandle_t)); + return metadata_->addLocalMemoryBuffer(desc, true); +} + +int IntraNodeNvlinkTransport::unregisterLocalMemory(void *addr, + bool update_metadata) { + return metadata_->removeLocalMemoryBuffer(addr, update_metadata); +} + +int IntraNodeNvlinkTransport::relocateSharedMemoryAddress(uint64_t &dest_addr, + uint64_t length, + uint64_t target_id) { + auto desc = metadata_->getSegmentDescByID(target_id); + int index = 0; + for (auto &entry : desc->buffers) { + if (!entry.shm_name.empty() && entry.addr <= dest_addr && + dest_addr + length <= entry.addr + entry.length) { + 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); + if (output_buffer.size() == sizeof(cudaIpcMemHandle_t)) { + cudaIpcMemHandle_t handle; + memcpy(&handle, output_buffer.data(), sizeof(handle)); + void *shm_addr = nullptr; + cudaError_t err = cudaIpcOpenMemHandle( + &shm_addr, handle, cudaIpcMemLazyEnablePeerAccess); + if (err != cudaSuccess) { + LOG(ERROR) << "IntraNodeNvlinkTransport: " + "cudaIpcOpenMemHandle failed: " + << cudaGetErrorString(err); + 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; + } else { + LOG(ERROR) << "Mismatched NVLink data transfer method"; + return -1; + } + } + 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; + } + index++; + } + LOG(ERROR) << "Requested address " << (void *)dest_addr << " to " + << (void *)(dest_addr + length) << " not found!"; + return ERR_INVALID_ARGUMENT; +} + +int IntraNodeNvlinkTransport::registerLocalMemoryBatch( + const std::vector &buffer_list, + const std::string &location) { + for (auto &buffer : buffer_list) + registerLocalMemory(buffer.addr, buffer.length, location, true, false); + return metadata_->updateLocalSegmentDesc(); +} + +int IntraNodeNvlinkTransport::unregisterLocalMemoryBatch( + const std::vector &addr_list) { + for (auto &addr : addr_list) unregisterLocalMemory(addr, false); + return metadata_->updateLocalSegmentDesc(); +} + +void *IntraNodeNvlinkTransport::allocatePinnedLocalMemory(size_t size) { + void *ptr = nullptr; + cudaError_t res = cudaMalloc(&ptr, size); + if (res == cudaSuccess) { + LOG(INFO) << "IntraNodeNvlinkTransport: Falling back to cudaMalloc for " + << size << " bytes (memory will NOT be exportable)"; + return ptr; + } else { + LOG(ERROR) + << "IntraNodeNvlinkTransport: cudaMalloc failed during fallback: " + << cudaGetErrorString(res); + return nullptr; + } +} + +void IntraNodeNvlinkTransport::freePinnedLocalMemory(void *ptr) { + cudaFree(ptr); + return; +} + +} // namespace mooncake