Skip to content

Commit e3a994e

Browse files
author
Stratum
committed
feat(ggml-rpc): bounded RDMA RX reassembly and RC link-param negotiation
- rx_reassembler caps frame copies to the requested size (RX_REASSEMBLER_CHUNK == RDMA_CHUNK) - carry local path MTU in the free top byte of the 24-bit PSN field and negotiate the common MTU - hop limit / traffic class via GGML_RDMA_HL and GGML_RDMA_TC - bump RPC_PROTO_PATCH_VERSION to 1 (GGML_OP_COUNT 101 -> 103) - unit tests: test-rpc-rx-reassembly, test-rpc-link-params
1 parent 144abd3 commit e3a994e

7 files changed

Lines changed: 333 additions & 19 deletions

File tree

‎ggml/include/ggml-rpc.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,10 @@ extern "C" {
88

99
#define RPC_PROTO_MAJOR_VERSION 8
1010
#define RPC_PROTO_MINOR_VERSION 0
11-
#define RPC_PROTO_PATCH_VERSION 0
11+
#define RPC_PROTO_PATCH_VERSION 1
1212

1313
#ifdef __cplusplus
14-
static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
14+
static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
1515
#endif
1616

1717
#define GGML_RPC_MAX_SERVERS 16
Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
#pragma once
2+
3+
// RC link parameter helpers for the RDMA transport: carry the local path MTU
4+
// in the free top byte of the 24-bit PSN field, and negotiate the common MTU.
5+
// Pure logic, no verbs types, so it is unit-testable with no NIC.
6+
7+
#include <cstdint>
8+
9+
// PSN is 24 bits (qp_num & 0xffffff), so the top byte is free. ibv_mtu is 1..5,
10+
// which fits in 3 bits. 0 in those bits means "peer advertised no MTU" (old peer).
11+
static constexpr uint32_t RDMA_PSN_MASK = 0x00ffffffu;
12+
static constexpr int RDMA_PSN_MTU_SHIFT = 24;
13+
static constexpr uint32_t RDMA_PSN_MTU_BITS = 0x7u;
14+
15+
// Pack a 24-bit PSN and an MTU enum (1..5) into one wire word.
16+
inline uint32_t rdma_psn_pack(uint32_t psn24, int mtu) {
17+
return (psn24 & RDMA_PSN_MASK) | ((uint32_t)(mtu & RDMA_PSN_MTU_BITS) << RDMA_PSN_MTU_SHIFT);
18+
}
19+
20+
// Recover the real 24-bit PSN for use as rq_psn.
21+
inline uint32_t rdma_psn_base(uint32_t psn) {
22+
return psn & RDMA_PSN_MASK;
23+
}
24+
25+
// Recover the peer MTU enum (0 if the peer advertised none).
26+
inline int rdma_psn_mtu(uint32_t psn) {
27+
return (int)((psn >> RDMA_PSN_MTU_SHIFT) & RDMA_PSN_MTU_BITS);
28+
}
29+
30+
// Common MTU: the smaller of the two, but a remote 0 means "unknown" so keep local.
31+
inline int rdma_mtu_min(int local_mtu, int remote_mtu) {
32+
if (remote_mtu == 0) {
33+
return local_mtu;
34+
}
35+
return remote_mtu < local_mtu ? remote_mtu : local_mtu;
36+
}
Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
#pragma once
2+
3+
// Turns completed fixed-size RECV frames into exact-size reads for the RDMA
4+
// transport. A frame longer than the pending read is kept; its tail serves the
5+
// next read. This bounds the copy that rdma_recv used to do unchecked.
6+
// Pure logic, no verbs types, so it is unit-testable with a plain buffer.
7+
8+
#include <cstddef>
9+
#include <cstdint>
10+
#include <cstring>
11+
12+
// A posted RECV buffer is RDMA_CHUNK bytes, so a frame can never be longer.
13+
// Kept as its own constant to decouple the header from transport.cpp internals.
14+
#ifndef RX_REASSEMBLER_CHUNK
15+
#define RX_REASSEMBLER_CHUNK (256 * 1024)
16+
#endif
17+
18+
struct rx_reassembler {
19+
bool has_pending() const { return slot >= 0; }
20+
21+
// Start a new frame. Returns false on a bad length (peer protocol error):
22+
// empty frame, or longer than a posted RECV buffer.
23+
bool set_frame(int frame_slot, size_t byte_len) {
24+
if (byte_len == 0 || byte_len > cap) {
25+
return false;
26+
}
27+
slot = frame_slot;
28+
off = 0;
29+
len = byte_len;
30+
return true;
31+
}
32+
33+
// Copy up to want bytes of the pending frame into dst. src is the pending
34+
// frame's RECV buffer. Returns the number of bytes copied (never > want,
35+
// never past the frame).
36+
size_t drain(const uint8_t * src, uint8_t * dst, size_t want) {
37+
size_t avail = len - off;
38+
size_t n = want < avail ? want : avail;
39+
memcpy(dst, src + off, n);
40+
off += n;
41+
return n;
42+
}
43+
44+
bool frame_done() const { return off == len; }
45+
46+
// Hand back the fully drained slot so the caller can repost it, clear state.
47+
int take_slot() {
48+
int s = slot;
49+
slot = -1;
50+
off = 0;
51+
len = 0;
52+
return s;
53+
}
54+
55+
int pending_slot() const { return slot; }
56+
57+
private:
58+
int slot = -1; // -1 = no pending frame
59+
size_t off = 0; // bytes of the frame already served
60+
size_t len = 0; // frame length
61+
size_t cap = RX_REASSEMBLER_CHUNK; // max valid frame length
62+
};

‎ggml/src/ggml-rpc/transport.cpp‎

Lines changed: 36 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -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)
5757
static 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");
5862
static 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
6064
static 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.
352357
bool 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 {

‎tests/test-rpc-link-params.cpp‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
// Unit test for rdma_link_params.h: PSN/MTU packing and MTU negotiation.
2+
// Pure logic, no NIC. Guards item 02.
3+
4+
#include "ggml-rpc/rdma_link_params.h"
5+
6+
#include <cstdint>
7+
#include <cstdio>
8+
9+
static int g_fail = 0;
10+
#define CHECK(cond) do { \
11+
if (!(cond)) { printf("FAIL %s:%d: %s\n", __FILE__, __LINE__, #cond); g_fail++; } \
12+
} while (0)
13+
14+
int main() {
15+
// 1. pack/unpack round-trip for every valid ibv_mtu enum (1..5).
16+
for (int mtu = 1; mtu <= 5; mtu++) {
17+
uint32_t psn = 0x00abcdef;
18+
uint32_t packed = rdma_psn_pack(psn, mtu);
19+
CHECK(rdma_psn_mtu(packed) == mtu);
20+
CHECK(rdma_psn_base(packed) == psn);
21+
}
22+
23+
// 2. old peer advertises top byte 0 -> no MTU.
24+
CHECK(rdma_psn_mtu(0x00abcdef) == 0);
25+
26+
// 3. MTU negotiation: remote 0 -> local, else the smaller.
27+
CHECK(rdma_mtu_min(3, 0) == 3);
28+
CHECK(rdma_mtu_min(3, 4) == 3);
29+
CHECK(rdma_mtu_min(4, 3) == 3);
30+
CHECK(rdma_mtu_min(5, 5) == 5);
31+
32+
// 4. base preserves the low 24 bits regardless of the carried byte.
33+
CHECK(rdma_psn_base(rdma_psn_pack(0x00ffffff, 5)) == 0x00ffffff);
34+
CHECK(rdma_psn_base(rdma_psn_pack(0x00000000, 1)) == 0x00000000);
35+
36+
// 5. a 24-bit base is never changed by packing then stripping.
37+
CHECK(rdma_psn_base(rdma_psn_pack(0x00123456, 2)) == 0x00123456);
38+
39+
if (g_fail == 0) {
40+
printf("test-rpc-link-params: all checks passed\n");
41+
return 0;
42+
}
43+
printf("test-rpc-link-params: %d checks failed\n", g_fail);
44+
return 1;
45+
}

0 commit comments

Comments
 (0)