From f4e884f2224a25612eaeaeac2a854c1dd330c144 Mon Sep 17 00:00:00 2001
From: Sheng Lin <linsh0@protonmail.com>
Date: Wed, 17 Dec 2025 17:52:58 +0800
Subject: [PATCH] [NIXL][Bugfix] Fix NIXL/RDMA registration failure over
 CuMemAllocator (#29569)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Somoku <linsh0@protonmail.com>
Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
---
 csrc/cumem_allocator.cpp | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/csrc/cumem_allocator.cpp b/csrc/cumem_allocator.cpp
index 78dc840a9..6c2c18a66 100644
--- a/csrc/cumem_allocator.cpp
+++ b/csrc/cumem_allocator.cpp
@@ -107,6 +107,16 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
   prop.location.id = device;
   prop.allocFlags.compressionType = CU_MEM_ALLOCATION_COMP_NONE;
 
+#ifndef USE_ROCM
+  int flag = 0;
+  CUDA_CHECK(cuDeviceGetAttribute(
+      &flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED,
+      device));
+  if (flag) {  // support GPUDirect RDMA if possible
+    prop.allocFlags.gpuDirectRDMACapable = 1;
+  }
+#endif
+
 #ifndef USE_ROCM
   // Allocate memory using cuMemCreate
   CUDA_CHECK(cuMemCreate(p_memHandle, size, &prop, 0));