wip

ggml/src/ggml-backend-impl.h

@@ -221,6 +221,9 @@ extern "C" {
         // (optional) get a pointer to a function in the backend
         // backends can add custom functions that are not part of the standard ggml-backend interface
         void * (*get_proc_address)(ggml_backend_reg_t reg, const char * name);
+
+        // (optional) free backend reg (before unloading the dynamic library)
+        void (*free)(ggml_backend_reg_t reg);
     };
 
     struct ggml_backend_reg {

ggml/src/ggml-backend-reg.cpp

@@ -106,6 +106,8 @@ static std::string path_str(const fs::path & path) {
 struct ggml_backend_reg_entry {
     ggml_backend_reg_t reg;
     dl_handle_ptr handle;
+    // we won't actually unload the backend because dlclose() may cause issues on some backends (e.g. CUDA), so we mark it here and reuse the handle later
+    bool unloaded = false;
 };
 
 struct ggml_backend_registry {
@@ -113,6 +115,10 @@ struct ggml_backend_registry {
     std::vector<ggml_backend_dev_t> devices;
 
     ggml_backend_registry() {
+        load_all_static();
+    }
+
+    void load_all_static() {
 #ifdef GGML_USE_CUDA
         register_backend(ggml_backend_cuda_reg());
 #endif
@@ -183,6 +189,13 @@ struct ggml_backend_registry {
 
         for (auto & entry : backends) {
             if (entry.reg == reg) {
+                if (entry.unloaded) {
+                    // reuse if needed
+                    entry.unloaded = false;
+                    for (size_t i = 0; i < ggml_backend_reg_dev_count(reg); i++) {
+                        register_device(ggml_backend_reg_dev_get(reg, i));
+                    }
+                }
                 return;
             }
         }
@@ -267,18 +280,28 @@ struct ggml_backend_registry {
             return;
         }
 
+        if (it->unloaded) {
+            return;
+        }
+
         if (!silent) {
             GGML_LOG_DEBUG("%s: unloading %s backend\n", __func__, ggml_backend_reg_name(reg));
         }
 
-        // remove devices
+        // remove devices from the registry's bookkeeping
+        // note: the ggml_backend_dev_t objects are kept so the backend can be resumed later
         devices.erase(
             std::remove_if(devices.begin(), devices.end(),
                             [reg](ggml_backend_dev_t dev) { return ggml_backend_dev_backend_reg(dev) == reg; }),
             devices.end());
 
-        // remove backend
-        backends.erase(it);
+        // if backend supports free(), use it
+        if (reg->iface.free) {
+            reg->iface.free(reg);
+        }
+
+        // mark as unloaded, so the handle can be reused later
+        it->unloaded = true;
     }
 };
 
@@ -553,6 +576,7 @@ static ggml_backend_reg_t ggml_backend_load_best(const char * name, bool silent,
 }
 
 void ggml_backend_load_all() {
+    get_reg().load_all_static();
     ggml_backend_load_all_from_path(nullptr);
 }
 

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -5722,17 +5722,33 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con
     return nullptr;
 }
 
+static bool ggml_backend_cuda_reg_initialized = false;
+
+static void ggml_backend_cuda_reg_free(ggml_backend_reg_t reg) {
+    ggml_backend_cuda_reg_context * ctx = (ggml_backend_cuda_reg_context *) reg->context;
+
+    for (ggml_backend_dev_t dev : ctx->devices) {
+        int device = ((ggml_backend_cuda_device_context *) dev->context)->device;
+
+        cudaError_t err = cudaSetDevice(device);
+        if (err == cudaSuccess) {
+            cudaDeviceReset();
+        }
+    }
+}
+
 static const ggml_backend_reg_i ggml_backend_cuda_reg_interface = {
     /* .get_name          = */ ggml_backend_cuda_reg_get_name,
     /* .get_device_count  = */ ggml_backend_cuda_reg_get_device_count,
     /* .get_device        = */ ggml_backend_cuda_reg_get_device,
     /* .get_proc_address  = */ ggml_backend_cuda_reg_get_proc_address,
+    /* .free              = */ ggml_backend_cuda_reg_free,
 };
 
 // backend registry
 ggml_backend_reg_t ggml_backend_cuda_reg() {
     static ggml_backend_reg reg;
-    static bool initialized = false;
+    bool & initialized = ggml_backend_cuda_reg_initialized;
 
     {
         static std::mutex mutex;

src/llama.cpp

@@ -96,9 +96,7 @@ void llama_backend_init(void) {
         ggml_free(ctx);
     }
 
-    if (!ggml_backend_reg_count()) {
-        ggml_backend_load_all();
-    }
+    ggml_backend_load_all();
 }
 
 void llama_numa_init(enum ggml_numa_strategy numa) {
@@ -115,6 +113,9 @@ void llama_numa_init(enum ggml_numa_strategy numa) {
 
 void llama_backend_free(void) {
     ggml_quantize_free();
+    for (size_t i = 0; i < ggml_backend_reg_count(); i++) {
+        ggml_backend_unload(ggml_backend_reg_get(i));
+    }
 }
 
 int64_t llama_time_us(void) {

tools/server/server-context.cpp

@@ -64,6 +64,16 @@ enum slot_state {
     SLOT_STATE_GENERATING,
 };
 
+struct server_backend {
+    server_backend(const common_params & params) {
+        llama_backend_init();
+        llama_numa_init(params.numa);
+    }
+    ~server_backend() {
+        llama_backend_free();
+    }
+};
+
 struct server_slot; // forward declaration
 
 struct server_batch {
@@ -937,6 +947,9 @@ private:
 
     int64_t t_last_load_progress_ms = 0;
 
+    // optional so that it's only initialized in non-router mode
+    std::optional<server_backend> backend;
+
     void destroy() {
         spec.reset();
         ctx_dft.reset();
@@ -949,6 +962,8 @@ private:
 
         mtmd_free(mctx);
         mctx = nullptr;
+
+        backend.reset();
     }
 
     void handle_sleeping_state(bool new_state) {
@@ -1004,6 +1019,8 @@ private:
         load_progress_data load_progress_mmproj(this, "mmproj_model");
         load_progress_data load_progress_spec  (this, "spec_model");
 
+        backend.emplace(params_base);
+
         const bool is_resume = sleeping;
 
         params_base = params;

tools/server/server.cpp

@@ -91,9 +91,6 @@ int llama_server(int argc, char ** argv) {
         return 1;
     }
 
-    llama_backend_init();
-    llama_numa_init(params.numa);
-
     common_models_handler models_handler;
     try {
         models_handler = common_models_handler_init(params, LLAMA_EXAMPLE_SERVER);
@@ -348,7 +345,6 @@ int llama_server(int argc, char ** argv) {
                 models_routes->stopping.store(true); // maybe redundant, but just to be safe
                 models_routes->models.unload_all();
             }
-            llama_backend_free();
         };
 
         if (!ctx_http.start()) {
@@ -374,7 +370,6 @@ int llama_server(int argc, char ** argv) {
             g_stream_sessions.stop_gc();
             ctx_http.stop();
             ctx_server.terminate();
-            llama_backend_free();
         };
 
         // start the HTTP server before loading the model to be able to serve /health requests