Repository navigation
Expand file tree
/
Copy pathds4_tp.h
More file actions
350 lines (320 loc) · 15.6 KB
/
Copy pathds4_tp.h
File metadata and controls
350 lines (320 loc) · 15.6 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
#ifndef DS4_TP_H
#define DS4_TP_H
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include "ds4.h"
/* Tensor-parallel transport and lockstep protocol.
*
* Two ranks run the same logical model, each with one contiguous half of the
* routed experts resident. Rank 0 (leader) is a normal frontend session that
* mirrors every ds4_session_sync()/ds4_session_eval() call to rank 1 (worker)
* over a TCP control socket, so both engines execute the identical graph
* sequence.
* Inside each decoded token, partial block outputs are exchanged through a
* registered memory slab: two-sided RDMA SEND/RECV when RDMA over
* Thunderbolt is available, or a full-duplex TCP exchange as fallback.
*
* Layering: ds4.c calls the session-mirroring and slab entry points;
* ds4_metal.m only ever sees ds4_tp_gate_exchange() through a callback
* registered with the GPU gate machinery. Nothing here touches tensors.
*/
typedef struct ds4_tp ds4_tp;
enum {
DS4_TP_GATE_ATTN = 0,
DS4_TP_GATE_FFN = 1,
DS4_TP_GATES_PER_LAYER = 2,
/* Includes the seed row as well as drafted tokens. */
DS4_TP_BATCH_MAX_ROWS = 8,
DS4_TP_BATCH_GATES_PER_LAYER = 2,
};
/* Engine identity exchanged in the hello so a mismatched pair aborts before
* any inference runs. */
typedef struct {
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;
/* Decode gate schedule, used to place RDMA recvs into the right slab
* slot. A nonempty mask lists the exact slots in firing order; otherwise
* slot(seq) = start + ((seq-1) % per_token) * step.
* per_token 0 falls back to the identity mapping over all slots
* (DS4: every layer fires ATTN then FFN). GLM fires one FFN gate per
* sparse layer in streaming mode. Hybrid GLM-5.3 uses the mask because
* DSA layers also fire ATTN while KDA layers do not. Exchanged in the
* hello; both sides must agree. */
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];
} ds4_tp_identity;
bool ds4_tp_enabled(const ds4_tp_options *opt);
typedef enum {
DS4_TP_CLI_ERROR = -1,
DS4_TP_CLI_NOT_MATCHED = 0,
DS4_TP_CLI_MATCHED = 1,
} ds4_tp_cli_parse_result;
/* CLI parsing, same contract as ds4_dist_parse_cli_arg(): returns 1 when the
* argument was consumed, 0 when not matched, -1 on error (err filled). */
int ds4_tp_parse_cli_arg(
const char *arg,
int *index,
int argc,
char **argv,
ds4_tp_options *opt,
char *err,
size_t errlen);
int ds4_tp_adopt_distributed_options(
ds4_tp_options *tp,
ds4_distributed_options *dist,
char *err,
size_t errlen);
void ds4_tp_usage(FILE *fp);
/* Validates option combinations that TP cannot run with (SSD streaming,
* distributed mode, MTP drafting, CPU backend). */
int ds4_tp_validate_engine_options(
const ds4_engine_options *opt,
char *err,
size_t errlen);
/* Connection bring-up. The leader listens and accepts one worker; the
* worker dials with retry. Both then exchange and validate identities.
* Blocking; call after the engine is loaded (identity needs the shape). */
int ds4_tp_create(
ds4_tp **out,
const ds4_tp_options *opt,
const ds4_tp_identity *id,
char *err,
size_t errlen);
void ds4_tp_free(ds4_tp *tp);
int ds4_tp_rank(const ds4_tp *tp);
bool ds4_tp_is_rdma(const ds4_tp *tp);
/* True once both peers agreed on the RDMA command mailbox (RC queue pairs,
* not disabled with DS4_TP_DISABLE_RDMA_MAILBOX on either rank). */
bool ds4_tp_command_mailbox_active(const ds4_tp *tp);
uint32_t ds4_tp_peer_ctx(const ds4_tp *tp);
enum {
DS4_TP_MODE_EXPERT_INTERMEDIATE_SPLIT = 1u,
DS4_TP_MODE_V41_SHARED_OWNER = 2u,
DS4_TP_MODE_V41_SHARED_BALANCE = 4u,
DS4_TP_MODE_V41_DSPARK = 8u,
};
/* Call on both engines before binding the GPU slab. Resident layouts and
* shared-expert ownership must match before either rank can reduce results.
* full_experts is a local capability on entry, the intersection on return. */
int ds4_tp_agree_execution_mode(ds4_tp *tp, uint32_t flags, bool *full_experts, char *err, size_t errlen);
bool ds4_tp_failed(const ds4_tp *tp);
void ds4_tp_mark_failed(ds4_tp *tp);
/* Gate slab. The engine allocates one shared GPU-visible block and hands
* its base VA here; ds4_tp registers it with the NIC (RDMA) and exchanges
* remote keys. Layout, all offsets from base, S = n_layer * 2 slots:
*
* out vectors S * vec_bytes written by local GPU kernels
* in vectors S * vec_bytes RDMA/TCP-written with the peer partials
* in seq flags S * 8 written strictly after each in vector
* token slot 16 {seq u64, token i32, pad} leader->worker
* (gpu flags, then batch out/in:
* n_layer * BATCH_GATES_PER_LAYER * BATCH_MAX_ROWS * vec_bytes
* each, row partials for the speculative verify-block gates)
*
* vec_bytes = n_embd * 4 (f32 partials, never quantized on the wire). */
uint64_t ds4_tp_slab_bytes(uint32_t n_layer, uint32_t n_embd);
uint64_t ds4_tp_slab_out_offset(const ds4_tp *tp, uint32_t layer, uint32_t gate);
uint64_t ds4_tp_slab_in_offset(const ds4_tp *tp, uint32_t layer, uint32_t gate);
uint64_t ds4_tp_slab_batch_out_offset(const ds4_tp *tp, uint32_t layer);
uint64_t ds4_tp_slab_batch_in_offset(const ds4_tp *tp, uint32_t layer);
uint64_t ds4_tp_slab_gpu_flags_offset(const ds4_tp *tp);
int ds4_tp_attach_slab(ds4_tp *tp, void *base, char *err, size_t errlen);
/* Stop the data plane before freeing its registered buffers. No more gates
* may be issued; free the transport after the engine has been unbound. */
void ds4_tp_detach_slab(ds4_tp *tp);
/* Exchange one gate: send out[layer][gate] to the peer's in[layer][gate]
* and wait until the peer's partial for `seq` has fully landed locally.
* Called from the GPU gate service thread. Returns 0 on failure. */
int ds4_tp_gate_exchange(ds4_tp *tp, uint32_t layer, uint32_t gate, uint64_t seq);
/* Verify-block batch gate: exchange `rows` row partials for one layer in one
* bulk RDMA transfer, with a symmetric TCP transfer as fallback. Called from
* the GPU gate service thread. */
int ds4_tp_batch_gate_exchange(ds4_tp *tp, uint32_t layer, uint32_t rows,
uint64_t seq);
/* Verify-block RDMA window (speculative decoding): call on both ranks right
* before/after a verify block with one batch gate per layer. */
int ds4_tp_batch_block_begin(ds4_tp *tp, uint32_t rows, uint32_t n_layers);
/* Negotiate one or two batch gates per layer before encoding the block.
* Both ranks receive the smaller offered count. Slots are then consecutive
* gate indices, not layer indices. Works with RDMA and TCP. */
int ds4_tp_batch_block_begin_gates(ds4_tp *tp, uint32_t rows, uint32_t n_layers,
uint32_t *gates_per_layer);
int ds4_tp_batch_block_end(ds4_tp *tp);
/* Prefill batch gate: arbitrary-size symmetric payload exchange over bulk
* RDMA, with interleaved 2MB TCP rounds as fallback (see ds4_tp.c). */
int ds4_tp_big_gate_exchange(ds4_tp *tp, uint32_t layer, uint64_t seq,
const void *out, void *in, uint64_t bytes);
/* Lockstep mirroring (leader side) and worker loop primitives. */
typedef struct {
uint64_t session_id;
int32_t token;
uint32_t reserved;
} ds4_tp_batch_item;
int ds4_tp_send_session_create(ds4_tp *tp, uint64_t session_id, int ctx_size);
int ds4_tp_send_session_destroy(ds4_tp *tp, uint64_t session_id);
int ds4_tp_send_sync(ds4_tp *tp, uint64_t session_id,
const int *tokens, uint32_t n_tokens);
int ds4_tp_send_sync_multimodal(ds4_tp *tp, uint64_t session_id,
const int *tokens, uint32_t n_tokens,
const ds4_vision_span *images,
uint32_t image_count);
int ds4_tp_send_eval(ds4_tp *tp, uint64_t session_id,
uint64_t seq, int token);
int ds4_tp_send_glm_mtp(ds4_tp *tp, uint64_t session_id,
uint64_t seq, int token, int limit);
int ds4_tp_send_rewind(ds4_tp *tp, uint64_t session_id, int pos);
/* State-only restore: ACK status 0 restored, 1 unavailable, -1 failed (closes).
* No TP collectives or logits reply. Await ds4_tp_wait_command_status() before
* mirroring replay EVALs; status 1 requires invalidation and full sync. */
int ds4_tp_send_spec_restore(ds4_tp *tp, uint64_t session_id, int pos);
int ds4_tp_send_invalidate(ds4_tp *tp, uint64_t session_id);
int ds4_tp_send_eval_batch(ds4_tp *tp, const ds4_tp_batch_item *items,
uint32_t count);
int ds4_tp_send_mixed_batch(ds4_tp *tp, uint64_t prefill_session_id,
const int *prompt, uint32_t prompt_count,
const ds4_tp_batch_item *items,
uint32_t count);
int ds4_tp_send_command_ack(ds4_tp *tp, uint64_t session_id, int status);
int ds4_tp_wait_command_ack(ds4_tp *tp, uint64_t session_id,
const char *operation, char *err, size_t errlen);
int ds4_tp_wait_command_status(ds4_tp *tp, uint64_t session_id, int *status,
const char *operation, char *err, size_t errlen);
/* Both ranks call at matching prefill boundaries. Cancellation is agreed
* here, never sampled independently while either rank is inside a GPU gate. */
int ds4_tp_sync_checkpoint(ds4_tp *tp, uint32_t point, int current, int total,
bool requested, bool *cancelled);
int ds4_tp_send_stop(ds4_tp *tp);
/* Worker: blocks for the next mirrored command. Frame types below; for
* DS4_TP_FRAME_SYNC the token array is returned in *tokens / *n_tokens
* (malloc'd, caller frees), for DS4_TP_FRAME_EVAL seq/token are filled. */
typedef enum {
DS4_TP_FRAME_ERROR = -1,
DS4_TP_FRAME_SYNC = 1,
DS4_TP_FRAME_EVAL = 2,
DS4_TP_FRAME_REWIND = 3,
DS4_TP_FRAME_INVALIDATE = 4,
DS4_TP_FRAME_STOP = 5,
DS4_TP_FRAME_HASH = 6,
DS4_TP_FRAME_RDMA_INFO = 7,
DS4_TP_FRAME_SYNC_ACK = 8,
DS4_TP_FRAME_RDMA_READY = 9,
DS4_TP_FRAME_LOGITS = 10,
DS4_TP_FRAME_VERIFY = 11,
DS4_TP_FRAME_VERIFY_COMMIT = 12,
DS4_TP_FRAME_SESSION_CREATE = 13,
DS4_TP_FRAME_SESSION_DESTROY = 14,
DS4_TP_FRAME_EVAL_BATCH = 15,
DS4_TP_FRAME_MIXED_BATCH = 16,
DS4_TP_FRAME_COMMAND_ACK = 17,
DS4_TP_FRAME_SYNC_MULTIMODAL = 18,
DS4_TP_FRAME_RDMA_WARM = 19,
DS4_TP_FRAME_RDMA_POSTED = 20,
DS4_TP_FRAME_GLM_MTP = 21,
DS4_TP_FRAME_SYNC_CHECKPOINT = 22,
DS4_TP_FRAME_DSPARK_DRAFT = 23,
DS4_TP_FRAME_SPEC_RESTORE = 24,
} ds4_tp_frame_type;
typedef struct {
ds4_tp_frame_type type;
uint64_t session_id;
uint64_t seq;
int value;
int limit;
int *tokens;
uint32_t n_tokens;
ds4_tp_batch_item *items;
uint32_t n_items;
ds4_vision_span *images;
uint32_t n_images;
} ds4_tp_command;
int ds4_tp_recv_command(
ds4_tp *tp,
ds4_tp_command *command,
char *err,
size_t errlen);
void ds4_tp_command_free(ds4_tp_command *command);
/* Debug lockstep check: both sides send their hidden-state hash for a token
* and compare. Returns 0 on transport failure, -1 on hash mismatch. */
int ds4_tp_hash_check(ds4_tp *tp, uint64_t seq, uint64_t hash, char *err, size_t errlen);
/* Vocab-split output head: the worker ships its logits half to the leader
* after every eval (and after a sync) on the control socket. */
int ds4_tp_send_logits_half(ds4_tp *tp, const float *half, uint32_t count);
int ds4_tp_recv_logits_half(ds4_tp *tp, float *half, uint32_t count);
/* Speculative verify mirroring. The leader announces a draft block right
* before both ranks run the expert-split batch verify; the worker then blocks
* on the commit frame. A full commit keeps the entire verified block, a prefix
* commit restores the matching verifier prefix on both ranks, and rollback
* restores the original frontier before replaying token_count tokens. */
typedef enum {
DS4_TP_VERIFY_ROLLBACK_REPLAY = 0,
DS4_TP_VERIFY_COMMIT_FULL = 1,
DS4_TP_VERIFY_COMMIT_PREFIX = 2,
} ds4_tp_verify_commit_mode;
int ds4_tp_send_verify(ds4_tp *tp, uint64_t session_id,
const int *drafts, uint32_t n);
/* VERIFY flags. DS4_TP_VERIFY_HEAD_SPLIT: greedy block, each rank computes
* the verify head on its vocabulary half and the ranks combine row top-1s
* (ds4_tp_small_exchange); the worker then sends its half of the one row
* the leader commits (ds4_tp_head_row_send/recv). */
#define DS4_TP_VERIFY_HEAD_SPLIT 1u
/* DS4_TP_VERIFY_DSPARK_TP: the block holds at least one row of the last
* lockstep DRAFT's proposal, which the worker also holds and checks (a block
* trimmed to its seed is unmarked); DS4_TP_VERIFY_DSPARK_SEED: the block
* starts with the seed, so the proposal rows are drafts[1..]. */
#define DS4_TP_VERIFY_DSPARK_TP 2u
#define DS4_TP_VERIFY_DSPARK_SEED 4u
#define DS4_TP_VERIFY_FLAGS \
(DS4_TP_VERIFY_HEAD_SPLIT | DS4_TP_VERIFY_DSPARK_TP | DS4_TP_VERIFY_DSPARK_SEED)
int ds4_tp_send_verify_flags(ds4_tp *tp, uint64_t session_id,
const int *drafts, uint32_t n, uint32_t flags);
/* Lockstep TP drafter. The leader sends DSPARK_DRAFT where it would prepare
* a DSpark proposal. DRAFT: both ranks run the proposal for (token, pos) on
* their halves of the drafter. MAINTAIN: both only bring the drafter history
* up to pos; the leader may still draft on its own with the full drafter.
* The worker receives the token in value, pos in limit and the mode in seq. */
enum {
DS4_TP_DSPARK_MAINTAIN = 1,
DS4_TP_DSPARK_DRAFT = 2,
};
int ds4_tp_send_dspark_draft(ds4_tp *tp, uint64_t session_id, int token,
int pos, uint32_t mode);
/* Startup agreement, after the execution mode: enabled only when both ranks
* opted in with byte-identical metadata. When only one opted in, both keep
* the leader-local drafter (*agreed false, *peer_on reports the peer);
* different metadata fails. */
#define DS4_TP_DSPARK_META_BYTES 64u
int ds4_tp_agree_dspark_split(ds4_tp *tp, bool mine_on, const void *meta,
bool *agreed, bool *peer_on,
char *err, size_t errlen);
/* Lockstep symmetric exchange of up to DS4_TP_SMALL_EXCHANGE_MAX bytes over
* polled slab records (RDMA command mailbox only; 0 without it). Both
* ranks must pass the same byte count, which the record carries and the
* receiver checks. Failure marks the context failed. */
#define DS4_TP_SMALL_EXCHANGE_MAX 192u
int ds4_tp_small_exchange(ds4_tp *tp, const void *mine, void *peer, uint32_t bytes);
/* Worker -> leader float row of up to DS4_TP_HEAD_ROW_MAX floats (mailbox
* only), sent after the current verify block's begin. The worker writes the
* row into ds4_tp_head_row_out() (host memory inside the registered slab),
* then sends; the leader receives it into dst. The record carries the
* verify block, row and float count and the leader requires all three, so
* a row sent for a block the leader then rolled back without reading is
* never mistaken for a later one. */
#define DS4_TP_HEAD_ROW_MAX 131072u
float *ds4_tp_head_row_out(ds4_tp *tp);
int ds4_tp_head_row_send(ds4_tp *tp, uint32_t row, uint32_t floats);
int ds4_tp_head_row_recv(ds4_tp *tp, uint32_t row, float *dst, uint32_t floats);
int ds4_tp_send_verify_commit(ds4_tp *tp, int32_t mode, int32_t token_count);
int ds4_tp_recv_verify_commit(ds4_tp *tp, int32_t *mode, int32_t *token_count);
/* Standalone worker mode entry. Loads nothing itself: the engine is already
* open. */
int ds4_tp_worker_run(ds4_engine *engine, const ds4_tp_options *opt);
#endif