fix: size layer-split ranges in proportion to device capacity (#2309)

src/core/layer_split_partition.cpp

@@ -73,10 +73,13 @@ namespace sd {
         return backend_index == 0 ? primary_backend_vram_limit : 0;
     }
 
+    // Capacity of a device that reports no free memory and has no VRAM limit.
+    static constexpr int64_t graph_cut_layer_split_unknown_capacity = std::numeric_limits<int64_t>::max() / 4;
+
     static std::vector<int64_t> graph_cut_layer_split_backend_capacities(const std::vector<ggml_backend_t>& backends,
                                                                          const std::vector<size_t>& backend_vram_limits,
                                                                          size_t primary_backend_vram_limit) {
-        std::vector<int64_t> capacities(backends.size(), std::numeric_limits<int64_t>::max() / 4);
+        std::vector<int64_t> capacities(backends.size(), graph_cut_layer_split_unknown_capacity);
         constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024;
         for (size_t i = 0; i < backends.size(); i++) {
             ggml_backend_dev_t dev = ggml_backend_get_device(backends[i]);
@@ -97,6 +100,57 @@ namespace sd {
         return capacities;
     }
 
+    // Whether filling the devices in order, each up to `fraction` of its capacity, places every
+    // segment. Mirrors the placement loop in partition_graph_cut_layer_split().
+    static bool graph_cut_layer_split_fits(const std::vector<int64_t>& segment_bytes,
+                                           const std::vector<int64_t>& capacities,
+                                           double fraction) {
+        size_t backend = 0;
+        int64_t used   = 0;
+        for (int64_t bytes : segment_bytes) {
+            while (backend + 1 < capacities.size() && bytes > 0 &&
+                   used + bytes > (int64_t)((double)capacities[backend] * fraction)) {
+                backend++;
+                used = 0;
+            }
+            if (bytes > 0 && used + bytes > (int64_t)((double)capacities[backend] * fraction)) {
+                return false;
+            }
+            used += bytes;
+        }
+        return true;
+    }
+
+    // Fill limits that size the contiguous ranges in proportion to the device capacities: the same
+    // fraction of every capacity, as small as possible while all segments still fit. Filling the
+    // first device up to its capacity would leave only the fixed headroom for compute, which is too
+    // small for large video graphs. A fraction of 1 is plain capacity-based placement, so every
+    // layout that fits the capacities is still found.
+    static std::vector<int64_t> graph_cut_layer_split_fill_limits(const std::vector<int64_t>& segment_bytes,
+                                                                  const std::vector<int64_t>& capacities) {
+        const bool capacities_known = std::all_of(capacities.begin(), capacities.end(), [](int64_t capacity) {
+            return capacity < graph_cut_layer_split_unknown_capacity;
+        });
+        if (!capacities_known || !graph_cut_layer_split_fits(segment_bytes, capacities, 1.0)) {
+            return capacities;
+        }
+        double low  = 0.0;
+        double high = 1.0;
+        for (int i = 0; i < 32; i++) {
+            const double mid = (low + high) / 2;
+            if (graph_cut_layer_split_fits(segment_bytes, capacities, mid)) {
+                high = mid;
+            } else {
+                low = mid;
+            }
+        }
+        std::vector<int64_t> limits(capacities.size());
+        for (size_t i = 0; i < capacities.size(); i++) {
+            limits[i] = (int64_t)((double)capacities[i] * high);
+        }
+        return limits;
+    }
+
     bool partition_graph_cut_layer_split(const char* desc,
                                          ggml_cgraph* gf,
                                          const sd::ggml_graph_cut::Plan& plan,
@@ -151,6 +205,11 @@ namespace sd {
             return param_assignments.count(param) != 0;
         });
 
+        // Move to the next device at the proportional fill limit; the capacity stays the hard limit.
+        const std::vector<int64_t> backend_limits = reuse_assignments
+                                                        ? backend_capacities
+                                                        : graph_cut_layer_split_fill_limits(segment_param_bytes, backend_capacities);
+
         std::vector<ggml_backend_t> backend_by_segment(plan.segments.size(), split_backends[0]);
         size_t current_backend = 0;
         int64_t current_used   = 0;
@@ -158,7 +217,7 @@ namespace sd {
             int64_t bytes = segment_param_bytes[seg_idx];
             while (!reuse_assignments && current_backend + 1 < split_backends.size() &&
                    bytes > 0 &&
-                   current_used + bytes > backend_capacities[current_backend]) {
+                   current_used + bytes > backend_limits[current_backend]) {
                 current_backend++;
                 current_used = 0;
             }