Repository navigation
Expand file tree
/
Copy pathds4_tp.c
More file actions
3960 lines (3778 loc) · 161 KB
/
Copy pathds4_tp.c
File metadata and controls
3960 lines (3778 loc) · 161 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
/* Tensor-parallel transport and lockstep protocol. See ds4_tp.h.
*
* Wire notes: supported ranks are little-endian; the hello magic
* doubles as a byte-order check. The control socket is a plain blocking
* TCP stream carrying framed commands. Gate traffic uses two-sided RDMA
* send/recv (Thunderbolt UC on Metal, RoCE RC on Linux) or a separate
* full-duplex TCP socket. */
#include <errno.h>
#include <fcntl.h>
#include <limits.h>
#include <netdb.h>
#include <stdarg.h>
#include <sys/uio.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <pthread.h>
#include <poll.h>
#include <stdatomic.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <time.h>
#include <unistd.h>
#include "ds4_tp.h"
#include "ds4_gpu.h"
#if (defined(__APPLE__) || defined(__linux__)) && defined(__has_include)
#if __has_include(<infiniband/verbs.h>)
#include <infiniband/verbs.h>
#include <dlfcn.h>
#define DS4_TP_HAVE_VERBS 1
#endif
#endif
#define DS4_TP_MAGIC UINT32_C(0x44533454) /* "DS4T" */
#define DS4_TP_BATCH_MAGIC UINT32_C(0x44533442) /* "DS4B" */
/* 17: RDMA command mailbox and verify-block begin records in the slab tail.
* 18: VERIFY flags (verify-head vocab split), small exchange and head-row
* records in the slab tail.
* 19: DSpark drafter agreement, draft commands and proposal-check flags.
* 20: VERIFY_COMMIT uses the command mailbox when negotiated.
* 21: SPEC_RESTORE restores a speculative prefix with a status-only ACK. */
#define DS4_TP_PROTOCOL_VERSION 21u
/* Leader -> worker command mailbox. Command frames are written with RDMA
* WRITE into a ring in the worker's slab and polled from host memory, so
* the latency-critical ones (VERIFY, EVAL) skip the kernel TCP receive path.
* A record holds a whole frame, or a marker saying the frame follows on the
* TCP control channel (large payloads), so command order is the ring order.
* Each record ends with its sequence number, written by a second RDMA WRITE
* after the body; the RC queue pair places the two in order. */
#define DS4_TP_MBOX_SLOTS 64u
#define DS4_TP_MBOX_SLOT_BYTES 4096u
#define DS4_TP_MBOX_SEQ_OFF (DS4_TP_MBOX_SLOT_BYTES - 8u)
#define DS4_TP_MBOX_FRAME_TCP UINT32_C(0x4d424f58) /* "MBOX": frame is on TCP */
#define DS4_TP_MBOX_ACK_EVERY 16u
#define DS4_TP_MBOX_RING_BYTES ((uint64_t)DS4_TP_MBOX_SLOTS * DS4_TP_MBOX_SLOT_BYTES)
/* Begin records: two parity slots per direction (in, out staging), 64 bytes,
* sequence at byte 56; then the credit word (in) and its staging (out). */
#define DS4_TP_BEGIN_REC_BYTES 64u
#define DS4_TP_BEGIN_SEQ_OFF 56u
#define DS4_TP_MBOX_TAIL_BYTES (2u * DS4_TP_MBOX_RING_BYTES + \
4u * DS4_TP_BEGIN_REC_BYTES + 2u * 64u)
/* Small symmetric exchange records (two parity slots per direction) and one
* worker -> leader float row (verify-head vocab split), also in the slab
* tail. Each ends with an 8-byte sequence word written after the body. */
#define DS4_TP_XCHG_REC_BYTES 256u
#define DS4_TP_XCHG_DATA_OFF 8u /* after a uint32_t byte count */
#define DS4_TP_XCHG_SEQ_OFF (DS4_TP_XCHG_REC_BYTES - 8u)
/* Head row record: {uint32_t floats, uint32_t row, uint64_t block} header,
* the floats, then the sequence word, which is the verify block id. */
#define DS4_TP_HEAD_ROW_DATA_OFF 16u
#define DS4_TP_HEAD_ROW_SEQ_OFF (DS4_TP_HEAD_ROW_DATA_OFF + (uint64_t)DS4_TP_HEAD_ROW_MAX * sizeof(float))
#define DS4_TP_HEAD_ROW_BYTES (DS4_TP_HEAD_ROW_SEQ_OFF + 64u)
#define DS4_TP_HEAD_TAIL_BYTES (4u * DS4_TP_XCHG_REC_BYTES + 2u * DS4_TP_HEAD_ROW_BYTES)
/* Spin this long waiting for a command before falling back to 1 ms polls. */
#define DS4_TP_MBOX_SPIN_MS 50u
#define DS4_TP_DEFAULT_TIMEOUT_SEC 300
/* Once both ranks enter a Metal gate, a live exchange normally completes in
* microseconds. Fail well before Metal's command-buffer watchdog if the peer
* stalls while keeping its sockets open. */
#define DS4_TP_DEFAULT_GATE_TIMEOUT_MS 750
#define DS4_TP_DEFAULT_BIG_GATE_ARRIVAL_MS 30000
typedef struct {
uint32_t magic;
uint32_t type;
uint32_t bytes;
} ds4_tp_frame_header;
typedef struct {
uint32_t magic; /* also detects byte-order mismatch */
uint32_t version;
uint32_t role;
uint32_t rdma_ok; /* this side has a usable verbs device */
uint64_t gguf_bytes;
uint32_t model_id;
uint32_t n_layer;
uint32_t n_embd;
uint32_t n_vocab;
uint32_t quant_bits;
uint32_t ctx_size;
uint32_t gate_slot_start;
uint32_t gate_slot_step;
uint32_t gates_per_token;
uint32_t rdma_max_msg; /* formerly zero padding: zero offers legacy 16 KiB */
uint64_t gate_slot_mask[DS4_TP_GATE_MASK_WORDS];
} ds4_tp_hello_fixed;
typedef struct {
uint64_t slab_base;
uint32_t rkey;
uint32_t qpn;
uint32_t psn;
uint32_t mtu;
uint16_t lid;
uint8_t gid[16];
uint8_t link_layer;
uint8_t reliable;
uint8_t mailbox; /* offers the slab-tail command mailbox */
} ds4_tp_rdma_info;
/* TCP gate frames carry a small header so a desynchronized pair fails loudly
* instead of silently mixing partials. */
typedef struct {
uint32_t magic;
uint16_t layer;
uint16_t gate;
uint64_t seq;
} ds4_tp_gate_header;
#ifdef DS4_TP_HAVE_VERBS
/* librdma is loaded at runtime so builds and machines without the RDMA
* stack (or with it disabled) fall back to TCP with no link-time cost.
* ibv_post_send()/ibv_poll_cq() are header inlines over context->ops, so
* only the setup entry points need dlsym. */
typedef struct {
void *handle;
struct ibv_device **(*get_device_list)(int *);
void (*free_device_list)(struct ibv_device **);
const char *(*get_device_name)(struct ibv_device *);
struct ibv_context *(*open_device)(struct ibv_device *);
int (*close_device)(struct ibv_context *);
int (*query_device)(struct ibv_context *, struct ibv_device_attr *);
int (*query_port)(struct ibv_context *, uint8_t, struct ibv_port_attr *);
int (*query_gid)(struct ibv_context *, uint8_t, int, union ibv_gid *);
#ifdef __linux__
int (*query_gid_ex)(struct ibv_context *, uint32_t, uint32_t,
struct ibv_gid_entry *, uint32_t, size_t);
#endif
struct ibv_pd *(*alloc_pd)(struct ibv_context *);
int (*dealloc_pd)(struct ibv_pd *);
struct ibv_mr *(*reg_mr)(struct ibv_pd *, void *, size_t, int);
int (*dereg_mr)(struct ibv_mr *);
struct ibv_cq *(*create_cq)(struct ibv_context *, int, void *, struct ibv_comp_channel *, int);
int (*destroy_cq)(struct ibv_cq *);
struct ibv_qp *(*create_qp)(struct ibv_pd *, struct ibv_qp_init_attr *);
int (*destroy_qp)(struct ibv_qp *);
int (*modify_qp)(struct ibv_qp *, struct ibv_qp_attr *, int);
int (*query_qp)(struct ibv_qp *, struct ibv_qp_attr *, int, struct ibv_qp_init_attr *);
} ds4_tp_verbs_api;
/* AppleThunderboltRDMA quirks (validated with scratchpad probes,
* 2026-07-06): only UC queue pairs exist (RC/UD: ENOTSUP); RDMA WRITE work
* requests are accepted but never execute, so the data plane is two-sided
* SEND/RECV like Apple's own JACCL; messages above 16KB are not delivered;
* RTR requires GRH addressing with the IPv4-mapped GID that appears only
* once the Thunderbolt member interface has an IPv4 address of its own.
* UC delivery is in-order and the gate sequence is globally deterministic
* (a model-fixed number of gates per token). After any initial bulk prefill,
* decode keeps a receive window posted by sequence number: recv for seq s
* lands in the slab in-slot (s-1) % slots and its completion is the arrival
* signal. The provider also reports CQEs for unsignaled sends: request and
* account for every send, rather than treating one CQE as a completed chain. */
#define DS4_TP_RDMA_MAX_MSG 16384
#define DS4_TP_RDMA_RECV_WINDOW 16
#define DS4_TP_RDMA_BULK_SLOTS 64
#define DS4_TP_RDMA_BULK_WR_TAG (UINT64_C(1) << 63)
#define DS4_TP_RDMA_BLOCK_WR_TAG (UINT64_C(1) << 61)
#define DS4_TP_RDMA_MBOX_WR_TAG (UINT64_C(1) << 60)
typedef struct {
ds4_tp_verbs_api api;
struct ibv_context *ctx;
struct ibv_pd *pd;
struct ibv_cq *cq;
struct ibv_qp *qp;
struct ibv_mr *mr;
struct ibv_port_attr port;
union ibv_gid gid;
int gid_index;
uint32_t max_inline;
uint32_t max_msg;
ds4_tp_rdma_info peer;
uint32_t send_outstanding; /* individual sends not yet reaped */
uint64_t recv_done; /* highest gate seq whose recv completed */
uint64_t last_gate_seq; /* last real decode receive consumed */
bool recv_window_active; /* decode recvs are queued ahead */
int warm_failed; /* last warm-up round failed (dead direction) */
uint32_t setup_attempt; /* queue pair recreations so far */
pthread_mutex_t post_lock;
bool post_lock_ready;
uint32_t recv_depth; /* queue pair receive depth actually granted */
uint32_t send_depth; /* queue pair send depth actually granted */
struct ibv_sge *win_sge; /* window work request arrays (recv_depth) */
struct ibv_recv_wr *win_rwr;
struct ibv_send_wr *win_swr;
/* Verify-block window (speculative decoding): one batch gate per layer,
* receives posted a layer ahead of every send, no per-gate control
* traffic (see ds4_tp_batch_block_begin). */
bool block_active;
uint32_t block_rows;
uint32_t block_layers;
uint32_t block_posted; /* layers whose receives are posted */
uint64_t block_recv_done; /* row messages received in this block */
/* Command mailbox and begin records (both peers offered them). */
bool mbox_active;
uint64_t mbox_sent; /* leader: records written */
uint64_t mbox_consumed; /* worker: records taken */
uint64_t mbox_acked; /* worker: last credit written back */
uint64_t begin_count; /* verify blocks begun (both ranks) */
uint64_t xchg_count; /* small exchanges (both ranks) */
} ds4_tp_rdma;
static uint32_t tp_rdma_negotiate_max_msg(uint32_t local, uint32_t peer) {
if (!local) local = DS4_TP_RDMA_MAX_MSG;
if (!peer) peer = DS4_TP_RDMA_MAX_MSG;
if (local < DS4_TP_RDMA_MAX_MSG || peer < DS4_TP_RDMA_MAX_MSG ||
(local & (local - 1u)) || (peer & (peer - 1u))) return 0;
return local < peer ? local : peer;
}
#endif
struct ds4_tp {
ds4_tp_options opt;
int rank; /* 0 leader, 1 worker */
int control_fd;
int data_fd; /* TCP fallback, headers, and verify gates */
bool rdma_active;
uint32_t peer_ctx;
uint32_t n_layer;
uint32_t n_embd;
uint64_t vec_bytes;
uint32_t n_slots;
/* Decode gate schedule (see ds4_tp_identity). */
uint32_t gate_slot_start;
uint32_t gate_slot_step;
uint32_t gates_per_token;
uint64_t gate_slot_mask[DS4_TP_GATE_MASK_WORDS];
uint8_t *slab;
uint64_t slab_bytes;
/* Slab regions, see ds4_tp.h layout comment. */
uint64_t out_off;
uint64_t in_off;
uint64_t in_flags_off;
uint64_t token_off;
uint64_t out_flags_off; /* local staging for RDMA flag writes */
uint64_t gpu_flags_off; /* GPU-written gate-ready flags (u32/slot) */
uint64_t batch_out_off; /* [layer][row] verify-block local partials */
uint64_t batch_in_off; /* [layer][row] verify-block peer partials */
uint64_t mbox_in_off; /* command ring written by the leader */
uint64_t mbox_out_off; /* local staging for outgoing records */
uint64_t begin_in_off; /* peer begin records (2 parity slots) */
uint64_t begin_out_off; /* local staging for begin records */
uint64_t credit_in_off; /* worker's consumed count (on the leader) */
uint64_t credit_out_off; /* local staging for the credit word */
uint64_t xchg_in_off; /* peer small-exchange records (2 parity) */
uint64_t xchg_out_off; /* local staging for them */
uint64_t head_row_in_off; /* leader: the worker's head row */
uint64_t head_row_out_off; /* worker: staging of its head row */
uint64_t timeout_sec;
uint64_t gate_timeout_ms;
atomic_bool failed;
uint64_t sync_checkpoint_seq;
#ifdef DS4_TP_HAVE_VERBS
ds4_tp_rdma rdma;
#endif
};
#ifdef DS4_TP_HAVE_VERBS
static uint32_t tp_rdma_max_msg(const ds4_tp *tp) {
return tp->rdma.max_msg ? tp->rdma.max_msg : DS4_TP_RDMA_MAX_MSG;
}
#endif
/* ------------------------------------------------------------------------
* Small socket helpers (same conventions as ds4_distributed.c).
* --------------------------------------------------------------------- */
static double tp_now_sec(void) {
struct timespec ts;
clock_gettime(CLOCK_MONOTONIC, &ts);
return (double)ts.tv_sec + (double)ts.tv_nsec / 1e9;
}
static void tp_set_err(char *err, size_t errlen, const char *fmt, ...) {
if (!err || !errlen) return;
va_list ap;
va_start(ap, fmt);
vsnprintf(err, errlen, fmt, ap);
va_end(ap);
}
static int tp_write_full(int fd, const void *buf, size_t len) {
const char *p = buf;
while (len) {
#ifdef MSG_NOSIGNAL
ssize_t w = send(fd, p, len, MSG_NOSIGNAL);
#else
ssize_t w = send(fd, p, len, 0);
#endif
if (w < 0) {
if (errno == EINTR) continue;
return 0;
}
if (w == 0) return 0;
p += w;
len -= (size_t)w;
}
return 1;
}
static int tp_read_full(int fd, void *buf, size_t len) {
char *p = buf;
while (len) {
ssize_t r = read(fd, p, len);
if (r < 0) {
if (errno == EINTR) continue;
return 0;
}
if (r == 0) return 0;
p += r;
len -= (size_t)r;
}
return 1;
}
static void tp_socket_tune(int fd) {
int >
#ifdef SO_NOSIGPIPE
setsockopt(fd, SOL_SOCKET, SO_NOSIGPIPE, &one, sizeof(one));
#endif
setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one));
/* Gate exchanges are latency-critical 16KB messages; large socket
* buffers only matter for the TCP fallback's pipelining. */
int sz = 4 * 1024 * 1024;
setsockopt(fd, SOL_SOCKET, SO_SNDBUF, &sz, sizeof(sz));
setsockopt(fd, SOL_SOCKET, SO_RCVBUF, &sz, sizeof(sz));
}
static int tp_socket_set_gate_timeout(int fd, uint64_t timeout_ms) {
struct timeval tv = {
.tv_sec = (time_t)(timeout_ms / 1000u),
.tv_usec = (suseconds_t)((timeout_ms % 1000u) * 1000u),
};
if (setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv)) != 0)
return 0;
if (setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)) != 0)
return 0;
return 1;
}
static void tp_iov_advance(struct msghdr *msg, size_t bytes) {
while (msg->msg_iovlen && bytes >= msg->msg_iov->iov_len) {
bytes -= msg->msg_iov->iov_len;
msg->msg_iov++;
msg->msg_iovlen--;
}
if (bytes) {
msg->msg_iov->iov_base = (char *)msg->msg_iov->iov_base + bytes;
msg->msg_iov->iov_len -= bytes;
}
}
/* Keep both directions moving even if the kernel clamps the socket buffers. */
static int tp_tcp_exchange_io(ds4_tp *tp, ds4_tp_gate_header *header,
ds4_tp_gate_header *peer_header,
const void *out, void *in, size_t bytes) {
struct iovec tx[] = {{header, sizeof(*header)}, {(void *)out, bytes}};
struct iovec rx[] = {{peer_header, sizeof(*peer_header)}, {in, bytes}};
struct msghdr send_msg = {.msg_iov = tx + !header, .msg_iovlen = header ? 2 : 1};
struct msghdr recv_msg = {.msg_iov = rx + !header, .msg_iovlen = header ? 2 : 1};
const double timeout = tp->gate_timeout_ms / 1000.0;
double deadline = tp_now_sec() + timeout;
while (send_msg.msg_iovlen || recv_msg.msg_iovlen) {
int progress = 0;
if (send_msg.msg_iovlen) {
int flags = MSG_DONTWAIT;
#ifdef MSG_NOSIGNAL
flags |= MSG_NOSIGNAL;
#endif
ssize_t n = sendmsg(tp->data_fd, &send_msg, flags);
if (n > 0) { tp_iov_advance(&send_msg, (size_t)n); progress = 1; }
else if (n == 0) { errno = EPIPE; return 0; }
else if (errno != EAGAIN && errno != EWOULDBLOCK && errno != EINTR) return 0;
}
if (recv_msg.msg_iovlen) {
ssize_t n = recvmsg(tp->data_fd, &recv_msg, MSG_DONTWAIT);
if (n > 0) { tp_iov_advance(&recv_msg, (size_t)n); progress = 1; }
else if (n == 0) { errno = ECONNRESET; return 0; }
else if (errno != EAGAIN && errno != EWOULDBLOCK && errno != EINTR) return 0;
}
if (progress) { deadline = tp_now_sec() + timeout; continue; }
const double remaining_ms = (deadline - tp_now_sec()) * 1000.0;
if (remaining_ms <= 0.0) { errno = ETIMEDOUT; return 0; }
struct pollfd pfd = {.fd = tp->data_fd,
.events = (send_msg.msg_iovlen ? POLLOUT : 0) |
(recv_msg.msg_iovlen ? POLLIN : 0)};
int wait_ms = remaining_ms >= INT_MAX ? INT_MAX : (int)remaining_ms + 1;
if (poll(&pfd, 1, wait_ms) < 0 && errno != EINTR) return 0;
if (pfd.revents & POLLNVAL) { errno = EBADF; return 0; }
}
return 1;
}
static int tp_tcp_exchange(ds4_tp *tp, ds4_tp_gate_header *header,
ds4_tp_gate_header *peer_header,
const void *out, void *in, size_t bytes) {
#ifdef __APPLE__
/* Darwin sendmsg can block despite MSG_DONTWAIT. The data socket has
* one exchange owner; restore its mode before blocking header I/O. */
const int flags = fcntl(tp->data_fd, F_GETFL);
if (flags < 0) return 0;
const bool changed = !(flags & O_NONBLOCK);
if (changed && fcntl(tp->data_fd, F_SETFL, flags | O_NONBLOCK) < 0) return 0;
#endif
const int ok = tp_tcp_exchange_io(tp, header, peer_header, out, in, bytes);
#ifdef __APPLE__
const int saved_errno = errno;
if (changed && fcntl(tp->data_fd, F_SETFL, flags) < 0) return 0;
errno = saved_errno;
#endif
return ok;
}
#ifdef DS4_TP_HAVE_VERBS
/* UC queue pairs do not report a dead remote reliably. The control socket
* does, so sample it while polling an RDMA completion and abort before the
* Metal command-buffer watchdog fires. */
static int tp_peer_closed(const ds4_tp *tp) {
char byte;
const ssize_t n = recv(tp->control_fd, &byte, 1,
MSG_PEEK | MSG_DONTWAIT);
if (n == 0) return 1;
if (n > 0) return 0;
return errno != EAGAIN && errno != EWOULDBLOCK && errno != EINTR;
}
#endif
static int tp_listen(const char *host, int port, char *err, size_t errlen) {
char portbuf[16];
snprintf(portbuf, sizeof(portbuf), "%d", port);
struct addrinfo hints = {0}, *res = NULL;
hints.ai_family = AF_UNSPEC;
hints.ai_socktype = SOCK_STREAM;
hints.ai_flags = AI_PASSIVE;
int rc = getaddrinfo(host && host[0] ? host : NULL, portbuf, &hints, &res);
if (rc != 0) {
tp_set_err(err, errlen, "tp listen resolve %s:%d: %s", host, port, gai_strerror(rc));
return -1;
}
int fd = -1;
for (struct addrinfo *ai = res; ai; ai = ai->ai_next) {
fd = socket(ai->ai_family, ai->ai_socktype, ai->ai_protocol);
if (fd < 0) continue;
int >
setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
if (bind(fd, ai->ai_addr, ai->ai_addrlen) == 0 && listen(fd, 2) == 0) break;
close(fd);
fd = -1;
}
freeaddrinfo(res);
if (fd < 0) tp_set_err(err, errlen, "tp listen %s:%d: %s", host, port, strerror(errno));
return fd;
}
static int tp_dial(const char *host, int port, double timeout_sec, char *err, size_t errlen) {
char portbuf[16];
snprintf(portbuf, sizeof(portbuf), "%d", port);
double deadline = tp_now_sec() + timeout_sec;
int last_errno = 0;
uint32_t attempts = 0;
do {
struct addrinfo hints = {0}, *res = NULL;
hints.ai_family = AF_UNSPEC;
hints.ai_socktype = SOCK_STREAM;
int gai = getaddrinfo(host, portbuf, &hints, &res);
if (gai == 0) {
for (struct addrinfo *ai = res; ai; ai = ai->ai_next) {
int fd = socket(ai->ai_family, ai->ai_socktype, ai->ai_protocol);
if (fd < 0) continue;
if (connect(fd, ai->ai_addr, ai->ai_addrlen) == 0) {
freeaddrinfo(res);
return fd;
}
last_errno = errno;
close(fd);
}
freeaddrinfo(res);
}
/* Retrying is normal while the peer loads its model; still say why
* every ~10s so a wrong address or a policy block is visible. */
if (attempts++ % 50 == 0) {
fprintf(stderr, "ds4-tp: connecting to %s:%d ... (%s)\n", host, port,
gai != 0 ? gai_strerror(gai) :
last_errno ? strerror(last_errno) : "no address worked");
}
usleep(200 * 1000);
} while (tp_now_sec() < deadline);
tp_set_err(err, errlen, "tp connect %s:%d: %s", host, port,
last_errno ? strerror(last_errno) : "unreachable");
return -1;
}
static int tp_send_frame(int fd, uint32_t type, const void *payload, uint32_t bytes) {
ds4_tp_frame_header h = { DS4_TP_MAGIC, type, bytes };
if (!tp_write_full(fd, &h, sizeof(h))) return 0;
if (bytes && !tp_write_full(fd, payload, bytes)) return 0;
return 1;
}
#ifdef DS4_TP_HAVE_VERBS
static int tp_mbox_active(const ds4_tp *tp);
static int tp_mbox_put(ds4_tp *tp, uint32_t type, const void *payload, uint32_t bytes);
static int tp_mbox_get(ds4_tp *tp, uint32_t *type, uint8_t **payload, uint32_t *bytes);
#endif
/* Leader -> worker frames consumed by the worker's command loop. With the
* RDMA mailbox a frame that fits travels in the ring; a larger one sends a
* ring marker and then the frame on the TCP control channel, so the worker
* still takes commands strictly in order. */
static int tp_send_command(ds4_tp *tp, uint32_t type, const void *payload, uint32_t bytes) {
if (!tp || ds4_tp_failed(tp)) return 0;
int ok;
#ifdef DS4_TP_HAVE_VERBS
if (tp_mbox_active(tp)) {
if (bytes <= DS4_TP_MBOX_SEQ_OFF - sizeof(ds4_tp_frame_header))
ok = tp_mbox_put(tp, type, payload, bytes);
else
ok = tp_mbox_put(tp, DS4_TP_MBOX_FRAME_TCP, NULL, 0) &&
tp_send_frame(tp->control_fd, type, payload, bytes);
} else
#endif
ok = tp_send_frame(tp->control_fd, type, payload, bytes);
if (!ok) ds4_tp_mark_failed(tp);
return ok;
}
static int tp_read_frame_header(int fd, uint32_t *type, uint32_t *bytes) {
ds4_tp_frame_header h;
if (!tp_read_full(fd, &h, sizeof(h))) return 0;
if (h.magic != DS4_TP_MAGIC) return 0;
*type = h.type;
*bytes = h.bytes;
return 1;
}
/* ------------------------------------------------------------------------
* Options and CLI.
* --------------------------------------------------------------------- */
bool ds4_tp_enabled(const ds4_tp_options *opt) {
return opt && opt->role != DS4_TP_NONE;
}
void ds4_tp_usage(FILE *fp) {
fprintf(fp,
"Tensor parallelism (two identical machines):\n"
" --tensor-parallel Use --role/--listen/--coordinator for a 50/50 TP pair.\n"
" --transport <auto|rdma|tcp> Gate transport (default auto).\n"
" --rdma-device <name> Select a verbs device such as rdma_en1.\n"
" --rdma-gid-index <n> Select the local verbs GID index.\n"
" --tensor-parallel-token-prefill\n"
" GLM diagnostic: prefill one token at a time.\n"
" --debug-hash <n> Cross-check hidden state every n tokens.\n");
}
int ds4_tp_parse_cli_arg(
cons
10BC8
t char *arg,
int *index,
int argc,
char **argv,
ds4_tp_options *opt,
char *err,
size_t errlen)
{
int i = *index;
if (!strcmp(arg, "--tensor-parallel")) {
opt->requested = true;
} else if (!strcmp(arg, "--transport")) {
if (i + 1 >= argc) goto missing;
const char *v = argv[++i];
if (!strcmp(v, "auto")) opt->transport = DS4_TP_TRANSPORT_AUTO;
else if (!strcmp(v, "rdma")) opt->transport = DS4_TP_TRANSPORT_RDMA;
else if (!strcmp(v, "tcp")) opt->transport = DS4_TP_TRANSPORT_TCP;
else {
tp_set_err(err, errlen, "invalid %s value: %s", arg, v);
return DS4_TP_CLI_ERROR;
}
} else if (!strcmp(arg, "--rdma-device")) {
if (i + 1 >= argc) goto missing;
opt->rdma_device = argv[++i];
} else if (!strcmp(arg, "--rdma-gid-index")) {
if (i + 1 >= argc) goto missing;
char *end = NULL;
errno = 0;
long value = strtol(argv[++i], &end, 10);
if (errno != 0 || !end || *end != '\0' || value < 0 || value > INT_MAX) {
tp_set_err(err, errlen, "invalid --rdma-gid-index %s", argv[i]);
return DS4_TP_CLI_ERROR;
}
opt->rdma_gid_index = (int)value;
opt->rdma_gid_index_set = true;
} else if (!strcmp(arg, "--tensor-parallel-token-prefill")) {
opt->glm_token_prefill = true;
} else if (!strcmp(arg, "--debug-hash")) {
if (i + 1 >= argc) goto missing;
opt->debug_hash = atoi(argv[++i]);
} else {
return DS4_TP_CLI_NOT_MATCHED;
}
*index = i;
return DS4_TP_CLI_MATCHED;
missing:
tp_set_err(err, errlen, "%s requires an argument", arg);
return DS4_TP_CLI_ERROR;
}
int ds4_tp_adopt_distributed_options(
ds4_tp_options *tp,
ds4_distributed_options *dist,
char *err,
size_t errlen)
{
if (!tp || !dist || !tp->requested) return 1;
if (tp->role != DS4_TP_NONE) {
tp_set_err(err, errlen,
"--tensor-parallel selects its role through --role");
return 0;
}
if (dist->role == DS4_DISTRIBUTED_NONE) {
tp_set_err(err, errlen,
"--tensor-parallel requires --role coordinator or --role worker");
return 0;
}
if (dist->layers.set) {
tp_set_err(err, errlen,
"tensor parallelism always uses one 50/50 worker; omit --layers");
return 0;
}
if (dist->prefill_chunk || dist->prefill_window || dist->activation_bits ||
dist->replay_check || dist->debug) {
tp_set_err(err, errlen,
"--dist-* and distributed debug options cannot be used with --tensor-parallel");
return 0;
}
if (dist->role == DS4_DISTRIBUTED_COORDINATOR) {
if (!dist->listen_host || dist->listen_port <= 0) {
tp_set_err(err, errlen,
"--role coordinator --tensor-parallel requires --listen HOST PORT");
return 0;
}
if (dist->coordinator_host || dist->coordinator_port) {
tp_set_err(err, errlen,
"--role coordinator must not use --coordinator");
return 0;
}
tp->role = DS4_TP_LEADER;
tp->listen_host = dist->listen_host;
tp->listen_port = dist->listen_port;
} else if (dist->role == DS4_DISTRIBUTED_WORKER) {
if (!dist->coordinator_host || dist->coordinator_port <= 0) {
tp_set_err(err, errlen,
"--role worker --tensor-parallel requires --coordinator HOST PORT");
return 0;
}
if (dist->listen_host || dist->listen_port) {
tp_set_err(err, errlen,
"--role worker --tensor-parallel must not use --listen");
return 0;
}
tp->role = DS4_TP_WORKER;
tp->leader_host = dist->coordinator_host;
tp->leader_port = dist->coordinator_port;
} else {
tp_set_err(err, errlen, "invalid tensor-parallel role");
return 0;
}
memset(dist, 0, sizeof(*dist));
return 1;
}
int ds4_tp_validate_engine_options(
const ds4_engine_options *opt,
char *err,
size_t errlen)
{
if (!ds4_tp_enabled(&opt->tp)) {
if (opt->tp.requested || opt->tp.transport != DS4_TP_TRANSPORT_AUTO ||
opt->tp.rdma_device || opt->tp.rdma_gid_index_set ||
opt->tp.glm_token_prefill || opt->tp.debug_hash != 0) {
tp_set_err(err, errlen,
"tensor-parallel options require --tensor-parallel and --role");
return 0;
}
return 1;
}
bool supported_backend = opt->backend == DS4_BACKEND_METAL;
#if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU)
supported_backend |= opt->backend == DS4_BACKEND_CUDA;
#endif
if (!supported_backend) {
tp_set_err(err, errlen, "network tensor parallelism requires Metal or supported CUDA models");
return 0;
}
if (opt->backend == DS4_BACKEND_CUDA && (opt->cuda_tensor_parallel || opt->ssd_streaming)) {
tp_set_err(err, errlen, "network CUDA TP requires one GPU per rank and resident expert shards");
return 0;
}
if (opt->distributed.role != DS4_DISTRIBUTED_NONE) {
tp_set_err(err, errlen, "tensor parallelism and --role distributed modes are exclusive");
return 0;
}
/* Speculative drafting (DSpark/MTP) is allowed on the leader: the
* verify block is mirrored to the worker via DS4_TP_FRAME_VERIFY and
* the legacy MTP path falls back to per-token decode under TP. */
if (opt->load_slice) {
tp_set_err(err, errlen, "tensor parallelism does not use distributed layer slices");
return 0;
}
return 1;
}
/* ------------------------------------------------------------------------
* Slab layout.
* --------------------------------------------------------------------- */
uint64_t ds4_tp_slab_bytes(uint32_t n_layer, uint32_t n_embd) {
uint64_t vec = (uint64_t)n_embd * sizeof(float);
uint64_t slots = (uint64_t)n_layer * DS4_TP_GATES_PER_LAYER;
return slots * vec * 2 + /* out + in vectors */
slots * 8 * 2 + /* in flags + out flag staging */
16 + /* token slot */
slots * 4 + /* GPU-written gate-ready flags */
(uint64_t)n_layer * DS4_TP_BATCH_GATES_PER_LAYER *
DS4_TP_BATCH_MAX_ROWS * vec * 2 + /* batch out+in */
DS4_TP_MBOX_TAIL_BYTES + /* command mailbox */
DS4_TP_HEAD_TAIL_BYTES; /* small exchange + head row */
}
static void tp_slab_layout(ds4_tp *tp) {
uint64_t vec = tp->vec_bytes;
uint64_t slots = tp->n_slots;
tp->out_off = 0;
tp->in_off = slots * vec;
tp->in_flags_off = tp->in_off + slots * vec;
tp->token_off = tp->in_flags_off + slots * 8;
tp->out_flags_off = tp->token_off + 16;
tp->gpu_flags_off = tp->out_flags_off + slots * 8;
tp->batch_out_off = tp->gpu_flags_off + slots * 4;
tp->batch_in_off = tp->batch_out_off +
(uint64_t)tp->n_layer * DS4_TP_BATCH_GATES_PER_LAYER *
DS4_TP_BATCH_MAX_ROWS * vec;
tp->mbox_in_off = tp->batch_in_off +
(uint64_t)tp->n_layer * DS4_TP_BATCH_GATES_PER_LAYER *
DS4_TP_BATCH_MAX_ROWS * vec;
tp->mbox_out_off = tp->mbox_in_off + DS4_TP_MBOX_RING_BYTES;
tp->begin_in_off = tp->mbox_out_off + DS4_TP_MBOX_RING_BYTES;
tp->begin_out_off = tp->begin_in_off + 2u * DS4_TP_BEGIN_REC_BYTES;
tp->credit_in_off = tp->begin_out_off + 2u * DS4_TP_BEGIN_REC_BYTES;
tp->credit_out_off = tp->credit_in_off + 64u;
tp->xchg_in_off = tp->credit_out_off + 64u;
tp->xchg_out_off = tp->xchg_in_off + 2u * DS4_TP_XCHG_REC_BYTES;
tp->head_row_in_off = tp->xchg_out_off + 2u * DS4_TP_XCHG_REC_BYTES;
tp->head_row_out_off = tp->head_row_in_off + DS4_TP_HEAD_ROW_BYTES;
tp->slab_bytes = tp->head_row_out_off + DS4_TP_HEAD_ROW_BYTES;
}
uint64_t ds4_tp_slab_gpu_flags_offset(const ds4_tp *tp) {
return tp->gpu_flags_off;
}
static uint32_t tp_slot(const ds4_tp *tp, uint32_t layer, uint32_t gate) {
(void)tp;
return layer * DS4_TP_GATES_PER_LAYER + gate;
}
uint64_t ds4_tp_slab_out_offset(const ds4_tp *tp, uint32_t layer, uint32_t gate) {
return tp->out_off + (uint64_t)tp_slot(tp, layer, gate) * tp->vec_bytes;
}
uint64_t ds4_tp_slab_in_offset(const ds4_tp *tp, uint32_t layer, uint32_t gate) {
return tp->in_off + (uint64_t)tp_slot(tp, layer, gate) * tp->vec_bytes;
}
uint64_t ds4_tp_slab_batch_out_offset(const ds4_tp *tp, uint32_t layer) {
return tp->batch_out_off +
(uint64_t)layer * DS4_TP_BATCH_MAX_ROWS * tp->vec_bytes;
}
uint64_t ds4_tp_slab_batch_in_offset(const ds4_tp *tp, uint32_t layer) {
return tp->batch_in_off +
(uint64_t)layer * DS4_TP_BATCH_MAX_ROWS * tp->vec_bytes;
}
static uint32_t tp_gate_mask_count(
const uint64_t mask[DS4_TP_GATE_MASK_WORDS]) {
uint32_t count = 0;
for (uint32_t i = 0; i < DS4_TP_GATE_MASK_WORDS; i++)
count += (uint32_t)__builtin_popcountll(mask[i]);
return count;
}
static int tp_gate_mask_fits(
const uint64_t mask[DS4_TP_GATE_MASK_WORDS],
uint32_t n_slots) {
for (uint32_t word = 0; word < DS4_TP_GATE_MASK_WORDS; word++) {
uint64_t bits = mask[word];
while (bits) {
const uint32_t slot =
word * 64u + (uint32_t)__builtin_ctzll(bits);
if (slot >= n_slots) return 0;
bits &= bits - 1u;
}
}
return 1;
}
/* ------------------------------------------------------------------------
* RDMA path.
* --------------------------------------------------------------------- */
#ifdef DS4_TP_HAVE_VERBS
static int tp_rdma_load_api(ds4_tp_verbs_api *api) {
if (api->handle) return 1;
#ifdef __APPLE__
void *h = dlopen("/usr/lib/librdma.dylib", RTLD_NOW | RTLD_LOCAL);
if (!h) h = dlopen("librdma.dylib", RTLD_NOW | RTLD_LOCAL);
#else
void *h = dlopen("libibverbs.so.1", RTLD_NOW | RTLD_LOCAL);
#endif
if (!h) return 0;
#define TP_SYM(field, name) \
do { \
api->field = (__typeof__(api->field))dlsym(h, name); \
if (!api->field) { dlclose(h); return 0; } \
} while (0)
TP_SYM(get_device_list, "ibv_get_device_list");
TP_SYM(free_device_list, "ibv_free_device_list");
TP_SYM(get_device_name, "ibv_get_device_name");
TP_SYM(open_device, "ibv_open_device");
TP_SYM(close_device, "ibv_close_device");
TP_SYM(query_device, "ibv_query_device");
TP_SYM(query_port, "ibv_query_port");
TP_SYM(query_gid, "ibv_query_gid");
#ifdef __linux__
TP_SYM(query_gid_ex, "_ibv_query_gid_ex");
#endif
TP_SYM(alloc_pd, "ibv_alloc_pd");
TP_SYM(dealloc_pd, "ibv_dealloc_pd");
TP_SYM(reg_mr, "ibv_reg_mr");
TP_SYM(dereg_mr, "ibv_dereg_mr");
TP_SYM(create_cq, "ibv_create_cq");
TP_SYM(destroy_cq, "ibv_destroy_cq");
TP_SYM(create_qp, "ibv_create_qp");
TP_SYM(destroy_qp, "ibv_destroy_qp");
TP_SYM(modify_qp, "ibv_modify_qp");
TP_SYM(query_qp, "ibv_query_qp");
#undef TP_SYM
api->handle = h;
return 1;
}
/* Probe only: does this machine expose a verbs device right now? */
static int tp_rdma_probe(ds4_tp_verbs_api *api) {
if (!tp_rdma_load_api(api)) return 0;
int num = 0;
struct ibv_device **devs = api->get_device_list(&num);
if (!devs) return 0;
api->free_device_list(devs);
return num > 0;
}
#ifdef __linux__
/* Prefer the RoCEv2 GID belonging to the control socket's direct-link address.
* Device order need not agree between hosts with several crossed NIC ports. */
static int tp_rdma_linux_gid(ds4_tp *tp, struct ibv_context *ctx,
const struct ibv_port_attr *port,
union ibv_gid *gid, int *index) {
struct sockaddr_storage local = {0};
socklen_t len = sizeof(local);
union ibv_gid address = {0};
bool have_address = getsockname(tp->control_fd, (struct sockaddr *)&local, &len) == 0;
if (have_address && local.ss_family == AF_INET) {
address.raw[10] = address.raw[11] = 0xff;
memcpy(address.raw + 12, &((struct sockaddr_in *)&local)->sin_addr, 4);
} else if (have_address && local.ss_family == AF_INET6) {
memcpy(address.raw, &((struct sockaddr_in6 *)&local)->sin6_addr, 16);
} else {
have_address = false;
}
int best = 0;
const int first = tp->opt.rdma_gid_index_set ? tp->opt.rdma_gid_index : 0;
if (first < 0 || first >= port->gid_tbl_len || first > UINT8_MAX) return 0;
const int end = tp->opt.rdma_gid_index_set ? first + 1 : port->gid_tbl_len;
for (int i = first; i < end && i <= UINT8_MAX; i++) {
struct ibv_gid_entry entry = {0};
if (tp->rdma.api.query_gid_ex(ctx, 1, (uint32_t)i, &entry, 0, sizeof(entry)) ||
entry.gid_type != IBV_GID_TYPE_ROCE_V2) continue;
const union ibv_gid zero = {0};
if (!memcmp(entry.gid.raw, zero.raw, sizeof(zero.raw))) continue;
const int score = have_address && !memcmp(entry.gid.raw, address.raw, 16) ? 2 : 1;
if (score > best) {
*gid = entry.gid;
*index = i;
best = score;
}
}
return best;
}
#endif
static enum ibv_qp_type tp_rdma_qp_type(void) {
#ifdef __APPLE__
return IBV_QPT_UC;
#else
return IBV_QPT_RC;
#endif
}
static int tp_rdma_open(ds4_tp *tp, char *err, size_t errlen) {
ds4_tp_rdma *r = &tp->rdma;
int num = 0;
struct ibv_device **devs = r->api.get_device_list(&num);
if (!devs || num == 0) {
tp_set_err(err, errlen, "tp rdma: no verbs devices");
if (devs) r->api.free_device_list(devs);
return 0;
}
/* One verbs device per Thunderbolt port (rdma_enN); pick the active one
* unless the caller selected a device explicitly. */
const char *want_name = tp->opt.rdma_device;
char states[256] = "";
#ifdef __linux__
int best = 0, candidates = 0;
for (int i = 0; i < num; i++) {
const char *name = r->api.get_device_name(devs[i]);
if (want_name && strcmp(want_name, name)) continue;
struct ibv_context *ctx = r->api.open_device(devs[i]);
if (!ctx) continue;
struct ibv_port_attr pa = {0};
union ibv_gid gid;
int index = -1;
const int score = !r->api.query_port(ctx, 1, &pa) && pa.state == IBV_PORT_ACTIVE &&
pa.link_layer == IBV_LINK_LAYER_ETHERNET ? tp_rdma_linux_gid(tp, ctx, &pa, &gid, &index) : 0;
if (score) candidates++;
if (score > best) {
if (r->ctx) r->api.close_device(r->ctx);
r->ctx = ctx;
r->port = pa;
r->gid = gid;
r->gid_index = index;
best = score;
snprintf(states, sizeof(states), "%s", name);
} else {
r->api.close_device(ctx);
}
}
r->api.free_device_list(devs);
if (!best || (!want_name && best == 1 && candidates > 1)) {
tp_set_err(err, errlen,
"tp rdma: %s; use the direct RoCE address for --coordinator, or select --rdma-device and --rdma-gid-index",
best ? "multiple active ports, none matches the control address" : "no active RoCEv2 GID for the requested device/index");
return 0;
}
fprintf(stderr, "ds4-tp: rdma device %s, RoCEv2 GID %d, RC\n", states, r->gid_index);
#else
for (int i = 0; i < num && !r->ctx; i++) {
const char *name = r->api.get_device_name(devs[i]);
if (want_name && strcmp(want_name, name) != 0) continue;
struct ibv_context *ctx = r->api.open_device(devs[i]);
if (!ctx) continue;
struct ibv_port_attr pa = {0};
if (r->api.query_port(ctx, 1, &pa) == 0 &&
(pa.state == IBV_PORT_ACTIVE || want_name)) {
r->ctx = ctx;
r->port = pa;
fprintf(stderr, "ds4-tp: rdma device %s (port state %d)\n", name, (int)pa.state);
break;
}
size_t off = strlen(states);
You can’t perform that action at this time.