char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
742347b2e
char/aliceai-t5-moe
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 { 44uint32_t magic ;// RDMA_SEG_MAGIC; a mismatch means the stream desynced 45uint32_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 { 50uint32_t qpn ; 51uint16_t lid ; 52uint16_t reserved ; 53uint8_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 { 59int fd = -1 ;// bootstrap TCP socket, kept as the liveness anchor 60 61struct ibv_context * ctx = nullptr ; 62struct ibv_pd * pd = nullptr ; 63struct ibv_cq * cq = nullptr ;// one CQ for both directions; RDMA_RECV_WR tags recv completions 64struct ibv_qp * qp = nullptr ; 65 66uint8_t * send_mem = nullptr ; 67struct ibv_mr * send_mr = nullptr ; 68uint8_t * recv_mem = nullptr ; 69struct ibv_mr * recv_mr = nullptr ; 70 71int 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 74struct {int buf ;uint32_t off ;uint32_t len ; }inq [RDMA_NBUF ]= {}; 75int inq_head = 0 ; 76int inq_count = 0 ; 77int pend_buf = -1 ; 78uint32_t pend_len = 0 ; 79bool broken = false; 80 81uint32_t qpn = 0 ; 82uint8_t port = 0 ; 83int gid_idx = 0 ; 84enum ibv_mtu path_mtu = IBV_MTU_1024 ; 85 86int progress (); 87bool acquire_pending (); 88bool post_pending (); 89 90bool post_recv (int i ) { 91struct ibv_sge sge = {}; 92sge .addr = (uintptr_t )(recv_mem + (size_t )i * RDMA_STRIDE ); 93sge .length = (uint32_t )RDMA_STRIDE ; 94sge .lkey = recv_mr -> lkey ; 95struct ibv_recv_wr wr = {},* bad = nullptr ; 96wr .wr_id = RDMA_RECV_WR | (uint64_t )i ; 97wr .sg_list = & sge ; 98wr .num_sge = 1 ; 99return ibv_post_recv (qp ,& wr ,& bad )== 0 ; 100 } 101 102bool post_send (int i ,size_t len ) { 103struct ibv_sge sge = {}; 104sge .addr = (uintptr_t )(send_mem + (size_t )i * RDMA_STRIDE ); 105sge .length = (uint32_t )len ; 106sge .lkey = send_mr -> lkey ; 107struct ibv_send_wr wr = {},* bad = nullptr ; 108wr .wr_id = (uint64_t )i ; 109wr .sg_list = & sge ; 110wr .num_sge = 1 ; 111wr .opcode = IBV_WR_SEND ; 112wr .send_flags = IBV_SEND_SIGNALED ; 113return ibv_post_send (qp ,& wr ,& bad )== 0 ; 114 } 115 116 ~impl () { 117broken = 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. 120if (qp )ibv_destroy_qp (qp ); 121if (send_mr )ibv_dereg_mr (send_mr ); 122if (recv_mr )ibv_dereg_mr (recv_mr ); 123free (send_mem ); 124free (recv_mem ); 125if (cq )ibv_destroy_cq (cq ); 126if (pd )ibv_dealloc_pd (pd ); 127if (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 { 136return 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 ) { 142ssize_t n ; 143do {n = ::send (fd ,& b ,sizeof (b ),0 ); }while (n < 0 && errno == EINTR ); 144return n == sizeof (b ); 145} 146 147static bool tcp_recv_byte (int fd ,uint8_t * b ) { 148ssize_t n ; 149do {n = ::recv (fd ,b ,sizeof (* b ),0 ); }while (n < 0 && errno == EINTR ); 150return 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 , 156const uint8_t * target ,union ibv_gid * out ) { 157for (int i = 0 ;i < gid_tbl_len ;i ++ ) { 158union ibv_gid g ; 159if (ibv_query_gid (ctx ,port ,i ,& g )!= 0 )continue ; 160if (memcmp (g .raw ,target ,RDMA_GID_SIZE )!= 0 )continue ; 161if (out )* out = g ; 162return i ; 163 } 164return -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 ) { 171struct ibv_device_attr da ; 172if (ibv_query_device (ctx ,& da )!= 0 )return 0 ; 173for (uint8_t p = 1 ;p <=da .phys_port_cnt ;p ++ ) { 174struct ibv_port_attr pa ; 175if (ibv_query_port (ctx ,p ,& pa )!= 0 )continue ; 176if (pa .state == IBV_PORT_ACTIVE ) {if (out )* out = pa ;return p ; } 177 } 178return 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 () { 184static const bool present = [] { 185void * handle = dlopen ("/usr/lib/librdma.dylib" ,RTLD_LAZY ); 186if (handle == nullptr ) { 187return false; 188 } 189dlclose (handle ); 190return true; 191 }(); 192return 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 ) { 200if (!rdma_library_present ()) { 201return nullptr ; 202 } 203int ndev = 0 ; 204ibv_device ** devs = ibv_get_device_list (& ndev ); 205if (!devs )return nullptr ; 206 207ibv_context * ctx = nullptr ; 208uint8_t port = 0 ; 209struct ibv_port_attr pa = {}; 210union ibv_gid gid = {}; 211int gid_idx = -1 ; 212 std::string matched ; 213for (int d = 0 ;d < ndev ;d ++ ) { 214ibv_context * c = ibv_open_device (devs [d ]); 215if (!c )continue ; 216struct ibv_port_attr p = {}; 217uint8_t pt = rdma_first_active_port (c ,& p ); 218int gi = pt ?rdma_match_gid (c ,pt ,p .gid_tbl_len ,target_gid ,& gid ) :-1 ; 219if (gi < 0 ) {ibv_close_device (c );continue ; } 220ctx = c ;port = pt ;pa = p ;gid_idx = gi ; 221const char * name = ibv_get_device_name (devs [d ]); 222matched = name ?name :"" ; 223break ; 224 } 225ibv_free_device_list (devs ); 226if (!ctx )return nullptr ; 227 228 std::unique_ptr < impl > c (new impl ()); 229c -> fd = fd ; 230c -> ctx = ctx ; 231c -> port = port ; 232c -> gid_idx = gid_idx ; 233c -> path_mtu = pa .active_mtu ; 234 235c -> pd = ibv_alloc_pd (ctx ); 236if (!c -> pd )return nullptr ; 237 238c -> cq = ibv_create_cq (ctx ,2 * RDMA_QP_WR + 1 ,nullptr ,nullptr ,0 ); 239if (!c -> cq )return nullptr ; 240 241ibv_qp_init_attr qia = {}; 242qia .send_cq = c -> cq ; 243qia .recv_cq = c -> cq ; 244qia .qp_type = IBV_QPT_UC ; 245qia .cap .max_send_wr = RDMA_QP_WR ; 246qia .cap .max_recv_wr = RDMA_QP_WR ; 247qia .cap .max_send_sge = 1 ; 248qia .cap .max_recv_sge = 1 ; 249c -> qp = ibv_create_qp (c -> pd ,& qia ); 250if (!c -> qp )return nullptr ; 251 252 { 253ibv_qp_attr a = {}; 254a .qp_state = IBV_QPS_INIT ; 255a .pkey_index = 0 ; 256a .port_num = port ; 257a .qp_access_flags = IBV_ACCESS_LOCAL_WRITE |IBV_ACCESS_REMOTE_READ |IBV_ACCESS_REMOTE_WRITE ; 258if (ibv_modify_qp (c -> qp ,& a , 259IBV_QP_STATE |IBV_QP_PKEY_INDEX |IBV_QP_PORT |IBV_QP_ACCESS_FLAGS )!= 0 ) { 260return nullptr ; 261 } 262 } 263 264long page = sysconf (_SC_PAGESIZE ); 265if (page <=0 )page = 4096 ; 266const size_t ring_bytes = (size_t )RDMA_NBUF * RDMA_STRIDE ; 267if (posix_memalign ((void ** )& c -> send_mem , (size_t )page ,ring_bytes )!= 0 )c -> send_mem = nullptr ; 268if (posix_memalign ((void ** )& c -> recv_mem , (size_t )page ,ring_bytes )!= 0 )c -> recv_mem = nullptr ; 269if (!c -> send_mem || !c -> recv_mem )return nullptr ; 270 271// Apple's provider rejects LOCAL_WRITE-only MRs even for two-sided SEND/RECV. 272const int mr_flags = IBV_ACCESS_LOCAL_WRITE |IBV_ACCESS_REMOTE_READ |IBV_ACCESS_REMOTE_WRITE ; 273c -> send_mr = ibv_reg_mr (c -> pd ,c -> send_mem ,ring_bytes ,mr_flags ); 274c -> recv_mr = ibv_reg_mr (c -> pd ,c -> recv_mem ,ring_bytes ,mr_flags ); 275if (!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 280c -> qpn = c -> qp -> qp_num ; 281 282apple_rdma_caps rc = {}; 283rc .qpn = c -> qpn ; 284rc .lid = pa .lid ; 285memcpy (rc .gid ,gid .raw ,RDMA_GID_SIZE ); 286memcpy (caps ,& rc ,sizeof (rc )); 287 288GGML_LOG_INFO ("RDMA(Apple/UC) probed: dev=%s port=%u gid=%d qpn=%u lid=%u mtu=%d ring=%d x %zu KiB\n" , 289matched .c_str (),port ,gid_idx ,c -> qpn , (unsigned )pa .lid ,128 <<c -> path_mtu , 290RDMA_NBUF ,RDMA_STRIDE /1024 ); 291return 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 ) { 297impl * c = pimpl .get (); 298 299apple_rdma_caps rc = {}; 300memcpy (& rc ,caps ,sizeof (rc )); 301 302bool ok = true; 303 { 304ibv_qp_attr a = {}; 305a .qp_state = IBV_QPS_RTR ; 306a .path_mtu = c -> path_mtu ; 307a .rq_psn = RDMA_PSN ; 308a .dest_qp_num = rc .qpn ; 309a .ah_attr .is_global = 1 ; 310a .ah_attr .port_num = c -> port ; 311a .ah_attr .sl = 0 ; 312a .ah_attr .src_path_bits = 0 ; 313a .ah_attr .dlid = rc .lid ; 314a .ah_attr .grh .hop_limit = 1 ; 315a .ah_attr .grh .sgid_index = (uint8_t )c -> gid_idx ; 316memcpy (& a .ah_attr .grh .dgid ,rc .gid ,RDMA_GID_SIZE ); 317if (ibv_modify_qp (c -> qp ,& a , 318IBV_QP_STATE |IBV_QP_AV |IBV_QP_PATH_MTU |IBV_QP_DEST_QPN |IBV_QP_RQ_PSN )!= 0 ) { 319GGML_LOG_ERROR ("RDMA(Apple/UC) RTR failed: %s\n" ,strerror (errno )); 320ok = false; 321 } 322 } 323if (ok ) { 324ibv_qp_attr a = {}; 325a .qp_state = IBV_QPS_RTS ; 326a .sq_psn = RDMA_PSN ; 327if (ibv_modify_qp (c -> qp ,& a ,IBV_QP_STATE |IBV_QP_SQ_PSN )!= 0 ) { 328GGML_LOG_ERROR ("RDMA(Apple/UC) RTS failed: %s\n" ,strerror (errno )); 329ok = false; 330 } 331 } 332 333// Recvs are posted only now: the controller starts processing them at RTR. 334for (int i = 0 ;ok && i < RDMA_NBUF ;i ++ ) { 335if (!c -> post_recv (i )) { 336GGML_LOG_ERROR ("RDMA(Apple/UC) post_recv %d/%d failed\n" ,i ,RDMA_NBUF ); 337ok = 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. 344uint8_t peer_ready = 0 ; 345if (!tcp_send_byte (c -> fd ,ok ?RDMA_SYNC_READY :0 )|| !tcp_recv_byte (c -> fd ,& peer_ready )) { 346return false; 347 } 348if (!ok || peer_ready != RDMA_SYNC_READY ) { 349return false; 350 } 351 352GGML_LOG_INFO ("RDMA(Apple/UC) activated: qpn=%u->%u mtu=%d rx_depth=%d\n" , 353c -> qpn ,rc .qpn ,128 <<c -> path_mtu ,RDMA_NBUF ); 354return 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 () { 360struct ibv_wc wc [RDMA_NBUF * 2 ]; 361int n = ibv_poll_cq (cq ,RDMA_NBUF * 2 ,wc ); 362if (n < 0 ) {GGML_LOG_ERROR ("RDMA(Apple/UC) poll_cq failed\n" );broken = true;return -1 ; } 363for (int j = 0 ;j < n ;j ++ ) { 364uint64_t id = wc [j ].wr_id ; 365bool is_recv = (id & RDMA_RECV_WR )!= 0 ; 366if (wc [j ].status != IBV_WC_SUCCESS ) { 367GGML_LOG_ERROR ("RDMA(Apple/UC) %s wc error: status=%d\n" ,is_recv ?"recv" :"send" ,wc [j ].status ); 368broken = true; 369return -1 ; 370 } 371if (is_recv ) { 372int b = (int )(id & RDMA_WR_IDX_MASK ); 373const rdma_seg_hdr * h = (const rdma_seg_hdr * )(recv_mem + (size_t )b * RDMA_STRIDE ); 374if (h -> magic != RDMA_SEG_MAGIC ) {GGML_LOG_ERROR ("RDMA(Apple/UC) bad frame magic\n" );broken = true;return -1 ; } 375if (h -> len > RDMA_PAYLOAD ) {GGML_LOG_ERROR ("RDMA(Apple/UC) frame len %u exceeds payload\n" ,h -> len );broken = true;return -1 ; } 376int slot = (inq_head + inq_count ) %RDMA_NBUF ; 377inq [slot ].buf = b ; 378inq [slot ].off = 0 ; 379inq [slot ].len = h -> len ; 380inq_count ++ ; 381 }else { 382send_busy [(int )(id & RDMA_WR_IDX_MASK )]= 0 ; 383 } 384 } 385return n ; 386} 387 388// Reserve a free send buffer to coalesce into, waiting on progress if none free. 389bool apple_rdma::impl::acquire_pending () { 390if (pend_buf >=0 )return true; 391for (;;) { 392if (broken )return false; 393for (int k = 0 ;k < RDMA_NBUF ;k ++ )if (!send_busy [k ]) {pend_buf = k ;pend_len = 0 ;return true; } 394if (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 () { 402if (pend_buf < 0 )return true; 403int i = pend_buf ; 404rdma_seg_hdr * h = (rdma_seg_hdr * )(send_mem + (size_t )i * RDMA_STRIDE ); 405h -> magic = RDMA_SEG_MAGIC ; 406h -> len = pend_len ; 407if (!post_send (i ,RDMA_STRIDE )) {broken = true;return false; } 408send_busy [i ]= 1 ; 409pend_buf = -1 ; 410pend_len = 0 ; 411return 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 ) { 417impl * c = pimpl .get (); 418const uint8_t * p = (const uint8_t * )data ; 419while (size > 0 ) { 420if (c -> broken )return false; 421if (!c -> acquire_pending ())return false; 422uint8_t * sb = c -> send_mem + (size_t )c -> pend_buf * RDMA_STRIDE ; 423size_t space = RDMA_PAYLOAD - c -> pend_len ; 424size_t chunk = size < space ?size :space ; 425memcpy (sb + sizeof (rdma_seg_hdr )+ c -> pend_len ,p ,chunk ); 426c -> pend_len += (uint32_t )chunk ; 427p += chunk ; 428size -= chunk ; 429if (c -> pend_len == RDMA_PAYLOAD ) {if (!c -> post_pending ())return false; } 430 } 431return true; 432} 433 434bool apple_rdma::recv (void * data ,size_t size ) { 435impl * c = pimpl .get (); 436uint8_t * p = (uint8_t * )data ; 437if (!c -> post_pending ())return false;// turnaround: flush the coalesced request 438unsigned idle = 0 ; 439while (size > 0 ) { 440if (c -> inq_count == 0 ) { 441if (c -> broken )return false; 442int n = c -> progress (); 443if (n < 0 )return false; 444if (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. 449if ((++ idle & 0xFFFFF )== 0 ) { 450struct pollfd pfd = {c -> fd ,POLLIN ,0 }; 451if (poll (& pfd ,1 ,0 )> 0 && 452 (pfd .revents & (POLLIN |POLLHUP |POLLERR |POLLNVAL ))) { 453return false; 454 } 455 } 456 }else { 457idle = 0 ; 458 } 459continue ; 460 } 461idle = 0 ; 462int slot = c -> inq_head ; 463int b = c -> inq [slot ].buf ; 464uint32_t avail = c -> inq [slot ].len - c -> inq [slot ].off ; 465uint32_t take = (size < (size_t )avail ) ? (uint32_t )size :avail ; 466memcpy (p ,c -> recv_mem + (size_t )b * RDMA_STRIDE + sizeof (rdma_seg_hdr )+ c -> inq [slot ].off ,take ); 467p += take ; 468size -= take ; 469c -> inq [slot ].off += take ; 470if (c -> inq [slot ].off == c -> inq [slot ].len ) { 471if (!c -> post_recv (b )) {c -> broken = true;return false; } 472c -> inq_head = (c -> inq_head + 1 ) %RDMA_NBUF ; 473c -> inq_count -- ; 474 } 475 } 476return true; 477} 478 479bool apple_rdma::flush () { 480return pimpl -> post_pending (); 481}