server: --pipeline-groups, run the slots over several contexts of one model

ggml/src/ggml-rpc/ggml-rpc.cpp

@@ -212,7 +212,7 @@ struct ggml_backend_rpc_device_context {
     uint32_t    device;
     std::string name;
     std::string description;
-    uint64_t    last_graph_uid;
+    // note: the uid of the last graph stored on the server is tracked per connection, see socket_t
 };
 
 struct ggml_backend_rpc_buffer_type_context {
@@ -300,7 +300,8 @@ static bool parse_endpoint(const std::string & endpoint, std::string & host, int
 
 // RPC request : | rpc_cmd (1 byte) | request_size (8 bytes) | request_data (request_size bytes) |
 // No response
-static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, size_t input_size) {
+// writes one whole message; the caller must hold sock->conn.mtx_send
+static bool send_rpc_cmd_locked(socket_ptr sock, enum rpc_cmd cmd, const void * input, size_t input_size) {
     uint8_t cmd_byte = cmd;
     if (!sock->send_data(&cmd_byte, sizeof(cmd_byte))) {
         return false;
@@ -314,12 +315,61 @@ static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input,
     return sock->flush();
 }
 
+static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, size_t input_size) {
+    std::lock_guard<std::mutex> lock(sock->conn.mtx_send);
+    return send_rpc_cmd_locked(sock, cmd, input, input_size);
+}
+
+// Reserves this thread's place in the response order of a connection. The server answers the
+// commands of one connection strictly in the order it received them, so the n-th response
+// belongs to the n-th response-bearing request that was written to the socket. The ticket is
+// taken while mtx_send is still held by the sender, and always released, so a failed send
+// cannot leave the later waiters stuck.
+struct rpc_response_ticket {
+    rpc_conn_state & conn;
+    uint64_t         seq;
+
+    // must be constructed with conn.mtx_send held
+    explicit rpc_response_ticket(rpc_conn_state & conn) : conn(conn) {
+        std::lock_guard<std::mutex> lock(conn.mtx_seq);
+        seq = conn.seq_next++;
+    }
+
+    void wait() {
+        std::unique_lock<std::mutex> lock(conn.mtx_seq);
+        conn.cv_seq.wait(lock, [this] { return conn.seq_serving == seq; });
+    }
+
+    ~rpc_response_ticket() {
+        std::lock_guard<std::mutex> lock(conn.mtx_seq);
+        conn.seq_serving = seq + 1;
+        conn.cv_seq.notify_all();
+    }
+};
+
 // RPC request : | rpc_cmd (1 byte) | request_size (8 bytes) | request_data (request_size bytes) |
 // RPC response: | response_size (8 bytes) | response_data (response_size bytes) |
 static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, size_t input_size, void * output, size_t output_size) {
-    if (!send_rpc_cmd(sock, cmd, input, input_size)) {
+    std::unique_ptr<rpc_response_ticket> ticket;
+    bool failed = false;
+    {
+        std::lock_guard<std::mutex> lock(sock->conn.mtx_send);
+        ticket.reset(new rpc_response_ticket(sock->conn));
+        if (!send_rpc_cmd_locked(sock, cmd, input, input_size)) {
+            // still take our turn, so the ticket is released in order and no later waiter is
+            // woken with a response that is not theirs
+            failed = true;
+        }
+    }
+
+    if (failed) {
+        ticket->wait();
         return false;
     }
+
+    // the response is read outside mtx_send, so the other threads can keep submitting
+    ticket->wait();
+
     uint64_t out_size;
     if (!sock->recv_data(&out_size, sizeof(out_size))) {
         return false;
@@ -731,21 +781,31 @@ static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, g
     ggml_backend_rpc_device_context * rpc_dev_ctx = (ggml_backend_rpc_device_context *)rpc_dev->context;
 
     GGML_ASSERT(cgraph->n_nodes > 0);
-    bool reuse = cgraph->uid != 0 && rpc_dev_ctx->last_graph_uid == cgraph->uid;
-    if (reuse) {
+    GGML_UNUSED(rpc_dev_ctx);
+
+    auto sock = get_socket(rpc_ctx->endpoint);
+
+    // The graph stored by RPC_CMD_GRAPH_COMPUTE lives on the server per connection and device,
+    // and one connection is shared by every backend of this endpoint - including the backends of
+    // other llama_contexts. So the uid of the last graph sent has to be tracked per connection,
+    // and the check has to happen under the same lock as the send, or a RECOMPUTE could re-run
+    // the graph another context stored in between.
+    std::unique_lock<std::mutex> lock(sock->conn.mtx_send);
+
+    auto & last_uid = sock->conn.last_graph_uid[rpc_ctx->device];
+    if (cgraph->uid != 0 && last_uid == cgraph->uid) {
         rpc_msg_graph_recompute_req request;
         request.device = rpc_ctx->device;
-        auto sock = get_socket(rpc_ctx->endpoint);
-        bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_RECOMPUTE, &request, sizeof(request));
-        RPC_STATUS_ASSERT(status);
-    } else {
-        rpc_dev_ctx->last_graph_uid = cgraph->uid;
-        std::vector<uint8_t> input;
-        serialize_graph(rpc_ctx->device, cgraph, input);
-        auto sock = get_socket(rpc_ctx->endpoint);
-        bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_COMPUTE, input.data(), input.size());
+        bool status = send_rpc_cmd_locked(sock, RPC_CMD_GRAPH_RECOMPUTE, &request, sizeof(request));
         RPC_STATUS_ASSERT(status);
+        return GGML_STATUS_SUCCESS;
     }
+
+    last_uid = cgraph->uid;
+    std::vector<uint8_t> input;
+    serialize_graph(rpc_ctx->device, cgraph, input);
+    bool status = send_rpc_cmd_locked(sock, RPC_CMD_GRAPH_COMPUTE, input.data(), input.size());
+    RPC_STATUS_ASSERT(status);
     return GGML_STATUS_SUCCESS;
 }
 
@@ -2044,7 +2104,6 @@ ggml_backend_reg_t ggml_backend_rpc_add_server(const char * endpoint) {
             /* .device      = */    ind,
             /* .name        = */    dev_name,
             /* .description = */    dev_desc,
-            /* .last_graph_uid = */ 0,
         };
 
         ggml_backend_dev_t dev = new ggml_backend_device {

ggml/src/ggml-rpc/transport.h

@@ -1,18 +1,42 @@
 #pragma once
 
+#include <condition_variable>
 #include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <mutex>
+#include <unordered_map>
 
 struct socket_t;
 typedef std::shared_ptr<socket_t> socket_ptr;
 
 static constexpr size_t MAX_CHUNK_SIZE = 1024ull * 1024ull * 1024ull; // 1 GiB
 static constexpr size_t RPC_CONN_CAPS_SIZE = 24;
 
+// State shared by every client thread that uses one connection. A connection is looked up by
+// endpoint and is therefore shared by all backends of that endpoint, including the backends of
+// different llama_contexts, so all of it has to be serialised:
+//   - mtx_send makes a whole RPC message atomic on the wire
+//   - seq_* hands the responses out in request order (the server answers strictly in order),
+//     without holding mtx_send while waiting, so another thread can keep submitting work
+//   - last_graph_uid mirrors the server's per-connection stored graph for a device, so that
+//     RPC_CMD_GRAPH_RECOMPUTE can never re-run a graph submitted by another context
+struct rpc_conn_state {
+    std::mutex              mtx_send;
+    std::mutex              mtx_seq;
+    std::condition_variable cv_seq;
+    uint64_t                seq_next    = 0;
+    uint64_t                seq_serving = 0;
+
+    std::unordered_map<uint32_t, uint64_t> last_graph_uid;
+};
+
 struct socket_t {
     ~socket_t();
 
+    // guarded by conn.mtx_send / conn.mtx_seq, see rpc_conn_state
+    rpc_conn_state conn;
+
     bool send_data(const void * data, size_t size);
     bool recv_data(void * data, size_t size);
     // Must be called at every message boundary: the RDMA transport coalesces

tools/server/README.md

@@ -2074,6 +2074,74 @@ Note that the following endpoints are exempt from being considered as incoming t
 - `GET /models`
 - `GET /metrics`
 
+## Pipeline groups
+
+`--pipeline-groups N` (default `1`) runs the server's slots over `N` independent `llama_context`
+objects created from the same model. Each group has its own batch, its own sampling and its own
+decode thread; the model weights, the task queue, the results queue and the HTTP layer are shared.
+
+This is meant for a layer split across two machines, e.g.
+
+```sh
+llama-server -m model.gguf -c 32768 --parallel 16 \
+    --rpc peer:50052 --device RPC0,CUDA0 -sm layer -ngl 99 \
+    --pipeline-groups 2
+```
+
+Note the device order: **list the remote device first and the local one last**. With `-sm layer`
+the devices are filled in the order they are given, so the last one holds the output layer. Put
+the local GPU last and the logits are produced locally, which removes a `n_vocab * n_rows * 4`
+byte transfer from every decode step (31.8 MB per step at 32 rows on a 248320-token vocabulary)
+and lets the sampler read them out of local memory. On a pair of DGX Sparks with
+Qwen3.8-27B-UD-Q4_K_XL at 32 concurrent clients this is worth more than the pipeline groups
+themselves, and the two compound:
+
+| device order | groups | tok/s | TPOT ms | GPU busy, local / remote |
+|---|---|---|---|---|
+| `CUDA0,RPC0` | 1 | 94.9 | 310 | 44 / 43 pc |
+| `CUDA0,RPC0` | 2 | 75.5 | 395 | 43 / 44 pc |
+| `RPC0,CUDA0` | 1 | 99.7 | 295 | 41 / 43 pc |
+| `RPC0,CUDA0` | 2 | **130.4** | 223 | 76 / 79 pc |
+
+With one context, a layer split is a two-stage pipeline that is fed one batch at a time, so each
+stage is idle while the other one computes. With two groups there are two batches in flight, so
+while group A is being computed on the second stage, group B is being computed on the first one.
+
+Details:
+
+- The slots are partitioned contiguously: with `--parallel P` and `--pipeline-groups N`, group `g`
+  owns slots `[g*P/N, (g+1)*P/N)`. `--parallel` must be a positive multiple of `--pipeline-groups`.
+- Each context is created with `n_seq_max = P/N` and `n_ctx = C/N`, so the per-slot context and the
+  total KV memory over all groups are the same as with a single context. `-c` must be given
+  explicitly and must be a multiple of `N`.
+- Slot selection for an incoming request still runs over *all* slots, so prompt cache similarity and
+  the slot save / restore endpoints work exactly as before: a returning conversation lands on the
+  slot that still holds its prefix, whichever group that slot belongs to.
+- Task processing briefly pauses the decode loops, so `/slots`, `/metrics` and cancellations are
+  answered after the in-flight decode of each group finishes rather than during it.
+- Speculative decoding works per group: every group owns a draft or MTP context bound to its own
+  target context and a `common_speculative` of its own, so `--spec-type draft-mtp` and
+  `--model-draft` combine with `N > 1` (with `--model-draft` the draft model is loaded once per
+  group). Slot save / restore, checkpoints and the prompt cache carry the draft state exactly as
+  with one group.
+- `N > 1` is refused at startup together with multimodal (`--mmproj`), `--control-vector` and
+  `--sleep-idle-seconds`.
+- With `N = 1` nothing changes: one context, one batch and one update loop on the main thread.
+- Each group samples its own rows over a small worker pool, because a serial pass over the slots
+  costs tens of milliseconds per step between the decode and the next submit and, at `N > 1`,
+  competes for memory bandwidth with the other group's GPU work. The total number of sampling
+  threads is the same however many groups there are; `LLAMA_SERVER_SAMPLE_THREADS=1` turns the
+  pool off. The sampled tokens do not depend on how the pass is scheduled.
+- Use `--cache-ram 0` with a layer split. The RAM prompt cache moves a whole slot state on every
+  slot handover, and on a split most of that state lives on the remote node, so it goes over the
+  wire; at 32 concurrent clients on the pair it costs about a quarter of the throughput with one
+  group and much more with two, because the handover runs on the single task thread and the other
+  group starves while it does.
+- `LLAMA_SERVER_PIPE_PROF=1` prints, every five seconds and per group, how long an iteration
+  spends waiting for the engine, in `pre_decode`, in `llama_decode`, in `llama_synchronize`, and
+  in `post_decode`, and splits `post_decode` per token into sampling, detokenization, stop-string
+  handling and the result queue. That is how the numbers above were found.
+
 ## More examples
 
 ### Interactive mode