rpc: direct server-to-server tensor transfer for a layer split - #196
rpc: direct server-to-server tensor transfer for a layer split#196danielhanchen wants to merge 11 commits into
Conversation
…ient A layer split over several RPC devices moved every hidden state through the host that runs the scheduler: ggml_backend_sched has no direct path between two RPC buffers on different endpoints, so it fell back to reading the tensor into the client's memory and writing it out again, two transfers and a synchronize per stage boundary. RPC_CMD_COPY_TENSOR_TO tells the server that holds the source tensor to write it into a tensor on another server. The source server opens a connection to the destination with the same HELLO negotiation a client uses (RDMA when both rails allow it, TCP otherwise), pushes the data as an ordinary RPC_CMD_SET_TENSOR so the destination applies its own tensor deserialization and buffer range checks, and waits for RPC_CMD_PEER_BARRIER before answering the client, so the destination cannot compute before the write has landed. Connections to other servers are cached per destination endpoint and closed when the client disconnects. Serving several connections at once is what this needs, so a server now runs one thread per connection over a shared buffer registry and a shared execution mutex; a session still owns and frees only the buffers it allocated. The client uses the command from the RPC backend's cpy_tensor_async when the source and the destination are RPC buffers on different endpoints and both servers report protocol minor 3 or higher; everything else, including two devices of one server, keeps its previous path. GGML_RPC_P2P=0 forces the old path. Protocol minor 2 -> 3, every existing command unchanged.
…/rpc-p2p # Conflicts: # ggml/src/ggml-rpc/ggml-rpc.cpp
Resolve the conflict in ggml/src/ggml-rpc/ggml-rpc.cpp by keeping this branch's peer to peer copy code and the base branch's comment wording. Verified with comment_tools check that only comments changed relative to 6e00db6.
|
@codex security review |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
get_peer_socket opens a connection from inside a server, and get_socket ends in negotiate_hello, whose failed exchange hits RPC_STATUS_ASSERT and aborts the process. A destination that accepts the connection and then disconnects, or returns a truncated HELLO, therefore took down the source server and every client it was serving. That also contradicted what copy_tensor_to documents: an unreachable or too old destination is a recoverable condition reported as result = 0, which puts the client back on its previous path. negotiate_hello and get_socket now take may_fail, set only for peer connections, and return failure instead of aborting. A client keeps aborting as before.
copy_tensor_to resized p2p_buf from hdr.size, which a client controls, before any of the checks against the source tensor and its buffer. A large value aborted the server on an uncaught bad_alloc, and a value near SIZE_MAX wrapped msg_size so that the two header copies ran past the end of an undersized vector. The checks already existed further down. The sizing and the header copies now happen after them, where hdr.size is known to fit both the buffer and the tensor.
rpc_flush_deferred_guarded only puts queued commands on the wire. It does not wait for them to be served, and the peer write arrives on a different connection, so the two are separated by nothing but the destination's execution mutex. An earlier GRAPH_COMPUTE or SET_TENSOR that has not yet taken that mutex could run after the peer's SET_TENSOR and read or overwrite a split input buffer that had just been reused. RPC_CMD_PEER_BARRIER already means that everything received earlier on a connection has been served, so it is now issued on dst_sock before the copy starts. A failure there is recoverable and falls back to routing the tensor through the client.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Raising RPC_PROTO_MINOR_VERSION to 3 locked out every already-deployed client: negotiate_hello rejects any server whose minor exceeds its own, so a 5.2 client could not connect to a 5.3 server for any operation at all, even though the new commands are appended compatibly. The version is back to 2 and support is advertised in the byte of the HELLO response that was previously pure padding. That byte is fixed size, already on the wire, and read as padding by every existing client, so an old client connects exactly as before and simply never sees the flag, while an old server sends zero and is never asked for a peer copy. The two gates that keyed on minor >= 3 now test the flag. Note on why the flag does not live in conn_caps, which would otherwise be the obvious place: update_caps treats any non-zero byte there as "the peer speaks RDMA", and rdma_caps already fills all 24 bytes, so a spare bit would both be unavailable and, if taken, make an RDMA-less peer look RDMA-capable.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Two ways a recoverable peer-copy failure could take down the source server. socket_t::connect() created a descriptor and then returned nullptr from three later failure paths without closing it: TCP_NODELAY, name resolution, and the connect() itself. Until now that cost at most one descriptor per process, because a client that cannot connect aborts anyway. get_peer_socket() reaches the same code once per tensor per boundary crossing, does not cache failed peers, and treats a refused connection as recoverable, so the same leak becomes unbounded and a long decode eventually exhausts the descriptor limit, at which point the server can no longer connect to peers or accept clients. socket_t::impl already closes the descriptor in its destructor, so the fix is to adopt it into the socket_ptr immediately rather than carrying it raw until the last line. accept() and create_server() had the same shape and are fixed the same way; accept()'s is per client connection. Separately, send() was called with no flags. On POSIX, writing to a socket whose peer has gone away raises SIGPIPE, and neither this library nor the rpc-server executable installs a handler, so the default action terminates the process. A destination that was restarted while a source held a cached connection to it is exactly the case that is supposed to end in result = 0; instead it killed the source server and every client it was serving. Sends now pass MSG_NOSIGNAL where it exists, and Apple, which has no MSG_NOSIGNAL, gets SO_NOSIGPIPE set once per descriptor instead. Windows has no SIGPIPE and is unaffected.
Moving the buffer registry into rpc_server_shared so a peer copy could resolve the destination buffer also widened every other command. deserialize_tensor() validated against shared.buffers, the set of all live buffers in the process, so any connection could name a buffer belonging to another connection and then read it with GET_TENSOR, overwrite it with SET_TENSOR, or reference it from a graph. Before the registry moved, buffers belonged to the connection's rpc_server and the same request simply did not resolve. deserialize_tensor() now validates against owned_buffers, this session's own allocations, and takes an allow_foreign flag for the one path that genuinely needs more. That path is SET_TENSOR: the destination of a peer copy receives it on the connection the source server opened rather than on the connection of the client that allocated the buffer, so it has to resolve a buffer another session owns. The write stays bounded by the existing buffer range checks. GET_TENSOR, GRAPH_COMPUTE, memset_tensor and the local copy path go back to being confined to the session, as do buffer_get_base() and buffer_clear(), which were checking process-wide liveness rather than ownership. free_buffer() already checked ownership. owned_buffers is maintained on allocation, on free, and on session teardown, so it is the same lifetime the pre-existing per-connection registry had. This restores the isolation that existed before peer copies for every command except the peer write itself. It is not authentication: the RPC server has none, and a caller that can open a connection can still address the destination buffer of a peer copy. Narrowing that further needs the destination to be told which writes to expect, which is a protocol change and not attempted here.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Two conflicts, both from the two branches having grown the same thing independently. srv_flags bit assignment. Batched get and peer copy were developed separately and each took bit 0. Resolving that by keeping one enum would have shipped a wire-level collision: a server that implements only batched get sets bit 0, and a client that knows about peer copy would read that same bit as peer-copy support and send it RPC_CMD_COPY_TENSOR_TO, which it does not implement. The two features now occupy distinct bits. Batched get keeps bit 0 because that value is already published on this branch; peer copy moves to bit 1. hello() advertises both, since this server implements both. rpc_conn_state::server_flags. Both branches added this member, so the textual auto-merge succeeded while producing a struct with the field declared twice, which does not compile. Kept one declaration alongside server_minor. Verified on CPU rather than by inspection: the merged rpc-server advertises flags 0x03 on the wire, and a llama-server generating over a real RPC socket against it completes normally. Both client-side gates survived the merge and remain keyed to their own flag.
|
Merged the base branch in to clear the conflict (999cdbc). This was not a mechanical merge, so recording what changed and why.
The one worth reviewing: The features now occupy distinct bits. Batched get keeps bit 0, because that value is already published on this branch and anything already built against it would otherwise be wrong; peer copy moves to bit 1. The second: Checked rather than assumed, on CPU with no GPU involved: the merged The head moved, so this needs another look. |
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…able peers, bound connections Three review items, each reproduced against this branch before anything was changed. All three are CPU reproducible, so the numbers below are measured rather than argued. Foreign writes. set_tensor() passed allow_foreign = true for every caller, because the destination of a peer copy is the one path that has to resolve a buffer belonging to another session. Ordinary client connections dispatch the same command, so any client could name any live buffer in the process and write it. Demonstrated with two plain connections to one server: client A allocates a buffer and fills it with 0x11, client B allocates nothing at all and sends SET_TENSOR naming A's buffer, and A reads its own buffer back. Before, 4096 of 4096 bytes were B's 0xAA and none were A's 0x11. After, 4096 of 4096 are still 0x11 and none are 0xAA. The widening now applies only to a link that declared itself with RPC_CMD_PEER_LINK, which a source server sends on the connection it opens to a destination and no ordinary client sends. To be clear about what this is: the RPC protocol has no authentication of any kind, so this is isolation rather than authorization and does not make the port safe to expose. What it changes is that the privilege belongs to the one connection that asked for it instead of to every connection. The control matters more than the fix here, because refusing too much would break peer copies silently: a rejected write surfaces as result = 0 and a fallback, not as an error. Driving RPC_CMD_COPY_TENSOR_TO between two servers, both before and after, gives result = 1 and 4096 of 4096 bytes of the source pattern at the destination. Identical on both sides of the change. Unreachable peers. A destination the coordinator can reach but the source server cannot, which is what asymmetric firewall or NAT rules produce, failed in get_peer_socket() on every single copy with nothing remembered. Where packets are dropped rather than refused that is a full TCP timeout per split boundary, which turns a mild fallback into a stall. Failures are now cached with a 1s, 2s, 4s backoff to a 60s cap, cleared on the first success, so a destination that is briefly down is retried promptly while one that is unreachable by routing stops costing anything measurable. Connection threads. Every accepted socket got a detached thread before HELLO was validated, and that thread can block in recv_data() forever because the transport sets no read timeout. Measured by opening 200 connections that connect and then say nothing: before, 200 accepted, 0 refused, thread count +200. After, 64 accepted, 136 refused, thread count +64, and back to baseline once the clients close. The limit is 64 by default and GGML_RPC_MAX_CONNECTIONS raises it. Not covered here: the backoff is reasoned from the code path rather than measured, because provoking a genuinely unreachable destination needs the two-node setup, and the peer path was exercised only between two servers on one host.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4f54a47d66
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if (src_is_rpc && dst_is_rpc) { | ||
| return ggml_backend_rpc_cpy_tensor_p2p(backend_src, backend_dst, src, dst); |
There was a problem hiding this comment.
When a cross-RPC transfer returns false (for example because the source server cannot reach the destination), ggml_backend_cpy_tensor_async_impl next invokes backend_src->iface.cpy_tensor_from_async; RPC registers this same function for both callbacks, so the entire destination barrier and COPY_TENSOR_TO exchange is attempted a second time before the generic fallback runs. On an unreachable peer this performs two failed direct-copy round trips—and currently two device-to-host staging reads—for every boundary copy. Use separate destination/source callback implementations, or make the source-side callback decline an RPC-to-RPC copy without retrying it.
Useful? React with 👍 / 👎.
| socket_ptr peer = get_peer_socket(endpoint); | ||
| if (peer == nullptr) { | ||
| return true; |
There was a problem hiding this comment.
When server-to-server routing is unavailable, get_peer_socket returns null—including immediately during its cached backoff—but it is called only after the source tensor has already been read into p2p_buf. The callback then returns false and the generic fallback reads the same activation from the source again, so every layer boundary continues paying a redundant device-to-host transfer and host allocation throughout the backoff period. Resolve the peer before staging so the intended fallback remains inexpensive when asymmetric firewall or NAT rules prevent direct connectivity.
Useful? React with 👍 / 👎.
| rpc_msg_peer_barrier_rsp barrier; | ||
| if (!send_rpc_cmd(peer, RPC_CMD_PEER_BARRIER, nullptr, 0, &barrier, sizeof(barrier))) { | ||
| GGML_LOG_ERROR("[%s] %s did not acknowledge the write\n", __func__, endpoint.c_str()); | ||
| peer_socks.erase(endpoint); |
There was a problem hiding this comment.
When multiple coordinator sessions have copied to the same destination, their peer_socks maps share the socket returned by the process-wide weak g_sockets cache. If the destination restarts, erasing only this session's reference leaves the failed socket alive while any other, possibly idle, session still holds it; subsequent retries retrieve that same dead socket from get_socket, fail the peer-link handshake, and can never establish a fresh connection until every other holder releases it. Explicitly invalidate the global cached socket, or give peer connections independent cache ownership, when a send or barrier fails.
Useful? React with 👍 / 👎.
Defect report, not a review of this PRIssues are disabled on this repository, so this is filed here because this PR is where it was What happensAt 32 concurrent streaming clients against Where it has been seenIt is not a property of any one branch, which is why this is an issue rather than a comment on a
Two unrelated branches, two harnesses, the same one-in-sixty-four. Nobody has picked it up. ReproductionAny 27B layer split over two nodes, or any server configuration that sustains 32 concurrent One of the 64 requests ends with zero streamed content chunks. The server's Why it mattersA serving deployment at this concurrency silently drops 1.6 percent of requests with a 200 and an What is not yet knownWhether it is the slot handover, the streaming path, or the HTTP layer. PR #187 also records a |
1 participant
What this does
A layer split over several RPC devices used to move every hidden state through the host that runs
the scheduler.
ggml_backend_sched_compute_splitshas no direct path between two RPC buffers ondifferent endpoints (
ggml_backend_rpc_buffer_cpy_tensorreturns false as soon as the socketsdiffer), so each stage boundary fell back to
ggml_backend_tensor_copy: aGET_TENSORinto ahost allocation on the coordinator, then a
SET_TENSORback out. Two transfers, a synchronize,and the hidden state of every boundary crossing the coordinator's memory.
With this change the coordinator still issues every control command, but for a boundary between
two remote servers it issues one command instead of two and the data goes straight from one
server to the other. P stages form a ring: the coordinator sends the embeddings to stage 1,
stage 1 sends its output to stage 2, and so on back to the device that holds the output layer.
Protocol
Minor version 2 -> 3. Every existing command byte and struct is unchanged, so old clients and
servers interoperate for everything else, and the new commands are only used after the HELLO
version check reports minor >= 3 on both endpoints.
RPC_CMD_COPY_TENSOR_TOis sent to the server that holds the source tensor. Payload is| src rpc_tensor | dst rpc_tensor | size | endpoint_len | destination endpoint |. The sourceserver validates the source region exactly as
RPC_CMD_GET_TENSORdoes (deserialize_tensorplus the buffer range check), reads it, opens or reuses a connection to the destination
endpoint using the same connect and HELLO negotiation a client uses (so RDMA when both rails
allow it and TCP otherwise), and pushes the data as an ordinary
RPC_CMD_SET_TENSOR. Nothingabout the destination pointer is trusted on the source side: the destination applies its own
deserialize_tensorand buffer range checks, the same ones it applies to a client.RPC_CMD_PEER_BARRIERis an empty request with a one byte response. A connection is servedstrictly in order, so its response proves the
SET_TENSORthat preceded it has completed.The source server waits for it before answering the client, which is the ordering guarantee:
the destination cannot compute before the write has landed, because the scheduler only sends
RPC_CMD_GRAPH_COMPUTEto the destination afterRPC_CMD_COPY_TENSOR_TOhas returned.Connections to other servers are cached per destination endpoint and closed when the client that
asked for them disconnects. Recoverable failures (destination unreachable, destination too old,
write rejected) come back as
result = 0and the client falls back to the previous path; only amalformed request or an out of bounds source closes the connection.
Serving several connections at once is what this needs, so a server now runs one thread per
connection over a shared buffer registry and a shared execution mutex; a session still owns and
frees only the buffers it allocated. The execution mutex is released before the peer connection
is used, so a ring of servers cannot deadlock on each other.
Client side, the RPC backend's
cpy_tensor_asynctakes the RPC to RPC case when the twoendpoints differ. Two devices of one server keep the server-local
RPC_CMD_COPY_TENSOR, andevery other combination returns false so the previous path runs unchanged.
GGML_RPC_P2P=0forces the old path for A/B.
Commands and bytes per decode step
Client side counters (
GGML_RPC_STATS=1), decode window, Qwen3.8-27B UD-Q4_K_XL, 32 rows,three stages (
--device RPC0,RPC1,CUDA0 -sm layer --tensor-split 1,1,1), so one RPC to RPCboundary per step. The hidden state is 32 x 5120 x 4 B = 655360 B.
One
GET_TENSORplus oneSET_TENSORbecome oneCOPY_TENSOR_TOper boundary, and 1.31 MB perstep stops crossing the coordinator: 2.69 MB of RPC traffic per step becomes 1.38 MB, a 49
percent reduction. For P stages the saving is (P - 2) boundaries per step, so it grows with the
number of nodes while the coordinator's work stays flat.
CPU-only harness (
tests/cpu/rpc_p2p.sh, three local rpc-servers on the CPU backend,--device RPC0,RPC1,RPC2 -sm layer, the process host doing the embedding lookup), single stream,exact per-step counts over the decode window:
Two RPC to RPC boundaries, two
COPY_TENSOR_TO, and the SET_TENSOR payload halves.Correctness
Greedy output over five prompts, 48 tokens, temperature 0, top_k 1, on the CPU-only harness:
md5
177dc61e0703eba3bdaf7bf1131f0458for all four arms - three stage and two stage,GGML_RPC_P2P=0andGGML_RPC_P2P=1- which is the same md5 the existing RPC harnesses record.tools/server/tests: 368 passed, 6 skipped, 199 deselected.Numbers on two GB10 nodes
Qwen3.8-27B UD-Q4_K_XL, 32 concurrent closed loop clients, npp 128 / ntg 256, 64 requests,
-c 16384(16386 for the three group cells),--cache-ram 0,-fa on,-t 6. Local CUDA0 holdsthe last third and the output layer; two
ggml-rpc-serverprocesses on the peer (ports 50052 and50053, one GPU between them) hold the first two thirds,
--tensor-split 1,1,1.Two stage reference on the same binaries,
--device RPC0,CUDA0, one rpc-server on the peer:Reading these honestly: throughput does not move. The peer's mean SM clock falls monotonically
through each round (2466, 1856, 1690 MHz), so the first cell of a round is always the fastest one,
and the ranking of the two arms flips when the order is reversed: hub is ahead by 16 percent when
it runs first (round a) and behind by 8 percent when the direct arm runs first (round b). Both
arms are inside that band, so this pair and this topology do not separate them on tok/s. That is
expected here, because the two remote stages share one GPU on one node: the boundary the change
removes from the coordinator becomes a loopback transfer on the peer, and nothing is freed on the
one link that matters. What is unambiguous is the protocol accounting above, which is what scales:
for P stages on P nodes the change removes (P - 2) hidden state round trips per step from the
coordinator's link and memory, so the coordinator stops being a per-boundary relay and the wire
cost of a stage boundary no longer depends on where the scheduler runs.
Also visible in the table: three stages over the same two GPUs is worse than two stages
(143.2 tok/s at two stages with two groups against 77 to 88 at three stages with three groups),
so a third stage should only be added with a third node.
Known limitation, not fixed here:
RPC_CMD_COPY_TENSOR_TOblocks the source server's connectionthread for the whole transfer plus the barrier, so with several pipeline groups over one
connection the groups serialise behind it. Moving the transfer to a per-destination worker and
letting the destination enforce the ordering with a sequence number, instead of the client
waiting for the acknowledgement, would remove that; it needs a second command on the destination
and another measurement window.
Non-RPC workloads
Every line of this change is in
ggml/src/ggml-rpc/:ggml-backend.cpp, the scheduler, and the CUDA, Metal, Vulkan, CPU and HIP backends are nottouched, so a build with
-DGGML_RPC=OFFcontains none of it andlibggml-baseandlibllamagain no new behaviour when the RPC backend is not loaded. Checks on one GB10:
-DGGML_RPC=OFF -DGGML_CUDA=ONconfigures and builds clean.test-backend-ops -b CUDA0: 13572/13572 tests passed, OK, on this branch and on the base, samecount.
Greedy generation on one GPU with no
--rpc, three prompts, 64 tokens, temperature 0, top_k 1:md5
1f99f8da109f7a456073901232e1f014on this branch and on the base, byte identical.llama-batched-benchon one GPU with no--rpc, npp 512, ntg 128, npl 1 / 8 / 32, run as abase / new / base bracket in one window (whole run tok/s):
The first run of the bracket is the fastest in every row (a cold GPU at the start of the
window); the new binary matches the second base run to within 0.2 percent at npl 1 and 8 and
sits between the two base runs at npl 32.
Scope: this does nothing on a two node pair
Worth stating plainly, because the byte table above is easy to read as a decode win and it is not one here. The saving is (P - 2) boundaries per step. On the two DGX Spark pair P = 2, so there is no RPC to RPC boundary at all and this change is inert: the one crossing goes to the device holding the output layer, which is the coordinator's own CUDA0. The three stage cells above were run with two rpc-servers sharing one GPU on the peer precisely because that is the only way to make a P = 3 boundary exist on two machines, and they show three stages over two GPUs is worse than two stages regardless of the arm.
There is also now a measurement that bounds what any interconnect change on the pair can be worth. A two node layer split moves 6.84 MB per decode step in total, which is 0.34 ms of a 220 ms step against the measured 20.31 GB/s NCCL bus bandwidth, under 0.005x of end to end speedup. Removing 1.31 MB of that would be worth under 0.001x even if the boundary existed. So this PR should be read as a three or more node capacity and scaling change, plus the per connection threaded rpc-server it introduces, and not as a throughput change for the pair. The honest reading of the tok/s table stands as written: the arms do not separate.