RPC: use RDMA completion channel to not spin (#29440)

* RPC: use RDMA completion queue to not spin

* add TODO for apple RDMA
This commit is contained in:
Aman Gupta
2026-09-27 17:28:32 +08:00
committed by GitHub
parent 7fb2b082ce
commit d7fb90e8e2
2 changed files with 64 additions and 4 deletions
+1
View File
@@ -26,6 +26,7 @@
// so every SEND posts a whole 128KiB stride over the wire, even when partially filled.
// (In testing 128KiB was the best performing among 32, 64, 128, 256)
//TODO: add mechanism similar to https://github.com/ggml-org/llama.cpp/pull/29440 to prevent idle CPU from spinning
static constexpr uint32_t RDMA_SEG_MAGIC = 0x52534547u; // "RSEG"
static constexpr int RDMA_NBUF = 16; // ring depth (frames per direction)
static constexpr size_t RDMA_FRAME = 4096; // Thunderbolt frame (fixed on Apple)
+63 -4
View File
@@ -25,6 +25,8 @@
#ifdef GGML_RPC_RDMA
# include <infiniband/verbs.h>
# include <array>
# include <cerrno>
# include <chrono>
# include <time.h>
# ifndef _WIN32
# include <poll.h>
@@ -54,6 +56,8 @@ using rdma_gid_t = std::array<uint8_t, RDMA_GID_SIZE>;
#if defined(GGML_RPC_RDMA) && !defined(GGML_RPC_RDMA_APPLE)
static constexpr size_t RDMA_CHUNK = 256 * 1024; // 256 KiB per send/recv (fits default 8 MiB memlock)
static constexpr int RDMA_RX_DEPTH = 24; // pre-posted recv ring: 24 × 256 KiB = 6 MiB
// keep polling the CQ for this long after the last activity, then sleep until the next completion
static constexpr auto RDMA_SPIN_TIME = std::chrono::milliseconds(100);
struct rdma_conn {
struct ibv_context * ctx = nullptr;
@@ -61,6 +65,9 @@ struct rdma_conn {
struct ibv_cq * scq = nullptr; // send completions
struct ibv_cq * rcq = nullptr; // recv completions
struct ibv_qp * qp = nullptr;
struct ibv_comp_channel * ch = nullptr; // CQ events, so an idle connection can sleep instead of spinning
std::chrono::steady_clock::time_point last_active; // last completion or posted send
void * tx_buf = nullptr;
struct ibv_mr * tx_mr = nullptr;
@@ -95,6 +102,7 @@ struct rdma_conn {
if (qp) ibv_destroy_qp(qp);
if (scq) ibv_destroy_cq(scq);
if (rcq) ibv_destroy_cq(rcq);
if (ch) ibv_destroy_comp_channel(ch);
if (pd) ibv_dealloc_pd(pd);
if (ctx) ibv_close_device(ctx);
}
@@ -142,6 +150,7 @@ struct socket_t::impl {
bool tcp_peer_closed();
bool rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, const uint8_t * remote_gid);
bool rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc);
bool rdma_wait_event();
std::unique_ptr<rdma_conn> rdma;
rdma_local_info rdma_local = {};
@@ -291,8 +300,10 @@ bool socket_t::impl::rdma_probe() {
rdma->pd = ibv_alloc_pd(ibctx);
if (!rdma->pd) return false;
rdma->scq = ibv_create_cq(ibctx, 16, nullptr, nullptr, 0);
rdma->rcq = ibv_create_cq(ibctx, RDMA_RX_DEPTH + 4, nullptr, nullptr, 0);
// without a completion channel rdma_poll() spins all the time, as before
rdma->ch = ibv_create_comp_channel(ibctx);
rdma->scq = ibv_create_cq(ibctx, 16, nullptr, rdma->ch, 0);
rdma->rcq = ibv_create_cq(ibctx, RDMA_RX_DEPTH + 4, nullptr, rdma->ch, 0);
if (!rdma->scq || !rdma->rcq) return false;
ibv_qp_init_attr qia = {};
@@ -394,15 +405,45 @@ bool socket_t::impl::rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, con
}
}
rdma->last_active = std::chrono::steady_clock::now();
GGML_LOG_INFO("RDMA activated: qpn=%u->%u mtu=%d rx_depth=%d\n",
rdma_local.qpn, remote_qpn, 128 << rdma_local.path_mtu, RDMA_RX_DEPTH);
return true;
}
// Sleep until the completion channel has an event or the TCP peer closes.
bool socket_t::impl::rdma_wait_event() {
rdma_conn * c = rdma.get();
// POLLHUP and POLLERR are always reported, the TCP socket carries no data after the RDMA upgrade
struct pollfd pfds[2] = {
{ c->ch->fd, POLLIN, 0 },
{ fd, POLLRDHUP, 0 },
};
if (poll(pfds, 2, -1) < 0) {
return errno == EINTR;
}
if (pfds[1].revents & (POLLHUP | POLLERR | POLLRDHUP)) {
return false;
}
if (pfds[0].revents & POLLIN) {
struct ibv_cq * ev_cq = nullptr;
void * ev_ctx = nullptr;
if (ibv_get_cq_event(c->ch, &ev_cq, &ev_ctx) != 0) {
return false;
}
ibv_ack_cq_events(ev_cq, 1);
}
return true;
}
bool socket_t::impl::rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc) {
for (uint64_t s = 0; ; s++) {
rdma_conn * c = rdma.get();
bool armed = false;
for (uint64_t s = 1; ; s++) {
int n = ibv_poll_cq(cq, 1, wc);
if (n > 0) {
c->last_active = std::chrono::steady_clock::now();
if (wc->status != IBV_WC_SUCCESS) {
GGML_LOG_ERROR("RDMA CQ wc error: status=%d (%s) vendor_err=0x%x\n",
wc->status, ibv_wc_status_str(wc->status), wc->vendor_err);
@@ -410,7 +451,24 @@ bool socket_t::impl::rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc) {
return wc->status == IBV_WC_SUCCESS;
}
if (n < 0) return false;
if ((s & 0xFFFFF) == 0 && s > 0) {
if (armed) {
// armed and still empty: sleep until the next completion
if (!rdma_wait_event()) {
return false;
}
armed = false;
continue;
}
// spin while the connection is busy, arm the CQ once it has been idle for RDMA_SPIN_TIME
// a completion that arrives before arming raises no event, so poll once more after arming
if (c->ch && (s & 0x3FF) == 0 && std::chrono::steady_clock::now() - c->last_active > RDMA_SPIN_TIME) {
if (ibv_req_notify_cq(cq, 0) != 0) {
return false;
}
armed = true;
continue;
}
if ((s & 0xFFFFF) == 0) {
if (tcp_peer_closed()) {
return false;
}
@@ -444,6 +502,7 @@ bool socket_t::impl::rdma_send(const void * data, size_t size) {
}
if (ibv_post_send(c->qp, &wr, &bad) != 0) return false;
c->last_active = std::chrono::steady_clock::now();
struct ibv_wc wc;
if (!rdma_poll(c->scq, &wc)) return false;