perf: optimize H3 inference and segmented weight prefetch (#2240)

src/core/compute_workspace.cpp

@@ -158,7 +158,7 @@ namespace sd {
         for (const auto& entry : copies) {
             const size_t source_index = ggml_hash_find(&graph->visited_hash_set, entry.first);
             if (source_index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, source_index)) {
-                const size_t copy_index = ggml_hash_insert(&copy_graph->visited_hash_set, entry.second);
+                const size_t copy_index            = ggml_hash_insert(&copy_graph->visited_hash_set, entry.second);
                 copy_graph->use_counts[copy_index] = graph->use_counts[source_index];
             }
         }

src/core/ggml_extend_backend.cpp

@@ -615,6 +615,32 @@ bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) {
     return false;
 }
 
+void sd_backend_set_buffer_cache_enabled(ggml_backend_t backend, bool enabled) {
+    auto dev = backend != nullptr ? ggml_backend_get_device(backend) : nullptr;
+    if (dev == nullptr) {
+        return;
+    }
+    auto reg                    = ggml_backend_dev_backend_reg(dev);
+    using SetBufferCacheEnabled = void (*)(ggml_backend_t, bool);
+    auto fn                     = reinterpret_cast<SetBufferCacheEnabled>(ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_buffer_cache_enabled"));
+    if (fn != nullptr) {
+        fn(backend, enabled);
+    }
+}
+
+void sd_backend_clear_buffer_cache(ggml_backend_t backend) {
+    auto dev = backend != nullptr ? ggml_backend_get_device(backend) : nullptr;
+    if (dev == nullptr) {
+        return;
+    }
+    auto reg               = ggml_backend_dev_backend_reg(dev);
+    using ClearBufferCache = void (*)(ggml_backend_t);
+    auto fn                = reinterpret_cast<ClearBufferCache>(ggml_backend_reg_get_proc_address(reg, "ggml_backend_clear_buffer_cache"));
+    if (fn != nullptr) {
+        fn(backend);
+    }
+}
+
 static ggml_cgraph sd_ggml_graph_view(ggml_cgraph* cgraph0, int i0, int i1) {
     ggml_cgraph cgraph = {
         /*.size             =*/0,

src/core/ggml_extend_backend.h

@@ -96,6 +96,9 @@ ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t dev
                                                           size_t max_tensor_size);
 ggml_backend_t sd_backend_cpu_init();
 bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
+void sd_backend_set_buffer_cache_enabled(ggml_backend_t backend, bool enabled);
+void sd_backend_clear_buffer_cache(ggml_backend_t backend);
+
 ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
                                                         ggml_cgraph* gf,
                                                         sd_graph_eval_callback_t callback_eval,

src/core/ggml_graph_cut.cpp

@@ -753,6 +753,14 @@ namespace sd::ggml_graph_cut {
         for (int node_idx : segment.internal_node_indices) {
             ggml_graph_add_node(segment_graph, ggml_graph_node(gf, node_idx));
         }
+        // Fusion checks need the original use counts, including consumers outside this segment.
+        for (ggml_tensor* tensor : represented) {
+            const size_t source_index = ggml_hash_find(&gf->visited_hash_set, tensor);
+            if (source_index != GGML_HASHSET_FULL && ggml_bitset_get(gf->visited_hash_set.used, source_index)) {
+                const size_t index               = ggml_hash_insert(&segment_graph->visited_hash_set, tensor);
+                segment_graph->use_counts[index] = gf->use_counts[source_index];
+            }
+        }
         *graph_ctx_out = graph_ctx;
         return segment_graph;
     }

src/core/ggml_runner.cpp

@@ -466,6 +466,10 @@ bool GGMLRunner::runner_start() {
     if (runner_started_) {
         return true;
     }
+    sd_backend_set_buffer_cache_enabled(runtime_backend, true);
+    for (auto backend : extra_runtime_backends) {
+        sd_backend_set_buffer_cache_enabled(backend, true);
+    }
     cache_.clear();
     workspace_.set_extra_backends(extra_runtime_backends);
     if (auto manager = residency_manager.lock()) {
@@ -490,6 +494,7 @@ void GGMLRunner::runner_end() {
     free_compute_ctx();
     compute_ctx_buffer_.reset();
     cache_.clear();
+    cut_cache_.clear();
     logged_compute_bytes_.clear();
     logged_segment_count_ = 0;
     if (auto manager = residency_manager.lock()) {
@@ -503,6 +508,10 @@ void GGMLRunner::runner_end() {
         manager->evict_compute_backend_params(tensors);
         manager->remove_runtime_owner(reinterpret_cast<uintptr_t>(this));
     }
+    sd_backend_clear_buffer_cache(runtime_backend);
+    for (auto backend : extra_runtime_backends) {
+        sd_backend_clear_buffer_cache(backend);
+    }
     runner_started_ = false;
 }
 
@@ -712,8 +721,11 @@ static size_t add_bytes(size_t a, size_t b) {
     return b > SIZE_MAX - a ? SIZE_MAX : a + b;
 }
 
-ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph) {
+ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, const std::unordered_set<const ggml_tensor*>* external_outputs) {
     auto external_backend = [&](const ggml_tensor* tensor) -> ggml_backend_t {
+        if (external_outputs != nullptr && external_outputs->count(tensor)) {
+            return runtime_backend;
+        }
         if (!params_tensor_set_.count(tensor)) {
             return nullptr;
         }
@@ -766,7 +778,7 @@ bool GGMLRunner::fits(const std::vector<DeviceMemoryRequest>& requests,
     return true;
 }
 
-bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) {
+bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads, const std::function<void()>& on_before_compute) {
     if (sd_backend_is_cpu(runtime_backend)) {
         sd_backend_cpu_set_n_threads(runtime_backend, n_threads);
     }
@@ -776,7 +788,27 @@ bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) {
     auto scheduler = workspace_.scheduler();
     GGML_ASSERT(scheduler != nullptr);
     ggml_backend_sched_set_eval_callback(scheduler, sd_get_backend_eval_callback(), sd_get_backend_eval_callback_data());
-    const auto status = ggml_backend_sched_graph_compute(scheduler, graph);
+#ifndef SD_USE_UPSTREAM_GGML
+    struct BeforeComputeCallback {
+        ggml_backend_t backend;
+        const std::function<void()>& callback;
+        bool invoked = false;
+    } before_compute{runtime_backend, on_before_compute};
+    ggml_backend_sched_set_compute_callback(
+        scheduler, [](ggml_backend_t backend, void* data) {
+            auto& hook = *static_cast<BeforeComputeCallback*>(data);
+            if (backend == hook.backend && !hook.invoked) {
+                hook.invoked = true;
+                hook.callback();
+            }
+        },
+        &before_compute);
+    const auto status = ggml_backend_sched_graph_compute_async(scheduler, graph);
+    ggml_backend_sched_set_compute_callback(scheduler, nullptr, nullptr);
+#else
+    on_before_compute();
+    const auto status = ggml_backend_sched_graph_compute_async(scheduler, graph);
+#endif
     workspace_.synchronize();
     if (status != GGML_STATUS_SUCCESS) {
         last_compute_status_ = status;
@@ -828,27 +860,31 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
     if (!assign_graph_cut_layer_split_backends(graph)) {
         return std::nullopt;
     }
-    const auto params           = collect_used_param_tensors(graph);
-    const auto& cached_plan     = resolve_graph_cut_plan(graph);
-    const auto full_measurement = measure(graph);
-    if (full_measurement.buffers.empty()) {
-        last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
-        return std::nullopt;
-    }
-    auto fits_monolithic = [&]() {
-        // Planning headroom absorbs allocation estimate drift; execution keeps the normal limits.
-        constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL;
-        auto requests                      = memory_requests(full_measurement.buffers, cache_.pending_bytes(graph));
-        for (auto& request : requests) {
-            request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom);
+    const auto params       = collect_used_param_tensors(graph);
+    const auto& cached_plan = resolve_graph_cut_plan(graph);
+    auto manager            = residency_manager.lock();
+    const bool can_segment  = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) &&
+                              manager != nullptr && manager->segmented_compute_enabled() &&
+                              cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1;
+    // An empty workspace request is a lower bound on the full graph's memory demand.
+    bool segmented = can_segment && !fits(memory_requests({{runtime_backend, 0}}, 0), params);
+    ComputeWorkspace::Measurement full_measurement;
+    if (!segmented) {
+        full_measurement = measure(graph);
+        if (full_measurement.buffers.empty()) {
+            last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+            return std::nullopt;
         }
-        return fits(requests, params);
-    };
-    auto manager         = residency_manager.lock();
-    const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) &&
-                           manager != nullptr && manager->segmented_compute_enabled() &&
-                           cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1 &&
-                           !fits_monolithic();
+        if (can_segment) {
+            // Planning headroom absorbs allocation estimate drift; execution keeps the normal limits.
+            constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL;
+            auto requests                      = memory_requests(full_measurement.buffers, cache_.pending_bytes(graph));
+            for (auto& request : requests) {
+                request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom);
+            }
+            segmented = !fits(requests, params);
+        }
+    }
     ggml_graph_cut::Plan monolithic_plan;
     if (!segmented) {
         monolithic_plan.segments.emplace_back();
@@ -922,14 +958,24 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
             }
         } segment_cleanup{*this, weights, bindings, segment_context};
 
-        auto measurement = segmented ? measure(segment_graph) : full_measurement;
+        std::unordered_set<const ggml_tensor*> cache_outputs;
+        if (segmented) {
+            for (int output_index : segment.output_node_indices) {
+                auto tensor = ggml_graph_node(graph, output_index);
+                if (ggml_graph_cut::is_graph_cut_tensor(tensor) && segment.future_cut_names.count(tensor->name) &&
+                    GraphCutTensorCache::can_bind_output(runtime_backend, tensor)) {
+                    cache_outputs.insert(tensor);
+                }
+            }
+        }
+        auto measurement = segmented ? measure(segment_graph, &cache_outputs) : full_measurement;
         if (!workspace_.prepare(measurement)) {
             last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
             return fail_segment("workspace preparation");
         }
-        const size_t cut_bytes       = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment);
-        const size_t new_cache_bytes = add_bytes(cut_bytes, cache_.pending_bytes(segment_graph));
-        auto ensure_capacity         = [&]() {
+        const size_t cut_bytes = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment);
+        size_t new_cache_bytes = add_bytes(cut_bytes, cache_.pending_bytes(segment_graph));
+        auto ensure_capacity   = [&]() {
             sync_runtime_residency();
             auto requests = memory_requests(measurement.buffers, new_cache_bytes);
             if (fits(requests, weights.params(index))) {
@@ -956,6 +1002,11 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
         if (!ensure_capacity()) {
             return fail_segment("workspace capacity check");
         }
+        if (!cut_cache_.bind_outputs(cache_outputs)) {
+            last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+            return fail_segment("output cache allocation");
+        }
+        new_cache_bytes = add_bytes(last ? 0 : cut_cache_.estimate_output_bytes(graph, segment), cache_.pending_bytes(segment_graph));
         if (!workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) {
                 pin_multi_device_nodes(scheduler, current);
             })) {
@@ -978,7 +1029,7 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
         }
         LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(),
                   index + 1, plan.segments.size(), segment.group_name.c_str());
-        if (!execute_segment(segment_graph, n_threads)) {
+        if (!execute_segment(segment_graph, n_threads, [&]() { weights.submit_next(); })) {
             return fail_segment("execution");
         }
         auto cache_status = cache_.capture(segment_graph);

src/core/ggml_runner.h

@@ -141,12 +141,12 @@ private:
     size_t logged_segment_count_     = 0;
     ggml_status last_compute_status_ = GGML_STATUS_SUCCESS;
 
-    sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph);
+    sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, const std::unordered_set<const ggml_tensor*>* external_outputs = nullptr);
     std::vector<DeviceMemoryRequest> memory_requests(const std::vector<sd::BackendBufferSize>& sizes,
                                                      size_t pending_cache_bytes) const;
     bool fits(const std::vector<DeviceMemoryRequest>& requests,
               const std::vector<ggml_tensor*>& params) const;
-    bool execute_segment(ggml_cgraph* graph, int n_threads);
+    bool execute_segment(ggml_cgraph* graph, int n_threads, const std::function<void()>& on_before_compute);
     std::optional<sd::Tensor<float>> execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs);
 
 protected:
@@ -162,8 +162,8 @@ protected:
     sd::ComputeWorkspace workspace_;
     ggml_context* compute_ctx = nullptr;
     std::unique_ptr<std::max_align_t[]> compute_ctx_buffer_;
-    bool runner_started_      = false;
-    bool graph_active_        = false;
+    bool runner_started_ = false;
+    bool graph_active_   = false;
 
     size_t max_graph_vram_bytes        = 0;
     bool graph_cut_layer_split_enabled = false;

src/core/runner_cache.cpp

@@ -24,15 +24,7 @@ namespace sd {
         ggml_free(context);
     }
 
-    std::unique_ptr<CachedTensor> CachedTensor::copy(ggml_backend_t backend,
-                                                     const std::string& name,
-                                                     ggml_tensor* source,
-                                                     ggml_status& status) {
-        status = GGML_STATUS_FAILED;
-        if (ggml_graph_cut::tensor_buffer(source) == nullptr) {
-            return nullptr;
-        }
-        status         = GGML_STATUS_ALLOC_FAILED;
+    std::unique_ptr<CachedTensor> CachedTensor::allocate(ggml_backend_t backend, const std::string& name, ggml_tensor* source) {
         auto entry     = std::make_unique<CachedTensor>();
         entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true});
         if (entry->context == nullptr) {
@@ -46,6 +38,22 @@ namespace sd {
         if (entry->buffer == nullptr) {
             return nullptr;
         }
+        return entry;
+    }
+
+    std::unique_ptr<CachedTensor> CachedTensor::copy(ggml_backend_t backend,
+                                                     const std::string& name,
+                                                     ggml_tensor* source,
+                                                     ggml_status& status) {
+        status = GGML_STATUS_FAILED;
+        if (ggml_graph_cut::tensor_buffer(source) == nullptr) {
+            return nullptr;
+        }
+        status     = GGML_STATUS_ALLOC_FAILED;
+        auto entry = allocate(backend, name, source);
+        if (entry == nullptr) {
+            return nullptr;
+        }
         if (source->view_src != nullptr || !ggml_is_contiguous(source) || source->buffer == nullptr) {
             std::vector<uint8_t> data(ggml_nbytes(source));
             ggml_backend_tensor_get(source, data.data(), 0, data.size());
@@ -155,6 +163,26 @@ namespace sd {
         return sd::resident_bytes(tensors_, device);
     }
 
+    bool GraphCutTensorCache::can_bind_output(ggml_backend_t backend, const ggml_tensor* tensor) {
+        return tensor != nullptr && tensor->view_src == nullptr && tensor->op != GGML_OP_NONE &&
+               ggml_is_contiguous(tensor) && ggml_backend_supports_op(backend, tensor);
+    }
+
+    bool GraphCutTensorCache::bind_outputs(const std::unordered_set<const ggml_tensor*>& outputs) {
+        for (const ggml_tensor* output : outputs) {
+            auto tensor = const_cast<ggml_tensor*>(output);
+            auto entry  = CachedTensor::allocate(backend_, tensor->name, tensor);
+            if (entry == nullptr) {
+                return false;
+            }
+            tensor->buffer         = entry->tensor->buffer;
+            tensor->data           = entry->tensor->data;
+            tensor->extra          = entry->tensor->extra;
+            tensors_[tensor->name] = std::move(entry);
+        }
+        return true;
+    }
+
     size_t GraphCutTensorCache::estimate_output_bytes(
         ggml_cgraph* graph,
         const ggml_graph_cut::Segment& segment) const {
@@ -171,6 +199,10 @@ namespace sd {
                 !segment.future_cut_names.count(output->name)) {
                 continue;
             }
+            const ggml_tensor* cached = get(output->name);
+            if (cached != nullptr && output->buffer == cached->buffer && output->data == cached->data) {
+                continue;
+            }
             ggml_tensor* source      = ggml_graph_cut::cache_source_tensor(output);
             const size_t tensor_size = GGML_PAD(
                 ggml_backend_buft_get_alloc_size(buffer_type, source), alignment);
@@ -196,6 +228,10 @@ namespace sd {
                 !segment.future_cut_names.count(output->name)) {
                 continue;
             }
+            const ggml_tensor* cached = get(output->name);
+            if (cached != nullptr && output->buffer == cached->buffer && output->data == cached->data) {
+                continue;
+            }
             ggml_status status;
             auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output), status);
             if (entry == nullptr) {

src/core/runner_cache.h

@@ -18,6 +18,7 @@ namespace sd {
         ggml_backend_buffer_t buffer = nullptr;
         ggml_tensor* tensor          = nullptr;
         ~CachedTensor();
+        static std::unique_ptr<CachedTensor> allocate(ggml_backend_t backend, const std::string& name, ggml_tensor* source);
         static std::unique_ptr<CachedTensor> copy(ggml_backend_t backend,
                                                   const std::string& name,
                                                   ggml_tensor* source,
@@ -59,6 +60,8 @@ namespace sd {
         size_t resident_bytes(ggml_backend_dev_t device) const;
         size_t estimate_output_bytes(ggml_cgraph* graph,
                                      const ggml_graph_cut::Segment& segment) const;
+        static bool can_bind_output(ggml_backend_t backend, const ggml_tensor* tensor);
+        bool bind_outputs(const std::unordered_set<const ggml_tensor*>& outputs);
         ggml_status capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc);
         void prune(const std::unordered_set<std::string>& keep_names);
         void clear() { tensors_.clear(); }

src/core/segment_weight_pipeline.cpp

@@ -178,7 +178,7 @@ namespace sd {
             disable();
             return;
         }
-        switch (manager->prefetch_params(owner_id_, params)) {
+        switch (manager->prepare_prefetch_params(owner_id_, params)) {
             case WeightPrefetchResult::Scheduled:
                 queued_params_  = std::move(params);
                 queued_segment_ = next_segment;
@@ -195,6 +195,14 @@ namespace sd {
         disable();
     }
 
+    void SegmentWeightPipeline::submit_next() {
+        if (queued_segment_ != SIZE_MAX) {
+            if (auto manager = residency_manager_.lock()) {
+                manager->submit_prefetched_params(owner_id_);
+            }
+        }
+    }
+
     void SegmentWeightPipeline::clear() {
         if (auto manager = residency_manager_.lock()) {
             manager->clear_prefetched_params(owner_id_);

src/core/segment_weight_pipeline.h

@@ -56,6 +56,7 @@ namespace sd {
         // Prefetch is best effort; segment_start falls back to synchronous loading.
         void enqueue_next(size_t segment_index,
                           const DeviceMemoryRequest& request);
+        void submit_next();
     };
 }