@@ -55,6 +55,10 @@ using rdma_gid_t = std::array<uint8_t, RDMA_GID_SIZE>;
5555
5656#if defined(GGML_RPC_RDMA) && !defined(GGML_RPC_RDMA_APPLE)
5757static constexpr size_t RDMA_CHUNK = 256 * 1024 ; // 256 KiB per send/recv (fits default 8 MiB memlock)
58+ #define RX_REASSEMBLER_CHUNK (256 * 1024 )
59+ #include " rdma_reassembler.h"
60+ #include " rdma_link_params.h"
61+ static_assert (RX_REASSEMBLER_CHUNK == RDMA_CHUNK , " reassembler cap must match RDMA_CHUNK" );
5862static constexpr int RDMA_RX_DEPTH = 24 ; // pre-posted recv ring: 24 × 256 KiB = 6 MiB
5963// keep polling the CQ for this long after the last activity, then sleep until the next completion
6064static constexpr auto RDMA_SPIN_TIME = std::chrono::milliseconds(100 );
@@ -75,6 +79,7 @@ struct rdma_conn {
7579 void * rx_buf = nullptr ; // RDMA_RX_DEPTH × RDMA_CHUNK contiguous
7680 struct ibv_mr * rx_mr = nullptr ;
7781 int rx_head = 0 ;
82+ rx_reassembler reasm; // bounds frame copies to the requested size
7883
7984 uint32_t max_inline = 0 ;
8085
@@ -350,6 +355,15 @@ bool socket_t::impl::rdma_probe() {
350355// Phase 2: Given remote QPN/PSN/GID, transition QP: RESET->INIT->pre-post->RTR->RTS.
351356// On success, the connection is live and ready for rdma_send/rdma_recv.
352357bool socket_t::impl::rdma_activate (uint32_t remote_qpn, uint32_t remote_psn, const uint8_t * remote_gid) {
358+ // the remote MTU rides in the top byte of the PSN field; 0 = old peer
359+ enum ibv_mtu mtu = (enum ibv_mtu)rdma_mtu_min (rdma_local.path_mtu , rdma_psn_mtu (remote_psn));
360+ uint32_t rq_psn = rdma_psn_base (remote_psn);
361+
362+ const char * hl_env = std::getenv (" GGML_RDMA_HL" );
363+ const char * tc_env = std::getenv (" GGML_RDMA_TC" );
364+ uint8_t hop_limit = hl_env ? (uint8_t )atoi (hl_env) : 64 ;
365+ uint8_t traffic_class = tc_env ? (uint8_t )atoi (tc_env) : 0 ;
366+
353367 // RESET -> INIT
354368 {
355369 struct ibv_qp_attr a = {};
@@ -371,14 +385,15 @@ bool socket_t::impl::rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, con
371385 {
372386 struct ibv_qp_attr a = {};
373387 a.qp_state = IBV_QPS_RTR ;
374- a.path_mtu = rdma_local. path_mtu ;
388+ a.path_mtu = mtu ;
375389 a.dest_qp_num = remote_qpn;
376- a.rq_psn = remote_psn ;
390+ a.rq_psn = rq_psn ;
377391 a.max_dest_rd_atomic = 1 ;
378392 a.min_rnr_timer = 1 ;
379393 a.ah_attr .is_global = 1 ;
380394 memcpy (&a.ah_attr .grh .dgid , remote_gid, RDMA_GID_SIZE );
381- a.ah_attr .grh .hop_limit = 1 ;
395+ a.ah_attr .grh .hop_limit = hop_limit;
396+ a.ah_attr .grh .traffic_class = traffic_class;
382397 a.ah_attr .grh .sgid_index = rdma_local.gid_idx ;
383398 a.ah_attr .dlid = 0 ;
384399 a.ah_attr .port_num = rdma_local.ib_port ;
@@ -407,8 +422,8 @@ bool socket_t::impl::rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, con
407422
408423 rdma->last_active = std::chrono::steady_clock::now ();
409424
410- GGML_LOG_INFO (" RDMA activated: qpn=%u->%u mtu=%d rx_depth=%d\n " ,
411- rdma_local.qpn , remote_qpn, 128 << rdma_local. path_mtu , RDMA_RX_DEPTH );
425+ GGML_LOG_INFO (" RDMA activated: qpn=%u->%u mtu=%d hop=%u tc=%u rx_depth=%d\n " ,
426+ rdma_local.qpn , remote_qpn, 128 << mtu, hop_limit, traffic_class , RDMA_RX_DEPTH );
412427 return true ;
413428}
414429
@@ -517,17 +532,21 @@ bool socket_t::impl::rdma_recv(void * data, size_t size) {
517532 uint8_t * dst = (uint8_t *)data;
518533 size_t rem = size;
519534 while (rem > 0 ) {
520- struct ibv_wc wc;
521- if (!rdma_poll (c->rcq , &wc)) return false ;
522-
523- int slot = (int )wc.wr_id ;
524- size_t got = wc.byte_len ;
525- memcpy (dst, c->rx_slot (slot), got);
526-
527- if (!c->post_rx (slot)) return false ;
528-
529- dst += got;
530- rem -= got;
535+ if (!c->reasm .has_pending ()) {
536+ struct ibv_wc wc;
537+ if (!rdma_poll (c->rcq , &wc)) return false ;
538+ if (!c->reasm .set_frame ((int )wc.wr_id , wc.byte_len )) {
539+ GGML_LOG_ERROR (" RDMA recv: bad frame len %u\n " , wc.byte_len );
540+ return false ;
541+ }
542+ }
543+ int slot = c->reasm .pending_slot ();
544+ size_t n = c->reasm .drain (c->rx_slot (slot), dst, rem);
545+ dst += n;
546+ rem -= n;
547+ if (c->reasm .frame_done () && !c->post_rx (c->reasm .take_slot ())) {
548+ return false ;
549+ }
531550 }
532551 return true ;
533552}
@@ -604,7 +623,7 @@ void socket_t::impl::get_caps(uint8_t * local_caps) {
604623 if (rdma_probe ()) {
605624 rdma_caps rc = {};
606625 rc.qpn = rdma_local.qpn ;
607- rc.psn = rdma_local.psn ;
626+ rc.psn = rdma_psn_pack ( rdma_local.psn , rdma_local. path_mtu ) ;
608627 memcpy (rc.gid , rdma_local.gid , RDMA_GID_SIZE );
609628 memcpy (local_caps, &rc, sizeof (rc));
610629 } else {
0 commit comments