Skip to content

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

Draft
danielhanchen wants to merge 20 commits into
masterfrom
feature/rpc-stage2
Draft

rpc: one tensor read per decode step and asynchronous split copies#193
danielhanchen wants to merge 20 commits into
masterfrom
feature/rpc-stage2

Conversation

@danielhanchen

@danielhanchen danielhanchen commented Sep 5, 2026

Copy link
Copy Markdown
Member

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

This PR targets master, and against master it is ten files, not four. An earlier version of
this section counted only the diff from the point in the pipeline-groups work that this branch was
cut from, which is not what a reviewer sees. Both numbers are given below, because the RPC-only
property is real but it is a property of the increment, not of the PR.

Against master, which is the base:

 .gitignore                      |   1 +
 ggml/include/ggml-rpc.h         |   2 +-
 ggml/src/ggml-backend-impl.h    |  14 +-
 ggml/src/ggml-backend.cpp       |  29 +-
 ggml/src/ggml-rpc/ggml-rpc.cpp  | 621 ++++++++++++++++++++++++++++++++--
 ggml/src/ggml-rpc/transport.h   |  39 +++
 tools/server/README.md          |  34 ++
 tools/server/server-context.cpp | 702 +++++++++++++++++++++++++++++------
 tools/server/server-context.h   |   3 +
 tools/server/server.cpp         |  44 +++
 10 files changed, 1305 insertions(+), 184 deletions(-)

The four tools/server/ files are the first pipeline-groups commits, which this branch is stacked
on and which are reviewed in #187. They are not part of what this PR proposes.

What this PR itself adds, on top of that stack, is confined to ggml/src/ggml-rpc/ apart from
ggml/src/ggml-backend.cpp and ggml/src/ggml-backend-impl.h, which gain the source-side async
copy entry described below. Nothing in src/llama-context.cpp, src/llama-sampling.cpp or
src/llama-graph.cpp is touched in either accounting.

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.

Hardening

Review found a defect in the GET_TENSORS handler that was reachable before this fix, and it is
worth stating precisely rather than as "unchecked allocation", because the allocation failure is
the milder half.

The handler summed the per-entry sizes off the wire and resized the response to that sum before any
validation ran. A single oversized entry therefore threw an uncaught std::bad_alloc and
terminated the server. But the sum accumulates uint64_t sizes into a size_t, so it can wrap: a
wrapped total allocates a small response, and the copy loop then still calls
ggml_backend_tensor_get for the full entries[i].size at out_offset. That is a heap overflow
write in a command handler on a network-facing server, not a clean failure.

The existing per-entry bounds check does not catch it, and the reason is the part worth
remembering: it constrains each region against its own source buffer and never against the
response. It looks like protection against exactly this and provides none.

Validation now runs in a pass before the allocation, and the total is accumulated with a checked
addition against a 4 GiB ceiling. The ceiling is deliberately kept even though validation now
precedes allocation: at that point each region is already bounded by a really allocated buffer, so
the ceiling is not what prevents the allocation failure. It prevents several individually valid
entries naming the same large buffer from summing into the terabytes on a machine with a large
unified pool, and it keeps the addition in a range where it cannot wrap. 4 GiB is orders of
magnitude above the activations and single tensor regions this path batches.

This should be read alongside the fact that the RPC server has no authentication of any kind. None
of the changes here add any, and none should be read as adding any.

Compatibility

An earlier revision of this branch bumped the protocol minor version from 1 to 2 and gated the
batched read on server_minor >= 2. That was a compatibility break, and in the direction that
matters most: negotiate_hello() rejects any server whose minor exceeds the client's own, so an
unchanged 5.1 client could not connect to this server at all. Upgrading the server alone is the
normal way this gets deployed, so the claim that older clients kept working was exactly backwards.

The general lesson, since this is the second branch where it has come up: a version bump carries two
meanings at once, "the wire format changed" and "this server has a new optional command", and only
the second was ever true here. Conflating them turns a purely additive feature into a compatibility
break, because an old client that would never send the new command is refused anyway. An optional
capability belongs in a capability flag in already-reserved space, not in the version.

So the minor stays at 1, and support is advertised as RPC_SRV_FLAG_BATCHED_GET in the byte of
rpc_msg_hello_rsp that was previously pure padding. That byte is already on the wire, is fixed
size, and is read as padding by existing clients, which see zero and simply never take the batched
path. rpc_supports_batched_get() tests the flag rather than the version. The request side could
not carry this: rpc_msg_hello_req is a fixed-size struct read as such by the server, so growing it
would lock out old clients in the same way the version bump did.

An older server keeps working, and an older client now genuinely 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.

A layer split over two nodes is a two-stage pipeline that a single
llama_context feeds one batch at a time, so each stage sits idle while
the other one computes. With --pipeline-groups N the server creates N
llama_contexts from the one model, partitions its slots between them and
gives each group its own batch and its own decode thread, so there are N
batches in flight and both stages have work.

Each context is created with n_seq_max = n_parallel / N and
n_ctx = n_ctx / N, so the per-slot context and the total KV memory are
unchanged. Slots are partitioned contiguously and carry the sequence id
they use inside their own context. Slot selection for a new task still
runs over all slots, so prompt cache similarity, the slot endpoints and
the KV prefix reuse behave exactly as before.

The model weights, the task queue, the results queue and the HTTP layer
are shared. Task processing pauses the decode loops for the moment it
looks at the slots. Speculative decoding, multimodal and idle sleeping
are refused with N > 1 rather than half supported.

With the default N = 1 there is one context, one batch and one update
loop on the main thread, no locks and no extra threads.
- the unlock around llama_decode is now RAII, so a throwing decode cannot
  leave a group marked busy (which would wedge every later task) nor
  return to the error handler without the engine lock
- n_cmpl is rejected when it exceeds the slots of one group, instead of
  being deferred forever: the child slots take their KV from the parent,
  so they have to live in the parent's context
- refuse --control-vector with more than one group, common_init_from_params
  only applies it to the context it creates
- the queued prompt stats and the empty batch kill switch move into the
  group, they were shared counters flushed per group
- post_decode uses the group's context, and the detokenize calls in the
  result path use the slot's own context
- free the contexts already created if a later one fails, and do not
  index groups[0] when no model is loaded
…onnection

One socket is cached per endpoint and is therefore shared by every backend of that endpoint,
including the backends of different llama_contexts. A message is written as three unlocked
send_data calls, so two contexts interleave their command streams and the server sees a
malformed request within seconds. Make a whole message atomic on the wire, and hand the
responses out in request order with a ticket, so a thread waiting for its response does not
hold the send lock and the other contexts can keep submitting.

last_graph_uid was kept per endpoint device while the graph it refers to is stored by the
server per connection, and it was read and written without a lock, so two contexts on one
connection could make RPC_CMD_GRAPH_RECOMPUTE re-run the other one's graph. Track it per
connection and device and check it under the send lock.

server: pause only the group that owns the slot a task touches

process_single_task stopped every pipeline group for every task and waited for all the
in-flight decodes. Holding the engine is already enough to keep the slot state stable, so
only wait for the group whose context the task touches: the owning group for completions,
cancel, control and slot save / restore / erase, every group for --cache-idle-slots and
SET_LORA, none for metrics, /slots and get-lora.
…chronously

The RPC backend was fully synchronous. Two consequences on a two node layer
split: a decode step with backend sampling read four small tensors per
sequence, one round trip each (about 130 at 32 sequences), and the hidden
state that crosses the split was staged through a host malloc after a full
synchronize of the producing device.

- get_tensor_async queues the read and the queue is drained as a single
  RPC_CMD_GET_TENSORS at the next flush point, so a step reads once.
- cpy_tensor_async takes the device to device copies. Device to RPC copies
  into pinned staging on the producing stream, records an event and sends
  from the staging once the event completed, without synchronizing the
  producing device. RPC to device reads into pinned staging and hands it to
  an asynchronous host to device copy on the consuming stream.
- ggml_backend now asks the source backend for an asynchronous copy when the
  destination cannot take it, which is what the RPC to device direction needs.
- GGML_RPC_STATS=1 prints the client side command counts and bytes.

Tensors are serialized when an operation is queued, not when it is flushed:
one llama_decode allocates and resets a graph per ubatch, so the pointer can
be gone by the next flush point.
The asynchronous entry points of a backend assert that the tensor lives in
that backend's default buffer type, so a tensor in a host buffer that the
same backend can also reach has to stay on the synchronous path.
…end declined

Two backends of the same type share one cpy_tensor_async, so asking the
source after the destination of the same type declined would just repeat the
same call with the same answer. Comparing the implementations keeps every
single-type setup, one GPU or several, on exactly the previous path.
A server serves the connections of a client one at a time, so opening a
second connection to an endpoint that already has a live one blocks until
the first closes. synchronize now looks the connection up without creating
one (with nothing connected there is nothing queued either), and the
asynchronous read and copy take the connection the tensor's buffer already
holds instead of asking for one by endpoint.
The RDMA transport is not a byte stream: a receive completion carries
exactly one send and recv_data copies all of it, so reading one sent message
back in several pieces overruns the first destination and then blocks for a
completion that never arrives. The batched read took the response apart with
one recv_data per entry, which worked over TCP and hung over RDMA as soon as
a step read more than one tensor, which is every step with backend sampling.
Receive the response once and scatter it in memory.
The measurement scripts write their cells and their nvidia-smi samples into
the worktree, and a blanket add swept them in. They are results, not source.
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@danielhanchen

Copy link
Copy Markdown
Member Author

@codex security review

@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Sep 9, 2026

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review Completed 2026-09-09T13:44:30.411567Z 6e2cbe7 Manual request
ℹ️ 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" or "@codex security review".

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 chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: ec74a6002f

ℹ️ 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".

Comment thread ggml/src/ggml-backend.cpp Outdated
Comment on lines +510 to +512
if (backend_src != NULL && backend_src->iface.cpy_tensor_async != NULL &&
(backend_dst == NULL || backend_src->iface.cpy_tensor_async != backend_dst->iface.cpy_tensor_async)) {
if (backend_src->iface.cpy_tensor_async(backend_src, backend_dst, src, dst)) {

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Restrict source-side dispatch to compatible backends

When the destination declines or has no async-copy callback, this now invokes every source backend's callback even though existing implementations were written as destination-side handlers. For example, a Vulkan-to-CPU copy reaches ggml_backend_vk_cpy_tensor_async, which interprets backend_dst->context as a Vulkan context before safely rejecting the destination; this can crash instead of using the synchronous fallback. Only dispatch source-side to implementations that explicitly support that role, or update each implementation to validate both backend types first.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 313019c. Agreed, and it is worse than a bad cast in one case: the dispatch guard also allowed backend_dst to be NULL, and ggml_backend_vk_cpy_tensor_async casts backend_dst->context on its first line, so that combination is a null dereference.

Rather than teach each implementation to validate both roles, the source side now has its own entry, cpy_tensor_from_async, appended to the end of ggml_backend_i so backends that initialise the members positionally are untouched. Only the RPC backend implements it, and its handler already checks ggml_backend_is_rpc on both sides, which is null-safe.

Comment thread ggml/src/ggml-rpc/ggml-rpc.cpp Outdated
Comment on lines +1152 to +1155
ggml_backend_event_t event = rpc_staging_event(st, other_dev);
if (event != nullptr) {
ggml_backend_event_record(event, backend_dst);
st.inflight = event;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Reuse events on RPC-to-local copies

On the RPC-to-local async-copy path, every copy increments events_used here, but that counter is reset only after rpc_flush_deferred() processes a nonempty deferred queue. The blocking ggml_backend_tensor_get() immediately above normally has no deferred operation, so CUDA/Vulkan split inference creates a new backend event per copied tensor and never frees it; a long-running server will continuously consume event and driver resources until allocation fails or the process becomes unstable.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in bb0e9ac. Confirmed the path you describe: rpc_flush_deferred_guarded returns early on an empty deferred queue, so the reset inside rpc_flush_deferred never ran for the RPC-to-local copy.

The reset now happens at the staging buffer wrap in rpc_staging_alloc instead. That is where the last recorded event has just been synchronized, and events go onto one backend in order, so the earlier ones are complete as well. It also bounds the pool by the number of copies per staging generation rather than by uptime.

break;
}

guard.wait_for(slot->id_group);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Wait for the group before probing its prompt cache

With multiple pipeline groups, this wait occurs only after get_available_slot(task), but that function may call prompt_save() and prompt_load(), which use llama_state_seq_* on the selected slot's context. If another slot in that group is currently inside llama_decode(), a new completion with prompt caching enabled can therefore access the same llama_context concurrently, risking corrupted state or a crash. Select the candidate without touching its context, wait for its group, and only then perform the cache save/load.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 433ef71. Confirmed: get_available_slot calls prompt_save and prompt_load, which use llama_state_seq_* on the selected slot context, and it ran before guard.wait_for(slot->id_group).

Done the way you suggest. Selection no longer touches the context, it only reports whether an update is due, and update_prompt_cache runs after the group has been waited for.

@unslothai unslothai deleted a comment from chatgpt-codex-connector Bot Sep 9, 2026
events_used was reset only inside rpc_flush_deferred, which
rpc_flush_deferred_guarded skips whenever the deferred queue is empty. The
RPC-to-local copy path blocks in ggml_backend_tensor_get with nothing deferred,
so that reset never ran and every copy created a backend event that was never
reused, until event or driver resources ran out.

The reset moves to the staging buffer wrap in rpc_staging_alloc. That is the
point where the last recorded event has just been synchronized, and events are
recorded on one backend in order, so the earlier ones are complete too.
… entry

The source-side dispatch reused cpy_tensor_async, but every implementation of
that is written as a destination-side handler and casts backend_dst to its own
context before deciding anything. ggml_backend_vk_cpy_tensor_async does it on
its first line, and the dispatch guard also allowed backend_dst to be NULL, so
a Vulkan source with a foreign or absent destination was a bad cast or a null
dereference rather than a clean fallback to the synchronous path.

Adds cpy_tensor_from_async, appended to the end of ggml_backend_i so backends
that initialise the members positionally are unaffected, and only implemented
where the source role is genuinely supported. The RPC backend qualifies: it
checks ggml_backend_is_rpc on both sides, which is null-safe, and declines
unless exactly one of them is an RPC backend.
get_available_slot calls prompt_save and prompt_load, which use
llama_state_seq_* on the selected slot's context, and it ran before
guard.wait_for(slot->id_group). With more than one pipeline group another slot
of that group can still be inside llama_decode at that moment, so a completion
with prompt caching enabled could touch a context that was in use.

Selection no longer touches the context: it reports whether an update is due
and update_prompt_cache does the work, after the caller has waited for the
group.
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

1 similar comment
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 433ef71971

ℹ️ 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".

// reinterpret a foreign backend_dst, or dereference a null one. Implement this only if
// the source role is genuinely supported. Appended last so backends that list the
// members positionally are unaffected.
bool (*cpy_tensor_from_async)(ggml_backend_t backend_src, ggml_backend_t backend_dst, const struct ggml_tensor * src, struct ggml_tensor * dst);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Bump the backend API version for the enlarged interface

When an executable loads a dynamic backend built against the previous interface, the loader still accepts it because GGML_BACKEND_API_VERSION remains 2, but adding this pointer enlarges the embedded ggml_backend_i and moves the core's expected device and context offsets. Accessing such a backend can therefore reinterpret its old fields or call a non-function address; increment the backend API version so stale plugins are rejected.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 35c4ca6. Correct and I should have done it in the same commit that added the member: a backend compiled against version 2 hands over a struct one member short, and reading past the end of it is precisely the failure the load-time check exists to prevent. Version is now 3, with a note on the line saying what changed.

Comment thread ggml/src/ggml-rpc/ggml-rpc.cpp Outdated
Comment on lines +1160 to +1161
ggml_backend_event_record(event, backend_dst);
st.inflight = event;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Wait for every stream before recycling RPC staging

With --pipeline-groups > 1, each llama_context has a distinct destination backend stream, yet all RPC-to-local copies for an endpoint share this staging arena. Overwriting st.inflight with only the latest stream's event means rpc_staging_alloc() later synchronizes that event and resets the entire arena/event pool even though an earlier copy on another stream may still be reading its region, causing copied activations to be overwritten; track and wait for all outstanding events or partition staging by backend stream.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed and fixed in 3b1f523. This one is right, and the code was documenting the very assumption that breaks.

rpc_staging_alloc() carried the comment "Events are recorded on one backend in order, so synchronizing the last one above means the earlier ones are done". That holds while one arena is written from one backend. The arena is keyed by socket_t *, so there is one per endpoint, and with --pipeline-groups > 1 each llama_context drives its own destination backend and therefore its own stream. Two groups copying over the same endpoint then overwrite each other's event in st.inflight, and the wrap path waits on whichever survived before resetting used and events_used, recycling the whole arena and the event pool while the other group's copy may still be reading its region. Events recorded on different streams have no ordering between them, so the surviving event says nothing about the discarded one.

Took the first of your two options. inflight becomes a vector of outstanding events and the wrap waits on all of them before recycling. Partitioning the arena by backend stream would also work, but it multiplies the pinned staging buffers by the number of groups, and on unified memory that is real capacity spent to avoid a wait that is usually already satisfied. Waiting on all outstanding events drops the single-stream assumption outright rather than patching around it, so it holds for any number of groups rather than for the two we happen to run.

The deferred SET path is unaffected: those events are drained by rpc_flush_deferred_guarded(), which the wrap check at the top of the non-RPC-source branch already calls before reaching the allocator.

Builds clean. This needs a two-node run with more than one pipeline group to be exercised properly rather than merely compiled, and that reservation is queued; I will report the result here rather than assume it.

Comment thread ggml/include/ggml-rpc.h Outdated

#define RPC_PROTO_MAJOR_VERSION 5
#define RPC_PROTO_MINOR_VERSION 1
#define RPC_PROTO_MINOR_VERSION 2

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Preserve compatibility with older RPC clients

When only the RPC server is upgraded, it now reports minor version 2 in HELLO, while an unchanged 5.1 client rejects any response whose minor version exceeds its own. Consequently an older client cannot connect to this newer server at all, despite the protocol change being presented as backward-compatible; negotiate an advertised/common minor version or otherwise keep the server acceptable to 5.1 clients.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed and fixed in 47aaac3. This is a real break and the framing in the PR was wrong: upgrading the server alone is the normal deployment, and that is exactly the case that stopped working.

negotiate_hello() rejects any response whose minor exceeds the client's own, so an unchanged 5.1 client refuses a 5.2 server outright, even though GET_TENSORS is purely additive and a 5.1 client would never send it. The version bump was carrying two meanings at once, "the wire format changed" and "this server has a new optional command", and only the second was true.

So the minor goes back to 1 and support is advertised in the byte of rpc_msg_hello_rsp that was pure padding. That byte is already on the wire, fixed size, and read as padding by existing clients, which see zero and simply never take the batched path. rpc_supports_batched_get() now tests the flag instead of server_minor >= 2.

Worth noting the same reasoning ruled out the obvious alternative of signalling on the request side: rpc_msg_hello_req is a fixed-size struct read as such by the server, so growing it would lock out existing clients in the same way the version bump did. Only the response had a spare byte, and only because it was already reserved.

Comment thread tools/server/server.cpp
Comment on lines +48 to +49
g_pipeline_groups = std::atoi(argv[++i]);
continue;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Propagate pipeline groups to router child servers

In router mode, this consumes and removes --pipeline-groups before server_models constructs its base preset from argc/argv. The value remains only in the router process's static variable, while each spawned model server receives arguments derived from the stripped preset and therefore starts with the default of one group, making the option silently ineffective for router-managed models.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Accepted, and confirmed by grep rather than by reading the diff: this branch has no LLAMA_ARG_PIPELINE_GROUPS handling in server-models.cpp or server.cpp at all, so router children genuinely never receive the setting.

Context worth having, since it affects how this should be fixed. The same defect was found and fixed on feature/pipeline-groups, twice: first by writing the variable into our own environment, which is wrong on Windows because common_set_env writes the CRT copy while get_environment reads the Win32 block and propagation between them is undocumented, and then by pushing it explicitly into child_env. A later round found that even that was incomplete, because base_env is a copy of the router environment, so an inherited value survives and wins: execve keeps duplicates and getenv returns the first, and an explicit --pipeline-groups 1 appended nothing at all. The landed form erases any inherited entry and then always sets the resolved value.

I tried merging feature/pipeline-groups here to bring that across and backed it out: it conflicts in five files, and since this PR targets master rather than stacking, the merge would also drag in scope that does not belong to it. So this needs the fix ported deliberately rather than merged, and I would rather port it correctly than quickly, given it took two attempts to get right on the other branch.

Marking agreed because the defect is real and the fix is understood, but the commit is not here yet.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correcting myself on the bookkeeping, not the verdict. I marked this agreed in the reply above while saying in the same breath that the commit is not here yet, which is exactly the thing I have been refusing to do on the other PRs: fixcheck treats a thumbs-up as a claim that a fix is on head, and one without a commit makes the gate report something untrue. I have removed the reaction. The analysis stands unchanged and I will re-react when the port lands.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed and fixed in 6ca4b11.

Reproduced first, on this branch, with stories260K on CPU. Router started with --pipeline-groups 2 --models-dir ...: the rendered child args contained no group setting at all, and the child logged 0 created llama_context for pipeline group lines. The value was parsed into a file-static in server.cpp and server_take_pipeline_groups() strips the flag from argv before server_models is constructed, so nothing survived to the child. Every model served at one group while the operator believed otherwise, silently.

The router now exports the resolved value before models_routes.emplace(...), which is what base_env(get_environment()) snapshots, and the child reads it back through an env fallback. That is the same channel LLAMA_ARG_HF_REPO and LLAMA_SERVER_ROUTER_PORT already use. An explicit flag on the child command line still wins over the inherited value, matching common/arg.cpp precedence.

One thing the reproduction turned up that the propagation alone would have made worse. Groups > 1 requires an explicit context size, and the router only renders --ctx-size into child args when the operator passed one. Propagating the setting by itself converted a silent downgrade into a load failure raised inside a subprocess on the first model request. The router now refuses at startup with an actionable message instead.

Measured, after the change:

configuration pipeline group contexts in child threadpool inits
router, no flag (control) 0 1
router --pipeline-groups 2 -c 512 1 2
router --pipeline-groups 2, no -c startup refusal, rc=1 n/a

Generation succeeds in every configuration that starts.

cparams.n_ctx = llama_n_ctx(ctx_tgt);
cparams.n_seq_max = llama_n_seq_max(ctx_tgt);

groups[g]->ctx = llama_init_from_model(model_tgt, cparams);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Initialize thread pools for every pipeline context

When a pipeline group performs any CPU-backed graph work, only group 0 has the reusable thread pools that common_init_from_params() attaches. These extra contexts are created directly with llama_init_from_model() and never receive a thread pool, so every CPU graph compute creates and destroys a disposable pool and ignores configured affinity, polling, and priority settings; create and attach a suitable pool for each concurrently executing context.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed and fixed in 6ca4b11.

Reproduced on this branch before changing anything. llama-server -m stories260K.gguf --pipeline-groups 2 -c 512 logged 1 llama threadpool init line while creating 2 group contexts. common_init_from_params builds and attaches a pool for ctx_tgt only; the contexts created by llama_init_from_model() for g >= 1 got none, so every extra group fell back to the disposable per-graph pool. That path ignores the configured CPU mask, priority and strict placement, which is the whole point of passing them.

Groups execute concurrently, so they cannot share one pool. Each server_group now owns a common_threadpools and initialises it for g >= 1. Group 0 keeps the pool it already had.

On destruction order, since that is the easy way to get this wrong: the existing teardown already calls llama_free(groups[g]->ctx) explicitly before groups.clear(), so each context is freed before the pool it was attached to. The error path during creation frees in the same order.

After the change the same command logs 2 threadpool inits for 2 groups. Controls that must not move, and did not: a plain server with no flag stays at 1 pool and 0 group contexts, and a router with no flag is unchanged. Generation succeeds in each case.

Comment on lines +1816 to +1820
size_t total = 0;
for (uint32_t i = 0; i < n_entries; i++) {
total += entries[i].size;
}
response.resize(total, 0);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Validate batched reads before allocating the response

Any client able to connect to the RPC server can put an arbitrary 64-bit size in a GET_TENSORS entry, and this resizes the response to the unchecked sum before any tensor or buffer-bound validation occurs. A tiny malformed request can therefore request a multi-terabyte allocation and throw an uncaught std::bad_alloc, terminating the RPC server; validate every entry, use checked addition, and impose the response-size limit before allocating.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed and fixed in 4ed1cd0. The unchecked sum is worse than an allocation failure, so I want to record the second half of it here.

You are right that response.resize(total) runs before any validation and that a single oversized entry throws an uncaught std::bad_alloc and takes the server down. The other case is that total is a size_t accumulating uint64_t sizes, so it can wrap. A wrapped total allocates a small response, and the copy loop then still calls ggml_backend_tensor_get(..., response.data() + out_offset, ..., entries[i].size) for the full requested size. The per-entry bounds check does not catch it, because it constrains a region against its own source buffer and never against the response, so that path is a heap overflow write rather than a clean failure.

The fix does all three things you listed. Validation moves into a first pass that runs before the allocation: each entry is deserialized and bounds checked there, and the running total is accumulated with a checked addition. On the response-size limit I want to be explicit about the reasoning rather than just picking a number. Once per-entry validation precedes the allocation, each region is bounded by a really allocated buffer, so an arbitrary cap is not what stops the bad_alloc case. What it stops is several individually valid entries naming the same large buffer, which on a machine with a large unified pool could still sum into the terabytes, and it keeps the addition in a range where it cannot wrap. It is set at 4 GiB, which is orders of magnitude above the activations and single tensor regions this path batches, so it should not be reachable by legitimate use.

ggml_backend_i gained cpy_tensor_from_async, so a dynamically loaded backend
compiled against version 2 supplies a struct one member short and ggml would
read past the end of it. That is exactly what the load-time version check
exists to catch, so the version moves with the layout.
The source-side dispatch permitted backend_dst == NULL. Any implementation of
the source role has to inspect the destination to decide whether it can help,
so a null destination pushes that dereference into every implementer.

The RPC handler was one of them. ggml_backend_is_rpc() tolerates null, so a
null destination against an RPC source passed the differing-kind test and then
reached ggml_backend_get_device(other) with other == nullptr. I had claimed
that handler was null-safe on the strength of the is_rpc calls alone, which was
wrong: the calls are safe, the code after them is not.

Guarded in both places, in the helper so the invariant lives in one spot rather
than depending on every implementer, and in the RPC handler so it does not rely
on its only caller.
The batched read summed the entry sizes straight off the wire and resized the
response to that sum, before any tensor or bounds validation ran. Two problems,
not one.

A single entry can name a size larger than anything allocatable, so a tiny
request threw an uncaught std::bad_alloc and terminated the server. And the sum
accumulates uint64_t sizes into a size_t, so it can wrap: a wrapped total
allocates a small response while the copy loop still writes entries[i].size
bytes at out_offset, running off the end of the block. The existing per-entry
bounds check does not catch that, because it constrains a region against its own
source buffer and never against the response.

Validation now happens in a first pass, before the allocation: every entry is
deserialized and bounds checked there, and the total is accumulated with a
checked addition against a ceiling. Per-entry validation alone would leave the
sum bounded only by really allocated buffers, which several entries naming the
same large buffer could still push into the terabytes, so the ceiling is what
keeps the addition in a range where it cannot wrap. At 4 GiB it is orders of
magnitude above the activations and single tensor regions this batches.
Bumping RPC_PROTO_MINOR_VERSION to 2 made this server unusable by every existing
client. negotiate_hello() rejects any response whose minor exceeds the client's
own, so an unchanged 5.1 client refuses to connect to a 5.2 server, even though
GET_TENSORS is purely additive and a 5.1 client would never send it. Upgrading
the server alone is the normal way this gets deployed, so presenting the change
as backward compatible while breaking exactly that case is the wrong way round.

The minor goes back to 1, and support is advertised in the byte of the HELLO
response that was pure padding. That byte is already on the wire, fixed size,
and read as padding by existing clients, which see zero and simply never use the
batched path. rpc_supports_batched_get() now tests the flag rather than the
version, so capability and version stop being conflated.
rpc_staging held a single inflight event, which assumed one arena is written
from one backend: the last event recorded would then be ordered after all the
earlier ones on the same stream, so waiting on it alone was enough. The comment
in rpc_staging_alloc() stated that assumption explicitly.

The arena is keyed by socket, so there is one per endpoint, and with
--pipeline-groups > 1 each llama_context drives its own destination backend and
therefore its own stream. Two groups copying over the same endpoint overwrote
each other's event in that slot, and the wrap path then waited on the survivor
before resetting used and events_used, recycling the whole arena and its event
pool while the other group's copy could still be reading its region. Events
recorded on different streams have no ordering between them, so the surviving
event says nothing about the discarded one, and the result is copied
activations overwritten under a copy still in progress.

The slot becomes a vector of outstanding events, and the wrap waits on all of
them before recycling. That drops the single-stream assumption rather than
patching around it, so it holds for any number of groups. The deferred SET path
is unaffected: its events are drained by rpc_flush_deferred_guarded(), which the
wrap check already calls before reaching here.
…p a thread pool

Two defects found by review and reproduced against this branch before changing anything,
using stories260K on CPU and counting log lines rather than inspecting code.

1. Router children never saw --pipeline-groups. server_take_pipeline_groups() strips the flag
   from argv before server_models is constructed, and the value lived only in a file-static in
   server.cpp, so the rendered child args carried nothing. A router started with
   --pipeline-groups 2 spawned children that logged 0 "created llama_context for pipeline group"
   lines: the operator's setting was silently discarded and every model served at one group.
   The router now exports the resolved value and the child reads it back, which is the channel
   LLAMA_ARG_HF_REPO and LLAMA_SERVER_ROUTER_PORT already use. An explicit flag still wins over
   the inherited value, matching common/arg.cpp precedence.

   Propagation on its own would have converted a silent downgrade into a per-request load
   failure, because groups > 1 needs a context size and the router only renders --ctx-size into
   child args when the operator passed one. The router now refuses at startup with an actionable
   message instead of failing inside a subprocess on the first model request.

2. Only group 0 had a thread pool. common_init_from_params attaches one to ctx_tgt; the contexts
   created by llama_init_from_model() for g >= 1 got none, so every extra group fell back to a
   disposable per-graph pool that ignores the configured CPU mask, priority and strict placement.
   With --pipeline-groups 2 the server logged 1 "llama threadpool init" line for 2 groups. Groups
   run concurrently and so cannot share a pool; each now builds its own. Pools are destroyed
   after the contexts that use them, which the existing teardown order already guarantees.

Controls, all on this build: plain server with no flag stays at 1 pool and 0 group contexts;
router with no flag is unchanged and still serves; env value alone applies; explicit flag
overrides the env value; generation succeeds in every configuration that starts.
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 6ca4b11763

ℹ️ 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".

Comment thread ggml/src/ggml-rpc/ggml-rpc.cpp Outdated
Comment on lines +1177 to +1180
std::lock_guard<std::mutex> lock(rpc_staging_mutex);
rpc_staging & st = rpc_staging_map[sock.get()];
staging = rpc_staging_alloc(st, host_buft, size);
}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Keep staging reservations live until the copy is registered

With multiple pipeline groups copying RPC tensors concurrently, this releases rpc_staging_mutex before the blocking read, async destination copy, and event registration. If this allocation fills the arena, another group can enter rpc_staging_alloc(), see an overflow, and reset st.used while this thread's region has no entry in st.outstanding yet; it can then reuse and overwrite the same bytes during the read or before the destination stream consumes them. The fresh evidence beyond the previous event-tracking issue is this unlocked allocation-to-registration gap; reserve/register the region atomically or partition staging per stream/group.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed and fixed in 6e2cbe7. The unlocked allocation-to-registration gap is exactly the hole, and the reasoning about st.used being reset while a region has no entry in st.outstanding holds as written.

Holding rpc_staging_mutex across the blocking read would serialize the pipeline groups the arena exists to keep concurrent, which is the thing the shared arena is for, so the region is reserved instead. rpc_staging_alloc() now takes the lock as a unique_lock, increments an in-flight count on success, and an overflow waits for that count to reach zero before it synchronizes outstanding and resets. The caller hands the region over with a commit once its event is recorded, or once the copy has been synchronized when no event was available. Every reservation is released without needing anything from the waiting thread, so the wait cannot deadlock against it.

Two things the fix had to get right beyond the literal report:

The deferred path had the same hole from the other side. The !src_is_rpc branch obtains its event under the lock but records it outside, and put it only in the deferred op, never in st.outstanding. So an overflow driven from the RPC-to-local path synchronized a set that did not include those events at all and reset the arena while a queued read was still filling it. That event is now registered in outstanding as well. What keeps those bytes alive until the deferred send has copied them out is still the flush-on-wrap, which I left as it was and also added to the second path, which did not have it.

Not the partition-per-stream option, deliberately. It removes the class rather than patching an instance, which is usually the better trade, but the deferred queue is per socket while the arenas would be per group, so a flush would be draining ops belonging to several arenas and the event-pool recycling at the end of rpc_flush_deferred would need to know which. That is a larger change than the reservation and I did not want to make it without being able to run it.

What I could not run, stated plainly. The staging path is behind set_tensor_async / get_tensor_async, which the CPU backend does not implement, so the whole function returns false before touching the arena and no CPU test reaches this code. Verifying the concurrent behaviour needs two nodes with real devices and --pipeline-groups 2, which is the configuration this arena work has never actually been exercised in. I have a two-node greedy-equivalence run queued behind the machine lock here: one group and two groups over the same RPC peer must produce byte-identical output at temperature 0, since staging corruption changes tokens rather than crashing. What is verified today is that it compiles, that ordinary RPC generation over a socket is unaffected, and that teardown is clean. I will post the two-node numbers when the lock frees.

Comment on lines +649 to +650
static std::mutex rpc_staging_mutex;
static std::unordered_map<socket_t *, rpc_staging> rpc_staging_map;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Release staging resources when an RPC socket expires

When the last RPC buffer for a connection is freed, the socket can expire from the weak socket cache, but this raw-pointer-keyed map never erases its entry and rpc_staging has no destructor that frees its backend buffer or events. A long-lived process that repeatedly creates and tears down RPC connections therefore permanently retains each connection's pinned staging allocation and driver events, potentially exhausting pinned memory; tie the staging entry to socket lifetime and free its buffer/events on removal.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed and fixed in 6e2cbe7.

Both halves of this were true: rpc_staging_map is keyed by a raw socket_t * and nothing ever erased an entry, and rpc_staging had no destructor, so neither the pinned host buffer nor the pool of driver events was ever released. The arena is only grown, never shrunk, so what a process retained was one pinned allocation per connection it had ever opened, for its whole lifetime.

rpc_staging now has a destructor that waits on anything still outstanding, frees the event pool and frees the buffer, and the map entry is erased from socket_t's destructor, which is the last moment the key is valid. Doing it there rather than at buffer-free time also closes the stale-arena case you would otherwise get if a later socket were allocated at the same address.

Ordering: the wait on outstanding happens before the events are freed, so an event is never destroyed while a copy might still reference it. The struct is made non-copyable so the map cannot silently duplicate and then double-free it; operator[] default-constructs in place and the container is node-based, so references stay valid.

Verified on CPU that connections open and close cleanly with the hook in place, including generation over a real RPC socket followed by teardown, with no aborts or asserts on either side. What I have not done is measure the retention itself: the arena is only ever created from the async staging paths, which a CPU backend never enters, so a CPU run has no entry to leak and RSS across connect and disconnect cycles would show nothing either way. The leak is clear from the code and the fix is clear, but the number is not measured, and I would rather say that than imply it was.

Comment on lines +1117 to +1121
params_ctx.n_parallel = n_seq_per_group;
params_ctx.n_ctx = params_base.n_ctx / n_groups;
}

llama_init = common_init_from_params(params_ctx);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Account for every pipeline context during memory fitting

With the default fit_params enabled, reducing params_ctx.n_ctx to 1/N before common_init_from_params() makes common_fit_params() budget device memory for only one reduced context. The code then creates N-1 additional contexts without rerunning or scaling the fit, so configurations that fit correctly as one context at the requested total context size can over-offload model layers and fail while constructing a later group because the other groups' KV and compute buffers were omitted from the estimate. Run fitting against the aggregate context memory required by all groups.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed and fixed in 6e2cbe7, and this one turned out to be executable after all rather than only arguable.

The accounting gap, with four groups at -c 262144: the fit reported projected to use 106 MiB of device memory and will leave 124503 >= 1024 MiB, no changes needed, while the four contexts that then get built actually need 40.00 MiB of KV plus 65.39 MiB of compute each, so 421.6 MiB. The shortfall is 316.2 MiB, exactly 3 x 105.39: the fit budgeted a quarter of the true device requirement, which is the N-1 omission stated directly.

The real over-offload was reproducible without a GPU, using an ggml-rpc-server on localhost as a genuine non-host device whose budget the fit will manage. With the budget set so about 200 MiB remains usable on that device: before, will leave 124503 >= 124409 MiB, no changes needed and offloaded 6/6 layers to GPU, with four contexts needing 421.6 MiB in a 200 MiB budget. After, the margin becomes 1024 + 316.18 = 1340 MiB, the fit walks 6/6 down to 0/6 and reports success.

The extra requirement is added to fit_params_target before common_init_from_params(), the same margin mechanism the existing mmproj estimate uses, measuring one group context and multiplying the context and compute parts by n_groups - 1. Model weights are excluded because the groups share them. A failure to measure degrades to a warning rather than refusing to start. The single-group control emits nothing new and is byte-identical: projected to use 58 MiB, will leave 124551 >= 124409, offloaded 6/6.

A separate pre-existing defect found while doing this, which I have deliberately not changed. The first version of this patch indexed fit_params_target by the global ggml_backend_dev_get() index, and the margin silently stayed at 1024 MiB; only the run exposed it, since the log line and the syntax check both looked correct. common_fit_params indexes margins by model-device order (common/fit.cpp:285). The existing mmproj block at server-context.cpp:1078 has that same mistake, writing fit_params_target[i] with a global device index, so its estimate lands on the wrong device whenever the model's devices are not the leading global devices. I could not construct a multimodal reproduction here, so I left it alone rather than change code I could not verify. Happy to fix it here or open it separately, whichever you prefer.

Comment on lines +2629 to 2632
std::vector<server_slot *> child_slots = get_free_slots(n_child_tasks, slot->id, slot->id_group);
if (child_slots.size() < n_child_tasks) {
SRV_DBG("not enough free slots for child tasks, n_free = %zu, n_children = %zu, defer task, id_task = %d\n", child_slots.size(), n_child_tasks, id_task);
queue_tasks.defer(std::move(task));

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Choose a group with enough free child slots

For n_cmpl > 1, slot selection happens globally before this group-local child-capacity check. If the selected idle slot is in a group whose sibling slots are busy while another group has enough idle slots, get_free_slots() returns too few and the request is repeatedly deferred even though sufficient compatible capacity exists; a long-running request in the preferred group can leave the other group idle indefinitely. Select the parent only among groups with enough free slots, or retry selection in another eligible group before deferring.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed and fixed in 6e2cbe7. Reproduced in both directions, because the outcome depends on slot iteration order and one direction alone would not have distinguished a real fix from simply preferring a higher-numbered group.

Two groups of three slots, --parallel 6, two slots occupied by long generations, then one n=3 completion:

latency before served by latency after served by
slots 0,1 busy in group 0, group 1 idle 10.51 s slots [0,1,2] 0.01 s slots [3,4,5]
slots 3,4 busy in group 1, group 0 idle 11.19 s slots [3,4,5] 0.00 s slots [0,1,2]

The not enough free slots for child tasks deferral goes from 1 to 0 in both. Selection now skips slots whose group cannot supply the whole request, in the similarity path as well as the LRU path, since both had the defect.

One ordering change that is load-bearing rather than cosmetic: the n_cmpl bound check moved ahead of selection. Left where it was, a request for more completions than any group has slots would find no eligible group and be deferred forever instead of rejected, which would have turned a clean 400 into a hang. There is a control for exactly that.

Controls, identical before and after: a single pipeline group with capacity still serves immediately; ordinary n = 1 with two groups is unaffected; a genuinely saturated single group still waits rather than the fix inventing capacity; and n greater than the slots per group is still rejected in 0.00 s with n_cmpl must not exceed the number of slots per pipeline group (3).

Worth recording, because it nearly produced a false pass: the first run of this reproduction reported PASS on the unfixed binary. --slot-prompt-similarity defaults to 0.1, so the LCP path selected a slot and masked the LRU path entirely. Setting it to 0 exposes the defect, and the LCP path was then found to have the same one and got the same guard. A test that passes because a different code path answered first is the failure mode I most want to avoid here.

… socket, fit and select for every group

Four review items. Each was reproduced against this branch before anything was changed, and each has
a control showing the single-group and non-RPC paths are unaffected.

Staging reservation gap. rpc_staging_alloc() released rpc_staging_mutex before the blocking read,
the async destination copy and the event registration, so between a region being handed out and its
event reaching `outstanding` the arena had no record of it. Another pipeline group entering on the
same socket could see an overflow, wait only on the events it could see, reset `used` to zero and be
handed the same bytes while the first group was still reading into them. Holding the mutex across
the copy would have serialized the groups this arena exists to keep concurrent, so the region is
reserved instead: an overflow now waits for every outstanding reservation to be handed over before
it synchronizes and resets. Every reservation is released without needing anything from the waiting
thread, so the wait cannot deadlock against it. The deferred path also registers its event in
`outstanding` now, which it never did, so a wrap driven from the other staging path no longer resets
the arena while a queued read is still filling it.

Staging lifetime. rpc_staging_map is keyed by a raw socket_t * and nothing ever erased an entry, and
rpc_staging had no destructor, so each connection's pinned host arena and its driver events were
retained for the life of the process. A long-running process that opens and closes RPC connections
therefore accumulated one pinned allocation per connection. The entry is now dropped from
socket_t's destructor, which is the last point the key is valid, and the arena frees its buffer and
events after waiting on anything still outstanding.

Memory fitting across groups. n_ctx is divided by N before common_init_from_params(), so
common_fit_params() budgeted model weights plus one reduced context and one compute buffer, while
N-1 further contexts were then created with a bare llama_init_from_model() that does no fitting.
Measured with four groups at -c 262144: the fit projected 106 MiB of device memory while the four
contexts actually need 4 x 105.39 = 421.6 MiB, a shortfall of exactly 3 x 105.39. Against a real
device with a tight budget the old code offloaded 6/6 layers into memory the later groups needed;
it now backs off to 0/6 and reports a successful fit. The extra requirement is added to
fit_params_target the same way the mmproj estimate already is, excluding model weights because all
groups share them. Single-group runs are byte-identical and emit nothing new.

Slot selection across groups. The parent slot was chosen globally and only then checked for free
sibling slots within its own group, so an n > 1 request whose global winner sat in a busy group was
deferred while another group sat idle, repeatedly, because the selection is deterministic. Selection
now skips slots whose group cannot supply the whole request, in both the similarity and the LRU
path. With two groups of three slots and two slots busy: latency 10.51 s before and 0.01 s after
when group 0 is the busy one, and 11.19 s before and 0.00 s after when group 1 is, with the served
slot ids moving to the idle group in each case. Both directions were driven because the outcome
depends on iteration order and one direction alone is also consistent with simply preferring a
higher group. The n_cmpl bound check moved ahead of selection, so an impossible request is still
rejected immediately rather than becoming a request no group can satisfy. Controls: a single group
with capacity still serves immediately, a genuinely saturated single group still waits rather than
inventing capacity, ordinary n = 1 is unaffected, and n greater than the slots per group is still
rejected.
@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 6e2cbe7a19

ℹ️ 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".

std::lock_guard<std::mutex> lock(rpc_staging_mutex);
rpc_staging & st = rpc_staging_map[sock.get()];
st.outstanding.push_back(event);
rpc_staging_commit(st);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Keep the reservation until the SET is queued

When two pipeline groups contend and this allocation fills the arena, this call drops in_flight before the corresponding deferred SET is appended. Another group can then observe a wrap, flush a queue that does not yet contain this SET, synchronize only the staging fill, reset used, and reuse these bytes; the subsequently queued SET transmits the overwritten data. The fresh evidence after the prior staging-gap fix is the remaining ordering: rpc_staging_commit() here still precedes deferred.push_back() at line 1239, so the reservation must remain live until the operation is registered.

Useful? React with 👍 / 👎.

std::lock_guard<std::mutex> lock(rpc_staging_mutex);
auto it = rpc_staging_map.find(sock.get());
if (it != rpc_staging_map.end()) {
it->second.events_used = 0;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Recycle only events whose copies have completed

When an RPC-to-local copy has recorded an event for a destination stream, an unrelated deferred flush can reach this unconditional reset while outstanding still contains that event and its stream may still be consuming staging. A following no-wrap allocation then reuses and re-records the same event object; on CUDA the new record replaces the state waited on by cudaEventSynchronize, so a later arena wrap may wait only for the newer stream and overwrite bytes still used by the original one. The fresh evidence after the event-tracking fixes is that this reset neither synchronizes nor clears outstanding; recycle the pool only after all associated events have completed.

Useful? React with 👍 / 👎.

Comment on lines +4454 to +4455
void metrics_post_decode(server_group & grp, int32_t off, int32_t n_tokens, bool has_output) {
auto & batch = grp.batch;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Count only the current group's slots per decode

With --pipeline-groups > 1, each invocation of this helper represents one llama_decode() on grp.ctx, but the loop at line 4458 still adds every processing slot in the server to metrics.n_busy_slots. For example, two groups with four active slots each add eight busy slots for each group's decode, so the exported n_busy_slots_per_decode reports eight even though every actual decode contains at most four slots. Iterate grp.slots here so the Prometheus metric retains its documented per-call meaning.

Useful? React with 👍 / 👎.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant