[TransferEngine] Add NCCL host and device transport backend

docs/source/getting_started/build.md

@@ -56,6 +56,7 @@ environment setup must be prepared separately.
 | Hardware / backend | Build option | External SDK / setup | Environment and notes |
 | --- | --- | --- | --- |
 | NVIDIA CUDA / GPUDirect | `-DUSE_CUDA=ON` | Install CUDA 12.1+ and enable `nvidia-fs` for cuFile builds. | Add CUDA libraries to `LIBRARY_PATH` and `LD_LIBRARY_PATH`, for example `/usr/local/cuda/lib64`. |
+| NVIDIA NCCL DeviceTransport | `-DUSE_NCCL_DEVICE=ON` | Install NCCL 2.30.4+ with `nccl_device.h`. Requires CUDA. | Set `NCCL_ROOT` when NCCL is outside the standard search paths. The compile-time NCCL headers must exactly match the runtime `libnccl`. |
 | NVIDIA Multi-Node NVLink | `-DUSE_MNNVL=ON` | Requires CUDA. | Also set `-DUSE_CUDA=ON`. Not used with MUSA, HIP, or MACA builds. |
 | Moore Threads MUSA | `-DUSE_MUSA=ON` | Install MUSA SDK and `mthreads-peermem` for GPUDirect RDMA. | Add `/usr/local/musa/lib` to `LIBRARY_PATH` and `LD_LIBRARY_PATH`. |
 | Cambricon MLU | `-DUSE_MLU=ON` | Install Cambricon Neuware SDK. | Set `NEUWARE_HOME`, or pass `-DNEUWARE_ROOT=/path/to/neuware`. Use `-DMLU_INCLUDE_DIR` and `-DMLU_LIB_DIR` for custom layouts. |
@@ -66,6 +67,17 @@ environment setup must be prepared separately.
 | Hygon DCU | `-DUSE_HYGON=ON` | Install DTK SDK. | Set `DTK_HOME`, or pass `-DDTK_ROOT=/path/to/dtk`. Use `-DDTK_INCLUDE_DIR` and `-DDTK_LIB_DIR` for custom layouts. |
 | Iluvatar CoreX | `-DUSE_COREX=ON` | Install CoreX SDK. | Set `COREX_HOME`, or pass `-DCOREX_ROOT=/path/to/corex`. Use `-DCOREX_INCLUDE_DIR` and `-DCOREX_LIB_DIR` for custom layouts. |
 
+```{admonition} NCCL DeviceTransport version contract
+:class: important
+Mooncake currently requires the NCCL headers used to compile its host and
+device code to exactly match the loaded runtime `libnccl.so`.
+`NcclTransport::initialize()` rejects a mismatch. After upgrading NCCL,
+rebuild Mooncake and every AOT CUDA kernel that includes
+`transport/device/nccl_device.cuh`. Invalidate and regenerate any cached NCCL
+Device API JIT kernels before running. This is required for GIN device kernels,
+which are not cross-version compatible.
+```
+
 ```{admonition} GPU-Direct RDMA
 :class: note
 Mooncake can use the DMA-BUF path for GPU-Direct RDMA, which does **not**
@@ -129,6 +141,7 @@ The following options can be passed to `cmake ..`.
 | Option | Default | Description |
 | --- | --- | --- |
 | `-DUSE_CUDA=ON/OFF` | `OFF` | Enable GPU memory support, including GPUDirect RDMA, NVMe-oF, and GPU-aware TCP transport. Required when transferring GPU memory, even when using TCP. |
+| `-DUSE_NCCL_DEVICE=ON/OFF` | `OFF` | Enable the experimental NCCL DeviceTransport backend. Requires CUDA and NCCL 2.30.4+; compile-time headers and runtime `libnccl` must match exactly. |
 | `-DUSE_MNNVL=ON/OFF` | `OFF` | Enable Multi-Node NVLink transport. Requires `-DUSE_CUDA=ON`; not used with MUSA, HIP, or MACA builds. |
 | `-DUSE_MUSA=ON/OFF` | `OFF` | Enable Moore Threads GPU support via MUSA. |
 | `-DUSE_MACA=ON/OFF` | `OFF` | Enable MetaX (Muxi) GPU support via MACA. |

mooncake-common/FindNCCLDevice.cmake

@@ -0,0 +1,91 @@
+# Locate NCCL 2.30.4+ with the experimental Device API.
+#
+# Prefer NCCL's exported package so version metadata, include paths, library
+# variants, CUDA, and thread dependencies come from NCCL itself. Keep a manual
+# fallback for tarball and pre-package installations.
+
+set(_NCCL_DEVICE_MIN_VERSION "2.30.4")
+set(_NCCL_DEVICE_USING_CONFIG FALSE)
+
+find_package(NCCL ${_NCCL_DEVICE_MIN_VERSION} CONFIG QUIET)
+if(NCCL_FOUND AND TARGET NCCL::nccl)
+  set(_NCCL_DEVICE_USING_CONFIG TRUE)
+  set(NCCLDevice_VERSION "${NCCL_VERSION}")
+  find_path(NCCL_DEVICE_INCLUDE_DIR
+    NAMES nccl_device.h
+    HINTS ${NCCL_INCLUDE_DIRS}
+    NO_DEFAULT_PATH)
+endif()
+
+if(NOT _NCCL_DEVICE_USING_CONFIG)
+  find_path(NCCL_DEVICE_INCLUDE_DIR
+    NAMES nccl_device.h
+    HINTS
+      ${NCCL_ROOT}
+      $ENV{NCCL_ROOT}
+      $ENV{NCCL_HOME}
+    PATH_SUFFIXES include build/include
+    PATHS
+      /usr/local/cuda
+      /usr/local
+      /usr)
+
+  find_library(NCCL_DEVICE_LIBRARY
+    NAMES nccl
+    HINTS
+      ${NCCL_ROOT}
+      $ENV{NCCL_ROOT}
+      $ENV{NCCL_HOME}
+    PATH_SUFFIXES lib lib64 build/lib build/lib64
+    PATHS
+      /usr/local/cuda
+      /usr/local
+      /usr)
+
+  if(NCCL_DEVICE_INCLUDE_DIR AND
+     EXISTS "${NCCL_DEVICE_INCLUDE_DIR}/nccl.h")
+    foreach(_component MAJOR MINOR PATCH)
+      file(STRINGS "${NCCL_DEVICE_INCLUDE_DIR}/nccl.h"
+        _nccl_${_component}_line
+        REGEX "^#define NCCL_${_component}[ \t]+[0-9]+"
+        LIMIT_COUNT 1)
+      string(REGEX MATCH "[0-9]+$" _nccl_${_component}
+        "${_nccl_${_component}_line}")
+    endforeach()
+    if(_nccl_MAJOR AND _nccl_MINOR AND _nccl_PATCH)
+      set(NCCLDevice_VERSION
+        "${_nccl_MAJOR}.${_nccl_MINOR}.${_nccl_PATCH}")
+    endif()
+  endif()
+endif()
+
+include(FindPackageHandleStandardArgs)
+if(_NCCL_DEVICE_USING_CONFIG)
+  find_package_handle_standard_args(NCCLDevice
+    REQUIRED_VARS NCCL_DEVICE_INCLUDE_DIR
+    VERSION_VAR NCCLDevice_VERSION)
+else()
+  find_package_handle_standard_args(NCCLDevice
+    REQUIRED_VARS NCCL_DEVICE_INCLUDE_DIR NCCL_DEVICE_LIBRARY
+    VERSION_VAR NCCLDevice_VERSION)
+endif()
+
+if(NCCLDevice_FOUND)
+  find_package(CUDAToolkit REQUIRED)
+
+  if(NOT TARGET NCCL::nccl)
+    add_library(NCCL::nccl UNKNOWN IMPORTED)
+    set_target_properties(NCCL::nccl PROPERTIES
+      IMPORTED_LOCATION "${NCCL_DEVICE_LIBRARY}"
+      INTERFACE_INCLUDE_DIRECTORIES "${NCCL_DEVICE_INCLUDE_DIR}"
+      INTERFACE_LINK_LIBRARIES "CUDA::cudart")
+  endif()
+
+  set_property(TARGET NCCL::nccl APPEND PROPERTY
+    INTERFACE_COMPILE_DEFINITIONS
+    NCCL_DEVICE_PERMIT_EXPERIMENTAL_CODE=1)
+endif()
+
+mark_as_advanced(NCCL_DEVICE_INCLUDE_DIR NCCL_DEVICE_LIBRARY)
+unset(_NCCL_DEVICE_MIN_VERSION)
+unset(_NCCL_DEVICE_USING_CONFIG)

mooncake-common/common.cmake

@@ -71,6 +71,8 @@ option(BUILD_EXAMPLES "Build examples" ON)
 option(BUILD_UNIT_TESTS "Build unit tests" ON)
 option(BUILD_BENCHMARK "Build benchmarks" ON)
 option(USE_CUDA "option for enabling gpu features for NVIDIA GPU" OFF)
+option(USE_NCCL_DEVICE "option for enabling the NCCL DeviceTransport backend" OFF)
+option(USE_NCCL_HOST "option for enabling the NCCL host RMA transport" OFF)
 option(USE_MLU "option for enabling Cambricon MLU features" OFF)
 option(USE_MUSA "option for enabling gpu features for MTHREADS GPU" OFF)
 option(USE_MACA "option for enabling gpu features for MUXI GPU with MACA" OFF)
@@ -206,6 +208,28 @@ if(USE_CUDA)
   link_directories(/usr/local/cuda/lib /usr/local/cuda/lib64)
 endif()
 
+if(USE_NCCL_DEVICE OR USE_NCCL_HOST)
+  if(NOT USE_CUDA)
+    message(FATAL_ERROR
+      "USE_NCCL_DEVICE and USE_NCCL_HOST require USE_CUDA=ON")
+  endif()
+  list(APPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_LIST_DIR})
+  find_package(NCCLDevice 2.30.4 REQUIRED MODULE)
+  add_compile_definitions(NCCL_DEVICE_PERMIT_EXPERIMENTAL_CODE=1)
+endif()
+
+if(USE_NCCL_DEVICE)
+  add_compile_definitions(USE_NCCL_DEVICE)
+  message(STATUS
+    "NCCL DeviceTransport support is enabled (NCCL ${NCCLDevice_VERSION})")
+endif()
+
+if(USE_NCCL_HOST)
+  add_compile_definitions(USE_NCCL_HOST)
+  message(STATUS
+    "NCCL host RMA transport is enabled (NCCL ${NCCLDevice_VERSION})")
+endif()
+
 if(USE_TPU)
   add_compile_definitions(USE_TPU)
   message(STATUS "TPU (PJRT) staging support is enabled")

mooncake-transfer-engine/example/CMakeLists.txt

@@ -106,3 +106,53 @@ if(USE_CUDA AND BUILD_DEVICE_TRANSPORT_EXAMPLE)
       CUDA_ARCHITECTURES "80;90")
   endif()
 endif()
+
+# NCCL DeviceTransport example. This is a manual two-rank validation because
+# communicator initialization and symmetric-window registration are collective.
+option(BUILD_NCCL_DEVICE_TRANSPORT_EXAMPLE
+       "Build the two-rank NCCL DeviceTransport LSA/GIN example"
+       OFF)
+if(USE_NCCL_DEVICE AND BUILD_NCCL_DEVICE_TRANSPORT_EXAMPLE)
+  enable_language(CUDA)
+  add_executable(nccl_device_transport_example
+                 ${WORKSPACE}/nccl_device_transport_example.cu)
+  target_include_directories(nccl_device_transport_example PRIVATE
+    ${CMAKE_SOURCE_DIR}/mooncake-transfer-engine/include)
+  target_link_libraries(nccl_device_transport_example PRIVATE
+    transfer_engine NCCL::nccl gflags::gflags glog::glog)
+  target_compile_options(nccl_device_transport_example PRIVATE
+    $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>)
+  set_target_properties(nccl_device_transport_example PROPERTIES
+    CUDA_STANDARD 20 CUDA_STANDARD_REQUIRED ON CUDA_EXTENSIONS OFF)
+
+  if(NOT CMAKE_CUDA_ARCHITECTURES)
+    if(TORCH_CUDA_ARCH_LIST)
+      set(_nccl_cuda_arch_list "")
+      foreach(_arch IN LISTS TORCH_CUDA_ARCH_LIST)
+        string(REPLACE "." "" _arch_clean "${_arch}")
+        list(APPEND _nccl_cuda_arch_list "${_arch_clean}")
+      endforeach()
+      set_target_properties(nccl_device_transport_example PROPERTIES
+        CUDA_ARCHITECTURES "${_nccl_cuda_arch_list}")
+    else()
+      set_target_properties(nccl_device_transport_example PROPERTIES
+        CUDA_ARCHITECTURES "80;90")
+    endif()
+  endif()
+endif()
+
+option(BUILD_NCCL_HOST_TRANSPORT_EXAMPLE
+       "Build the two-GPU NCCL host RMA Transfer Engine example"
+       OFF)
+if(USE_NCCL_HOST AND BUILD_NCCL_HOST_TRANSPORT_EXAMPLE)
+  add_executable(nccl_host_transport_example
+                 ${WORKSPACE}/nccl_host_transport_example.cpp)
+  target_include_directories(nccl_host_transport_example PRIVATE
+    ${CMAKE_SOURCE_DIR}/mooncake-transfer-engine/include)
+  target_link_libraries(nccl_host_transport_example PRIVATE
+    transfer_engine NCCL::nccl glog::glog gflags::gflags)
+  add_executable(nccl_host_transport_distributed_example
+                 ${WORKSPACE}/nccl_host_transport_distributed_example.cpp)
+  target_link_libraries(nccl_host_transport_distributed_example PRIVATE
+    transfer_engine NCCL::nccl glog::glog gflags::gflags)
+endif()