rpc : reuse compute graphs

rgerganov · rgerganov · commit 537b2376e266 · 2025-08-18T22:50:20.000+03:00
Store compute graphs on the server side and reuse them when possible. Compute
graphs are kept in a ring buffer with fixed size, so we can avoid serializing
and deserializing the same graph every time. Add two new commands:
 * RPC_CMD_GRAPH_COMPUTE_AND_STORE -- store the graph, compute it and return its ID
 * RPC_CMD_GRAPH_RECOMPUTE -- recompute the graph with the given ID

Currently there is no good way to associate an ID with `ggml_cgraph`,
so we abuse `tensor-&gt;extra` of the first node for this purpose.
diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h
@@ -8,7 +8,7 @@ extern "C" {
 #endif
 
 #define RPC_PROTO_MAJOR_VERSION    2
-#define RPC_PROTO_MINOR_VERSION    0
+#define RPC_PROTO_MINOR_VERSION    5
 #define RPC_PROTO_PATCH_VERSION    0
 #define GGML_RPC_MAX_SERVERS       16
 
diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp
@@ -99,12 +99,17 @@ enum rpc_cmd {
     RPC_CMD_INIT_TENSOR,
     RPC_CMD_GET_ALLOC_SIZE,
     RPC_CMD_HELLO,
+    RPC_CMD_GRAPH_COMPUTE_AND_STORE,
+    RPC_CMD_GRAPH_RECOMPUTE,
     RPC_CMD_COUNT,
 };
 
 // Try RPC_CMD_SET_TENSOR_HASH first when data size is larger than this threshold
 const size_t HASH_THRESHOLD = 10 * 1024 * 1024;
 
+const int     MAX_STORED_GRAPHS = 64;
+const int64_t INVALID_GRAPH_ID = -1;
+
 struct rpc_msg_hello_rsp {
     uint8_t major;
     uint8_t minor;
@@ -186,6 +191,19 @@ struct rpc_msg_graph_compute_rsp {
     uint8_t result;
 };
 
+struct rpc_msg_graph_compute_and_store_rsp {
+    uint8_t result;
+    int32_t graph_id;
+};
+
+struct rpc_msg_graph_recompute_req {
+    int32_t graph_id;
+};
+
+struct rpc_msg_graph_recompute_rsp {
+    uint8_t result;
+};
+
 struct rpc_msg_get_device_memory_rsp {
     uint64_t free_mem;
     uint64_t total_mem;
@@ -209,6 +227,7 @@ struct ggml_backend_rpc_buffer_type_context {
 struct ggml_backend_rpc_context {
     std::string endpoint;
     std::string name;
+    int32_t curr_graph_id;
 };
 
 struct ggml_backend_rpc_buffer_context {
@@ -563,6 +582,8 @@ static enum ggml_status ggml_backend_rpc_buffer_init_tensor(ggml_backend_buffer_
         bool status = send_rpc_cmd(ctx->sock, RPC_CMD_INIT_TENSOR, &request, sizeof(request), nullptr, 0);
         RPC_STATUS_ASSERT(status);
     }
+    // HACK: use the extra field for storing the graph ID
+    tensor->extra = reinterpret_cast<void*>(INVALID_GRAPH_ID);
     return GGML_STATUS_SUCCESS;
 }
 
@@ -772,13 +793,29 @@ static void serialize_graph(const ggml_cgraph * cgraph, std::vector<uint8_t> & o
 
 static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) {
     ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *)backend->context;
-    std::vector<uint8_t> input;
-    serialize_graph(cgraph, input);
-    rpc_msg_graph_compute_rsp response;
-    auto sock = get_socket(rpc_ctx->endpoint);
-    bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_COMPUTE, input.data(), input.size(), &response, sizeof(response));
-    RPC_STATUS_ASSERT(status);
-    return (enum ggml_status)response.result;
+    GGML_ASSERT(cgraph->n_nodes > 0);
+    // HACK: we store the graph ID in the first node's extra field
+    int64_t stored_graph_id = reinterpret_cast<int64_t>(cgraph->nodes[0]->extra);
+    bool reuse_graph = stored_graph_id != INVALID_GRAPH_ID && (stored_graph_id + MAX_STORED_GRAPHS > rpc_ctx->curr_graph_id);
+    if (reuse_graph) {
+        rpc_msg_graph_recompute_req request;
+        request.graph_id = stored_graph_id;
+        rpc_msg_graph_recompute_rsp response;
+        auto sock = get_socket(rpc_ctx->endpoint);
+        bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_RECOMPUTE, &request, sizeof(request), &response, sizeof(response));
+        RPC_STATUS_ASSERT(status);
+        return (enum ggml_status)response.result;
+    } else {
+        std::vector<uint8_t> input;
+        serialize_graph(cgraph, input);
+        rpc_msg_graph_compute_and_store_rsp response;
+        auto sock = get_socket(rpc_ctx->endpoint);
+        bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_COMPUTE_AND_STORE, input.data(), input.size(), &response, sizeof(response));
+        RPC_STATUS_ASSERT(status);
+        rpc_ctx->curr_graph_id = response.graph_id;
+        cgraph->nodes[0]->extra = reinterpret_cast<void*>(response.graph_id);
+        return (enum ggml_status)response.result;
+    }
 }
 
 static ggml_backend_i ggml_backend_rpc_interface = {
@@ -831,8 +868,9 @@ ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint) {
 
 ggml_backend_t ggml_backend_rpc_init(const char * endpoint) {
     ggml_backend_rpc_context * ctx = new ggml_backend_rpc_context {
-        /* .endpoint  = */ endpoint,
-        /* .name      = */ "RPC[" + std::string(endpoint) + "]",
+        /* .endpoint      = */ endpoint,
+        /* .name          = */ "RPC[" + std::string(endpoint) + "]",
+        /*. curr_graph_id = */ 0,
     };
 
     ggml_backend_t backend = new ggml_backend {
@@ -871,7 +909,8 @@ void ggml_backend_rpc_get_device_memory(const char * endpoint, size_t * free, si
 class rpc_server {
 public:
     rpc_server(ggml_backend_t backend, const char * cache_dir)
-        : backend(backend), cache_dir(cache_dir) {
+        : backend(backend), cache_dir(cache_dir), curr_graph_id(0) {
+        stored_graphs.resize(MAX_STORED_GRAPHS);
     }
     ~rpc_server();
 
@@ -887,21 +926,31 @@ class rpc_server {
     bool get_tensor(const rpc_msg_get_tensor_req & request, std::vector<uint8_t> & response);
     bool copy_tensor(const rpc_msg_copy_tensor_req & request, rpc_msg_copy_tensor_rsp & response);
     bool graph_compute(const std::vector<uint8_t> & input, rpc_msg_graph_compute_rsp & response);
+    bool graph_compute_and_store(const std::vector<uint8_t> & input, rpc_msg_graph_compute_and_store_rsp & response);
+    bool graph_recompute(const rpc_msg_graph_recompute_req & request, rpc_msg_graph_recompute_rsp & response);
     bool init_tensor(const rpc_msg_init_tensor_req & request);
     bool get_alloc_size(const rpc_msg_get_alloc_size_req & request, rpc_msg_get_alloc_size_rsp & response);
 
+    struct stored_graph {
+        ggml_context_ptr ctx_ptr;
+        ggml_cgraph * graph;
+    };
+
 private:
     bool get_cached_file(uint64_t hash, std::vector<uint8_t> & data);
     ggml_tensor * deserialize_tensor(struct ggml_context * ctx, const rpc_tensor * tensor);
     ggml_tensor * create_node(uint64_t id,
                               struct ggml_context * ctx,
                               const std::unordered_map<uint64_t, const rpc_tensor*> & tensor_ptrs,
                               std::unordered_map<uint64_t, struct ggml_tensor*> & tensor_map);
-
+    bool store_graph(const std::vector<uint8_t> & input, stored_graph & sg);
 
     ggml_backend_t backend;
     const char * cache_dir;
     std::unordered_set<ggml_backend_buffer_t> buffers;
+    int64_t curr_graph_id;
+    // ring buffer for storing graphs
+    std::vector<stored_graph> stored_graphs;
 };
 
 void rpc_server::hello(rpc_msg_hello_rsp & response) {
@@ -1323,7 +1372,7 @@ ggml_tensor * rpc_server::create_node(uint64_t id,
     return result;
 }
 
-bool rpc_server::graph_compute(const std::vector<uint8_t> & input, rpc_msg_graph_compute_rsp & response) {
+bool rpc_server::store_graph(const std::vector<uint8_t> & input, stored_graph & sg) {
     // serialization format:
     // | n_nodes (4 bytes) | nodes (n_nodes * sizeof(uint64_t) | n_tensors (4 bytes) | tensors (n_tensors * sizeof(rpc_tensor)) |
     if (input.size() < sizeof(uint32_t)) {
@@ -1373,7 +1422,42 @@ bool rpc_server::graph_compute(const std::vector<uint8_t> & input, rpc_msg_graph
             return false;
         }
     }
-    ggml_status status = ggml_backend_graph_compute(backend, graph);
+    sg.ctx_ptr.swap(ctx_ptr);
+    sg.graph = graph;
+    return true;
+}
+
+bool rpc_server::graph_compute(const std::vector<uint8_t> & input, rpc_msg_graph_compute_rsp & response) {
+    stored_graph sg;
+    if (!store_graph(input, sg)) {
+        return false;
+    }
+    ggml_status status = ggml_backend_graph_compute(backend, sg.graph);
+    response.result = status;
+    return true;
+}
+
+bool rpc_server::graph_compute_and_store(const std::vector<uint8_t> & input, rpc_msg_graph_compute_and_store_rsp & response) {
+    int graph_slot = curr_graph_id % MAX_STORED_GRAPHS;
+    if (!store_graph(input, stored_graphs[graph_slot])) {
+        return false;
+    }
+    ggml_status status = ggml_backend_graph_compute(backend, stored_graphs[graph_slot].graph);
+    response.result = status;
+    response.graph_id = curr_graph_id;
+    curr_graph_id++;
+    return true;
+}
+
+bool rpc_server::graph_recompute(const rpc_msg_graph_recompute_req & request, rpc_msg_graph_recompute_rsp & response) {
+    if (request.graph_id < 0) {
+        return false;
+    }
+    int graph_slot = request.graph_id % MAX_STORED_GRAPHS;
+    if (stored_graphs[graph_slot].graph == nullptr) {
+        return false;
+    }
+    ggml_status status = ggml_backend_graph_compute(backend, stored_graphs[graph_slot].graph);
     response.result = status;
     return true;
 }
@@ -1585,6 +1669,34 @@ static void rpc_serve_client(ggml_backend_t backend, const char * cache_dir,
                 }
                 break;
             }
+            case RPC_CMD_GRAPH_COMPUTE_AND_STORE: {
+                std::vector<uint8_t> input;
+                if (!recv_msg(sockfd, input)) {
+                    return;
+                }
+                rpc_msg_graph_compute_and_store_rsp response;
+                if (!server.graph_compute_and_store(input, response)) {
+                    return;
+                }
+                if (!send_msg(sockfd, &response, sizeof(response))) {
+                    return;
+                }
+                break;
+            }
+            case RPC_CMD_GRAPH_RECOMPUTE: {
+                rpc_msg_graph_recompute_req request;
+                if (!recv_msg(sockfd, &request, sizeof(request))) {
+                    return;
+                }
+                rpc_msg_graph_recompute_rsp response;
+                if (!server.graph_recompute(request, response)) {
+                    return;
+                }
+                if (!send_msg(sockfd, &response, sizeof(response))) {
+                    return;
+                }
+                break;
+            }
             case RPC_CMD_GET_DEVICE_MEMORY: {
                 if (!recv_msg(sockfd, nullptr, 0)) {
                     return;
