fix: make parameter loading backend-aware (#1828)

src/core/ggml_extend.hpp

@@ -1753,7 +1753,7 @@ protected:
     std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
 
     std::vector<ggml_backend_t> extra_runtime_backends;  // borrowed (SDBackendManager-owned)
-    ggml_backend_sched_t sched             = nullptr;    // owned, multi-device only
+    ggml_backend_sched_t sched             = nullptr;    // owned
     ggml_backend_t cpu_fallback_backend    = nullptr;    // owned, sched requires a trailing CPU backend
     bool multi_device_eval_callback_warned = false;
 
@@ -2147,8 +2147,22 @@ protected:
         return !extra_runtime_backends.empty();
     }
 
+    bool graph_requires_backend_fallback(ggml_cgraph* gf) const {
+        if (gf == nullptr || sd_backend_is_cpu(runtime_backend)) {
+            return false;
+        }
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        for (int i = 0; i < n_nodes; ++i) {
+            ggml_tensor* node = ggml_graph_node(gf, i);
+            if (node != nullptr && !ggml_backend_supports_op(runtime_backend, node)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     bool alloc_compute_buffer(ggml_cgraph* gf) {
-        if (is_multi_device()) {
+        if (sched != nullptr || is_multi_device() || graph_requires_backend_fallback(gf)) {
             // The sched replaces the gallocr. Do NOT ggml_backend_sched_reserve
             // the graph here: reserve runs split_graph, which rewires the
             // graph's src pointers to sched-internal copy tensors, and the
@@ -2156,6 +2170,10 @@ protected:
             // rewired graph, silently corrupting every cross-backend input. A
             // graph must be split at most once; the alloc in execute_graph
             // performs the real allocation.
+            if (compute_allocr != nullptr) {
+                ggml_gallocr_free(compute_allocr);
+                compute_allocr = nullptr;
+            }
             return ensure_sched(gf);
         }
         if (compute_allocr != nullptr) {
@@ -2753,7 +2771,7 @@ protected:
         };
         ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer);
 
-        if (is_multi_device()) {
+        if (sched != nullptr) {
             ggml_backend_sched_reset(sched);
             pin_multi_device_nodes(gf);  // reset clears the pins; re-apply before alloc
             if (!ggml_backend_sched_alloc_graph(sched, gf)) {
@@ -2774,9 +2792,9 @@ protected:
         }
 
         ggml_status status;
-        if (is_multi_device()) {
+        if (sched != nullptr) {
             if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) {
-                LOG_WARN("%s: eval callback is not supported with multiple runtime backends; ignoring",
+                LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring",
                          get_desc().c_str());
                 multi_device_eval_callback_warned = true;
             }
@@ -3018,12 +3036,9 @@ public:
 
     // do copy after alloc graph
     void set_backend_tensor_data(ggml_tensor* tensor, const void* data) {
-        if (is_multi_device()) {
-            // The sched only assigns a backend (and thus a buffer) to tensors
-            // that participate in the graph; flag standalone data tensors as
-            // inputs so they get one.
-            ggml_set_input(tensor);
-        }
+        // The scheduler only allocates standalone data tensors when they are
+        // marked as graph inputs. The flag is harmless for single-backend graphs.
+        ggml_set_input(tensor);
         backend_tensor_data_map[tensor] = data;
     }
 
@@ -3240,6 +3255,11 @@ protected:
 
     virtual void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {}
 
+    virtual enum ggml_op param_usage_op(const std::string& name) const {
+        (void)name;
+        return GGML_OP_NONE;
+    }
+
 public:
     void init(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") {
         if (prefix.size() > 0) {
@@ -3290,6 +3310,18 @@ public:
         }
     }
 
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
+        for (auto& pair : blocks) {
+            pair.second->get_param_tensor_ops(tensor_ops);
+        }
+        for (auto& pair : params) {
+            enum ggml_op op = param_usage_op(pair.first);
+            if (op != GGML_OP_NONE) {
+                tensor_ops[pair.second] = op;
+            }
+        }
+    }
+
     virtual std::string get_desc() {
         return "GGMLBlock";
     }
@@ -3417,6 +3449,10 @@ protected:
         params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings);
     }
 
+    enum ggml_op param_usage_op(const std::string& name) const override {
+        return name == "weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE;
+    }
+
 public:
     Embedding(int64_t num_embeddings, int64_t embedding_dim)
         : embedding_dim(embedding_dim),