feat: drive layer split from graph-cut segments (#1762)

src/conditioning/conditioner.hpp

@@ -118,6 +118,8 @@ public:
     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) {}
+    virtual void set_graph_cut_layer_split_enabled(bool enabled) {}
+    virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {}
     virtual void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {}
     virtual void set_flash_attention_enabled(bool enabled) = 0;
     virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
@@ -181,6 +183,27 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         }
     }
 
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        text_model->set_runtime_backends(backends);
+        if (sd_version_is_sdxl(version)) {
+            text_model2->set_runtime_backends(backends);
+        }
+    }
+
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        text_model->set_graph_cut_layer_split_enabled(enabled);
+        if (sd_version_is_sdxl(version)) {
+            text_model2->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        text_model->set_graph_cut_layer_split_backend_vram_limits(limits);
+        if (sd_version_is_sdxl(version)) {
+            text_model2->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         text_model->set_flash_attention_enabled(enabled);
         if (sd_version_is_sdxl(version)) {
@@ -639,11 +662,41 @@ struct SD3CLIPEmbedder : public Conditioner {
     }
 
     void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        if (clip_l) {
+            clip_l->set_runtime_backends(backends);
+        }
+        if (clip_g) {
+            clip_g->set_runtime_backends(backends);
+        }
         if (t5) {
             t5->set_runtime_backends(backends);
         }
     }
 
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        if (clip_l) {
+            clip_l->set_graph_cut_layer_split_enabled(enabled);
+        }
+        if (clip_g) {
+            clip_g->set_graph_cut_layer_split_enabled(enabled);
+        }
+        if (t5) {
+            t5->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        if (clip_l) {
+            clip_l->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+        if (clip_g) {
+            clip_g->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+        if (t5) {
+            t5->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
     void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         if (t5) {
             t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
@@ -1010,11 +1063,32 @@ struct FluxCLIPEmbedder : public Conditioner {
     }
 
     void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        if (clip_l) {
+            clip_l->set_runtime_backends(backends);
+        }
         if (t5) {
             t5->set_runtime_backends(backends);
         }
     }
 
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        if (clip_l) {
+            clip_l->set_graph_cut_layer_split_enabled(enabled);
+        }
+        if (t5) {
+            t5->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        if (clip_l) {
+            clip_l->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+        if (t5) {
+            t5->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
     void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         if (t5) {
             t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
@@ -1278,6 +1352,18 @@ struct T5CLIPEmbedder : public Conditioner {
         }
     }
 
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        if (t5) {
+            t5->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        if (t5) {
+            t5->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
     void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         if (t5) {
             t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
@@ -1482,6 +1568,18 @@ struct MiniT2IConditioner : public Conditioner {
         }
     }
 
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        if (t5) {
+            t5->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        if (t5) {
+            t5->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
     void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         if (t5) {
             t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
@@ -1576,6 +1674,14 @@ struct AnimaConditioner : public Conditioner {
         llm->set_runtime_backends(backends);
     }
 
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        llm->set_graph_cut_layer_split_enabled(enabled);
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        llm->set_graph_cut_layer_split_backend_vram_limits(limits);
+    }
+
     void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         llm->get_param_tensors(tensors, "text_encoders.llm");
     }
@@ -1729,6 +1835,18 @@ struct LLMEmbedder : public Conditioner {
         llm->set_runtime_backends(backends);
     }
 
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        if (llm) {
+            llm->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        if (llm) {
+            llm->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
     void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         llm->get_param_tensors(tensors, "text_encoders.llm");
     }
@@ -2406,6 +2524,14 @@ struct LTXAVEmbedder : public Conditioner {
         llm->set_runtime_backends(backends);
     }
 
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        llm->set_graph_cut_layer_split_enabled(enabled);
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        llm->set_graph_cut_layer_split_backend_vram_limits(limits);
+    }
+
     void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         llm->get_param_tensors(tensors, "text_encoders.llm");
     }

src/core/ggml_extend.hpp

@@ -21,10 +21,12 @@
 #include <sstream>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 #include "core/ggml_extend_backend.h"
 #include "core/ggml_graph_cut.h"
+#include "core/layer_split_partition.h"
 #include "ggml-alloc.h"
 #include "ggml-backend.h"
 #include "ggml.h"
@@ -1745,6 +1747,8 @@ protected:
     size_t max_graph_vram_bytes           = 0;
     bool stream_layers_enabled            = false;
     size_t observed_max_effective_budget_ = 0;
+    bool graph_cut_layer_split_enabled    = false;
+    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
@@ -1776,6 +1780,9 @@ protected:
 
     sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_;
     std::unordered_set<const ggml_tensor*> params_tensor_set_;
+    std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_assignments_;
+    std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_node_assignments_;
+    bool graph_cut_layer_split_primary_notice_logged_ = false;
 
     template <typename T>
     static sd::Tensor<T> take_or_empty(std::optional<sd::Tensor<T>> tensor) {
@@ -1874,6 +1881,20 @@ protected:
         params_tensor_set_dirty_ = false;
     }
 
+    ggml_tensor* canonical_param_tensor(ggml_tensor* tensor) {
+        if (tensor == nullptr) {
+            return nullptr;
+        }
+        if (params_tensor_set_.find(tensor) != params_tensor_set_.end()) {
+            return tensor;
+        }
+        if (tensor->view_src != nullptr &&
+            params_tensor_set_.find(tensor->view_src) != params_tensor_set_.end()) {
+            return tensor->view_src;
+        }
+        return nullptr;
+    }
+
     std::vector<ggml_tensor*> collect_used_param_tensors(ggml_cgraph* gf) {
         std::vector<ggml_tensor*> used_params;
         rebuild_params_tensor_set();
@@ -1886,12 +1907,8 @@ protected:
         seen_params.reserve(static_cast<size_t>(n_leafs));
         for (int i = 0; i < n_leafs; ++i) {
             ggml_tensor* leaf       = sd::ggml_graph_cut::leaf_tensor(gf, i);
-            ggml_tensor* param_leaf = leaf;
+            ggml_tensor* param_leaf = canonical_param_tensor(leaf);
-            if (param_leaf != nullptr && params_tensor_set_.find(param_leaf) == params_tensor_set_.end()) {
-                param_leaf = param_leaf->view_src;
-            }
             if (param_leaf != nullptr &&
-                params_tensor_set_.find(param_leaf) != params_tensor_set_.end() &&
                 seen_params.insert(param_leaf).second) {
                 used_params.push_back(param_leaf);
             }
@@ -2102,12 +2119,18 @@ protected:
         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);
+            auto node_assignment = graph_cut_layer_split_node_assignments_.find(node);
+            if (node_assignment != graph_cut_layer_split_node_assignments_.end()) {
+                current = node_assignment->second;
+            }
             for (int s = 0; s < GGML_MAX_SRC; s++) {
                 ggml_backend_t weight_backend = backend_for_weight(node->src[s]);
                 if (weight_backend != nullptr) {
+                    if (node_assignment == graph_cut_layer_split_node_assignments_.end()) {
                         current = weight_backend;
                     }
                 }
+            }
             if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
                 node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) {
                 continue;
@@ -2435,6 +2458,123 @@ protected:
         return true;
     }
 
+    bool resolve_graph_cut_layer_split_plan(ggml_cgraph* gf,
+                                            GraphCutPlan* plan_out) {
+        GGML_ASSERT(plan_out != nullptr);
+        GGML_ASSERT(gf != nullptr);
+        *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend,
+                                                     gf,
+                                                     &graph_cut_plan_cache_,
+                                                     0,
+                                                     params_tensor_set_,
+                                                     get_desc().c_str());
+        return true;
+    }
+
+    bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
+        graph_cut_layer_split_node_assignments_.clear();
+        if (!graph_cut_layer_split_enabled) {
+            return true;
+        }
+        if (!is_multi_device()) {
+            LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", get_desc().c_str());
+            return false;
+        }
+
+        GraphCutPlan plan;
+        if (!resolve_graph_cut_layer_split_plan(gf, &plan)) {
+            return false;
+        }
+        if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) {
+            auto manager = weight_manager.lock();
+            if (manager == nullptr) {
+                LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
+                return false;
+            }
+            std::vector<ggml_tensor*> graph_params = collect_used_param_tensors(gf);
+            if (!graph_params.empty() &&
+                !manager->assign_compute_backend(graph_params, runtime_backend)) {
+                LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s",
+                          get_desc().c_str(),
+                          sd::layer_split_backend_device_display_name(runtime_backend).c_str());
+                return false;
+            }
+            for (ggml_tensor* param : graph_params) {
+                if (param != nullptr) {
+                    graph_cut_layer_split_assignments_[param] = runtime_backend;
+                }
+            }
+            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) {
+                    graph_cut_layer_split_node_assignments_[node] = runtime_backend;
+                }
+            }
+            if (!graph_cut_layer_split_primary_notice_logged_) {
+                LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
+                         get_desc().c_str(),
+                         sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
+                         graph_params.size());
+                graph_cut_layer_split_primary_notice_logged_ = true;
+            } else {
+                LOG_DEBUG("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
+                          get_desc().c_str(),
+                          sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
+                          graph_params.size());
+            }
+            return true;
+        }
+
+        std::vector<ggml_backend_t> split_backends;
+        split_backends.reserve(extra_runtime_backends.size() + 1);
+        split_backends.push_back(runtime_backend);
+        for (ggml_backend_t backend : extra_runtime_backends) {
+            if (backend != nullptr) {
+                split_backends.push_back(backend);
+            }
+        }
+
+        auto manager = weight_manager.lock();
+        if (manager == nullptr) {
+            LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
+            return false;
+        }
+
+        sd::GraphCutLayerSplitAssignment assignment;
+        auto canonicalize_param = [this](ggml_tensor* tensor) {
+            return canonical_param_tensor(tensor);
+        };
+        if (!sd::partition_graph_cut_layer_split(get_desc().c_str(),
+                                                 gf,
+                                                 plan,
+                                                 split_backends,
+                                                 graph_cut_layer_split_backend_vram_limits_,
+                                                 max_graph_vram_bytes,
+                                                 graph_cut_layer_split_assignments_,
+                                                 canonicalize_param,
+                                                 &assignment)) {
+            return false;
+        }
+
+        for (size_t i = 0; i < split_backends.size(); i++) {
+            if (assignment.tensors_by_backend[i].empty()) {
+                continue;
+            }
+            if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) {
+                LOG_ERROR("%s graph-cut layer split failed to assign params to %s",
+                          get_desc().c_str(),
+                          sd::layer_split_backend_device_display_name(split_backends[i]).c_str());
+                return false;
+            }
+        }
+
+        graph_cut_layer_split_node_assignments_ = std::move(assignment.node_assignments);
+        sd::log_graph_cut_layer_split_assignment(get_desc().c_str(), split_backends, assignment);
+
+        return true;
+    }
+
     struct PersistentExternalBinding {
         ggml_backend_buffer_t buffer = nullptr;
         void* data                   = nullptr;
@@ -2972,6 +3112,11 @@ public:
         GGML_ASSERT(gf != nullptr);
         rebuild_params_tensor_set();
 
+        if (!assign_graph_cut_layer_split_backends(gf)) {
+            free_compute_ctx();
+            return std::nullopt;
+        }
+
         if (can_attempt_graph_cut_segmented_compute()) {
             GraphCutPlan plan;
             if (!resolve_graph_cut_plan(gf, &plan)) {
@@ -3025,6 +3170,22 @@ public:
         stream_layers_enabled = enabled;
     }
 
+    void set_graph_cut_layer_split_enabled(bool enabled) {
+        graph_cut_layer_split_enabled = enabled;
+        if (!enabled) {
+            graph_cut_layer_split_assignments_.clear();
+            graph_cut_layer_split_node_assignments_.clear();
+            graph_cut_layer_split_primary_notice_logged_ = false;
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {
+        graph_cut_layer_split_backend_vram_limits_ = limits;
+        graph_cut_layer_split_assignments_.clear();
+        graph_cut_layer_split_node_assignments_.clear();
+        graph_cut_layer_split_primary_notice_logged_ = false;
+    }
+
     void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {
         extra_runtime_backends.clear();
         for (ggml_backend_t backend : backends) {
@@ -3036,6 +3197,9 @@ public:
                 extra_runtime_backends.push_back(backend);
             }
         }
+        graph_cut_layer_split_assignments_.clear();
+        graph_cut_layer_split_node_assignments_.clear();
+        graph_cut_layer_split_primary_notice_logged_ = false;
         if (is_multi_device() && stream_layers_enabled) {
             LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring",
                      get_desc().c_str());

src/core/layer_split_partition.cpp

@@ -1,9 +1,11 @@
 #include "core/layer_split_partition.h"
 
 #include <algorithm>
-#include <cstdint>
 #include <cstdlib>
 #include <cstring>
+#include <limits>
+#include <unordered_set>
+#include <utility>
 
 #include "core/util.h"
 
@@ -62,160 +64,194 @@ namespace sd {
         return name != nullptr ? name : "unknown";
     }
 
-    static bool layer_split_backend_supports_tensor(ggml_backend_t backend, const ggml_tensor* tensor) {
+    static size_t graph_cut_layer_split_backend_vram_limit(const std::vector<size_t>& backend_vram_limits,
-        return backend != nullptr && tensor != nullptr && ggml_backend_supports_op(backend, tensor);
+                                                           size_t backend_index,
+                                                           size_t primary_backend_vram_limit) {
+        if (backend_index < backend_vram_limits.size()) {
+            return backend_vram_limits[backend_index];
+        }
+        return backend_index == 0 ? primary_backend_vram_limit : 0;
     }
 
-    static size_t layer_split_supported_target(const std::string& desc,
+    static std::vector<int64_t> graph_cut_layer_split_backend_capacities(const std::vector<ggml_backend_t>& backends,
-                                               const std::string& tensor_name,
+                                                                         const std::vector<size_t>& backend_vram_limits,
-                                               const ggml_tensor* tensor,
+                                                                         size_t primary_backend_vram_limit) {
-                                               const std::vector<ggml_backend_t>& backends,
+        std::vector<int64_t> capacities(backends.size(), std::numeric_limits<int64_t>::max() / 4);
-                                               size_t preferred) {
-        if (tensor == nullptr || backends.empty()) {
-            return preferred;
-        }
-        size_t preferred_safe = std::min(preferred, backends.size() - 1);
-        if (layer_split_backend_supports_tensor(backends[preferred_safe], tensor)) {
-            return preferred_safe;
-        }
-        for (size_t i = 0; i < backends.size(); i++) {
-            if (layer_split_backend_supports_tensor(backends[i], tensor)) {
-                LOG_WARN("%s layer split: moving tensor '%s' from %s to %s because the preferred backend cannot run op=%s type=%s nbytes=%.2f MB",
-                         desc.c_str(),
-                         tensor_name.c_str(),
-                         layer_split_backend_device_display_name(backends[preferred_safe]).c_str(),
-                         layer_split_backend_device_display_name(backends[i]).c_str(),
-                         ggml_op_name(tensor->op),
-                         ggml_type_name(tensor->type),
-                         ggml_nbytes(tensor) / (1024.0 * 1024.0));
-                return i;
-            }
-        }
-        LOG_WARN("%s layer split: tensor '%s' is not supported by any split backend: op=%s type=%s nbytes=%.2f MB",
-                 desc.c_str(),
-                 tensor_name.c_str(),
-                 ggml_op_name(tensor->op),
-                 ggml_type_name(tensor->type),
-                 ggml_nbytes(tensor) / (1024.0 * 1024.0));
-        return preferred_safe;
-    }
-
-    std::vector<std::map<std::string, ggml_tensor*>> partition_layer_split_tensors(
-        const std::string& desc,
-        const std::map<std::string, ggml_tensor*>& tensors,
-        const std::map<std::string, ggml_tensor*>& split_tensors,
-        const std::vector<ggml_backend_t>& backends) {
-        std::vector<std::map<std::string, ggml_tensor*>> partitions(backends.size());
-        if (backends.empty()) {
-            LOG_WARN("%s: no backend available for a layer split", desc.c_str());
-            return partitions;
-        }
-
-        std::map<int, int64_t> block_bytes;
-        std::map<std::string, size_t> non_block_targets;
-        std::vector<int64_t> other_bytes_by_backend(backends.size(), 0);
-        int64_t total_block_bytes = 0;
-        int64_t total_other_bytes = 0;
-        int n_blocks              = 0;
-        for (const auto& kv : tensors) {
-            int64_t bytes = (int64_t)ggml_nbytes(kv.second);
-            int idx       = split_tensors.count(kv.first) != 0 ? layer_split_tensor_block_index(kv.first) : -1;
-            if (idx >= 0) {
-                block_bytes[idx] += bytes;
-                total_block_bytes += bytes;
-                n_blocks = std::max(n_blocks, idx + 1);
-            } else {
-                size_t target               = layer_split_supported_target(desc, kv.first, kv.second, backends, 0);
-                non_block_targets[kv.first] = target;
-                other_bytes_by_backend[target] += bytes;
-                total_other_bytes += bytes;
-            }
-        }
-        if (n_blocks == 0) {
-            LOG_WARN("%s: no transformer blocks found for a layer split; keeping tensors on compatible backends starting from %s",
-                     desc.c_str(),
-                     layer_split_backend_device_display_name(backends[0]).c_str());
-            for (const auto& kv : tensors) {
-                size_t target  = 0;
-                auto target_it = non_block_targets.find(kv.first);
-                if (target_it != non_block_targets.end()) {
-                    target = target_it->second;
-                }
-                partitions[target][kv.first] = kv.second;
-            }
-            return partitions;
-        }
-
-        // Reserve compute headroom and subtract each device's actual non-block
-        // bytes from its block budget.
         constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024;
-        std::vector<double> device_weights(backends.size(), 1.0);
-        double weight_sum = 0.0;
         for (size_t i = 0; i < backends.size(); i++) {
             ggml_backend_dev_t dev = ggml_backend_get_device(backends[i]);
             size_t free_bytes = 0, total_bytes = 0;
             if (dev != nullptr) {
                 ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
             }
-            // Keep a small share even for tight devices instead of dropping them.
+            if (free_bytes > 0) {
-            int64_t usable_bytes = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes,
+                capacities[i] = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes, 0);
-                                                     (int64_t)free_bytes / 8);
+            }
-            device_weights[i]    = usable_bytes > 0 ? (double)usable_bytes : 1.0;
+            size_t limit_bytes = graph_cut_layer_split_backend_vram_limit(backend_vram_limits,
-            weight_sum += device_weights[i];
+                                                                          i,
+                                                                          primary_backend_vram_limit);
+            if (limit_bytes > 0) {
+                capacities[i] = std::min<int64_t>(capacities[i], (int64_t)limit_bytes);
+            }
+        }
+        return capacities;
     }
 
-        std::vector<int64_t> block_budgets(backends.size(), 0);
+    bool partition_graph_cut_layer_split(const char* desc,
-        const int64_t total_bytes = total_block_bytes + total_other_bytes;
+                                         ggml_cgraph* gf,
-        for (size_t i = 0; i < backends.size(); i++) {
+                                         const sd::ggml_graph_cut::Plan& plan,
-            int64_t budget   = (int64_t)((double)total_bytes * device_weights[i] / weight_sum);
+                                         const std::vector<ggml_backend_t>& split_backends,
-            budget           = std::max<int64_t>(budget - other_bytes_by_backend[i], 0);
+                                         const std::vector<size_t>& backend_vram_limits,
-            block_budgets[i] = budget;
+                                         size_t primary_backend_vram_limit,
+                                         std::unordered_map<const ggml_tensor*, ggml_backend_t>& param_assignments,
+                                         const std::function<ggml_tensor*(ggml_tensor*)>& canonical_param_tensor,
+                                         GraphCutLayerSplitAssignment* assignment_out) {
+        GGML_ASSERT(gf != nullptr);
+        GGML_ASSERT(assignment_out != nullptr);
+        GGML_ASSERT(canonical_param_tensor != nullptr);
+        GGML_ASSERT(!split_backends.empty());
+
+        GraphCutLayerSplitAssignment assignment;
+        assignment.segment_count = plan.segments.size();
+        assignment.tensors_by_backend.resize(split_backends.size());
+        assignment.bytes_by_backend.resize(split_backends.size(), 0);
+        assignment.first_segment_by_backend.resize(split_backends.size(), plan.segments.size());
+        assignment.last_segment_by_backend.resize(split_backends.size(), 0);
+
+        std::vector<std::vector<ggml_tensor*>> segment_params(plan.segments.size());
+        std::vector<int64_t> segment_param_bytes(plan.segments.size(), 0);
+        std::unordered_set<ggml_tensor*> seen_params;
+        for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
+            std::vector<ggml_tensor*> params = sd::ggml_graph_cut::param_tensors(gf, plan.segments[seg_idx]);
+            for (ggml_tensor* raw_param : params) {
+                ggml_tensor* param = canonical_param_tensor(raw_param);
+                if (param == nullptr || !seen_params.insert(param).second) {
+                    continue;
+                }
+                segment_params[seg_idx].push_back(param);
+                segment_param_bytes[seg_idx] += (int64_t)ggml_nbytes(param);
+            }
         }
 
-        std::vector<int> boundaries(backends.size(), n_blocks);
+        int64_t total_param_bytes = 0;
-        size_t current = 0;
+        for (int64_t bytes : segment_param_bytes) {
-        int64_t used   = 0;
+            total_param_bytes += bytes;
-        for (int b = 0; b < n_blocks; b++) {
-            int64_t bytes = block_bytes.count(b) != 0 ? block_bytes[b] : 0;
-            if (current + 1 < backends.size() && used > 0 && used + bytes > block_budgets[current]) {
-                boundaries[current] = b;
-                current++;
-                used = 0;
         }
-            used += bytes;
+        if (total_param_bytes <= 0) {
+            LOG_ERROR("%s graph-cut layer split found no graph params to assign", desc);
+            return false;
         }
 
-        for (const auto& kv : tensors) {
+        std::vector<int64_t> backend_capacities = graph_cut_layer_split_backend_capacities(split_backends,
-            size_t target = 0;
+                                                                                           backend_vram_limits,
-            int idx       = split_tensors.count(kv.first) != 0 ? layer_split_tensor_block_index(kv.first) : -1;
+                                                                                           primary_backend_vram_limit);
-            if (idx >= 0) {
+
-                while (target < boundaries.size() && idx >= boundaries[target]) {
+        std::vector<ggml_backend_t> backend_by_segment(plan.segments.size(), split_backends[0]);
-                    target++;
+        size_t current_backend = 0;
+        int64_t current_used   = 0;
+        for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
+            int64_t bytes = segment_param_bytes[seg_idx];
+            while (current_backend + 1 < split_backends.size() &&
+                   bytes > 0 &&
+                   current_used + bytes > backend_capacities[current_backend]) {
+                current_backend++;
+                current_used = 0;
             }
-                target = std::min(target, backends.size() - 1);
+            if (bytes > 0 && current_used + bytes > backend_capacities[current_backend]) {
-                target = layer_split_supported_target(desc, kv.first, kv.second, backends, target);
+                LOG_ERROR("%s graph-cut layer split: segment %zu needs %.1f MB on %s, but only %.1f MB is available under current VRAM limits",
+                          desc,
+                          seg_idx,
+                          (current_used + bytes) / (1024.0 * 1024.0),
+                          layer_split_backend_device_display_name(split_backends[current_backend]).c_str(),
+                          backend_capacities[current_backend] / (1024.0 * 1024.0));
+                return false;
+            }
+            current_used += bytes;
+            backend_by_segment[seg_idx] = split_backends[current_backend];
+
+            for (ggml_tensor* param : segment_params[seg_idx]) {
+                ggml_backend_t target_backend = split_backends[current_backend];
+                auto assigned_it              = param_assignments.find(param);
+                if (assigned_it == param_assignments.end()) {
+                    param_assignments[param]            = target_backend;
+                    assignment.has_new_param_assignment = true;
                 } else {
-                auto target_it = non_block_targets.find(kv.first);
+                    target_backend = assigned_it->second;
-                if (target_it != non_block_targets.end()) {
-                    target = target_it->second;
-                }
-            }
-            partitions[target][kv.first] = kv.second;
                 }
 
-        int range_start = 0;
+                auto backend_it = std::find(split_backends.begin(), split_backends.end(), target_backend);
-        for (size_t i = 0; i < backends.size(); i++) {
+                if (backend_it == split_backends.end()) {
-            int range_end                = boundaries[i];
+                    LOG_ERROR("%s graph-cut layer split tensor '%s' is assigned to an unavailable backend",
-            const char* non_block_suffix = other_bytes_by_backend[i] > 0 ? " + non-block tensors" : "";
+                              desc,
-            LOG_INFO("%s layer split: %s <- blocks [%d, %d)%s",
+                              ggml_get_name(param));
-                     desc.c_str(),
+                    return false;
-                     layer_split_backend_device_display_name(backends[i]).c_str(),
+                }
-                     range_start,
+                size_t backend_idx                               = (size_t)std::distance(split_backends.begin(), backend_it);
-                     range_end,
+                assignment.first_segment_by_backend[backend_idx] = std::min(assignment.first_segment_by_backend[backend_idx], seg_idx);
-                     non_block_suffix);
+                assignment.last_segment_by_backend[backend_idx]  = std::max(assignment.last_segment_by_backend[backend_idx], seg_idx + 1);
-            range_start = range_end;
+                assignment.tensors_by_backend[backend_idx].push_back(param);
+                assignment.bytes_by_backend[backend_idx] += (int64_t)ggml_nbytes(param);
+            }
+        }
+
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
+            ggml_backend_t backend = backend_by_segment[seg_idx];
+            const auto& segment    = plan.segments[seg_idx];
+            for (int node_index : segment.internal_node_indices) {
+                if (node_index < 0 || node_index >= n_nodes) {
+                    continue;
+                }
+                ggml_tensor* node = ggml_graph_node(gf, node_index);
+                if (node != nullptr) {
+                    assignment.node_assignments[node] = backend;
+                }
+            }
+            for (int node_index : segment.output_node_indices) {
+                if (node_index < 0 || node_index >= n_nodes) {
+                    continue;
+                }
+                ggml_tensor* node = ggml_graph_node(gf, node_index);
+                if (node != nullptr) {
+                    assignment.node_assignments[node] = backend;
+                }
+            }
+        }
+
+        *assignment_out = std::move(assignment);
+        return true;
+    }
+
+    void log_graph_cut_layer_split_assignment(const char* desc,
+                                              const std::vector<ggml_backend_t>& split_backends,
+                                              const GraphCutLayerSplitAssignment& assignment) {
+        for (size_t i = 0; i < split_backends.size(); i++) {
+            if (i >= assignment.tensors_by_backend.size() ||
+                assignment.tensors_by_backend[i].empty()) {
+                continue;
+            }
+            size_t first_segment = assignment.first_segment_by_backend[i] == assignment.segment_count
+                                       ? 0
+                                       : assignment.first_segment_by_backend[i];
+            size_t last_segment  = assignment.last_segment_by_backend[i];
+            if (assignment.has_new_param_assignment) {
+                LOG_INFO("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
+                         desc,
+                         layer_split_backend_device_display_name(split_backends[i]).c_str(),
+                         first_segment,
+                         last_segment,
+                         assignment.tensors_by_backend[i].size(),
+                         assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
+            } else {
+                LOG_DEBUG("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
+                          desc,
+                          layer_split_backend_device_display_name(split_backends[i]).c_str(),
+                          first_segment,
+                          last_segment,
+                          assignment.tensors_by_backend[i].size(),
+                          assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
+            }
         }
-        return partitions;
     }
 
 }  // namespace sd

src/core/layer_split_partition.h

@@ -1,23 +1,43 @@
 #ifndef __SD_CORE_LAYER_SPLIT_PARTITION_H__
 #define __SD_CORE_LAYER_SPLIT_PARTITION_H__
 
-#include <map>
+#include <cstdint>
+#include <functional>
 #include <string>
+#include <unordered_map>
 #include <vector>
 
 #include "ggml-backend.h"
 #include "ggml.h"
 
+#include "core/ggml_graph_cut.h"
+
 namespace sd {
 
+    struct GraphCutLayerSplitAssignment {
+        std::vector<std::vector<ggml_tensor*>> tensors_by_backend;
+        std::vector<int64_t> bytes_by_backend;
+        std::vector<size_t> first_segment_by_backend;
+        std::vector<size_t> last_segment_by_backend;
+        std::unordered_map<const ggml_tensor*, ggml_backend_t> node_assignments;
+        size_t segment_count          = 0;
+        bool has_new_param_assignment = false;
+    };
+
     std::string layer_split_backend_device_display_name(ggml_backend_t backend);
     int layer_split_tensor_block_index(const std::string& name);
-
+    bool partition_graph_cut_layer_split(const char* desc,
-    std::vector<std::map<std::string, ggml_tensor*>> partition_layer_split_tensors(
+                                         ggml_cgraph* gf,
-        const std::string& desc,
+                                         const sd::ggml_graph_cut::Plan& plan,
-        const std::map<std::string, ggml_tensor*>& tensors,
+                                         const std::vector<ggml_backend_t>& split_backends,
-        const std::map<std::string, ggml_tensor*>& split_tensors,
+                                         const std::vector<size_t>& backend_vram_limits,
-        const std::vector<ggml_backend_t>& backends);
+                                         size_t primary_backend_vram_limit,
+                                         std::unordered_map<const ggml_tensor*, ggml_backend_t>& param_assignments,
+                                         const std::function<ggml_tensor*(ggml_tensor*)>& canonical_param_tensor,
+                                         GraphCutLayerSplitAssignment* assignment_out);
+    void log_graph_cut_layer_split_assignment(const char* desc,
+                                              const std::vector<ggml_backend_t>& split_backends,
+                                              const GraphCutLayerSplitAssignment& assignment);
 
 }  // namespace sd
 

src/model_manager.cpp

@@ -134,7 +134,8 @@ bool ModelManager::register_param_tensors(const std::string& desc,
                                           ggml_backend_t compute_backend,
                                           ggml_backend_t params_backend,
                                           size_t* registered_tensor_size,
-                                          bool allow_split_buffer) {
+                                          bool allow_split_buffer,
+                                          bool params_follow_compute_backend) {
     if (desc.empty()) {
         LOG_ERROR("model manager tensor desc is empty");
         return false;
@@ -166,6 +167,7 @@ bool ModelManager::register_param_tensors(const std::string& desc,
         state->compute_backend               = compute_backend;
         state->params_backend                = params_backend;
         state->allow_split_buffer            = allow_split_buffer;
+        state->params_follow_compute_backend = params_follow_compute_backend;
         new_states.push_back(std::move(state));
     }
 
@@ -919,6 +921,54 @@ bool ModelManager::resolve_required_tensor_states(const std::vector<ggml_tensor*
     return true;
 }
 
+bool ModelManager::assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
+                                          ggml_backend_t compute_backend) {
+    if (tensors.empty()) {
+        return true;
+    }
+    if (compute_backend == nullptr) {
+        LOG_ERROR("model manager cannot assign tensors to a null compute backend");
+        return false;
+    }
+
+    std::vector<TensorState*> required_states;
+    if (!resolve_required_tensor_states(tensors, required_states)) {
+        return false;
+    }
+
+    for (TensorState* state : required_states) {
+        if (state == nullptr || state->tensor == nullptr) {
+            continue;
+        }
+
+        const bool params_follow_compute = state->params_follow_compute_backend ||
+                                           state->residency_mode == ResidencyMode::Disk;
+        const bool compute_changes = state->compute_backend != compute_backend;
+        const bool params_changes  = params_follow_compute && state->params_backend != compute_backend;
+        if (!compute_changes && !params_changes) {
+            continue;
+        }
+
+        if (state->active_prepare_count > 0 || state->staged_to_compute_backend) {
+            LOG_ERROR("model manager cannot move active tensor '%s' to another compute backend",
+                      state->name.c_str());
+            return false;
+        }
+        if (params_changes && state->loaded_to_params_backend) {
+            LOG_ERROR("model manager cannot move loaded tensor '%s' to another params backend",
+                      state->name.c_str());
+            return false;
+        }
+
+        state->compute_backend = compute_backend;
+        if (params_follow_compute) {
+            state->params_backend = compute_backend;
+        }
+    }
+
+    return true;
+}
+
 bool ModelManager::prepare_params(const std::vector<ggml_tensor*>& tensors) {
     if (tensors.empty()) {
         return true;

src/model_manager.h

@@ -37,6 +37,7 @@ private:
         ggml_backend_t compute_backend     = nullptr;
         ggml_backend_t params_backend      = nullptr;
         bool allow_split_buffer            = false;
+        bool params_follow_compute_backend = false;
         bool metadata_validated            = false;
 
         int active_prepare_count = 0;
@@ -130,7 +131,8 @@ public:
                                 ggml_backend_t compute_backend,
                                 ggml_backend_t params_backend,
                                 size_t* registered_tensor_size     = nullptr,
-                                bool allow_split_buffer        = false);
+                                bool allow_split_buffer            = false,
+                                bool params_follow_compute_backend = false);
 
     template <typename Runner>
     bool register_runner_params(const std::string& desc,
@@ -170,6 +172,8 @@ public:
     bool validate_registered_tensors();
     bool load_all_params_eagerly();
 
+    bool assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
+                                ggml_backend_t compute_backend) override;
     bool prepare_params(const std::vector<ggml_tensor*>& tensors) override;
     void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) override;
     void release_params_backend_params(const std::vector<ggml_tensor*>& tensors) override;

src/stable-diffusion.cpp

@@ -268,6 +268,15 @@ public:
         return max_vram_assignment.bytes_for_backend(backend_for(module));
     }
 
+    std::vector<size_t> layer_split_vram_limits_for_backends(const std::vector<ggml_backend_t>& backends) {
+        std::vector<size_t> limits;
+        limits.reserve(backends.size());
+        for (ggml_backend_t backend : backends) {
+            limits.push_back(max_vram_assignment.bytes_for_backend(backend));
+        }
+        return limits;
+    }
+
     bool ensure_backend_pair(SDBackendModule module) {
         if (backend_for(module) == nullptr) {
             return false;
@@ -427,8 +436,9 @@ public:
                                                      params_mem_size);
     }
 
-    // Register each layer-split partition with its compute backend; the
+    // Register graph-cut layer-split tensors on the primary backend first.
-    // ModelManager handles allocation, staging, and LoRA by backend.
+    // The first real graph assigns each param tensor to a runtime backend
+    // before weights are loaded or staged.
     template <typename T>
     bool register_layer_split_runner_params(const std::string& desc,
                                             const std::shared_ptr<T>& model,
@@ -459,52 +469,29 @@ public:
                                                          params_mem_size);
         }
 
-        std::map<std::string, ggml_tensor*> split_tensors;
+        model->set_runtime_backends(module_backends);
-        if constexpr (std::is_base_of_v<Conditioner, T>) {
+        model->set_graph_cut_layer_split_backend_vram_limits(layer_split_vram_limits_for_backends(module_backends));
-            model->get_layer_split_param_tensors(split_tensors);
+        model->set_graph_cut_layer_split_enabled(true);
-        } else {
+        const bool params_follow_runtime = backend_manager.params_backend_follows_runtime(module) ||
-            split_tensors = group_tensors;
+                                           backend_manager.params_backend_is_disk(module);
+        ggml_backend_t initial_params_backend = params_follow_runtime ? module_backends[0] : params_backend_for(module);
+        if (initial_params_backend == nullptr) {
+            return false;
         }
 
-        auto partitions = sd::partition_layer_split_tensors(desc, group_tensors, split_tensors, module_backends);
+        LOG_INFO("%s graph-cut layer split: deferring %zu tensors across %zu runtime backends until first graph",
-        bool is_split   = false;
+                 desc.c_str(),
-        for (size_t i = 1; i < partitions.size(); i++) {
+                 group_tensors.size(),
-            if (!partitions[i].empty()) {
+                 module_backends.size());
-                is_split = true;
+
-                break;
-            }
-        }
-        if (!is_split) {
         return model_manager->register_param_tensors(desc,
                                                      std::move(group_tensors),
                                                      residency_mode,
                                                      module_backends[0],
-                                                         params_backend_for(module),
+                                                     initial_params_backend,
-                                                         params_mem_size);
+                                                     params_mem_size,
-        }
+                                                     false,
-
+                                                     params_follow_runtime);
-        model->set_runtime_backends(module_backends);
-        const bool params_follow_runtime = backend_manager.params_backend_follows_runtime(module) ||
-                                           backend_manager.params_backend_is_disk(module);
-        for (size_t i = 0; i < module_backends.size(); i++) {
-            if (partitions[i].empty()) {
-                continue;
-            }
-            ggml_backend_t partition_params_backend =
-                params_follow_runtime ? module_backends[i] : params_backend_for(module);
-            if (partition_params_backend == nullptr) {
-                return false;
-            }
-            if (!model_manager->register_param_tensors(desc,
-                                                       std::move(partitions[i]),
-                                                       residency_mode,
-                                                       module_backends[i],
-                                                       partition_params_backend,
-                                                       params_mem_size)) {
-                return false;
-            }
-        }
-        return true;
     }
 
     bool init_backend() {
@@ -529,6 +516,16 @@ public:
         return false;
     }
 
+    bool graph_cut_layer_split_active() {
+        for (SDBackendModule module : {SDBackendModule::DIFFUSION, SDBackendModule::TE}) {
+            if (backend_manager.split_mode(module) == SDSplitMode::LAYER &&
+                backend_manager.runtime_backends(module).size() > 1) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     std::shared_ptr<RNG> get_rng(rng_type_t rng_type) {
         if (rng_type == STD_DEFAULT_RNG) {
             return std::make_shared<STDDefaultRNG>();
@@ -785,6 +782,10 @@ public:
             LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring");
             stream_layers = false;
         }
+        if (eager_load && graph_cut_layer_split_active()) {
+            LOG_WARN("--eager-load is not supported with graph-cut layer split; weights will be prepared lazily");
+            eager_load = false;
+        }
 
         std::map<ggml_type, uint32_t> wtype_stat                 = model_loader.get_wtype_stat();
         std::map<ggml_type, uint32_t> conditioner_wtype_stat     = model_loader.get_conditioner_wtype_stat();

src/weight_manager.h

@@ -3,10 +3,14 @@
 
 #include <vector>
 
+#include "ggml-backend.h"
+
 struct ggml_tensor;
 
 struct RunnerWeightManager {
     virtual ~RunnerWeightManager()                                                        = default;
+    virtual bool assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
+                                        ggml_backend_t compute_backend)                   = 0;
     virtual bool prepare_params(const std::vector<ggml_tensor*>& tensors)                 = 0;
     virtual void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;
     virtual void release_params_backend_params(const std::vector<ggml_tensor*>& tensors)  = 0;