rpc: one tensor read per decode step and asynchronous split copies - #193
rpc: one tensor read per decode step and asynchronous split copies#193danielhanchen wants to merge 20 commits into
Conversation
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.
|
@codex review |
|
@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. |
There was a problem hiding this comment.
💡 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".
| 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)) { |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
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.
| ggml_backend_event_t event = rpc_staging_event(st, other_dev); | ||
| if (event != nullptr) { | ||
| ggml_backend_event_record(event, backend_dst); | ||
| st.inflight = event; |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
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); |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
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.
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.
|
@codex review |
1 similar comment
|
@codex review |
There was a problem hiding this comment.
💡 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); |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
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.
| ggml_backend_event_record(event, backend_dst); | ||
| st.inflight = event; |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
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.
|
|
||
| #define RPC_PROTO_MAJOR_VERSION 5 | ||
| #define RPC_PROTO_MINOR_VERSION 1 | ||
| #define RPC_PROTO_MINOR_VERSION 2 |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
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.
| g_pipeline_groups = std::atoi(argv[++i]); | ||
| continue; |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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); |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
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.
| size_t total = 0; | ||
| for (uint32_t i = 0; i < n_entries; i++) { | ||
| total += entries[i].size; | ||
| } | ||
| response.resize(total, 0); |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
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.
|
@codex review |
There was a problem hiding this comment.
💡 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".
| 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); | ||
| } |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
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.
| static std::mutex rpc_staging_mutex; | ||
| static std::unordered_map<socket_t *, rpc_staging> rpc_staging_map; |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
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.
| 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); |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
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.
| 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)); |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
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.
|
@codex review |
There was a problem hiding this comment.
💡 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); |
There was a problem hiding this comment.
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; |
There was a problem hiding this comment.
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 👍 / 👎.
| void metrics_post_decode(server_group & grp, int32_t off, int32_t n_tokens, bool has_output) { | ||
| auto & batch = grp.batch; |
There was a problem hiding this comment.
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 👍 / 👎.
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-samplingbuilds one sampler subgraph per output row, andllm_graph_context::build_samplingfills four per-row vectors (t_sampled,t_sampled_probs,t_sampled_logits,t_candidates).llama_context::decodethen walks each vector withcopy_tensor_async_rows, which issues oneggml_backend_tensor_get_asyncper row and pervector. The RPC backend had
get_tensor_async = NULL, so every one of those became its ownsynchronous
RPC_CMD_GET_TENSORround 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 asone new
RPC_CMD_GET_TENSORSat the next flush point, which is asynchronize, a graphcompute, 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_decodeallocates 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-> blockingggml_backend_tensor_get-> synchronousSET_TENSOR, becausecpy_tensor_asyncwasNULLandggml_backend_sched_compute_splitsfalls back to that.cpy_tensor_asyncis now implemented for both directions:producing device's host buffer type, an event recorded on the producing stream, and the
SET_TENSORsent by the dispatcher once that event has completed. The producing device isnever fully synchronized and the host thread serialises the graph while the copy is still in
flight.
set_tensor_asyncon the consumingstream, so the consuming device is not synchronized either.
ggml_backend_tensor_copy_asyncand the scheduler now ask the source backend for the copywhen 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=1prints the client side command counts and bytes (period fromGGML_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=1andGGML_RPC_NO_ASYNC_COPY=1turn each half off for A/B testing.RPC commands per decode step
CPU only harness: two local
ggml-rpc-serverinstances on the CPU backend,llama-serverwith--device RPC0,RPC1 -sm layer, 8 concurrent, npp 128, ntg 256, counted over a decode-onlywindow (prefill excluded).
The read count is now one per step at any batch size. What is left is the per-row
SET_TENSORthat each
distsampler 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: md5177dc61e0703eba3bdaf7bf1131f0458for the default CPU sampling path at--pipeline-groups1 and 2, and the same md5 for--backend-samplingon both the base treeand 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,RPC0so the output layer, the logits and the sampler are on the remote nodeand 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.
Qwen3.5-4B UD-Q4_K_XL, same split and workload with ntg 128, all ten arms:
CUDA0,RPC0base, N=1 / N=2CUDA0,RPC0async copy only, N=1 / N=2CUDA0,RPC0backend sampling, N=1 / N=2RPC0,CUDA0both features, N=1 / N=2RPC0,CUDA0async copy only, N=1 / N=2What the numbers say:
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.
Each
distsampler 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.
2.4 percent with two (138.09 against 134.87). It removes a full
ggml_backend_synchronizeofthe 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.
get_socketconnects when the cachedconnection has expired, and an
rpc-serverserves 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
synchronizelooks up without creating.recv_datacopies all of it, so a sent message must be read back in onerecv_dataof thesame 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 againstmasterit is ten files, not four. An earlier version ofthis 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:The four
tools/server/files are the first pipeline-groups commits, which this branch is stackedon 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 fromggml/src/ggml-backend.cppandggml/src/ggml-backend-impl.h, which gain the source-side asynccopy entry described below. Nothing in
src/llama-context.cpp,src/llama-sampling.cpporsrc/llama-graph.cppis touched in either accounting.The
ggml-backend.cppchange factors the existing "ask the destination backend for anasynchronous copy" into a helper that then also asks the source backend, but only when the two
backends have different
cpy_tensor_asyncimplementations. Two backends of the same type shareone 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=OFFconfigures and builds clean.temperature 0,top_k 1,seed 1234): md5177dc61e0703eba3bdaf7bf1131f0458on the branch base, on this branch, and onthis branch built with
GGML_RPC=OFF. On a single GPU with the 27B and no--rpc, md574926c4ef135f3cc89ad20cd5ec7e445on the branch base and on this branch.pytest -q -m "not slow"run serially: branch base 368 passed, 6 skipped, 199deselected; this branch 368 passed, 6 skipped, 199 deselected. Same set, no new failures.
llama-batched-benchbracket 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_TENSORShandler that was reachable before this fix, and it isworth 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_allocandterminated the server. But the sum accumulates
uint64_tsizes into asize_t, so it can wrap: awrapped total allocates a small response, and the copy loop then still calls
ggml_backend_tensor_getfor the fullentries[i].sizeatout_offset. That is a heap overflowwrite 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 thatmatters most:
negotiate_hello()rejects any server whose minor exceeds the client's own, so anunchanged 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_GETin the byte ofrpc_msg_hello_rspthat was previously pure padding. That byte is already on the wire, is fixedsize, 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 couldnot carry this:
rpc_msg_hello_reqis a fixed-size struct read as such by the server, so growing itwould 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=1andGGML_RPC_NO_ASYNC_COPY=1turn each half off,which is how the arms above were taken on one binary.