Introduce RdmaProvider and ERDMA support

Signed-off-by: Guangguan Wang <[email protected]>

Author: GuangguanWang

src/rdma_provider.h

@@ -0,0 +1,116 @@
+// Copyright 2025 Alibaba Group. or its affiliates. All Rights Reserved.
+//
+// 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.
+// =============================================================================
+
+#ifndef RDMA_PROVIDER_H_
+#define RDMA_PROVIDER_H_
+
+#ifdef DMLC_USE_RDMA
+
+#include <string>
+#include <rdma/rdma_cma.h>
+
+#include "dmlc/logging.h"
+#include "./ibvwarp.h"
+
+namespace ps {
+
+class RdmaProvider {
+ public:
+  static RdmaProvider *GetProvider(struct ibv_context *context);
+
+  virtual int InlineSize() = 0;
+  virtual int SetQPLag(struct ibv_qp* qp, int port_num) = 0;
+  virtual int ErrIgnore(struct ibv_wc *wc) = 0;
+
+ protected:
+  RdmaProvider() {}
+};
+
+class ErdmaProvider: public RdmaProvider {
+ public:
+  static inline RdmaProvider *Get() {
+    static ErdmaProvider *inst_ptr = new ErdmaProvider();
+    return inst_ptr;
+  }
+
+  static inline const char *DevPrefix() { return "erdma"; }
+  int InlineSize() override { return 96; }
+  int SetQPLag(struct ibv_qp *qp, int port_num) override { return 1; }
+  int ErrIgnore(struct ibv_wc *wc) override {
+    if (wc->status == IBV_WC_SUCCESS ||
+        (wc->status == IBV_WC_WR_FLUSH_ERR && !wc->vendor_err)) {
+      return 1;
+    }
+    return 0;
+  }
+
+ private:
+  explicit ErdmaProvider() {}
+};
+
+class Mlx5Provider: public RdmaProvider {
+ public:
+
+  static inline RdmaProvider *Get() {
+    static RdmaProvider *inst_ptr = new Mlx5Provider();
+    return inst_ptr;
+  }
+
+  static inline const char *DevPrefix() { return "mlx5"; }
+  int InlineSize() override { return 512; }
+
+  int SetQPLag(struct ibv_qp *qp, int port_num) override {
+    int ret = wrap_mlx5dv_modify_qp_lag_port(qp, port_num);
+    if (ret == 1) {
+      uint8_t set_port = 0xff, act_port = 0xff;
+      wrap_mlx5dv_query_qp_lag_port(qp, &set_port, &act_port);
+      PS_LOG(INFO) << "QP LAG Port: QP: " << qp->qp_num
+                   << ", Modify Port: " << port_num
+                   << ", Set to Port: " << static_cast<int>(set_port)
+                   << ", Active Port: " << static_cast<int>(act_port);
+    }
+    return ret;
+  }
+
+  int ErrIgnore(struct ibv_wc *wc) override {
+    return wc->status == IBV_WC_SUCCESS;
+  }
+
+ private:
+  explicit Mlx5Provider() {
+    if (wrap_ibv_symbols() != 1) {
+      PS_LOG(WARNING) << "Load mlx5 symbols fails.";
+    }
+  }
+};
+
+RdmaProvider *RdmaProvider::GetProvider(struct ibv_context *context) {
+  const char *dev_name = ibv_get_device_name(context->device);
+  if (strstr(dev_name, ErdmaProvider::DevPrefix())) {
+    return ErdmaProvider::Get();
+  } else {
+    if (!strstr(dev_name, Mlx5Provider::DevPrefix())) {
+      PS_LOG(WARNING) << "rdma device(" << dev_name
+                      << ") with unknow provider, use mlx5 as default, maybe "
+                         "not compatible.";
+    }
+    return Mlx5Provider::Get();
+  }
+}
+
+}  // namespace ps
+
+#endif  // DMLC_USE_RDMA
+#endif  // RDMA_PROVIDER_H_

src/rdma_transport.h

@@ -27,7 +27,7 @@
 #include <unordered_map>
 #include <vector>
 
-#include "./ibvwarp.h"
+#include "./rdma_provider.h"
 #include "./rdma_utils.h"
 #include "dmlc/logging.h"
 #include "ps/internal/multi_qp.h"
@@ -54,19 +54,19 @@ struct Endpoint {
   int kStartDepth = 128;
   int kRxDepth = 256;
   int kReplyDepth = kRxDepth;
+  int kMaxInlineSize;
   WRContext *rx_ctx;
   WRContext *start_ctx;
   WRContext *reply_ctx;
 
   ThreadsafeQueue<WRContext *> free_start_ctx;
   ThreadsafeQueue<WRContext *> free_reply_ctx;
 
+  RdmaProvider *rdma_provider = nullptr;
+
   uint8_t inited = 0;
 
   Endpoint() : node_id(Node::kEmpty), rx_ctx() {
-    if (wrap_ibv_symbols() != 1) {
-      PS_LOG(WARNING) << "Load mlx5 symbols fails.";
-    }
     FOR_QPS {
       cm_ids[qpIndex] = nullptr;
       status_list[qpIndex] = IDLE;
@@ -171,7 +171,7 @@ struct Endpoint {
     }
   }
 
-  void Init(struct ibv_cq *cq, struct ibv_pd *pd, rdma_cm_id *id = nullptr) {
+  void Init(struct ibv_cq *cq, struct ibv_pd *pd, RdmaProvider *provider, rdma_cm_id *id = nullptr) {
     struct ibv_qp_init_attr attr;
     memset(&attr, 0, sizeof(ibv_qp_init_attr));
     attr.send_cq = cq;
@@ -180,16 +180,18 @@ struct Endpoint {
     attr.cap.max_recv_wr = kRxDepth;
     attr.cap.max_send_sge = kSGEntry;
     attr.cap.max_recv_sge = kSGEntry;
-    attr.cap.max_inline_data = 256;
+    attr.cap.max_inline_data = provider->InlineSize();
     attr.qp_type = IBV_QPT_RC;
     attr.sq_sig_all = 0;
     PS_CHECK_EQ(rdma_create_qp(id, pd, &attr), 0)
         << "Create RDMA queue pair failed: " << strerror(errno);
     id->pd = pd;
+    kMaxInlineSize = attr.cap.max_inline_data;
 
     PS_LOG(TRACE) << "qp created: pd=" << pd << " , cq=" << cq
-                  << ", qp=" << id->qp->qp_num;
+                  << ", qp=" << id->qp->qp_num << ", maxInline=" << kMaxInlineSize;
     if (inited == 0) {
+      rdma_provider = provider;
       InitSendContextHelper(pd, start_ctx, &free_start_ctx, kStartDepth,
                             kRendezvousStartContext);
       InitSendContextHelper(pd, reply_ctx, &free_reply_ctx, kReplyDepth,
@@ -227,21 +229,14 @@ struct Endpoint {
 
     if (val == 1) {
       multi_qp_ = true;
+      PS_CHECK(rdma_provider);
       FOR_QPS {
         int lag = 1 + qpIndex % 2;
-        int ret = wrap_mlx5dv_modify_qp_lag_port(cm_ids[qpIndex]->qp, lag);
+        int ret = rdma_provider->SetQPLag(cm_ids[qpIndex]->qp, lag);
         if (ret != 1) {
-          PS_LOG(INFO) << "Failed to mlx5dv_modify_qp_lag_port qp ["
+          PS_LOG(INFO) << "Failed to SetQPLag qp ["
                        << cm_ids[qpIndex]->qp->qp_num << "] to port: " << lag
                        << ", qp type: " << cm_ids[qpIndex]->qp->qp_type;
-        } else {
-          uint8_t set_port = 0xff, act_port = 0xff;
-          wrap_mlx5dv_query_qp_lag_port(cm_ids[qpIndex]->qp, &set_port,
-                                        &act_port);
-          PS_LOG(INFO) << "QP LAG Port: QP: " << cm_ids[qpIndex]->qp->qp_num
-                       << ", Modify Port: " << lag
-                       << ", Set to Port: " << static_cast<int>(set_port)
-                       << ", Active Port: " << static_cast<int>(act_port);
         }
       }
     }
@@ -310,6 +305,9 @@ class RDMATransport : public Transport {
     allocator_ = PS_CHECK_NOTNULL(allocator);
     pagesize_ = sysconf(_SC_PAGESIZE);
 
+    PS_CHECK_GT(endpoint_->kMaxInlineSize, 0);
+    max_inline_size_ = endpoint_->kMaxInlineSize;
+
     postoffice_ = postoffice;
     is_server_ = postoffice_->is_server();
 #ifdef STEPMESH_USE_GDR
@@ -342,32 +340,51 @@ class RDMATransport : public Transport {
                                 uint32_t rkey, uint32_t idx,
                                 bool inline_write = false,
                                 struct ibv_send_wr *prev_wr = nullptr) {
-    struct ibv_sge sge;
-    sge.addr = reinterpret_cast<uint64_t>(msg_buf->inline_buf);
-    sge.length = msg_buf->inline_len;
-    sge.lkey = allocator_->LocalKey(msg_buf->inline_buf);
-
-    struct ibv_send_wr wr = {}, *bad_wr = nullptr;
-    memset(&wr, 0, sizeof(wr));
-    wr.wr_id = reinterpret_cast<uint64_t>(msg_buf);
-    wr.opcode = IBV_WR_RDMA_WRITE_WITH_IMM;
-    wr.next = nullptr;
-    wr.imm_data = idx;
-    wr.send_flags = IBV_SEND_SIGNALED;
-    wr.sg_list = &sge;
-    wr.num_sge = 1;
-    wr.wr.rdma.remote_addr = remote_addr;
-    wr.wr.rdma.rkey = rkey;
+    struct ibv_send_wr wr[kRdmaMaxWRs], *bad_wr = nullptr;
+    struct ibv_sge sge[kRdmaMaxWRs];
+    uint32_t lkey = allocator_->LocalKey(msg_buf->inline_buf);
+    size_t last_len = 0, offset = 0, len = 0;
+    int num_wr = 1, last_flags = 0;
 
     if (inline_write) {
-      wr.send_flags |= IBV_SEND_INLINE;
+      num_wr = DivUp(msg_buf->inline_len, max_inline_size_);
+      last_len = msg_buf->inline_len % max_inline_size_;
+      last_flags = IBV_SEND_INLINE | IBV_SEND_SIGNALED;
+    } else {
+      num_wr = 1;
+      last_len = msg_buf->inline_len;
+      last_flags = IBV_SEND_SIGNALED;
+    }
+
+    PS_CHECK_LE(num_wr, kRdmaMaxWRs)
+        << "too many wrs, send_len: " << msg_buf->inline_len
+        << " , max_inline: " << max_inline_size_;
+    memset(wr, 0, sizeof(struct ibv_send_wr) * num_wr);
+
+    for (int i = 0; i < num_wr; ++i) {
+      bool is_last = (i == (num_wr - 1));
+      len = is_last ? last_len : max_inline_size_;
+      sge[i].addr = reinterpret_cast<uint64_t>(msg_buf->inline_buf + offset);
+      sge[i].length = len;
+      sge[i].lkey = lkey;
+
+      wr[i].wr_id = reinterpret_cast<uint64_t>(msg_buf);
+      wr[i].opcode = is_last ? IBV_WR_RDMA_WRITE_WITH_IMM : IBV_WR_RDMA_WRITE;
+      wr[i].next = is_last ? nullptr : &wr[i + 1];
+      wr[i].imm_data = is_last ? idx : 0;
+      wr[i].send_flags = is_last ? last_flags : 0;
+      wr[i].sg_list = &sge[i];
+      wr[i].num_sge = 1;
+      wr[i].wr.rdma.remote_addr = remote_addr + offset;
+      wr[i].wr.rdma.rkey = rkey;
+      offset += len;
     }
 
     if (prev_wr == nullptr) {
-      PS_CHECK_EQ(ibv_post_send(endpoint_->cm_ids[0]->qp, &wr, &bad_wr), 0)
+      PS_CHECK_EQ(ibv_post_send(endpoint_->cm_ids[0]->qp, wr, &bad_wr), 0)
           << "ibv_post_send failed.";
     } else {
-      prev_wr->next = &wr;
+      prev_wr->next = &wr[0];
       PS_CHECK_EQ(ibv_post_send(endpoint_->cm_ids[0]->qp, prev_wr, &bad_wr), 0)
           << "ibv_post_send failed.";
     }
@@ -839,6 +856,7 @@ class RDMATransport : public Transport {
 
  protected:
   size_t pagesize_ = 4096;
+  size_t max_inline_size_ = 0;
   Endpoint *endpoint_;
   MemoryAllocator *allocator_;
   BackendMemoryAllocator *mem_allocator_;

src/rdma_utils.h

@@ -62,6 +62,7 @@ static const int kSGEntry = 1;
 static const int kTimeoutms = 1000;
 static const int kRdmaListenBacklog = 128;
 static const int kMaxHostnameLength = 16;
+static const int kRdmaMaxWRs = 12;
 
 // should have the same prefix with BytePS shared memory
 // for pcie reduce:  BytePS_Pcie_{pcie_id}_ShM_{JOB_ID}_{BYTEPS_KEY}

src/rdma_van.h

@@ -732,6 +732,9 @@ class RDMAVan : public Van {
     context_ = context;
     PS_CHECK(context_) << "ibv_context* empty";
 
+    provider_ = RdmaProvider::GetProvider(context_);
+    PS_CHECK(provider_) << "failed to get rdma provider";
+
     pd_ = ibv_alloc_pd(context_);
     PS_CHECK(pd_) << "Failed to allocate protection domain";
 
@@ -774,13 +777,18 @@ class RDMAVan : public Van {
 
       PS_CHECK_GE(ne, 0);
       for (int i = 0; i < ne; ++i) {
-        PS_CHECK(wc[i].status == IBV_WC_SUCCESS)
-            << "Failed status \n"
-            << ibv_wc_status_str(wc[i].status) << " " << wc[i].status << " "
-            << static_cast<uint64_t>(wc[i].wr_id) << " " << wc[i].vendor_err
-            << " " << wc[i].opcode << " "
-            << (wc[i].opcode == IBV_WC_RECV ? "RECV" : "OTHER")
-            << " postoffice ptr: " << reinterpret_cast<void *>(postoffice_);
+        if (wc[i].status != IBV_WC_SUCCESS) {
+          if (provider_->ErrIgnore(&wc[i])) {
+            continue;
+          }
+          PS_LOG(FATAL) << "Failed status \n"
+                        << ibv_wc_status_str(wc[i].status) << " " << wc[i].status << " "
+                        << static_cast<uint64_t>(wc[i].wr_id) << " " << wc[i].vendor_err
+                        << " " << wc[i].opcode << " "
+                        << (wc[i].opcode == IBV_WC_RECV ? "RECV" : "OTHER")
+                        << " postoffice ptr: " << reinterpret_cast<void *>(postoffice_);
+        }
+
 
         // IBV_WC_RDMA_WRITE use msg_buf as the wr_id
         // so there won't be context and endpoint for this op
@@ -991,7 +999,7 @@ class RDMAVan : public Van {
       InitContext(id->verbs);
     }
 
-    endpoint->Init(cq_, pd_, id);
+    endpoint->Init(cq_, pd_, provider_, id);
 
     bool is_local_node =
         disable_ipc_
@@ -1045,7 +1053,7 @@ class RDMAVan : public Van {
     if (context_ == nullptr) {
       InitContext(id->verbs);
     }
-    endpoint->Init(cq_, pd_, id);
+    endpoint->Init(cq_, pd_, provider_, id);
     endpoint->inComingCount++;
     RequestContext ctx;
     ctx.node = static_cast<uint32_t>(my_node_.id);
@@ -1122,6 +1130,7 @@ class RDMAVan : public Van {
 
   struct rdma_event_channel *event_channel_ = nullptr;
   struct ibv_context *context_ = nullptr;
+  RdmaProvider *provider_ = nullptr;
 
   // ibverbs protection domain
   struct ibv_pd *pd_ = nullptr;