#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

#define DS4_TP_MAGIC UINT32_C(0x44533454) /* "DS4T" */

#define DS4_TP_BATCH_MAGIC UINT32_C(0x44533442) /* "DS4B" */

#define DS4_TP_PROTOCOL_VERSION 21u

#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)

#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)

#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)

#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)

#define DS4_TP_MBOX_SPIN_MS 50u

#define DS4_TP_DEFAULT_TIMEOUT_SEC 300

#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;

typedef struct {

uint32_t magic;

uint16_t layer;

uint16_t gate;

uint64_t seq;

} ds4_tp_gate_header;

#ifdef DS4_TP_HAVE_VERBS

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);

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;

#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,

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;

}

struct ds4_tp {

};

#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;

}

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);

ssize_t w = send(fd, p, len, 0);

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 one = 1;

#ifdef SO_NOSIGPIPE

setsockopt(fd, SOL_SOCKET, SO_NOSIGPIPE, &one, sizeof(one));

setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one));

/* Gate exchanges are latency-critical 16KB messages; large socket

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;

}

}

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;

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

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;

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;

return ok;

}

#ifdef DS4_TP_HAVE_VERBS

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;

}

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 one = 1;

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

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);

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

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;

}

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(

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(

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;

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

if (opt->load_slice) {

tp_set_err(err, errlen, "tensor parallelism does not use distributed layer slices");

return 0;

}

return 1;

}

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;

}

#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);

void *h = dlopen("libibverbs.so.1", RTLD_NOW | RTLD_LOCAL);

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");

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;

}

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__

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;

}

static enum ibv_qp_type tp_rdma_qp_type(void) {

#ifdef __APPLE__

return IBV_QPT_UC;

return IBV_QPT_RC;

}

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

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);

for (int i = 0; i < num && !r->ctx; i++) {