[X] working with rpc, but slow

Author: LeaveNhA

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

@@ -33,6 +33,9 @@
 
 namespace fs = std::filesystem;
 
+// Forward declaration for device map access
+static std::unordered_map<std::string, ggml_backend_dev_t>& get_rpc_dev_map();
+
 static constexpr size_t MAX_CHUNK_SIZE = 1024ull * 1024ull * 1024ull; // 1 GiB
 
 #ifdef _WIN32
@@ -1760,16 +1763,33 @@ static const char * ggml_backend_rpc_reg_get_name(ggml_backend_reg_t reg) {
 }
 
 static size_t ggml_backend_rpc_reg_get_device_count(ggml_backend_reg_t reg) {
-    return 0;
+    const auto& dev_map = get_rpc_dev_map();
+    return dev_map.size();
 
     GGML_UNUSED(reg);
 }
 
 static ggml_backend_dev_t ggml_backend_rpc_reg_get_device(ggml_backend_reg_t reg, size_t index) {
-    GGML_ABORT("The RPC backend does not have enumerated devices - use ggml_backend_add_device instead");
+    const auto& dev_map = get_rpc_dev_map();
+    
+    if (index >= dev_map.size()) {
+        return nullptr;
+    }
+    
+    // Convert unordered_map to vector to access by index
+    std::vector<ggml_backend_dev_t> devices;
+    devices.reserve(dev_map.size());
+    for (const auto& pair : dev_map) {
+        devices.push_back(pair.second);
+    }
+    
+    if (index < devices.size()) {
+        return devices[index];
+    }
+    
+    return nullptr;
 
     GGML_UNUSED(reg);
-    GGML_UNUSED(index);
 }
 
 static ggml_backend_buffer_type_t ggml_backend_rpc_split_buffer_type(int main_device, const float * tensor_split) {
@@ -1818,8 +1838,14 @@ ggml_backend_reg_t ggml_backend_rpc_reg(void) {
     return &ggml_backend_rpc_reg;
 }
 
-ggml_backend_dev_t ggml_backend_rpc_add_device(const char * endpoint) {
+// Expose the device map for enumeration
+static std::unordered_map<std::string, ggml_backend_dev_t>& get_rpc_dev_map() {
     static std::unordered_map<std::string, ggml_backend_dev_t> dev_map;
+    return dev_map;
+}
+
+ggml_backend_dev_t ggml_backend_rpc_add_device(const char * endpoint) {
+    auto& dev_map = get_rpc_dev_map();
 
     static std::mutex mutex;
     std::lock_guard<std::mutex> lock(mutex);

src/llama-model.cpp

@@ -377,8 +377,13 @@ static buft_list_t make_gpu_buft_list(ggml_backend_dev_t dev, llama_split_mode s
         if (ggml_backend_split_buffer_type_fn) {
             size_t dev_index = [&]() {
                 auto * reg = ggml_backend_dev_backend_reg(dev);
-                for (size_t i = 0; i < ggml_backend_reg_dev_count(reg); ++i) {
-                    if (ggml_backend_reg_dev_get(reg, i) == dev) {
+                size_t reg_dev_count = ggml_backend_reg_dev_count(reg);
+                LLAMA_LOG_DEBUG("%s: device %s, reg %s, device count %zu\n", __func__, ggml_backend_dev_name(dev), ggml_backend_reg_name(reg), reg_dev_count);
+                for (size_t i = 0; i < reg_dev_count; ++i) {
+                    ggml_backend_dev_t reg_dev = ggml_backend_reg_dev_get(reg, i);
+                    LLAMA_LOG_DEBUG("%s: comparing device %s with reg device %s at index %zu\n", __func__, ggml_backend_dev_name(dev), ggml_backend_dev_name(reg_dev), i);
+                    if (reg_dev == dev) {
+                        LLAMA_LOG_DEBUG("%s: found device %s at index %zu\n", __func__, ggml_backend_dev_name(dev), i);
                         return i;
                     }
                 }

tools/llama-bench/llama-bench.cpp

@@ -20,6 +20,7 @@
 
 #include "common.h"
 #include "ggml.h"
+#include "ggml-rpc.h"
 #include "llama.h"
 
 #ifdef _WIN32
@@ -1827,6 +1828,39 @@ int main(int argc, char ** argv) {
 
     cmd_params params = parse_cmd_params(argc, argv);
 
+    // Register RPC devices if specified
+    for (const auto& rpc_servers_str : params.rpc_servers) {
+        if (!rpc_servers_str.empty()) {
+            auto rpc_servers = string_split<std::string>(rpc_servers_str, ',');
+            if (!rpc_servers.empty()) {
+                ggml_backend_reg_t rpc_reg = ggml_backend_reg_by_name("RPC");
+                if (!rpc_reg) {
+                    fprintf(stderr, "%s: failed to find RPC backend\n", __func__);
+                    return 1;
+                }
+
+                typedef ggml_backend_dev_t (*ggml_backend_rpc_add_device_t)(const char * endpoint);
+                ggml_backend_rpc_add_device_t ggml_backend_rpc_add_device_fn = 
+                    (ggml_backend_rpc_add_device_t) ggml_backend_reg_get_proc_address(rpc_reg, "ggml_backend_rpc_add_device");
+                if (!ggml_backend_rpc_add_device_fn) {
+                    fprintf(stderr, "%s: failed to find RPC device add function\n", __func__);
+                    return 1;
+                }
+
+                // Register each RPC device
+                for (const std::string & server : rpc_servers) {
+                    ggml_backend_dev_t dev = ggml_backend_rpc_add_device_fn(server.c_str());
+                    if (dev) {
+                        ggml_backend_device_register(dev);
+                    } else {
+                        fprintf(stderr, "%s: failed to add RPC device for server '%s'\n", __func__, server.c_str());
+                        return 1;
+                    }
+                }
+            }
+        }
+    }
+
     auto * cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
     if (!cpu_dev) {
         fprintf(stderr, "%s: error: CPU backend is not loaded\n", __func__);