[TransferEngine][ROCm] Add HIP dmabuf MR registration for AMD GPUs (fixes #751)

.typos.toml

@@ -1,11 +1,14 @@
 [default]
-extend-ignore-words = ["CANN", "ASO", "fre", "wqs"]
+extend-ignore-words = ["CANN", "ASO", "fre", "wqs", "hsa"]
 
 [default.extend-words]
 CANN = "CANN"
 ASO = "ASO"
 fre = "fre"
 wqs = "wqs"
+# AMD HSA runtime symbol prefix (hsa_*, hsaRes, hsaErr, etc.) — used by the
+# ROCm dmabuf MR registration path.
+hsa = "hsa"
 
 [files]
 extend-exclude = [

mooncake-transfer-engine/src/CMakeLists.txt

@@ -96,6 +96,22 @@ endif()
 if(USE_HIP)
   target_include_directories(transfer_engine PRIVATE ${HIP_INCLUDE_DIRS})
   target_link_libraries(transfer_engine PUBLIC hip::host rt)
+
+  # Optional dmabuf MR registration path (requires hsa-runtime64 and a
+  # kernel with CONFIG_PCI_P2PDMA + CONFIG_DMABUF_MOVE_NOTIFY).
+  option(USE_HIP_DMABUF "Enable HIP dmabuf RDMA MR registration" ON)
+  if(USE_HIP_DMABUF)
+    find_package(hsa-runtime64 CONFIG)
+    if(hsa-runtime64_FOUND)
+      target_compile_definitions(transfer_engine PRIVATE USE_HIP_DMABUF)
+      target_link_libraries(transfer_engine PUBLIC hsa-runtime64::hsa-runtime64)
+      message(STATUS "HIP dmabuf MR registration enabled (hsa-runtime64 found)")
+    else()
+      message(STATUS "HIP dmabuf MR registration disabled (hsa-runtime64 not found)")
+    endif()
+  else()
+    message(STATUS "HIP dmabuf MR registration disabled (USE_HIP_DMABUF=OFF)")
+  endif()
 endif()
 
 if(USE_MLU)

mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp

@@ -30,6 +30,14 @@
 #include "config.h"
 #include "cuda_alike.h"
 #include "environ.h"
+#if defined(USE_HIP_DMABUF)
+#include <sys/utsname.h>
+
+#include <vector>
+
+#include <hsa/hsa.h>
+#include <hsa/hsa_ext_amd.h>
+#endif
 #include "transport/rdma_transport/endpoint_store.h"
 #include "transport/rdma_transport/rdma_endpoint.h"
 #include "transport/rdma_transport/rdma_transport.h"
@@ -50,6 +58,89 @@ bool containsAddress(const MemoryRegionMeta &region, uintptr_t addr) {
     const auto region_length = static_cast<uintptr_t>(region.mr->length);
     return region_start <= addr && addr - region_start < region_length;
 }
+
+#if defined(USE_HIP_DMABUF)
+// Returns true when the kernel has CONFIG_PCI_P2PDMA and
+// CONFIG_DMABUF_MOVE_NOTIFY enabled, which are required for ibv_reg_dmabuf_mr
+// to produce working RDMA transfers (not just successful registration).
+// Result is cached after the first call.
+bool isKernelDmabufSupported() {
+    static const bool supported = []() {
+        if (const char *env = std::getenv("MOONCAKE_DISABLE_HIP_DMABUF")) {
+            if (std::string(env) != "0") {
+                LOG(INFO)
+                    << "HIP dmabuf disabled via MOONCAKE_DISABLE_HIP_DMABUF";
+                return false;
+            }
+        }
+        struct utsname uts{};
+        std::string release;
+        if (uname(&uts) == 0) release = uts.release;
+
+        const char *needles[] = {"CONFIG_PCI_P2PDMA=y",
+                                 "CONFIG_DMABUF_MOVE_NOTIFY=y"};
+        std::vector<bool> found(2, false);
+
+        const std::vector<std::string> candidates = {
+            "/proc/config.gz",  // compressed; checked below via gzip-magic skip
+            "/boot/config-" + release,
+            "/usr/src/linux-" + release + "/.config",
+            "/usr/src/linux/.config",
+            "/usr/lib/modules/" + release + "/config",
+            "/usr/lib/ostree-boot/config-" + release,
+            "/usr/lib/kernel/config-" + release,
+            "/usr/src/linux-headers-" + release + "/.config",
+            "/lib/modules/" + release + "/build/.config",
+        };
+
+        for (const auto &path : candidates) {
+            std::ifstream f(path);
+            if (!f.good()) continue;
+            // /proc/config.gz is gzipped; we skip it here (the kallsyms
+            // fallback below covers that case in practice).
+            if (path.find(".gz") != std::string::npos) continue;
+            std::string line;
+            while (std::getline(f, line)) {
+                for (size_t i = 0; i < 2; ++i) {
+                    if (!found[i] && line.find(needles[i]) != std::string::npos)
+                        found[i] = true;
+                }
+                if (found[0] && found[1]) break;
+            }
+            if (found[0] && found[1]) break;
+        }
+
+        // Fallback: probe /proc/kallsyms for the corresponding kernel symbols.
+        if (!found[0] || !found[1]) {
+            std::ifstream f("/proc/kallsyms");
+            if (f.good()) {
+                std::string line;
+                while (std::getline(f, line)) {
+                    if (!found[0] &&
+                        line.find("pci_p2pdma") != std::string::npos)
+                        found[0] = true;
+                    if (!found[1] &&
+                        line.find("dma_buf_move_notify") != std::string::npos)
+                        found[1] = true;
+                    if (found[0] && found[1]) break;
+                }
+            }
+        }
+
+        bool ok = found[0] && found[1];
+        if (!ok) {
+            LOG(WARNING)
+                << "Kernel lacks CONFIG_PCI_P2PDMA / CONFIG_DMABUF_MOVE_NOTIFY "
+                << "(p2pdma=" << found[0] << " move_notify=" << found[1]
+                << "); HIP dmabuf MR registration disabled, falling back to "
+                << "ibv_reg_mr() (which requires an amdgpu peermem driver). "
+                << "Rebuild kernel with both options for GPU-direct RDMA.";
+        }
+        return ok;
+    }();
+    return supported;
+}
+#endif  // USE_HIP_DMABUF
 }  // namespace
 
 RdmaContext::RdmaContext(RdmaTransport &engine, const std::string &device_name)
@@ -338,6 +429,99 @@ int RdmaContext::registerMemoryRegionInternal(void *addr, size_t length,
         cuDevicePrimaryCtxRelease(cuDev);
 #endif
     }
+#elif defined(USE_HIP_DMABUF)
+    hipPointerAttribute_t hipAttr{};
+    hipError_t hipRes = hipPointerGetAttributes(&hipAttr, addr);
+
+    if (hipRes != hipSuccess || hipAttr.type == hipMemoryTypeHost ||
+        hipAttr.type == hipMemoryTypeUnregistered) {
+        // Host memory — standard ibv_reg_mr() path.
+        mrMeta.addr = addr;
+        mrMeta.mr = ibv_reg_mr(pd_, addr, length, access);
+    } else if (hipAttr.type == hipMemoryTypeManaged) {
+        // Managed (unified) memory pages can migrate between host and device;
+        // hsa_amd_portable_export_dmabuf captures the device-side handle at
+        // export time only, making the dmabuf fd stale after migration. Fall
+        // back to ibv_reg_mr() for safety.
+        LOG(WARNING) << "HIP managed memory at " << (uintptr_t)addr
+                     << " — dmabuf export skipped (pages may migrate); "
+                        "falling back to ibv_reg_mr";
+        mrMeta.addr = addr;
+        mrMeta.mr = ibv_reg_mr(pd_, addr, length, access);
+    } else if (hipAttr.type == hipMemoryTypeDevice &&
+               !isKernelDmabufSupported()) {
+        // Kernel lacks CONFIG_PCI_P2PDMA / CONFIG_DMABUF_MOVE_NOTIFY —
+        // ibv_reg_dmabuf_mr may succeed but transfers will silently fail.
+        // Fail at registration time instead.
+        mrMeta.addr = addr;
+        mrMeta.mr = ibv_reg_mr(pd_, addr, length, access);
+    } else if (hipAttr.type == hipMemoryTypeDevice) {
+        // Device memory + kernel support — export dmabuf fd and register.
+        // Pin to the owning device for the duration of the export calls.
+        struct HipDeviceGuard {
+            int prev_device = 0;
+            bool need_restore = false;
+            bool set_ok = false;
+            explicit HipDeviceGuard(int target_device) {
+                if (hipGetDevice(&prev_device) == hipSuccess) {
+                    need_restore = (prev_device != target_device);
+                }
+                set_ok = (hipSetDevice(target_device) == hipSuccess);
+            }
+            ~HipDeviceGuard() {
+                if (need_restore) {
+                    (void)hipSetDevice(prev_device);
+                }
+            }
+        } dev_guard(hipAttr.device);
+        if (!dev_guard.set_ok) {
+            LOG(ERROR) << "Failed to set HIP device to " << hipAttr.device
+                       << " for dmabuf export of " << (uintptr_t)addr;
+            return ERR_CONTEXT;
+        }
+
+        // Get the allocation base + size, since `addr` may sit at an offset
+        // within a larger hipMalloc block (caching allocators pack tensors).
+        hipDeviceptr_t allocBase = nullptr;
+        size_t allocSize = 0;
+        hipRes = hipMemGetAddressRange(&allocBase, &allocSize,
+                                       reinterpret_cast<hipDeviceptr_t>(addr));
+        if (hipRes != hipSuccess) {
+            LOG(ERROR) << "Failed to call hipMemGetAddressRange for "
+                       << (uintptr_t)addr
+                       << " hip error=" << hipGetErrorString(hipRes);
+            return ERR_CONTEXT;
+        }
+
+        int dmabuf_fd = -1;
+        uint64_t hsa_dmabuf_offset = 0;
+        hsa_status_t hsaRes = hsa_amd_portable_export_dmabuf(
+            allocBase, allocSize, &dmabuf_fd, &hsa_dmabuf_offset);
+        if (hsaRes != HSA_STATUS_SUCCESS) {
+            const char *hsaErr = nullptr;
+            hsa_status_string(hsaRes, &hsaErr);
+            LOG(ERROR) << "Failed to retrieve dmabuf for " << (uintptr_t)addr
+                       << " base=" << (uintptr_t)allocBase
+                       << " size=" << allocSize
+                       << " hsa error=" << (hsaErr ? hsaErr : "unknown");
+            return ERR_CONTEXT;
+        }
+
+        mrMeta.addr = addr;
+        // Offset within the dmabuf-backed region: distance from the
+        // allocation base, plus any offset hsa returned for the export.
+        uint64_t reg_offset =
+            (uintptr_t)addr - (uintptr_t)allocBase + hsa_dmabuf_offset;
+        mrMeta.mr = ibv_reg_dmabuf_mr(pd_, reg_offset, length, (uintptr_t)addr,
+                                      dmabuf_fd, access);
+        const int regErrno = errno;
+        if (close(dmabuf_fd) != 0) {
+            PLOG(WARNING) << "Failed to close dmabuf fd";
+        }
+        if (!mrMeta.mr) {
+            errno = regErrno;
+        }
+    }
 #else
     mrMeta.addr = addr;
     mrMeta.mr = ibv_reg_mr(pd_, addr, length, access);