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

src/conditioning/conditioner.hpp

@@ -117,6 +117,7 @@ public:
     virtual SDCondition get_learned_condition(int n_threads,
                                               const ConditionerParams& conditioner_params) = 0;
     virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors)           = 0;
+    virtual void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {}
     virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
     virtual void set_stream_layers_enabled(bool enabled) {}
     virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {}
@@ -1664,6 +1665,10 @@ struct AnimaConditioner : public Conditioner {
         llm->get_param_tensors(tensors, "text_encoders.llm");
     }
 
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
+        llm->get_param_tensor_ops(tensor_ops);
+    }
+
     void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
         llm->set_max_graph_vram_bytes(max_vram_bytes);
     }
@@ -1847,6 +1852,10 @@ struct LLMEmbedder : public Conditioner {
         }
     }
 
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
+        llm->get_param_tensor_ops(tensor_ops);
+    }
+
     void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
         llm->set_max_graph_vram_bytes(max_vram_bytes);
         if (byt5) {
@@ -2828,6 +2837,10 @@ struct LTXAVEmbedder : public Conditioner {
         projector->get_param_tensors(tensors, "text_embedding_projection");
     }
 
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
+        llm->get_param_tensor_ops(tensor_ops);
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         llm->set_flash_attention_enabled(enabled);
         projector->set_flash_attention_enabled(enabled);

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 scheduler only allocates standalone data tensors when they are
-            // The sched only assigns a backend (and thus a buffer) to tensors
+        // marked as graph inputs. The flag is harmless for single-backend graphs.
-            // that participate in the graph; flag standalone data tensors as
-            // inputs so they get one.
         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),

src/model/te/llm.hpp

@@ -1657,6 +1657,10 @@ namespace LLM {
             model.get_param_tensors(tensors, prefix);
         }
 
+        void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
+            model.get_param_tensor_ops(tensor_ops);
+        }
+
         ggml_tensor* forward(GGMLRunnerContext* ctx,
                              ggml_tensor* input_ids,
                              ggml_tensor* input_pos,

src/model_loader.cpp

@@ -1053,7 +1053,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
         std::atomic<size_t> tensor_idx(0);
         std::atomic<bool> failed(false);
         std::vector<std::thread> workers;
-        std::mutex rpc_backend_mutex;
+        std::mutex backend_tensor_set_mutex;
 
         for (int i = 0; i < n_threads; ++i) {
             workers.emplace_back([&, file_path, is_zip]() {
@@ -1214,17 +1214,8 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                     if (dst_tensor->buffer != nullptr && !ggml_backend_buffer_is_host(dst_tensor->buffer)) {
                         t0 = ggml_time_ms();
 
-                        // RPC backends require serialized access to prevent concurrency issues
+                        std::lock_guard<std::mutex> lock(backend_tensor_set_mutex);
-                        const char* buffer_type_name = ggml_backend_buft_name(ggml_backend_buffer_get_type(dst_tensor->buffer));
-                        bool is_rpc_buffer           = buffer_type_name != nullptr &&
-                                             std::string(buffer_type_name).find("RPC") != std::string::npos;
-
-                        if (is_rpc_buffer) {
-                            std::lock_guard<std::mutex> lock(rpc_backend_mutex);
                         ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
-                        } else {
-                            ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
-                        }
 
                         t1 = ggml_time_ms();
                         copy_to_backend_time_ms.fetch_add(t1 - t0);

src/model_manager.cpp

@@ -53,6 +53,48 @@ static bool backend_supports_host_buffer(ggml_backend_t backend) {
     return props.caps.buffer_from_host_ptr;
 }
 
+static bool device_supports_param_op(ggml_backend_dev_t device,
+                                      ggml_tensor* weight,
+                                      enum ggml_op op,
+                                      ggml_backend_buffer_type_t buft) {
+    if (op == GGML_OP_NONE) {
+        return true;
+    }
+    if (device == nullptr || weight == nullptr || buft == nullptr || weight->buffer != nullptr) {
+        return false;
+    }
+
+    ggml_init_params params;
+    params.mem_size   = ggml_tensor_overhead() * 2;
+    params.mem_buffer = nullptr;
+    params.no_alloc   = true;
+    ggml_context* ctx = ggml_init(params);
+    if (ctx == nullptr) {
+        return false;
+    }
+
+    ggml_tensor* op_tensor = nullptr;
+    if (op == GGML_OP_GET_ROWS) {
+        ggml_tensor* indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1);
+        op_tensor            = ggml_get_rows(ctx, weight, indices);
+    }
+    if (op_tensor == nullptr) {
+        ggml_free(ctx);
+        return false;
+    }
+
+    weight->buffer = ggml_backend_buft_alloc_buffer(buft, 0);
+    if (weight->buffer == nullptr) {
+        ggml_free(ctx);
+        return false;
+    }
+    bool supported = ggml_backend_dev_supports_op(device, op_tensor);
+    ggml_backend_buffer_free(weight->buffer);
+    weight->buffer = nullptr;
+    ggml_free(ctx);
+    return supported;
+}
+
 ModelManager::~ModelManager() {
     release_all();
 }
@@ -135,7 +177,8 @@ bool ModelManager::register_param_tensors(const std::string& desc,
                                           ggml_backend_t params_backend,
                                           size_t* registered_tensor_size,
                                           bool allow_split_buffer,
-                                          bool params_follow_compute_backend) {
+                                          bool params_follow_compute_backend,
+                                          const std::map<ggml_tensor*, enum ggml_op>* tensor_ops) {
     if (desc.empty()) {
         LOG_ERROR("model manager tensor desc is empty");
         return false;
@@ -168,6 +211,12 @@ bool ModelManager::register_param_tensors(const std::string& desc,
         state->params_backend                = params_backend;
         state->allow_split_buffer            = allow_split_buffer;
         state->params_follow_compute_backend = params_follow_compute_backend;
+        if (tensor_ops != nullptr) {
+            auto op_it = tensor_ops->find(tensor);
+            if (op_it != tensor_ops->end()) {
+                state->usage_op = op_it->second;
+            }
+        }
         new_states.push_back(std::move(state));
     }
 
@@ -844,6 +893,22 @@ ggml_backend_buffer_type_t ModelManager::params_buffer_type_for(const TensorStat
     if (params_buft == nullptr) {
         params_buft = ggml_backend_get_default_buffer_type(state.params_backend);
     }
+    if (state.usage_op != GGML_OP_NONE &&
+        state.compute_backend != nullptr) {
+        ggml_backend_dev_t compute_dev = ggml_backend_get_device(state.compute_backend);
+        if (device_supports_param_op(compute_dev, state.tensor, state.usage_op, params_buft)) {
+            return params_buft;
+        }
+
+        ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
+        params_buft                = cpu_dev != nullptr ? ggml_backend_dev_buffer_type(cpu_dev) : nullptr;
+        if (!device_supports_param_op(cpu_dev, state.tensor, state.usage_op, params_buft)) {
+            LOG_ERROR("model manager has no compatible buffer for tensor '%s' used by %s",
+                      state.name.c_str(),
+                      ggml_op_name(state.usage_op));
+            return nullptr;
+        }
+    }
     return params_buft;
 }
 

src/model_manager.h

@@ -39,6 +39,7 @@ private:
         bool allow_split_buffer            = false;
         bool params_follow_compute_backend = false;
         bool metadata_validated            = false;
+        enum ggml_op usage_op              = GGML_OP_NONE;
 
         int active_prepare_count = 0;
 
@@ -132,7 +133,8 @@ public:
                                 ggml_backend_t params_backend,
                                 size_t* registered_tensor_size     = nullptr,
                                 bool allow_split_buffer            = false,
-                                bool params_follow_compute_backend = false);
+                                bool params_follow_compute_backend = false,
+                                const std::map<ggml_tensor*, enum ggml_op>* tensor_ops = nullptr);
 
     bool unregister_param_tensors(const std::string& desc,
                                   size_t* registered_tensor_size = nullptr);

src/stable-diffusion.cpp

@@ -320,7 +320,11 @@ public:
             return true;
         }
         std::map<std::string, ggml_tensor*> group_tensors;
+        std::map<ggml_tensor*, enum ggml_op> tensor_ops;
         model->get_param_tensors(group_tensors);
+        if constexpr (std::is_base_of_v<Conditioner, T>) {
+            model->get_param_tensor_ops(tensor_ops);
+        }
         if (model_manager == nullptr) {
             return true;
         }
@@ -337,6 +341,7 @@ public:
                                                                 module,
                                                                 module_backends,
                                                                 std::move(group_tensors),
+                                                                tensor_ops,
                                                                 residency_mode,
                                                                 params_mem_size);
                     }
@@ -345,6 +350,7 @@ public:
                                                               module,
                                                               module_backends,
                                                               std::move(group_tensors),
+                                                              tensor_ops,
                                                               residency_mode,
                                                               params_mem_size);
                 }
@@ -358,7 +364,10 @@ public:
                                                      residency_mode,
                                                      backend_for(module),
                                                      params_backend_for(module),
-                                                     params_mem_size);
+                                                     params_mem_size,
+                                                     false,
+                                                     false,
+                                                     &tensor_ops);
     }
 
     template <typename T>
@@ -367,6 +376,7 @@ public:
                                           SDBackendModule module,
                                           const std::vector<ggml_backend_t>& module_backends,
                                           std::map<std::string, ggml_tensor*> group_tensors,
+                                          const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
                                           ModelManager::ResidencyMode residency_mode,
                                           size_t* params_mem_size) {
         ggml_backend_t main_backend = module_backends[0];
@@ -378,6 +388,7 @@ public:
                                                       module,
                                                       module_backends,
                                                       std::move(group_tensors),
+                                                      tensor_ops,
                                                       residency_mode,
                                                       params_mem_size);
         };
@@ -452,7 +463,9 @@ public:
                                                    main_backend,
                                                    params_backend_for(module),
                                                    params_mem_size,
-                                                   /*allow_split_buffer=*/true)) {
+                                                   /*allow_split_buffer=*/true,
+                                                   false,
+                                                   &tensor_ops)) {
             return false;
         }
         return model_manager->register_param_tensors(desc,
@@ -460,7 +473,10 @@ public:
                                                      residency_mode,
                                                      main_backend,
                                                      params_backend_for(module),
-                                                     params_mem_size);
+                                                     params_mem_size,
+                                                     false,
+                                                     false,
+                                                     &tensor_ops);
     }
 
     // Register graph-cut layer-split tensors on the primary backend first.
@@ -472,6 +488,7 @@ public:
                                             SDBackendModule module,
                                             const std::vector<ggml_backend_t>& module_backends,
                                             std::map<std::string, ggml_tensor*> group_tensors,
+                                            const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
                                             ModelManager::ResidencyMode residency_mode,
                                             size_t* params_mem_size) {
         bool has_cpu_device = false;
@@ -493,7 +510,10 @@ public:
                                                          residency_mode,
                                                          module_backends[0],
                                                          params_backend_for(module),
-                                                         params_mem_size);
+                                                         params_mem_size,
+                                                         false,
+                                                         false,
+                                                         &tensor_ops);
         }
 
         model->set_runtime_backends(module_backends);
@@ -518,7 +538,8 @@ public:
                                                      initial_params_backend,
                                                      params_mem_size,
                                                      false,
-                                                     params_follow_runtime);
+                                                     params_follow_runtime,
+                                                     &tensor_ops);
     }
 
     bool unload_control_net() {