wylab/llama.cpp
0
0
Fork 1
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-14 18:02:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

move graph_uids to rpc_dispatcher

Browse Source
This commit is contained in:
Aman Gupta
2026-09-14 13:12:33 +03:00
committed by ggerganov
parent d3b97d0821
commit f9b423302d
1 changed files with 23 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-rpc/ggml-rpc.cpp
+23 -21
View File
@@ -258,8 +258,6 @@ struct ggml_backend_rpc_device_context {
uint32_t device;
std::string name;
std::string description;
// uids of graphs cached by the server for this device
std::unordered_set<uint64_t> graph_uids;
};
struct ggml_backend_rpc_buffer_type_context {
@@ -480,6 +478,7 @@ public:
void synchronize();
void busy_spin_acquire();
void busy_spin_release();
void graph_compute(uint32_t device, const ggml_cgraph * cgraph);
void start(const std::string & endpoint);
void work();
@@ -501,6 +500,8 @@ private:
rpc_msg_ptr msg;
std::shared_future<void> sf;
};
std::mutex graph_mutex;
std::unordered_map<uint32_t, std::unordered_set<uint64_t>> graph_uids;
rpc_msg_queue queue;
socket_ptr sock;
std::atomic_uint busy_spin_users = 0;
@@ -702,7 +703,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, const std::shared_ptr<rpc_dispatcher> & dispatcher = nullptr) {
static rpc_tensor serialize_tensor(const ggml_tensor * tensor, const rpc_dispatcher * dispatcher = nullptr) {
rpc_tensor result;
if (!tensor) {
memset(&result, 0, sizeof(result));
@@ -715,7 +716,7 @@ static rpc_tensor serialize_tensor(const ggml_tensor * tensor, const std::shared
ggml_backend_buffer_t buffer = tensor->buffer;
ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
// ref: https://github.com/ggml-org/llama.cpp/pull/26500
if (ctx != nullptr && (dispatcher == nullptr || ctx->dispatcher == dispatcher)) {
if (ctx != nullptr && (dispatcher == nullptr || ctx->dispatcher.get() == dispatcher)) {
result.buffer = ctx->remote_ptr;
result.data = reinterpret_cast<uint64_t>(tensor->data);
} else {
@@ -1081,7 +1082,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, const std::shared_ptr<rpc_dispatcher> & dispatcher, std::vector<rpc_tensor> & tensors, std::unordered_set<ggml_tensor*> & visited) {
static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, const rpc_dispatcher * dispatcher, std::vector<rpc_tensor> & tensors, std::unordered_set<ggml_tensor*> & visited) {
if (tensor == nullptr) {
return;
}
@@ -1101,7 +1102,7 @@ static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, const s
tensors.push_back(result);
}
static uint8_t * serialize_graph(uint32_t device, const ggml_cgraph * cgraph, const std::shared_ptr<rpc_dispatcher> & dispatcher, size_t * output_size) {
static uint8_t * serialize_graph(uint32_t device, const ggml_cgraph * cgraph, const 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;
@@ -1131,31 +1132,33 @@ static uint8_t * serialize_graph(uint32_t device, const ggml_cgraph * cgraph, co
return output;
}
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;
ggml_backend_dev_t rpc_dev = ggml_backend_get_device(backend);
ggml_backend_rpc_device_context * rpc_dev_ctx = (ggml_backend_rpc_device_context *)rpc_dev->context;
void rpc_dispatcher::graph_compute(uint32_t device, const ggml_cgraph * cgraph) {
std::lock_guard<std::mutex> lock(graph_mutex);
GGML_ASSERT(cgraph->n_nodes > 0);
auto & graph_uids = rpc_dev_ctx->graph_uids;
bool reuse = cgraph->uid != 0 && graph_uids.count(cgraph->uid) > 0;
auto & device_graph_uids = graph_uids[device];
bool reuse = cgraph->uid != 0 && device_graph_uids.count(cgraph->uid) > 0;
if (reuse) {
auto request = std::make_shared<rpc_msg_graph_recompute_req>();
request->device = rpc_ctx->device;
request->device = device;
request->uid = cgraph->uid;
rpc_ctx->dispatcher->send_async(RPC_CMD_GRAPH_RECOMPUTE, request, sizeof(*request));
send_async(RPC_CMD_GRAPH_RECOMPUTE, request, sizeof(*request));
} else {
if (cgraph->uid != 0) {
if (graph_uids.size() >= GRAPH_CACHE_MAX) {
graph_uids.clear();
if (device_graph_uids.size() >= GRAPH_CACHE_MAX) {
device_graph_uids.clear();
}
graph_uids.insert(cgraph->uid);
device_graph_uids.insert(cgraph->uid);
}
size_t input_size = 0;
uint8_t * input = serialize_graph(rpc_ctx->device, cgraph, rpc_ctx->dispatcher, &input_size);
uint8_t * input = serialize_graph(device, cgraph, this, &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);
send_async(RPC_CMD_GRAPH_COMPUTE, input_ptr, input_size);
}
}
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;
rpc_ctx->dispatcher->graph_compute(rpc_ctx->device, cgraph);
return GGML_STATUS_SUCCESS;
}
@@ -3016,7 +3019,6 @@ ggml_backend_reg_t ggml_backend_rpc_add_server(const char * endpoint) {
/* .device = */ ind,
/* .name = */ dev_name,
/* .description = */ dev_desc,
/* .graph_uids = */ {},
};
ggml_backend_dev_t dev = new ggml_backend_device {