rpc: one tensor read per decode step and asynchronous split copies

[danielhanchen]

What this changes

On a two node layer split the RPC backend was fully synchronous, and that cost two things per
decode step.

Reads. --backend-sampling builds one sampler subgraph per output row, and
llm_graph_context::build_sampling fills four per-row vectors (t_sampled, t_sampled_probs,
t_sampled_logits, t_candidates). llama_context::decode then walks each vector with
copy_tensor_async_rows, which issues one ggml_backend_tensor_get_async per row and per
vector. The RPC backend had get_tensor_async = NULL, so every one of those became its own
synchronous RPC_CMD_GET_TENSOR round trip: four per sequence, about 130 at 32 sequences.
That is why keeping the sampler on the remote device was slower than shipping a megabyte of
logits per row back to the client and sampling on the CPU.

The RPC backend now implements get_tensor_async. It queues the read and drains the queue as
one new RPC_CMD_GET_TENSORS at the next flush point, which is a synchronize, a graph
compute, or any other command that has to keep its place in the wire order. One decode step of
any batch size is one read.

The queued operation carries the serialised tensor, not the pointer: one llama_decode
allocates and resets a graph per ubatch, so by the time a later flush point is reached the
tensor of an earlier ubatch can be gone.

Writes. The hidden state that crosses the split went ggml_backend_synchronize(CUDA0) ->
host malloc -> blocking ggml_backend_tensor_get -> synchronous SET_TENSOR, because
cpy_tensor_async was NULL and ggml_backend_sched_compute_splits falls back to that.

cpy_tensor_async is now implemented for both directions:

• device -> RPC: an asynchronous device to host copy into pinned staging taken from the
  producing device's host buffer type, an event recorded on the producing stream, and the
  SET_TENSOR sent by the dispatcher once that event has completed. The producing device is
  never fully synchronized and the host thread serialises the graph while the copy is still in
  flight.
• RPC -> device: read into pinned staging and hand it to set_tensor_async on the consuming
  stream, so the consuming device is not synchronized either.

ggml_backend_tensor_copy_async and the scheduler now ask the source backend for the copy
when the destination declines. Only the destination used to be asked, which left the
RPC -> device direction on the synchronous fallback even though the RPC backend can accelerate
reads out of itself.

GGML_RPC_STATS=1 prints the client side command counts and bytes (period from
GGML_RPC_STATS_MS, default 5000).

Compatibility

The protocol minor version goes 1 -> 2 and the client only sends the batched read to a server
that reported minor >= 2 in HELLO, so an older server keeps working. TCP is unchanged.
GGML_RPC_NO_BATCHED_GET=1 and GGML_RPC_NO_ASYNC_COPY=1 turn each half off for A/B testing.

RPC commands per decode step

CPU only harness: two local ggml-rpc-server instances on the CPU backend, llama-server with
--device RPC0,RPC1 -sm layer, 8 concurrent, npp 128, ntg 256, counted over a decode-only
window (prefill excluded).

arm	GET_TENSOR	GET_TENSORS	SET_TENSOR	total commands per step
CPU sampling, before	2.00	-	12.96	16.96
CPU sampling, after	1.00	1.00	13.00	17.00
backend sampling, before	34.16	-	36.08	79.92
backend sampling, after	0.96	1.04	38.00	42.00

The read count is now one per step at any batch size. What is left is the per-row SET_TENSOR
that each dist sampler uses for its own four byte uniform input; those are one-way messages,
not round trips, and merging them needs a single uniform input tensor shared by the sampler
chains, which is follow-up work.

Correctness

Greedy equivalence on the CPU two-RPC harness, five prompts, temperature 0, top_k 1,
seed 1234: md5 177dc61e0703eba3bdaf7bf1131f0458 for the default CPU sampling path at
--pipeline-groups 1 and 2, and the same md5 for --backend-sampling on both the base tree
and this branch, so backend sampled greedy is byte-identical to the CPU sampler here.

Numbers on two DGX Sparks

Qwen3.8-27B UD-Q4_K_XL layer split over two DGX Sparks, 32 concurrent, npp 128, ntg 256,
--cache-ram 0, RDMA transport, one clock state (local 2388 to 2394 MHz, peer 2393 to 2398 MHz),
device order CUDA0,RPC0 so the output layer, the logits and the sampler are on the remote node
and only token ids come back. Whole-cell tok/s from a closed-loop client, best of the cells taken
in one window. "base" is this branch with both features switched off, which is the branch base's
behaviour on the same binary.

arm	tok/s	TPOT ms	TTFT s	util local/peer	reads/step	writes/step	cmds/step	bytes/step
base, N=1	96.90	302	7.40	44/45	1.00	7.00	9.00	692 K
base, N=2	134.87	211	7.28	-	1.00	7.00	9.00	1031 K
async copy only, N=1	96.47	305	6.79	45/45	1.00	6.88	15.00	690 K
async copy only, N=2	138.09	205	6.31	83/79	1.00	7.00	9.00	1007 K
backend sampling, N=1	101.35	288	6.91	47/45	1.00 batched	83.14	124.43	1001 K
backend sampling, N=2	93.68	314	6.47	54/59	1.27 batched	52.82	66.18	476 K

Qwen3.5-4B UD-Q4_K_XL, same split and workload with ntg 128, all ten arms:

arm	32 concurrent	8 concurrent
CUDA0,RPC0 base, N=1 / N=2	322.45 / 462.21	217.97 / 302.06
CUDA0,RPC0 async copy only, N=1 / N=2	317.13 / 444.07	216.39 / 295.98
CUDA0,RPC0 backend sampling, N=1 / N=2	366.65 / 328.24	236.38 / 226.03
RPC0,CUDA0 both features, N=1 / N=2	345.55 / 503.45	229.92 / 306.61
RPC0,CUDA0 async copy only, N=1 / N=2	346.33 / 483.87	227.34 / 308.26

What the numbers say:

• The batched read turns backend sampling from a loss into a win with one context. On the same
  pair and workload it used to cost 16 percent (80.0 against 95.2 tok/s); it is now 4.6 percent
  faster than CPU sampling on the 27B (101.35 against 96.90) and 13.7 percent faster on the 4B
  (366.65 against 322.45), with the logits never crossing the wire.
• It still loses with two contexts, and the reason is in the table: 53 to 83 writes per step.
  Each dist sampler sets its own four byte uniform input, one write per output row per group,
  and each carries a ~380 byte serialised tensor header. The reads are fixed, the writes are not.
  The fix is a single uniform input tensor shared by the sampler chains and viewed per row, which
  needs a small change to the sampler backend interface and is left as follow-up. Until then,
  backend sampling belongs at one context.
• The asynchronous copy on its own is worth nothing with one context (96.47 against 96.90) and
  2.4 percent with two (138.09 against 134.87). It removes a full ggml_backend_synchronize of
  the producing device and a malloc per split, but with one context the step is a serial chain,
  so there is little for the saved time to overlap with. There is no GPUDirect RDMA on GB10 and
  this transport is a send/receive byte stream over its own registered frames, so the path is
  device to pinned host to wire; on GB10 that pinned staging is physically the same memory the
  GPU uses, which is why it costs so little to begin with.

Two hazards this uncovered

Both are in the notes because anything added to this protocol will hit them.

• Never open a connection from an asynchronous path. get_socket connects when the cached
  connection has expired, and an rpc-server serves the connections of a client one at a time,
  so a second connection blocks until the first closes. The asynchronous read and copy take the
  connection the tensor's buffer already holds, and synchronize looks up without creating.
• The RDMA transport is not a byte stream. A receive completion carries exactly one send and
  recv_data copies all of it, so a sent message must be read back in one recv_data of the
  same size. The first version of the batched read scattered the response straight into its
  destinations with one receive per entry. That is correct over TCP and hangs over RDMA as soon
  as a step reads more than one tensor, which is every step with backend sampling, with both ends
  spinning on their completion queues. The response is now received once and scattered in memory.

Impact on non-RPC workloads

The change is in ggml/src/ggml-rpc/ except for 33 lines in ggml/src/ggml-backend.cpp. Nothing
in src/llama-context.cpp, src/llama-sampling.cpp, src/llama-graph.cpp or tools/server/ is
touched; the diff against the branch base is four files:

 .gitignore                     |   1 +
 ggml/include/ggml-rpc.h        |   2 +-
 ggml/src/ggml-backend.cpp      |  33 ++-
 ggml/src/ggml-rpc/ggml-rpc.cpp | 574 ++++++++++++++++++++++-
 ggml/src/ggml-rpc/transport.h  |  29 +++

The ggml-backend.cpp change factors the existing "ask the destination backend for an
asynchronous copy" into a helper that then also asks the source backend, but only when the two
backends have different cpy_tensor_async implementations. Two backends of the same type share
one implementation, which already saw the pair and declined, so on any single-type setup (one GPU,
several GPUs of the same type, CPU only) the helper makes exactly the calls the old code made, in
the same order. The extra call can only happen when two different backend types meet, which on
this tree means an RPC backend on one side.

Evidence:

• cmake -DGGML_RPC=OFF -DGGML_CUDA=OFF configures and builds clean.
• Greedy output with no RPC backend present (five prompts, temperature 0, top_k 1,
  seed 1234): md5 177dc61e0703eba3bdaf7bf1131f0458 on the branch base, on this branch, and on
  this branch built with GGML_RPC=OFF. On a single GPU with the 27B and no --rpc, md5
  74926c4ef135f3cc89ad20cd5ec7e445 on the branch base and on this branch.
• Server test set, pytest -q -m "not slow" run serially: branch base 368 passed, 6 skipped, 199
  deselected; this branch 368 passed, 6 skipped, 199 deselected. Same set, no new failures.
• Single GPU llama-batched-bench bracket base/new/base on the 27B, npp 512, ntg 128, npl 1/8/32,
  S t/s: 55.23 / 54.43 / 53.56, then 231.91 / 228.22 / 211.38, then 367.34 / 334.29 / 336.27. The
  GPU cooled between passes, so the two base passes differ by up to 8.5 percent; the new binary
  sits inside that spread on every row and within 0.6 percent of the second base pass.

Compatibility

The protocol minor version goes 1 to 2 and the client only sends the batched read to a server that
reported minor 2 or higher in HELLO, so an older server keeps working and an older client keeps
working against a new server. TCP is unchanged. GGML_RPC_NO_BATCHED_GET=1 and
GGML_RPC_NO_ASYNC_COPY=1 turn each half off, which is how the arms above were taken on one
binary.