rpc: avoid serializing buffers from other servers (llama/26500)

* rpc: avoid serializing buffers from other servers

Only include remote buffer pointers when the buffer belongs to the RPC dispatcher receiving the graph. Add a two-server regression test for cross-server tensor serialization.

Assisted-by: Codex

* cont : add ref

---------

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
This commit is contained in:
hmirin
2026-09-04 13:39:35 +03:00
committed by Georgi Gerganov
co-authored by Georgi Gerganov
parent 749683d30f
commit 4089fa628a
+16 -10
View File
@@ -625,7 +625,7 @@ static bool ggml_backend_buffer_is_rpc(ggml_backend_buffer_t buffer) {
return buffer->iface.free_buffer == ggml_backend_rpc_buffer_free_buffer; return buffer->iface.free_buffer == ggml_backend_rpc_buffer_free_buffer;
} }
static rpc_tensor serialize_tensor(const ggml_tensor * tensor) { static rpc_tensor serialize_tensor(const ggml_tensor * tensor, const std::shared_ptr<rpc_dispatcher> & dispatcher = nullptr) {
rpc_tensor result; rpc_tensor result;
if (!tensor) { if (!tensor) {
memset(&result, 0, sizeof(result)); memset(&result, 0, sizeof(result));
@@ -637,8 +637,14 @@ static rpc_tensor serialize_tensor(const ggml_tensor * tensor) {
if (tensor->buffer && ggml_backend_buffer_is_rpc(tensor->buffer)) { if (tensor->buffer && ggml_backend_buffer_is_rpc(tensor->buffer)) {
ggml_backend_buffer_t buffer = tensor->buffer; ggml_backend_buffer_t buffer = tensor->buffer;
ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context; ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
result.buffer = ctx != nullptr ? ctx->remote_ptr : 0; // ref: https://github.com/ggml-org/llama.cpp/pull/26500
result.data = reinterpret_cast<uint64_t>(tensor->data); if (ctx != nullptr && (dispatcher == nullptr || ctx->dispatcher == dispatcher)) {
result.buffer = ctx->remote_ptr;
result.data = reinterpret_cast<uint64_t>(tensor->data);
} else {
result.buffer = 0;
result.data = 0;
}
} else { } else {
result.buffer = 0; result.buffer = 0;
result.data = 0; result.data = 0;
@@ -958,7 +964,7 @@ static void ggml_backend_rpc_synchronize(ggml_backend_t backend) {
rpc_ctx->dispatcher->synchronize(); rpc_ctx->dispatcher->synchronize();
} }
static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::vector<rpc_tensor> & tensors, std::unordered_set<ggml_tensor*> & visited) { static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, const std::shared_ptr<rpc_dispatcher> & dispatcher, std::vector<rpc_tensor> & tensors, std::unordered_set<ggml_tensor*> & visited) {
if (tensor == nullptr) { if (tensor == nullptr) {
return; return;
} }
@@ -967,10 +973,10 @@ static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::ve
} }
visited.insert(tensor); visited.insert(tensor);
for (int i = 0; i < GGML_MAX_SRC; i++) { for (int i = 0; i < GGML_MAX_SRC; i++) {
add_tensor(tensor->src[i], cgraph, tensors, visited); add_tensor(tensor->src[i], cgraph, dispatcher, tensors, visited);
} }
add_tensor(tensor->view_src, cgraph, tensors, visited); add_tensor(tensor->view_src, cgraph, dispatcher, tensors, visited);
rpc_tensor result = serialize_tensor(tensor); rpc_tensor result = serialize_tensor(tensor, dispatcher);
const size_t hash_pos = ggml_hash_find(&cgraph->visited_hash_set, tensor); const size_t hash_pos = ggml_hash_find(&cgraph->visited_hash_set, tensor);
if (hash_pos != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, hash_pos)) { if (hash_pos != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, hash_pos)) {
result.use_count = cgraph->use_counts[hash_pos]; result.use_count = cgraph->use_counts[hash_pos];
@@ -978,12 +984,12 @@ static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::ve
tensors.push_back(result); tensors.push_back(result);
} }
static uint8_t * serialize_graph(uint32_t device, const ggml_cgraph * cgraph, size_t * output_size) { static uint8_t * serialize_graph(uint32_t device, const ggml_cgraph * cgraph, const std::shared_ptr<rpc_dispatcher> & dispatcher, size_t * output_size) {
uint32_t n_nodes = cgraph->n_nodes; uint32_t n_nodes = cgraph->n_nodes;
std::vector<rpc_tensor> tensors; std::vector<rpc_tensor> tensors;
std::unordered_set<ggml_tensor*> visited; std::unordered_set<ggml_tensor*> visited;
for (uint32_t i = 0; i < n_nodes; i++) { for (uint32_t i = 0; i < n_nodes; i++) {
add_tensor(cgraph->nodes[i], cgraph, tensors, visited); add_tensor(cgraph->nodes[i], cgraph, dispatcher, tensors, visited);
} }
// serialization format: // serialization format:
// | device (4 bytes) | n_nodes (4 bytes) | nodes (n_nodes * sizeof(uint64_t) | n_tensors (4 bytes) | tensors (n_tensors * sizeof(rpc_tensor)) | // | device (4 bytes) | n_nodes (4 bytes) | nodes (n_nodes * sizeof(uint64_t) | n_tensors (4 bytes) | tensors (n_tensors * sizeof(rpc_tensor)) |
@@ -1020,7 +1026,7 @@ static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, g
} else { } else {
rpc_dev_ctx->last_graph_uid = cgraph->uid; rpc_dev_ctx->last_graph_uid = cgraph->uid;
size_t input_size = 0; size_t input_size = 0;
uint8_t * input = serialize_graph(rpc_ctx->device, cgraph, &input_size); uint8_t * input = serialize_graph(rpc_ctx->device, cgraph, rpc_ctx->dispatcher, &input_size);
std::shared_ptr<uint8_t> input_ptr(input, std::default_delete<uint8_t[]>()); std::shared_ptr<uint8_t> input_ptr(input, std::default_delete<uint8_t[]>());
rpc_ctx->dispatcher->send_async(RPC_CMD_GRAPH_COMPUTE, input_ptr, input_size); rpc_ctx->dispatcher->send_async(RPC_CMD_GRAPH_COMPUTE, input_ptr, input_size);
} }