char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Ryan Crpc: fix apple rdma error spew on teardown (#27908)742347b2e

char/aliceai-t5-moe
18.3 KiB481 linesraw
1#include "transport-apple.h"
2#include "transport.h"
3#include "ggml-impl.h"
4
5#include <infiniband/verbs.h>
6
7#include <cerrno>
8#include <cstdlib>
9#include <cstring>
10#include <string>
11#include <dlfcn.h>
12#include <poll.h>
13#include <sys/socket.h>
14#include <unistd.h>
15
16// Apple RDMA-over-Thunderbolt (see Apple TN3205).
17//
18// Apple's RDMA is quite different from what's supported in Linux - deserving of its own transport implementation.
19// see https://developer.apple.com/documentation/technotes/tn3205-low-latency-communication-with-rdma-over-thunderbolt for details
20// at a high level the main differences are:
21// UC(unreliable connection) on Apple vs RC(reliable connection) QP transport types on Linux (though in practice UC on Apple is still lossless)
22// fixed 128KiB stride on Apple vs variable chunk size on Linux
23// relying on Apple's hardware credit based flow control vs RNR NAKs + retries on Linux
24//
25// on Apple a SEND and its corresponding RECV must cover the same number of 4 KiB Thunderbolt frames,
26// so every SEND posts a whole 128KiB stride over the wire, even when partially filled.
27// (In testing 128KiB was the best performing among 32, 64, 128, 256)
28
29static constexpr uint32_t RDMA_SEG_MAGIC   = 0x52534547u; // "RSEG"
30static constexpr int      RDMA_NBUF        = 16;          // ring depth (frames per direction)
31static constexpr size_t   RDMA_FRAME       = 4096;        // Thunderbolt frame (fixed on Apple)
32static constexpr size_t   RDMA_STRIDE      = 128 * 1024;  // 32 Thunderbolt frames; NBUF x this = 2 MiB pinned per direction
33static constexpr uint32_t RDMA_PSN         = 0;           // any value works if both sides match: UC has no retransmit
34static constexpr size_t   RDMA_GID_SIZE    = 16;
35
36static_assert(RDMA_STRIDE % RDMA_FRAME == 0, "RDMA_STRIDE must be a whole number of frames");
37// TN3205 counts queue depth in Thunderbolt frames, not work requests.
38static constexpr uint32_t RDMA_QP_WR       = (uint32_t)RDMA_NBUF * (RDMA_STRIDE / RDMA_FRAME);
39static constexpr uint64_t RDMA_RECV_WR     = 1ull << 20;  // wr_id bit tagging recv completions
40static constexpr uint64_t RDMA_WR_IDX_MASK = 0xffff;      // buffer index in the low bits of wr_id
41static constexpr uint8_t  RDMA_SYNC_READY  = 0x2A;        // readiness-handshake byte (peer activated)
42
43struct rdma_seg_hdr {
44    uint32_t magic; // RDMA_SEG_MAGIC; a mismatch means the stream desynced
45    uint32_t len;   // payload bytes in this frame; the rest of the stride is padding
46};
47static constexpr size_t RDMA_PAYLOAD = RDMA_STRIDE - sizeof(rdma_seg_hdr);
48
49struct apple_rdma_caps {
50    uint32_t qpn;
51    uint16_t lid;
52    uint16_t reserved;
53    uint8_t  gid[RDMA_GID_SIZE];
54};
55
56static_assert(sizeof(apple_rdma_caps) == RPC_CONN_CAPS_SIZE, "apple_rdma_caps must match conn_caps size");
57
58struct apple_rdma::impl {
59    int fd = -1;                      // bootstrap TCP socket, kept as the liveness anchor
60
61    struct ibv_context * ctx = nullptr;
62    struct ibv_pd * pd = nullptr;
63    struct ibv_cq * cq = nullptr;     // one CQ for both directions; RDMA_RECV_WR tags recv completions
64    struct ibv_qp * qp = nullptr;
65
66    uint8_t       * send_mem = nullptr;
67    struct ibv_mr * send_mr  = nullptr;
68    uint8_t       * recv_mem = nullptr;
69    struct ibv_mr * recv_mr  = nullptr;
70
71    int      send_busy[RDMA_NBUF] = {};  // 1 while this buffer has a send in flight
72    // completed recv frames, oldest first: ring index, bytes already handed to
73    // the reader, and total payload length
74    struct { int buf; uint32_t off; uint32_t len; } inq[RDMA_NBUF] = {};
75    int      inq_head = 0;
76    int      inq_count = 0;
77    int      pend_buf = -1;
78    uint32_t pend_len = 0;
79    bool     broken = false;
80
81    uint32_t     qpn = 0;
82    uint8_t      port = 0;
83    int          gid_idx = 0;
84    enum ibv_mtu path_mtu = IBV_MTU_1024;
85
86    int  progress();
87    bool acquire_pending();
88    bool post_pending();
89
90    bool post_recv(int i) {
91        struct ibv_sge sge = {};
92        sge.addr   = (uintptr_t)(recv_mem + (size_t)i * RDMA_STRIDE);
93        sge.length = (uint32_t)RDMA_STRIDE;
94        sge.lkey   = recv_mr->lkey;
95        struct ibv_recv_wr wr = {}, * bad = nullptr;
96        wr.wr_id   = RDMA_RECV_WR | (uint64_t)i;
97        wr.sg_list = &sge;
98        wr.num_sge = 1;
99        return ibv_post_recv(qp, &wr, &bad) == 0;
100    }
101
102    bool post_send(int i, size_t len) {
103        struct ibv_sge sge = {};
104        sge.addr   = (uintptr_t)(send_mem + (size_t)i * RDMA_STRIDE);
105        sge.length = (uint32_t)len;
106        sge.lkey   = send_mr->lkey;
107        struct ibv_send_wr wr = {}, * bad = nullptr;
108        wr.wr_id   = (uint64_t)i;
109        wr.sg_list = &sge;
110        wr.num_sge = 1;
111        wr.opcode  = IBV_WR_SEND;
112        wr.send_flags = IBV_SEND_SIGNALED;
113        return ibv_post_send(qp, &wr, &bad) == 0;
114    }
115
116    ~impl() {
117        broken = true;
118        // destroy the QP first: it can still write to the rings until it is gone.
119        // no IBV_QPS_ERR before it - Apple's provider then fails every region unmap.
120        if (qp) ibv_destroy_qp(qp);
121        if (send_mr) ibv_dereg_mr(send_mr);
122        if (recv_mr) ibv_dereg_mr(recv_mr);
123        free(send_mem);
124        free(recv_mem);
125        if (cq)  ibv_destroy_cq(cq);
126        if (pd)  ibv_dealloc_pd(pd);
127        if (ctx) ibv_close_device(ctx);
128    }
129};
130
131apple_rdma::apple_rdma(std::unique_ptr<impl> p) : pimpl(std::move(p)) {}
132
133apple_rdma::~apple_rdma() = default;
134
135bool apple_rdma::broken() const {
136    return pimpl->broken;
137}
138
139// The readiness handshake below still runs over the bootstrap socket, one byte
140// each way, before the transport is declared live.
141static bool tcp_send_byte(int fd, uint8_t b) {
142    ssize_t n;
143    do { n = ::send(fd, &b, sizeof(b), 0); } while (n < 0 && errno == EINTR);
144    return n == sizeof(b);
145}
146
147static bool tcp_recv_byte(int fd, uint8_t * b) {
148    ssize_t n;
149    do { n = ::recv(fd, b, sizeof(*b), 0); } while (n < 0 && errno == EINTR);
150    return n == (ssize_t)sizeof(*b);
151}
152
153// Index of the GID on this port equal to the target, or -1. Thunderbolt GIDs are
154// RoCEv2 IPv4-mapped (::ffff:a.b.c.d), so this matches the local TCP address.
155static int rdma_match_gid(struct ibv_context * ctx, uint8_t port, int gid_tbl_len,
156                          const uint8_t * target, union ibv_gid * out) {
157    for (int i = 0; i < gid_tbl_len; i++) {
158        union ibv_gid g;
159        if (ibv_query_gid(ctx, port, i, &g) != 0) continue;
160        if (memcmp(g.raw, target, RDMA_GID_SIZE) != 0) continue;
161        if (out) *out = g;
162        return i;
163    }
164    return -1;
165}
166
167// First ACTIVE port on the device. Only a cabled, up Thunderbolt link reports
168// ACTIVE, and it is not always port 1, so the port cannot be hardcoded the way
169// the Linux path does. Returns 0 if none.
170static uint8_t rdma_first_active_port(struct ibv_context * ctx, struct ibv_port_attr * out) {
171    struct ibv_device_attr da;
172    if (ibv_query_device(ctx, &da) != 0) return 0;
173    for (uint8_t p = 1; p <= da.phys_port_cnt; p++) {
174        struct ibv_port_attr pa;
175        if (ibv_query_port(ctx, p, &pa) != 0) continue;
176        if (pa.state == IBV_PORT_ACTIVE) { if (out) *out = pa; return p; }
177    }
178    return 0;
179}
180
181// librdma.dylib is weak-linked, so its symbols are null when it is absent. Nothing may
182// call one before this has returned true.
183static bool rdma_library_present() {
184    static const bool present = [] {
185        void * handle = dlopen("/usr/lib/librdma.dylib", RTLD_LAZY);
186        if (handle == nullptr) {
187            return false;
188        }
189        dlclose(handle);
190        return true;
191    }();
192    return present;
193}
194
195// Called before the endpoints are exchanged: pick the local device facing this
196// peer, create a UC QP and register the frame rings. RDMA is point-to-point, so
197// the device is the one whose GID equals the bootstrap connection's local
198// address, i.e. the one cabled to the peer.
199std::unique_ptr<apple_rdma> apple_rdma::probe(int fd, const uint8_t * target_gid, uint8_t * caps) {
200    if (!rdma_library_present()) {
201        return nullptr;
202    }
203    int ndev = 0;
204    ibv_device ** devs = ibv_get_device_list(&ndev);
205    if (!devs) return nullptr;
206
207    ibv_context * ctx = nullptr;
208    uint8_t port = 0;
209    struct ibv_port_attr pa = {};
210    union ibv_gid gid = {};
211    int gid_idx = -1;
212    std::string matched;
213    for (int d = 0; d < ndev; d++) {
214        ibv_context * c = ibv_open_device(devs[d]);
215        if (!c) continue;
216        struct ibv_port_attr p = {};
217        uint8_t pt = rdma_first_active_port(c, &p);
218        int gi = pt ? rdma_match_gid(c, pt, p.gid_tbl_len, target_gid, &gid) : -1;
219        if (gi < 0) { ibv_close_device(c); continue; }
220        ctx = c; port = pt; pa = p; gid_idx = gi;
221        const char * name = ibv_get_device_name(devs[d]);
222        matched = name ? name : "";
223        break;
224    }
225    ibv_free_device_list(devs);
226    if (!ctx) return nullptr;
227
228    std::unique_ptr<impl> c(new impl());
229    c->fd  = fd;
230    c->ctx = ctx;
231    c->port = port;
232    c->gid_idx = gid_idx;
233    c->path_mtu = pa.active_mtu;
234
235    c->pd = ibv_alloc_pd(ctx);
236    if (!c->pd) return nullptr;
237
238    c->cq = ibv_create_cq(ctx, 2 * RDMA_QP_WR + 1, nullptr, nullptr, 0);
239    if (!c->cq) return nullptr;
240
241    ibv_qp_init_attr qia = {};
242    qia.send_cq = c->cq;
243    qia.recv_cq = c->cq;
244    qia.qp_type = IBV_QPT_UC;
245    qia.cap.max_send_wr  = RDMA_QP_WR;
246    qia.cap.max_recv_wr  = RDMA_QP_WR;
247    qia.cap.max_send_sge = 1;
248    qia.cap.max_recv_sge = 1;
249    c->qp = ibv_create_qp(c->pd, &qia);
250    if (!c->qp) return nullptr;
251
252    {
253        ibv_qp_attr a = {};
254        a.qp_state = IBV_QPS_INIT;
255        a.pkey_index = 0;
256        a.port_num = port;
257        a.qp_access_flags = IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_WRITE;
258        if (ibv_modify_qp(c->qp, &a,
259                IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS) != 0) {
260            return nullptr;
261        }
262    }
263
264    long page = sysconf(_SC_PAGESIZE);
265    if (page <= 0) page = 4096;
266    const size_t ring_bytes = (size_t)RDMA_NBUF * RDMA_STRIDE;
267    if (posix_memalign((void **)&c->send_mem, (size_t)page, ring_bytes) != 0) c->send_mem = nullptr;
268    if (posix_memalign((void **)&c->recv_mem, (size_t)page, ring_bytes) != 0) c->recv_mem = nullptr;
269    if (!c->send_mem || !c->recv_mem) return nullptr;
270
271    // Apple's provider rejects LOCAL_WRITE-only MRs even for two-sided SEND/RECV.
272    const int mr_flags = IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_WRITE;
273    c->send_mr = ibv_reg_mr(c->pd, c->send_mem, ring_bytes, mr_flags);
274    c->recv_mr = ibv_reg_mr(c->pd, c->recv_mem, ring_bytes, mr_flags);
275    if (!c->send_mr || !c->recv_mr) return nullptr;
276
277    // Recvs are posted in activate() after the RTS transition, not here: Apple's
278    // provider rejects ibv_post_recv on a QP that has not reached RTS.
279
280    c->qpn = c->qp->qp_num;
281
282    apple_rdma_caps rc = {};
283    rc.qpn = c->qpn;
284    rc.lid = pa.lid;
285    memcpy(rc.gid, gid.raw, RDMA_GID_SIZE);
286    memcpy(caps, &rc, sizeof(rc));
287
288    GGML_LOG_INFO("RDMA(Apple/UC) probed: dev=%s port=%u gid=%d qpn=%u lid=%u mtu=%d ring=%d x %zu KiB\n",
289                  matched.c_str(), port, gid_idx, c->qpn, (unsigned)pa.lid, 128 << c->path_mtu,
290                  RDMA_NBUF, RDMA_STRIDE / 1024);
291    return std::unique_ptr<apple_rdma>(new apple_rdma(std::move(c)));
292}
293
294// Called once the peer's endpoint has arrived: INIT -> RTR -> RTS (UC: GID/GRH
295// addressing, no timeout/retry/rnr/rd_atomic), then the readiness handshake.
296bool apple_rdma::activate(const uint8_t * caps) {
297    impl * c = pimpl.get();
298
299    apple_rdma_caps rc = {};
300    memcpy(&rc, caps, sizeof(rc));
301
302    bool ok = true;
303    {
304        ibv_qp_attr a = {};
305        a.qp_state   = IBV_QPS_RTR;
306        a.path_mtu   = c->path_mtu;
307        a.rq_psn     = RDMA_PSN;
308        a.dest_qp_num = rc.qpn;
309        a.ah_attr.is_global     = 1;
310        a.ah_attr.port_num      = c->port;
311        a.ah_attr.sl            = 0;
312        a.ah_attr.src_path_bits = 0;
313        a.ah_attr.dlid          = rc.lid;
314        a.ah_attr.grh.hop_limit  = 1;
315        a.ah_attr.grh.sgid_index = (uint8_t)c->gid_idx;
316        memcpy(&a.ah_attr.grh.dgid, rc.gid, RDMA_GID_SIZE);
317        if (ibv_modify_qp(c->qp, &a,
318                IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN | IBV_QP_RQ_PSN) != 0) {
319            GGML_LOG_ERROR("RDMA(Apple/UC) RTR failed: %s\n", strerror(errno));
320            ok = false;
321        }
322    }
323    if (ok) {
324        ibv_qp_attr a = {};
325        a.qp_state = IBV_QPS_RTS;
326        a.sq_psn   = RDMA_PSN;
327        if (ibv_modify_qp(c->qp, &a, IBV_QP_STATE | IBV_QP_SQ_PSN) != 0) {
328            GGML_LOG_ERROR("RDMA(Apple/UC) RTS failed: %s\n", strerror(errno));
329            ok = false;
330        }
331    }
332
333    // Recvs are posted only now: the controller starts processing them at RTR.
334    for (int i = 0; ok && i < RDMA_NBUF; i++) {
335        if (!c->post_recv(i)) {
336            GGML_LOG_ERROR("RDMA(Apple/UC) post_recv %d/%d failed\n", i, RDMA_NBUF);
337            ok = false;
338        }
339    }
340
341    // A queue pair processes receives only after RTR and the transitions above can
342    // fail on one side alone, so neither peer sends a frame until both report their
343    // recvs posted.
344    uint8_t peer_ready = 0;
345    if (!tcp_send_byte(c->fd, ok ? RDMA_SYNC_READY : 0) || !tcp_recv_byte(c->fd, &peer_ready)) {
346        return false;
347    }
348    if (!ok || peer_ready != RDMA_SYNC_READY) {
349        return false;
350    }
351
352    GGML_LOG_INFO("RDMA(Apple/UC) activated: qpn=%u->%u mtu=%d rx_depth=%d\n",
353                  c->qpn, rc.qpn, 128 << c->path_mtu, RDMA_NBUF);
354    return true;
355}
356
357// Drain the CQ: release completed send buffers, queue completed recv frames for
358// the reader. Returns the number of completions reaped, or -1 on error.
359int apple_rdma::impl::progress() {
360    struct ibv_wc wc[RDMA_NBUF * 2];
361    int n = ibv_poll_cq(cq, RDMA_NBUF * 2, wc);
362    if (n < 0) { GGML_LOG_ERROR("RDMA(Apple/UC) poll_cq failed\n"); broken = true; return -1; }
363    for (int j = 0; j < n; j++) {
364        uint64_t id = wc[j].wr_id;
365        bool is_recv = (id & RDMA_RECV_WR) != 0;
366        if (wc[j].status != IBV_WC_SUCCESS) {
367            GGML_LOG_ERROR("RDMA(Apple/UC) %s wc error: status=%d\n", is_recv ? "recv" : "send", wc[j].status);
368            broken = true;
369            return -1;
370        }
371        if (is_recv) {
372            int b = (int)(id & RDMA_WR_IDX_MASK);
373            const rdma_seg_hdr * h = (const rdma_seg_hdr *)(recv_mem + (size_t)b * RDMA_STRIDE);
374            if (h->magic != RDMA_SEG_MAGIC) { GGML_LOG_ERROR("RDMA(Apple/UC) bad frame magic\n"); broken = true; return -1; }
375            if (h->len > RDMA_PAYLOAD) { GGML_LOG_ERROR("RDMA(Apple/UC) frame len %u exceeds payload\n", h->len); broken = true; return -1; }
376            int slot = (inq_head + inq_count) % RDMA_NBUF;
377            inq[slot].buf  = b;
378            inq[slot].off  = 0;
379            inq[slot].len  = h->len;
380            inq_count++;
381        } else {
382            send_busy[(int)(id & RDMA_WR_IDX_MASK)] = 0;
383        }
384    }
385    return n;
386}
387
388// Reserve a free send buffer to coalesce into, waiting on progress if none free.
389bool apple_rdma::impl::acquire_pending() {
390    if (pend_buf >= 0) return true;
391    for (;;) {
392        if (broken) return false;
393        for (int k = 0; k < RDMA_NBUF; k++) if (!send_busy[k]) { pend_buf = k; pend_len = 0; return true; }
394        if (progress() < 0) return false;
395    }
396}
397
398// Post the pending frame. The whole STRIDE goes out even when only partly filled:
399// TN3205 requires a SEND and its matching RECV to cover the same number of
400// Thunderbolt frames, so a short send would fail the peer's receive.
401bool apple_rdma::impl::post_pending() {
402    if (pend_buf < 0) return true;
403    int i = pend_buf;
404    rdma_seg_hdr * h = (rdma_seg_hdr *)(send_mem + (size_t)i * RDMA_STRIDE);
405    h->magic = RDMA_SEG_MAGIC;
406    h->len   = pend_len;
407    if (!post_send(i, RDMA_STRIDE)) { broken = true; return false; }
408    send_busy[i] = 1;
409    pend_buf = -1;
410    pend_len = 0;
411    return true;
412}
413
414// Coalescing write: append into the pending frame, posting a full frame when it
415// fills. The trailing partial is posted by flush() at each message boundary.
416bool apple_rdma::send(const void * data, size_t size) {
417    impl * c = pimpl.get();
418    const uint8_t * p = (const uint8_t *)data;
419    while (size > 0) {
420        if (c->broken) return false;
421        if (!c->acquire_pending()) return false;
422        uint8_t * sb = c->send_mem + (size_t)c->pend_buf * RDMA_STRIDE;
423        size_t space = RDMA_PAYLOAD - c->pend_len;
424        size_t chunk = size < space ? size : space;
425        memcpy(sb + sizeof(rdma_seg_hdr) + c->pend_len, p, chunk);
426        c->pend_len += (uint32_t)chunk;
427        p += chunk;
428        size -= chunk;
429        if (c->pend_len == RDMA_PAYLOAD) { if (!c->post_pending()) return false; }
430    }
431    return true;
432}
433
434bool apple_rdma::recv(void * data, size_t size) {
435    impl * c = pimpl.get();
436    uint8_t * p = (uint8_t *)data;
437    if (!c->post_pending()) return false;   // turnaround: flush the coalesced request
438    unsigned idle = 0;
439    while (size > 0) {
440        if (c->inq_count == 0) {
441            if (c->broken) return false;
442            int n = c->progress();
443            if (n < 0) return false;
444            if (n == 0) {
445                // UC gives no disconnect notification, so the bootstrap TCP fd is
446                // the liveness anchor: nothing crosses it once RDMA is up, so any
447                // readability means the peer's FIN (macOS has no POLLRDHUP).
448                // Same idle interval as the Linux path.
449                if ((++idle & 0xFFFFF) == 0) {
450                    struct pollfd pfd = { c->fd, POLLIN, 0 };
451                    if (poll(&pfd, 1, 0) > 0 &&
452                        (pfd.revents & (POLLIN | POLLHUP | POLLERR | POLLNVAL))) {
453                        return false;
454                    }
455                }
456            } else {
457                idle = 0;
458            }
459            continue;
460        }
461        idle = 0;
462        int slot = c->inq_head;
463        int b = c->inq[slot].buf;
464        uint32_t avail = c->inq[slot].len - c->inq[slot].off;
465        uint32_t take = (size < (size_t)avail) ? (uint32_t)size : avail;
466        memcpy(p, c->recv_mem + (size_t)b * RDMA_STRIDE + sizeof(rdma_seg_hdr) + c->inq[slot].off, take);
467        p += take;
468        size -= take;
469        c->inq[slot].off += take;
470        if (c->inq[slot].off == c->inq[slot].len) {
471            if (!c->post_recv(b)) { c->broken = true; return false; }
472            c->inq_head = (c->inq_head + 1) % RDMA_NBUF;
473            c->inq_count--;
474        }
475    }
476    return true;
477}
478
479bool apple_rdma::flush() {
480    return pimpl->post_pending();
481}