Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

rpc: avoid serializing buffers from other servers (#26500)

Browse Source 
* 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-08-30 20:26:16 +03:00
committed by GitHub
co-authored by Georgi Gerganov
parent 6d1479c148
commit a7cc83bbae
4 changed files with 118 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-rpc/ggml-rpc.cpp
+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;
}
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;
if (!tensor) {
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)) {
ggml_backend_buffer_t buffer = tensor->buffer;
ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
result.buffer = ctx != nullptr ? ctx->remote_ptr : 0;
result.data = reinterpret_cast<uint64_t>(tensor->data);
// ref: https://github.com/ggml-org/llama.cpp/pull/26500
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 {
result.buffer = 0;
result.data = 0;
@@ -958,7 +964,7 @@ static void ggml_backend_rpc_synchronize(ggml_backend_t backend) {
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) {
return;
}
@@ -967,10 +973,10 @@ static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::ve
}
visited.insert(tensor);
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);
rpc_tensor result = serialize_tensor(tensor);
add_tensor(tensor->view_src, cgraph, dispatcher, tensors, visited);
rpc_tensor result = serialize_tensor(tensor, dispatcher);
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)) {
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);
}
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;
std::vector<rpc_tensor> tensors;
std::unordered_set<ggml_tensor*> visited;
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:
// | 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 {
rpc_dev_ctx->last_graph_uid = cgraph->uid;
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[]>());
rpc_ctx->dispatcher->send_async(RPC_CMD_GRAPH_COMPUTE, input_ptr, input_size);
}