Fallback when mlx5dv is not supported. (#1665)

Summary:

This change adds fallback support when mlx5dv (Mellanox device-specific extensions) is not available for RDMA operations. It modifies the queue pair creation logic to conditionally use either extended mlx5dv-based queue pairs (when supported) or standard ibverbs queue pairs (as fallback). The pt_cuda_alloc flag is updated to require mlx5dv support since it's necessary for merging memory segments when using PyTorch's CUDA allocator. The change adds a new `is_extended` parameter to control whether to create extended or standard queue pairs at runtime.

Adds an env variable `MONARCH_RDMA_MLX5DV_DISABLED` to test the new code path on dev machine.

Differential Revision: D85504061

Author: casteryh

monarch_rdma/src/ibverbs_primitives.rs

@@ -650,6 +650,10 @@ pub fn mlx5dv_supported() -> bool {
 }
 
 fn mlx5dv_supported_impl() -> bool {
+    if std::env::var("MONARCH_RDMA_DEBUG_MLX5DV_DISABLED_DO_NOT_USE").is_ok() {
+        eprintln!("mlx5dv support disabled by MONARCH_RDMA_DEBUG_MLX5DV_DISABLED_DO_NOT_USE");
+        return false;
+    }
     // SAFETY: We are calling C functions from libibverbs and libmlx5.
     unsafe {
         let mut num_devices = 0;
@@ -733,7 +737,7 @@ fn ibverbs_supported_impl() -> bool {
 ///
 /// `true` if both ibverbs devices and mlx5dv extensions are available, `false` otherwise.
 pub fn rdma_supported() -> bool {
-    ibverbs_supported() && mlx5dv_supported()
+    ibverbs_supported()
 }
 
 /// Represents a view of a memory region that can be registered with an RDMA device.
@@ -1108,24 +1112,4 @@ mod tests {
         let mlx5dv_support = mlx5dv_supported();
         println!("mlx5dv_supported: {}", mlx5dv_support);
     }
-
-    #[test]
-    fn test_rdma_supported_combines_checks() {
-        // This test verifies that rdma_supported() properly combines both checks
-        let ibverbs_support = ibverbs_supported();
-        let mlx5dv_support = mlx5dv_supported();
-        let rdma_support = rdma_supported();
-
-        // rdma_supported should be true only if both checks pass
-        assert_eq!(
-            rdma_support,
-            ibverbs_support && mlx5dv_support,
-            "rdma_supported should equal (ibverbs_supported && mlx5dv_supported)"
-        );
-
-        println!(
-            "ibverbs_supported: {}, mlx5dv_supported: {}, rdma_supported: {}",
-            ibverbs_support, mlx5dv_support, rdma_support
-        );
-    }
 }

monarch_rdma/src/rdma_components.rs

@@ -66,6 +66,7 @@ use crate::ibverbs_primitives::IbvWc;
 use crate::ibverbs_primitives::IbverbsConfig;
 use crate::ibverbs_primitives::RdmaOperation;
 use crate::ibverbs_primitives::RdmaQpInfo;
+use crate::ibverbs_primitives::mlx5dv_supported;
 
 #[derive(Debug, Named, Clone, Serialize, Deserialize)]
 pub struct DoorBell {
@@ -530,6 +531,7 @@ impl RdmaQueuePair {
                 config.max_recv_wr.try_into().unwrap(),
                 config.max_send_sge.try_into().unwrap(),
                 config.max_recv_sge.try_into().unwrap(),
+                mlx5dv_supported() as u8,
             );
 
             if qp.is_null() {

monarch_rdma/src/rdma_manager_actor.rs

@@ -48,6 +48,7 @@ use crate::ibverbs_primitives::IbverbsConfig;
 use crate::ibverbs_primitives::RdmaMemoryRegionView;
 use crate::ibverbs_primitives::RdmaQpInfo;
 use crate::ibverbs_primitives::ibverbs_supported;
+use crate::ibverbs_primitives::mlx5dv_supported;
 use crate::rdma_components::RdmaBuffer;
 use crate::rdma_components::RdmaDomain;
 use crate::rdma_components::RdmaQueuePair;
@@ -120,7 +121,7 @@ pub enum RdmaManagerMessage {
         reply: OncePortRef<RdmaQpInfo>,
     },
     ReleaseQueuePair {
-        /// `other` - The ActorId to release queue pair for  
+        /// `other` - The ActorId to release queue pair for
         other: ActorRef<RdmaManagerActor>,
         self_device: String,
         other_device: String,
@@ -150,6 +151,8 @@ pub struct RdmaManagerActor {
     // True if both C10 CUDA allocator is enabled AND expandable segments are enabled
     pt_cuda_alloc: bool,
 
+    mlx5dv_supported: bool,
+
     // Map of unique RdmaMemoryRegionView to ibv_mr*.  In case of cuda w/ pytorch its -1
     // since its managed independently.  Only used for registration/deregistration purposes
     mr_map: HashMap<usize, usize>,
@@ -248,7 +251,7 @@ impl Drop for RdmaManagerActor {
         }
 
         // 4. Deregister all CUDA segments (if using PyTorch CUDA allocator)
-        if self.pt_cuda_alloc {
+        if self.cuda_pt_alloc_enabled() {
             unsafe {
                 let result = rdmaxcel_sys::deregister_segments();
                 if result != 0 {
@@ -265,6 +268,11 @@ impl Drop for RdmaManagerActor {
 }
 
 impl RdmaManagerActor {
+    /// Whether to register all memory regions allocated by the PyTorch CUDA allocator
+    /// True if both `pt_cuda_alloc` and `mlx5dv_supported` are true
+    fn cuda_pt_alloc_enabled(&self) -> bool {
+        self.pt_cuda_alloc && self.mlx5dv_supported
+    }
     /// Get or create a domain and loopback QP for the specified RDMA device
     fn get_or_create_device_domain(
         &mut self,
@@ -420,10 +428,7 @@ impl RdmaManagerActor {
             let mut mr: *mut rdmaxcel_sys::ibv_mr = std::ptr::null_mut();
             let mrv;
 
-            // Copy pt_cuda_alloc to avoid borrowing issues
-            let pt_cuda_alloc = self.pt_cuda_alloc;
-
-            if is_cuda && pt_cuda_alloc {
+            if is_cuda && self.cuda_pt_alloc_enabled() {
                 // Get registered segments and check if our memory range is covered
                 let mut maybe_mrv = self.find_cuda_segment_for_address(addr, size);
                 // not found, lets re-sync with caching allocator  and retry
@@ -529,6 +534,8 @@ impl Actor for RdmaManagerActor {
 
         let pt_cuda_alloc = crate::rdma_components::pt_cuda_allocator_compatibility();
 
+        let mlx5dv_supported = mlx5dv_supported();
+
         // check config and hardware support align
         if config.use_gpu_direct {
             match validate_execution_context().await {
@@ -557,6 +564,7 @@ impl Actor for RdmaManagerActor {
             device_domains: HashMap::new(),
             config,
             pt_cuda_alloc,
+            mlx5dv_supported,
             mr_map: HashMap::new(),
             mrv_id: 0,
             pci_to_device,

rdmaxcel-sys/src/rdmaxcel.c

@@ -28,7 +28,8 @@ struct ibv_qp* create_qp(
     int max_send_wr,
     int max_recv_wr,
     int max_send_sge,
-    int max_recv_sge) {
+    int max_recv_sge,
+    uint8_t is_extended) {
   // Create separate completion queues for send and receive operations
   struct ibv_cq* send_cq = ibv_create_cq(context, cq_entries, NULL, NULL, 0);
   if (!send_cq) {
@@ -43,44 +44,77 @@ struct ibv_qp* create_qp(
     return NULL;
   }
 
-  // Initialize extended queue pair attributes
-  struct ibv_qp_init_attr_ex qp_init_attr_ex = {
-      .qp_context = NULL,
-      .send_cq = send_cq,
-      .recv_cq = recv_cq,
-      .srq = NULL,
-      .cap =
-          {
-              .max_send_wr = max_send_wr,
-              .max_recv_wr = max_recv_wr,
-              .max_send_sge = max_send_sge,
-              .max_recv_sge = max_recv_sge,
-              .max_inline_data = 0,
-          },
-      .qp_type = IBV_QPT_RC,
-      .sq_sig_all = 0,
-      .pd = pd,
-      .comp_mask = IBV_QP_INIT_ATTR_PD | IBV_QP_INIT_ATTR_SEND_OPS_FLAGS,
-      .send_ops_flags = IBV_QP_EX_WITH_RDMA_WRITE | IBV_QP_EX_WITH_RDMA_READ |
-          IBV_QP_EX_WITH_SEND,
-      .create_flags = 0,
-  };
-
-  struct mlx5dv_qp_init_attr mlx5dv_attr = {};
-  mlx5dv_attr.comp_mask |= MLX5DV_QP_INIT_ATTR_MASK_SEND_OPS_FLAGS;
-  mlx5dv_attr.send_ops_flags =
-      MLX5DV_QP_EX_WITH_MKEY_CONFIGURE | MLX5DV_QP_EX_WITH_MR_LIST;
-
-  // Create extended queue pair
-  struct ibv_qp* qp = mlx5dv_create_qp(context, &qp_init_attr_ex, &mlx5dv_attr);
-  if (!qp) {
-    perror("failed to create extended queue pair (QP)");
-    ibv_destroy_cq(send_cq);
-    ibv_destroy_cq(recv_cq);
-    return NULL;
+  if (is_extended) {
+    // Initialize extended queue pair attributes
+    struct ibv_qp_init_attr_ex qp_init_attr_ex = {
+        .qp_context = NULL,
+        .send_cq = send_cq,
+        .recv_cq = recv_cq,
+        .srq = NULL,
+        .cap =
+            {
+                .max_send_wr = max_send_wr,
+                .max_recv_wr = max_recv_wr,
+                .max_send_sge = max_send_sge,
+                .max_recv_sge = max_recv_sge,
+                .max_inline_data = 0,
+            },
+        .qp_type = IBV_QPT_RC,
+        .sq_sig_all = 0,
+        .pd = pd,
+        .comp_mask = IBV_QP_INIT_ATTR_PD | IBV_QP_INIT_ATTR_SEND_OPS_FLAGS,
+        .send_ops_flags = IBV_QP_EX_WITH_RDMA_WRITE | IBV_QP_EX_WITH_RDMA_READ |
+            IBV_QP_EX_WITH_SEND,
+        .create_flags = 0,
+    };
+
+    struct mlx5dv_qp_init_attr mlx5dv_attr = {};
+    mlx5dv_attr.comp_mask |= MLX5DV_QP_INIT_ATTR_MASK_SEND_OPS_FLAGS;
+    mlx5dv_attr.send_ops_flags =
+        MLX5DV_QP_EX_WITH_MKEY_CONFIGURE | MLX5DV_QP_EX_WITH_MR_LIST;
+
+    // Create extended queue pair
+    struct ibv_qp* qp =
+        mlx5dv_create_qp(context, &qp_init_attr_ex, &mlx5dv_attr);
+    if (!qp) {
+      perror("failed to create extended queue pair (QP)");
+      ibv_destroy_cq(send_cq);
+      ibv_destroy_cq(recv_cq);
+      return NULL;
+    }
+
+    return qp;
+
+  } else {
+    // Initialize queue pair attributes
+    struct ibv_qp_init_attr qp_init_attr = {
+        .qp_context = NULL,
+        .send_cq = send_cq,
+        .recv_cq = recv_cq,
+        .srq = NULL,
+        .cap =
+            {
+                .max_send_wr = max_send_wr,
+                .max_recv_wr = max_recv_wr,
+                .max_send_sge = max_send_sge,
+                .max_recv_sge = max_recv_sge,
+                .max_inline_data = 0,
+            },
+        .qp_type = IBV_QPT_RC,
+        .sq_sig_all = 0,
+    };
+
+    // Create queue pair
+    struct ibv_qp* qp = ibv_create_qp(pd, &qp_init_attr);
+    if (!qp) {
+      perror("failed to create queue pair (QP)");
+      ibv_destroy_cq(send_cq);
+      ibv_destroy_cq(recv_cq);
+      return NULL;
+    }
+
+    return qp;
   }
-
-  return qp;
 }
 
 struct mlx5dv_qp* create_mlx5dv_qp(struct ibv_qp* qp) {

rdmaxcel-sys/src/rdmaxcel.h

@@ -69,7 +69,8 @@ struct ibv_qp* create_qp(
     int max_send_wr,
     int max_recv_wr,
     int max_send_sge,
-    int max_recv_sge);
+    int max_recv_sge,
+    uint8_t is_extended);
 
 struct mlx5dv_qp* create_mlx5dv_qp(struct ibv_qp* qp);