refactor: unify runner execution and weight residency (#2321)

docs/backend.md

@@ -160,10 +160,12 @@ the runner's graph-cut capacity checks.
 Runtime capacity checks also leave 512 MiB of currently free device memory for
 backend scratch buffers and pipelines, including with explicit backend assignments.
 They cap free-memory reports by the device's total memory minus tracked
-resident allocations. Vulkan reports exceeding total memory are rejected because
-its heap-budget subtraction can underflow. Other backends use the cap instead of
-treating such reports as zero free memory. Failed checks log the reported free and
-total memory alongside tracked weight and runtime allocations.
+resident allocations. Vulkan reports exceeding total memory are rejected on
+discrete GPUs because its heap-budget subtraction can underflow. On integrated
+GPUs, a wrapped report triggers reclamation before the tracked estimate can be
+accepted, subject to the byte checks and configured budgets. Other backends use
+the cap instead of treating such reports as zero free memory. Failed checks log
+the reported free and total memory alongside tracked weight and runtime allocations.
 With `--mmap`, device-backed mappings count toward these budgets at their full
 mapped-file size, once per device buffer even when multiple parameter blocks
 share it. Mappings retained in the loader cache continue to count.
@@ -181,8 +183,18 @@ GPU cache space follows the same component priority. Before a lower-priority
 component can become permanently resident, the planner leaves room for the full
 weights and estimated compute space of higher-priority offloaded components.
 If offloaded diffusion already needs the entire main GPU budget, TE and VAE also
-use offloaded parameters. Their GPU copies can then be released after their
-phases, leaving more room to reuse diffusion weights across sampling steps.
+use offloaded parameters. GPU copies remain cached across sampling steps while
+a component's execution is active and are released under memory pressure.
+Ending a component's execution releases its staged GPU copies, disk-backed
+weight buffers, temporary compute buffers and execution caches. Parameters
+assigned permanently to a backend remain resident. Hires model upscaling shares
+the runtime budget and can
+reclaim idle weight copies from other components.
+Reclamation releases idle compute buffers first, then cached weight allocations
+until the current work fits. Weights used by active computation or pending
+transfers stay protected.
+For CPU computation with disk-backed parameters, cached weights are released
+when execution ends so they do not accumulate in RAM across components.
 CPU parameter residency allows GPU weight caching; it does not force every
 weight to be copied again at every step.
 

src/core/execution_engine.cpp

@@ -0,0 +1,992 @@
+#include "core/execution_engine.h"
+
+#include <algorithm>
+#include <chrono>
+#include <exception>
+#include <thread>
+#include <unordered_set>
+
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/layer_split_partition.h"
+#include "core/segment_graph_bindings.h"
+#include "core/segment_weight_pipeline.h"
+#include "core/tensor_ggml.hpp"
+#include "core/util.h"
+
+std::shared_ptr<sd::ExecutionEngine> DeviceResidencyManager::execution_engine() {
+    if (execution_engine_ == nullptr) {
+        execution_engine_ = std::make_shared<sd::ExecutionEngine>(this);
+    }
+    return execution_engine_;
+}
+
+namespace sd {
+    class GraphExecution {
+        ExecutionEngine& engine_;
+        ExecutionSession& session_;
+        const ComputeTask& task_;
+        std::shared_ptr<DeviceResidencyManager> manager_;
+        std::unordered_set<const ggml_tensor*> params_tensor_set_;
+        ggml_status last_compute_status_ = GGML_STATUS_FAILED;
+        Tensor<float> output_;
+
+        using GraphCutPlan = ggml_graph_cut::Plan;
+        ggml_tensor* canonical_param_tensor(ggml_tensor* tensor);
+        std::vector<ggml_tensor*> collect_used_param_tensors(ggml_cgraph* graph);
+        ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const;
+        // Weightless ops follow the most recent weight device. Views must stay
+        // unpinned so the scheduler preserves cross-device copies for consumers.
+        void pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* graph, ggml_cgraph* original_graph = nullptr);
+        Tensor<float> read_graph_tensor(ggml_tensor* tensor, const char* label);
+        void copy_data_to_backend_tensor(ggml_cgraph* graph);
+        const GraphCutPlan& resolve_graph_cut_plan(ggml_cgraph* graph);
+        bool assign_graph_cut_layer_split_backends(ggml_cgraph* graph);
+        ComputeWorkspace::Measurement measure(ggml_cgraph* graph, const std::unordered_set<const ggml_tensor*>* external_outputs = nullptr);
+        std::vector<DeviceMemoryRequest> memory_requests(const std::vector<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, const std::function<void()>& on_before_compute);
+        bool execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs);
+        bool is_multi_device() const { return !session_.extra_backends_.empty(); }
+        size_t retained_runtime_buffer_bytes(ggml_backend_t backend) const { return session_.resident_bytes(backend); }
+        void sync_runtime_residency() { engine_.sync_residency(session_); }
+
+    public:
+        GraphExecution(ExecutionEngine& engine, ExecutionSession& session, const ComputeTask& task)
+            : engine_(engine), session_(session), task_(task), manager_(session.residency_manager_.lock()) {
+            params_tensor_set_.insert(task.weights.begin(), task.weights.end());
+            if (manager_ != nullptr) {
+                for (int i = 0; i < ggml_graph_cut::leaf_count(task.graph); ++i) {
+                    auto* parameter = manager_->resolve_param_tensor(ggml_graph_cut::leaf_tensor(task.graph, i));
+                    if (parameter != nullptr) {
+                        params_tensor_set_.insert(parameter);
+                    }
+                }
+                for (auto* tensor : params_tensor_set_) {
+                    auto* parameter = manager_->resolve_param_tensor(const_cast<ggml_tensor*>(tensor));
+                    if (parameter != nullptr) {
+                        session_.params_.insert(parameter);
+                    }
+                }
+            }
+        }
+
+        ComputeResult run() {
+            if (!execute_graph(task_.graph, task_.n_threads, task_.no_return, task_.read_outputs)) {
+                return {last_compute_status_, {}};
+            }
+            return {GGML_STATUS_SUCCESS, std::move(output_)};
+        }
+    };
+
+    static bool backend_can_allocate_op(ggml_backend_t backend, const ggml_tensor* node) {
+        // Vulkan reports a preferred allocation chunk size, not a hard tensor limit.
+        if (sd_backend_is(backend, "Vulkan")) {
+            return true;
+        }
+        auto buft        = ggml_backend_get_default_buffer_type(backend);
+        size_t alignment = ggml_backend_buft_get_alignment(buft);
+        size_t max_size  = ggml_backend_buft_get_max_size(buft);
+        max_size -= max_size % alignment;
+        if (ggml_backend_buft_get_alloc_size(buft, node) > max_size) {
+            return false;
+        }
+        // Cross-backend copies must also fit the destination's allocation limit.
+        for (auto src : node->src) {
+            if (src != nullptr && ggml_backend_buft_get_alloc_size(buft, src) > max_size) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    static size_t add_bytes(size_t a, size_t b) {
+        return b > SIZE_MAX - a ? SIZE_MAX : a + b;
+    }
+
+    ggml_tensor* GraphExecution::canonical_param_tensor(ggml_tensor* tensor) {
+        for (auto* current = tensor; current != nullptr; current = current->view_src) {
+            if (params_tensor_set_.count(current) != 0)
+                return current;
+        }
+        return nullptr;
+    }
+
+    std::vector<ggml_tensor*> GraphExecution::collect_used_param_tensors(ggml_cgraph* gf) {
+        std::vector<ggml_tensor*> used_params;
+        if (gf == nullptr || params_tensor_set_.empty()) {
+            return used_params;
+        }
+
+        std::unordered_set<const ggml_tensor*> seen_params;
+        const int n_leafs = sd::ggml_graph_cut::leaf_count(gf);
+        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 = canonical_param_tensor(leaf);
+            if (param_leaf != nullptr &&
+                seen_params.insert(param_leaf).second) {
+                used_params.push_back(param_leaf);
+            }
+        }
+        return used_params;
+    }
+
+    ggml_backend_t GraphExecution::backend_for_weight(const ggml_tensor* tensor) const {
+        if (tensor == nullptr || tensor->buffer == nullptr) {
+            return nullptr;
+        }
+        if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS ||
+            ggml_backend_buffer_is_host(tensor->buffer)) {
+            return nullptr;
+        }
+        ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer));
+        if (dev == nullptr) {
+            return nullptr;
+        }
+        if (ggml_backend_get_device(session_.backend_) == dev) {
+            return session_.backend_;
+        }
+        for (ggml_backend_t backend : session_.extra_backends_) {
+            if (ggml_backend_get_device(backend) == dev) {
+                return backend;
+            }
+        }
+        return nullptr;
+    }
+
+    void GraphExecution::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph) {
+        if (sched == nullptr || gf == nullptr) {
+            return;
+        }
+        ggml_backend_t current = session_.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);
+            auto node_assignment = session_.node_assignments_.find(original_graph == nullptr ? node : ggml_graph_node(original_graph, i));
+            if (node_assignment != session_.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 == session_.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;
+            }
+            if (!backend_can_allocate_op(current, node)) {
+                auto cpu = session_.workspace_.cpu_backend();
+                if (cpu != nullptr && ggml_backend_supports_op(cpu, node)) {
+                    ggml_backend_sched_set_tensor_backend(sched, node, cpu);
+                }
+            } else if (ggml_backend_supports_op(current, node)) {
+                ggml_backend_sched_set_tensor_backend(sched, node, current);
+            }
+        }
+    }
+
+    Tensor<float> GraphExecution::read_graph_tensor(ggml_tensor* tensor, const char* label) {
+        if (tensor == nullptr) {
+            LOG_ERROR("%s %s tensor is null", task_.description.c_str(), label);
+            return {};
+        }
+        if (tensor->type != GGML_TYPE_F32) {
+            LOG_ERROR("%s %s tensor type mismatch: got %s",
+                      task_.description.c_str(),
+                      label,
+                      ggml_type_name(tensor->type));
+            return {};
+        }
+        ggml_backend_buffer_t buf = sd::ggml_graph_cut::tensor_buffer(tensor);
+        if (buf == nullptr) {
+            LOG_ERROR("%s %s tensor buffer missing: name=%s op=%s buffer=%p view_src=%p view_src_buffer=%p data=%p",
+                      task_.description.c_str(),
+                      label,
+                      tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
+                      ggml_op_name(tensor->op),
+                      tensor->buffer,
+                      tensor->view_src,
+                      tensor->view_src ? tensor->view_src->buffer : nullptr,
+                      tensor->data);
+            return {};
+        }
+
+        return sd::make_sd_tensor_from_ggml<float>(tensor);
+    }
+
+    void GraphExecution::copy_data_to_backend_tensor(ggml_cgraph* gf) {
+        GGML_ASSERT(gf != nullptr);
+        std::unordered_set<const ggml_tensor*> graph_tensor_set;
+        const int n_leafs = sd::ggml_graph_cut::leaf_count(gf);
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        graph_tensor_set.reserve(static_cast<size_t>(n_leafs + n_nodes));
+        for (int i = 0; i < n_leafs; ++i) {
+            graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i));
+        }
+        for (int i = 0; i < n_nodes; ++i) {
+            graph_tensor_set.insert(ggml_graph_node(gf, i));
+        }
+
+        for (auto& kv : task_.inputs) {
+            auto tensor = kv.first;
+            auto data   = kv.second;
+            if (tensor == nullptr || data == nullptr) {
+                continue;
+            }
+            const char* name = ggml_get_name(tensor);
+            if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) {
+                continue;
+            }
+            if (tensor->buffer == nullptr) {
+                LOG_WARN("%s skip backend tensor copy: tensor buffer not set, name='%s', ne=[%lld,%lld,%lld,%lld], type=%s",
+                         task_.description.c_str(),
+                         name != nullptr ? name : "",
+                         (long long)tensor->ne[0],
+                         (long long)tensor->ne[1],
+                         (long long)tensor->ne[2],
+                         (long long)tensor->ne[3],
+                         ggml_type_name(tensor->type));
+                continue;
+            }
+
+            ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
+            if (buf == nullptr) {
+                LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p",
+                         task_.description.c_str(),
+                         tensor && tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
+                         tensor ? ggml_op_name(tensor->op) : "<null>",
+                         data,
+                         tensor ? tensor->view_src : nullptr,
+                         (tensor && tensor->view_src) ? tensor->view_src->buffer : nullptr);
+                continue;
+            }
+
+            ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor));
+        }
+    }
+
+    const GraphExecution::GraphCutPlan& GraphExecution::resolve_graph_cut_plan(ggml_cgraph* gf) {
+        GGML_ASSERT(gf != nullptr);
+        return sd::ggml_graph_cut::resolve_plan(session_.backend_,
+                                                gf,
+                                                &session_.plan_cache_,
+                                                params_tensor_set_,
+                                                task_.description.c_str());
+    }
+
+    bool GraphExecution::assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
+        session_.node_assignments_.clear();
+        if (!session_.layer_split_enabled_) {
+            return true;
+        }
+        if (!is_multi_device()) {
+            LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", task_.description.c_str());
+            return false;
+        }
+
+        const auto& plan = resolve_graph_cut_plan(gf);
+        if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) {
+            auto manager = manager_;
+            if (manager == nullptr) {
+                LOG_ERROR("%s weight manager is not set for graph-cut layer split", task_.description.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, session_.backend_)) {
+                LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s",
+                          task_.description.c_str(),
+                          sd::layer_split_backend_device_display_name(session_.backend_).c_str());
+                return false;
+            }
+            for (ggml_tensor* param : graph_params) {
+                if (param != nullptr) {
+                    session_.weight_assignments_[param] = session_.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) {
+                    session_.node_assignments_[node] = session_.backend_;
+                }
+            }
+            if (!session_.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",
+                         task_.description.c_str(),
+                         sd::layer_split_backend_device_display_name(session_.backend_).c_str(),
+                         graph_params.size());
+                session_.primary_notice_logged_ = true;
+            } else {
+                LOG_VERBOSE("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
+                            task_.description.c_str(),
+                            sd::layer_split_backend_device_display_name(session_.backend_).c_str(),
+                            graph_params.size());
+            }
+            return true;
+        }
+
+        std::vector<ggml_backend_t> split_backends;
+        split_backends.reserve(session_.extra_backends_.size() + 1);
+        split_backends.push_back(session_.backend_);
+        for (ggml_backend_t backend : session_.extra_backends_) {
+            if (backend != nullptr) {
+                split_backends.push_back(backend);
+            }
+        }
+
+        auto manager = manager_;
+        if (manager == nullptr) {
+            LOG_ERROR("%s weight manager is not set for graph-cut layer split", task_.description.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(task_.description.c_str(),
+                                                 gf,
+                                                 plan,
+                                                 split_backends,
+                                                 session_.backend_limits_,
+                                                 session_.max_backend_bytes_,
+                                                 session_.weight_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",
+                          task_.description.c_str(),
+                          sd::layer_split_backend_device_display_name(split_backends[i]).c_str());
+                return false;
+            }
+        }
+
+        session_.node_assignments_ = std::move(assignment.node_assignments);
+        sd::log_graph_cut_layer_split_assignment(task_.description.c_str(), split_backends, assignment);
+
+        return true;
+    }
+
+    ComputeWorkspace::Measurement GraphExecution::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 session_.backend_;
+            }
+            if (!params_tensor_set_.count(tensor)) {
+                return nullptr;
+            }
+            auto placement = session_.weight_assignments_.find(tensor);
+            return placement == session_.weight_assignments_.end() ? session_.backend_ : placement->second;
+        };
+        auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) {
+            pin_multi_device_nodes(scheduler, copy, graph);
+        };
+        return session_.workspace_.measure(graph, external_backend, assign_nodes, !is_multi_device());
+    }
+
+    std::vector<DeviceMemoryRequest> GraphExecution::memory_requests(
+        const std::vector<BackendBufferSize>& sizes,
+        size_t pending_cache_bytes) const {
+        std::vector<DeviceMemoryRequest> requests;
+        for (const auto& size : sizes) {
+            const size_t retained    = retained_runtime_buffer_bytes(size.backend);
+            const size_t reusable    = session_.workspace_.bytes(size.backend);
+            const size_t cache_bytes = size.backend == session_.backend_ ? pending_cache_bytes : 0;
+            const size_t pending     = add_bytes(size.bytes > reusable ? size.bytes - reusable : 0, cache_bytes);
+            size_t limit             = session_.max_backend_bytes_;
+            if (is_multi_device()) {
+                size_t index = 0;
+                if (size.backend != session_.backend_) {
+                    auto position = std::find(session_.extra_backends_.begin(), session_.extra_backends_.end(), size.backend);
+                    index         = static_cast<size_t>(position - session_.extra_backends_.begin()) + 1;
+                }
+                if (index < session_.backend_limits_.size()) {
+                    limit = session_.backend_limits_[index];
+                }
+            }
+            requests.push_back({size.backend, session_.owner_id(), pending,
+                                retained, limit});
+        }
+        return requests;
+    }
+
+    bool GraphExecution::fits(const std::vector<DeviceMemoryRequest>& requests,
+                              const std::vector<ggml_tensor*>& params) const {
+        auto manager = manager_;
+        if (manager == nullptr) {
+            return params.empty();
+        }
+        for (const auto& request : requests) {
+            if (!manager->fits_compute_backend_capacity(request, params)) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    bool GraphExecution::execute_segment(ggml_cgraph* graph, int n_threads, const std::function<void()>& on_before_compute) {
+        if (sd_backend_is_cpu(session_.backend_)) {
+            sd_backend_cpu_set_n_threads(session_.backend_, n_threads);
+        }
+        if (session_.workspace_.cpu_backend() != nullptr) {
+            sd_backend_cpu_set_n_threads(session_.workspace_.cpu_backend(), n_threads);
+        }
+        auto scheduler = session_.workspace_.scheduler();
+        GGML_ASSERT(scheduler != nullptr);
+        ggml_backend_sched_set_eval_callback(scheduler, sd_get_backend_eval_callback(), sd_get_backend_eval_callback_data());
+#ifndef SD_USE_UPSTREAM_GGML
+        struct BeforeComputeCallback {
+            ggml_backend_t backend;
+            const std::function<void()>& callback;
+            bool invoked = false;
+        } before_compute{session_.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
+        session_.workspace_.synchronize();
+        if (status != GGML_STATUS_SUCCESS) {
+            last_compute_status_ = status;
+            LOG_ERROR("%s compute failed: %s", task_.description.c_str(), ggml_status_to_string(status));
+            return false;
+        }
+        const std::string description = task_.description;
+        if (!task_.debug_tensors.empty()) {
+            std::unordered_set<const ggml_tensor*> graph_tensors;
+            const int leaf_count = ggml_graph_cut::leaf_count(graph);
+            const int node_count = ggml_graph_n_nodes(graph);
+            graph_tensors.reserve(static_cast<size_t>(leaf_count + node_count));
+            for (int index = 0; index < leaf_count; ++index) {
+                graph_tensors.insert(ggml_graph_cut::leaf_tensor(graph, index));
+            }
+            for (int index = 0; index < node_count; ++index) {
+                graph_tensors.insert(ggml_graph_node(graph, index));
+            }
+
+            for (const auto& entry : task_.debug_tensors) {
+                ggml_tensor* tensor = entry.first;
+                if (tensor == nullptr || graph_tensors.find(tensor) == graph_tensors.end()) {
+                    continue;
+                }
+                ggml_backend_buffer_t buffer =
+                    tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer;
+                if (buffer == nullptr) {
+                    LOG_WARN("%s skip debug tensor '%s': tensor buffer not set",
+                             description.c_str(),
+                             entry.second.c_str());
+                    continue;
+                }
+                if (tensor->type != GGML_TYPE_F32) {
+                    LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s",
+                             description.c_str(),
+                             entry.second.c_str(),
+                             ggml_type_name(tensor->type));
+                    continue;
+                }
+                auto debug_tensor = make_sd_tensor_from_ggml<float>(tensor);
+                print_sd_tensor(debug_tensor, false, entry.second.c_str());
+            }
+        }
+
+        return true;
+    }
+
+    bool GraphExecution::execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs) {
+        if (!assign_graph_cut_layer_split_backends(graph)) {
+            return false;
+        }
+        const auto params       = collect_used_param_tensors(graph);
+        const auto& cached_plan = resolve_graph_cut_plan(graph);
+        auto manager            = manager_;
+        const bool can_segment  = !is_multi_device() && !sd_backend_is_cpu(session_.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({{session_.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 false;
+            }
+            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, session_.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();
+            auto& segment               = monolithic_plan.segments.back();
+            segment.group_name          = "graph";
+            segment.compute_buffer_size = cached_plan.compute_buffer_size;
+            segment.internal_node_indices.reserve(ggml_graph_n_nodes(graph));
+            segment.input_refs.reserve(ggml_graph_cut::leaf_count(graph));
+            for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
+                segment.internal_node_indices.push_back(i);
+            }
+            for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) {
+                auto tensor = ggml_graph_cut::leaf_tensor(graph, i);
+                ggml_graph_cut::Segment::InputRef input;
+                input.leaf_index = i;
+                input.type       = canonical_param_tensor(tensor) != nullptr
+                                       ? ggml_graph_cut::Segment::INPUT_PARAM
+                                       : ggml_graph_cut::Segment::INPUT_EXTERNAL;
+                segment.input_refs.push_back(input);
+            }
+        }
+        const auto& plan            = segmented ? cached_plan : monolithic_plan;
+        const bool segments_changed = plan.segments.size() != session_.logged_segment_count_;
+        if (segments_changed && (segmented || session_.logged_segment_count_ > 1)) {
+            LOG_VERBOSE("%s using %zu segment%s", task_.description.c_str(),
+                        plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
+        }
+        SegmentGraphBindings bindings(session_.cut_cache_, plan, graph);
+        SegmentWeightPipeline weights(manager, session_.backend_, session_.owner_id(),
+                                      graph, plan, params_tensor_set_,
+                                      segmented && manager != nullptr && manager->prefetch_enabled());
+
+        std::map<ggml_backend_t, size_t> peak_compute_bytes;
+        auto track_compute_buffer = [&](ggml_backend_t backend) {
+            if (backend != nullptr) {
+                auto& peak = peak_compute_bytes[backend];
+                peak       = std::max(peak, session_.workspace_.bytes(backend));
+            }
+        };
+        for (size_t index = 0; index < plan.segments.size(); ++index) {
+            const auto& segment = plan.segments[index];
+            const bool last     = index + 1 == plan.segments.size();
+            auto fail_segment   = [&](const char* phase) {
+                LOG_ERROR("%s segment %zu/%zu (%s) failed during %s", task_.description.c_str(),
+                          index + 1, plan.segments.size(), segment.group_name.c_str(), phase);
+                return false;
+            };
+            session_.cut_cache_.prune(segment.live_cut_names);
+            bindings.reset(segment);
+            if (!bindings.bind_cached_inputs(segment, task_.description.c_str())) {
+                return fail_segment("input binding");
+            }
+            ggml_context* segment_context = nullptr;
+            auto segment_graph            = segmented
+                                                ? ggml_graph_cut::build_segment_graph(graph, segment, &segment_context)
+                                                : graph;
+            struct SegmentCleanup {
+                GraphExecution& runner;
+                SegmentWeightPipeline& weights;
+                SegmentGraphBindings& bindings;
+                ggml_context* context;
+                ~SegmentCleanup() {
+                    if (!runner.session_.workspace_.segment_end()) {
+                        runner.last_compute_status_ = GGML_STATUS_FAILED;
+                    }
+                    bindings.restore();
+                    weights.segment_end();
+                    ggml_free(context);
+                    runner.sync_runtime_residency();
+                }
+            } segment_cleanup{*this, weights, bindings, segment_context};
+
+            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(session_.backend_, tensor)) {
+                        cache_outputs.insert(tensor);
+                    }
+                }
+            }
+            auto measurement = segmented ? measure(segment_graph, &cache_outputs) : full_measurement;
+            if (!session_.workspace_.prepare(measurement)) {
+                last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+                return fail_segment("workspace preparation");
+            }
+            const size_t cut_bytes = last ? 0 : session_.cut_cache_.estimate_output_bytes(graph, segment);
+            size_t new_cache_bytes = add_bytes(cut_bytes, session_.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))) {
+                    return true;
+                }
+                if (session_.workspace_.release_excess(measurement)) {
+                    sync_runtime_residency();
+                    requests = memory_requests(measurement.buffers, new_cache_bytes);
+                }
+                const bool ready = weights.ensure_segment_capacity(index, requests);
+                if (!ready && manager != nullptr) {
+                    last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+                }
+                return ready;
+            };
+            if (!weights.segment_start(index, ensure_capacity)) {
+                return fail_segment("weight preparation");
+            }
+            // Preparing weights can execute LoRA graphs and reclaim an idle workspace.
+            if (!session_.workspace_.prepare(measurement)) {
+                last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+                return fail_segment("workspace preparation");
+            }
+            if (!ensure_capacity()) {
+                return fail_segment("workspace capacity check");
+            }
+            if (!session_.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 : session_.cut_cache_.estimate_output_bytes(graph, segment), session_.cache_.pending_bytes(segment_graph));
+            if (!session_.workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) {
+                    pin_multi_device_nodes(scheduler, current);
+                })) {
+                last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+                return fail_segment("workspace allocation");
+            }
+            for (const auto& size : measurement.buffers) {
+                track_compute_buffer(size.backend);
+            }
+            if (session_.workspace_.scheduler() != nullptr) {
+                track_compute_buffer(session_.workspace_.cpu_backend());
+            }
+            if (!ensure_capacity()) {
+                return fail_segment("allocated capacity check");
+            }
+            copy_data_to_backend_tensor(segment_graph);
+            auto prefetch_requests = memory_requests(measurement.buffers, new_cache_bytes);
+            if (!prefetch_requests.empty()) {
+                weights.enqueue_next(index, prefetch_requests.front());
+            }
+            LOG_DEBUG("%s executing segment %zu/%zu: %s", task_.description.c_str(),
+                      index + 1, plan.segments.size(), segment.group_name.c_str());
+            if (!execute_segment(segment_graph, n_threads, [&]() { weights.submit_next(); })) {
+                return fail_segment("execution");
+            }
+            auto cache_status = session_.cache_.capture(segment_graph);
+            if (cache_status == GGML_STATUS_SUCCESS) {
+                cache_status = session_.cut_cache_.capture(graph, segment, task_.description.c_str());
+            }
+            if (cache_status != GGML_STATUS_SUCCESS) {
+                last_compute_status_ = cache_status;
+                return fail_segment("output caching");
+            }
+            sync_runtime_residency();
+            if (last) {
+                if (read_outputs && !read_outputs()) {
+                    return fail_segment("output finalization");
+                }
+                if (!no_return) {
+                    output_ = read_graph_tensor(task_.output, "output");
+                    if (output_.empty()) {
+                        return fail_segment("output readback");
+                    }
+                }
+            }
+            if (!session_.workspace_.segment_end()) {
+                last_compute_status_ = GGML_STATUS_FAILED;
+                return fail_segment("workspace synchronization");
+            }
+            // Final outputs and their callbacks may still be views of consumed cuts.
+            session_.cut_cache_.prune(segment.future_cut_names);
+        }
+        if (segments_changed || peak_compute_bytes != session_.logged_compute_bytes_) {
+            for (const auto& entry : peak_compute_bytes) {
+                LOG_VERBOSE("%s compute buffer size: %.2f MB(%s) on %s (peak across %zu segment%s)",
+                            task_.description.c_str(), entry.second / (1024.0 * 1024.0),
+                            sd_backend_is_cpu(entry.first) ? "RAM" : "VRAM", ggml_backend_name(entry.first),
+                            plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
+            }
+            session_.logged_compute_bytes_ = std::move(peak_compute_bytes);
+            session_.logged_segment_count_ = plan.segments.size();
+        }
+        return true;
+    }
+
+    ExecutionSession::ExecutionSession(ggml_backend_t backend, const std::shared_ptr<DeviceResidencyManager>& manager)
+        : backend_(backend), residency_manager_(manager), cache_(backend), cut_cache_(backend), workspace_(backend) {
+        GGML_ASSERT(backend_ != nullptr);
+    }
+
+    size_t ExecutionSession::resident_bytes(ggml_backend_t backend) const {
+        size_t bytes = workspace_.bytes(backend);
+        if (backend == backend_) {
+            const auto add = [](size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; };
+            bytes          = add(bytes, cache_.resident_bytes(ggml_backend_get_device(backend)));
+            bytes          = add(bytes, cut_cache_.resident_bytes(ggml_backend_get_device(backend)));
+        }
+        return bytes;
+    }
+
+    void ExecutionSession::set_layer_split_enabled(bool enabled) {
+        layer_split_enabled_ = enabled;
+        if (!enabled) {
+            weight_assignments_.clear();
+            node_assignments_.clear();
+            primary_notice_logged_ = false;
+        }
+    }
+
+    void ExecutionSession::set_backend_limits(const std::vector<size_t>& limits) {
+        backend_limits_ = limits;
+        weight_assignments_.clear();
+        node_assignments_.clear();
+        primary_notice_logged_ = false;
+    }
+
+    void ExecutionSession::set_runtime_backends(const std::vector<ggml_backend_t>& backends) {
+        extra_backends_.clear();
+        for (auto backend : backends) {
+            if (backend != nullptr && backend != backend_ &&
+                std::find(extra_backends_.begin(), extra_backends_.end(), backend) == extra_backends_.end()) {
+                extra_backends_.push_back(backend);
+            }
+        }
+        workspace_.set_extra_backends(extra_backends_);
+        weight_assignments_.clear();
+        node_assignments_.clear();
+        primary_notice_logged_ = false;
+    }
+
+    ExecutionEngine::~ExecutionEngine() {
+        for (auto& entry : sessions_) {
+            end(*entry.second);
+        }
+    }
+
+    ExecutionSession& ExecutionEngine::create_session(ggml_backend_t backend, const std::shared_ptr<DeviceResidencyManager>& manager) {
+        GGML_ASSERT(manager.get() == manager_);
+        auto session = std::make_unique<ExecutionSession>(backend, manager);
+        auto& result = *session;
+        sessions_.emplace(session->owner_id(), std::move(session));
+        return result;
+    }
+
+    void ExecutionEngine::release_session(ExecutionSession& session) {
+        end(session);
+        sessions_.erase(session.owner_id());
+    }
+
+    bool ExecutionEngine::start(ExecutionSession& session) {
+        if (session.started_) {
+            return true;
+        }
+        sd_backend_set_buffer_cache_enabled(session.backend_, true);
+        for (auto backend : session.extra_backends_) {
+            sd_backend_set_buffer_cache_enabled(backend, true);
+        }
+        session.workspace_.set_extra_backends(session.extra_backends_);
+        if (auto manager = session.residency_manager_.lock()) {
+            manager->register_runtime_owner(session.owner_id());
+        }
+        session.started_ = true;
+        return true;
+    }
+
+    void ExecutionEngine::end(ExecutionSession& session) {
+        GGML_ASSERT(!session.executing_);
+        if (!session.started_) {
+            return;
+        }
+        session.workspace_.release();
+        session.cache_.clear();
+        session.cut_cache_.clear();
+        session.logged_compute_bytes_.clear();
+        session.logged_segment_count_ = 0;
+        if (auto manager = session.residency_manager_.lock()) {
+            manager->clear_prefetched_params(session.owner_id());
+            std::vector<ggml_tensor*> params(session.params_.begin(), session.params_.end());
+            manager->evict_compute_backend_params(params);
+            manager->remove_runtime_owner(session.owner_id());
+        }
+        session.params_.clear();
+        sd_backend_clear_buffer_cache(session.backend_);
+        for (auto backend : session.extra_backends_) {
+            sd_backend_clear_buffer_cache(backend);
+        }
+        session.started_ = false;
+    }
+
+    void ExecutionEngine::sync_residency(ExecutionSession& session) {
+        if (auto manager = session.residency_manager_.lock()) {
+            manager->update_runtime_residency(session.owner_id(), session.backend_, session.resident_bytes(session.backend_));
+            for (auto backend : session.extra_backends_) {
+                manager->update_runtime_residency(session.owner_id(), backend, session.resident_bytes(backend));
+            }
+        }
+    }
+
+    bool ExecutionEngine::ensure_capacity(
+        DeviceResidencyManager& manager,
+        const DeviceMemoryRequest& request,
+        const std::vector<ggml_tensor*>& required_params,
+        const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
+        const std::vector<ggml_tensor*>& protected_params) {
+        GGML_ASSERT(&manager == manager_);
+        auto capacity      = [&]() { return manager.inspect_compute_backend_capacity(request, required_params); };
+        const auto initial = capacity();
+        if (!initial.valid) {
+            return false;
+        }
+        if (initial.fits()) {
+            return true;
+        }
+        for (auto& entry : sessions_) {
+            auto& session = *entry.second;
+            if (session.owner_id() != request.owner_id && session.workspace_.release()) {
+                sync_residency(session);
+                if (capacity().fits(true)) {
+                    return true;
+                }
+            }
+        }
+        auto protected_weights = protected_params;
+        protected_weights.insert(protected_weights.end(), required_params.begin(), required_params.end());
+        auto candidates = manager.evictable_compute_backend_blocks(request.compute_backend, protected_weights);
+        std::unordered_map<const ggml_tensor*, size_t> preferred_positions;
+        for (size_t index = 0; index < preferred_eviction_order.size(); ++index) {
+            for (auto* tensor : preferred_eviction_order[index]) {
+                auto* weight = manager.resolve_param_tensor(tensor);
+                preferred_positions.emplace(weight != nullptr ? weight : tensor, index);
+            }
+        }
+        auto preferred_position = [&](const WeightResidencyBlock& block) {
+            size_t position = 0;
+            // A shared allocation follows the last of its members in the eviction order.
+            for (auto* tensor : block.tensors) {
+                const auto found = preferred_positions.find(tensor);
+                if (found == preferred_positions.end()) {
+                    return SIZE_MAX;
+                }
+                position = std::max(position, found->second);
+            }
+            return position;
+        };
+        std::stable_sort(candidates.begin(), candidates.end(), [&](const auto& lhs, const auto& rhs) {
+            if (lhs.requires_demotion != rhs.requires_demotion) {
+                return !lhs.requires_demotion;
+            }
+            if (!lhs.requires_demotion) {
+                const auto left = preferred_position(lhs), right = preferred_position(rhs);
+                if (left != right) {
+                    return left < right;
+                }
+                if (left != SIZE_MAX) {
+                    return false;
+                }
+            }
+            return lhs.last_use_epoch < rhs.last_use_epoch;
+        });
+        for (const auto& block : candidates) {
+            if (manager.evict_compute_backend_block(block) && capacity().fits(true)) {
+                return true;
+            }
+        }
+        auto available = capacity();
+        if (available.fits(true)) {
+            return true;
+        }
+        if (sd_backend_is(request.compute_backend, "Vulkan") &&
+            available.required_budget_bytes <= available.available_budget_bytes) {
+            ggml_backend_synchronize(request.compute_backend);
+            sd_backend_clear_buffer_cache(request.compute_backend);
+            const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(100);
+            do {
+                available = capacity();
+                if (available.fits(true)) {
+                    return true;
+                }
+                if (available.required_budget_bytes > available.available_budget_bytes) {
+                    break;
+                }
+                std::this_thread::sleep_for(std::chrono::milliseconds(10));
+            } while (std::chrono::steady_clock::now() < deadline);
+        }
+        available               = manager.inspect_compute_backend_capacity(request, required_params, true);
+        const auto device_bytes = available.available_device_bytes == SIZE_MAX
+                                      ? std::string("unknown")
+                                      : sd_format("%.2f MB", available.available_device_bytes / (1024.0 * 1024.0));
+        const auto budget_bytes = available.available_budget_bytes == SIZE_MAX
+                                      ? std::string("unlimited")
+                                      : sd_format("%.2f MB", available.available_budget_bytes / (1024.0 * 1024.0));
+        LOG_WARN("execution engine cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %s device / %s budget",
+                 ggml_backend_name(request.compute_backend),
+                 available.required_device_bytes / (1024.0 * 1024.0),
+                 available.required_budget_bytes / (1024.0 * 1024.0), device_bytes.c_str(), budget_bytes.c_str());
+        return false;
+    }
+
+    void ExecutionEngine::clear_cache(ExecutionSession& session) {
+        session.cache_.clear();
+        sync_residency(session);
+    }
+
+    void ExecutionEngine::discard_graph(ExecutionSession& session) {
+        GGML_ASSERT(!session.executing_);
+        session.cache_.graph_end(false);
+        session.cut_cache_.clear();
+        session.workspace_.release();
+        sync_residency(session);
+    }
+
+    ComputeResult ExecutionEngine::execute(ExecutionSession& session, const ComputeTask& task) {
+        if (session.executing_ || task.graph == nullptr || !start(session)) {
+            return {};
+        }
+        session.executing_ = true;
+        ComputeResult result;
+        try {
+            GraphExecution execution(*this, session, task);
+            result = execution.run();
+        } catch (const std::bad_alloc&) {
+            result.status = GGML_STATUS_ALLOC_FAILED;
+            LOG_ERROR("%s graph allocation failed", task.description.c_str());
+        } catch (const std::exception& error) {
+            LOG_ERROR("%s graph execution failed on %s: %s", task.description.c_str(),
+                      ggml_backend_name(session.backend_), error.what());
+        } catch (...) {
+            LOG_ERROR("%s graph execution failed on %s: unknown exception", task.description.c_str(),
+                      ggml_backend_name(session.backend_));
+        }
+        if (!session.workspace_.segment_end()) {
+            result.status = GGML_STATUS_FAILED;
+        }
+        const bool success = result.status == GGML_STATUS_SUCCESS;
+        session.cache_.graph_end(success);
+        session.cut_cache_.clear();
+        if (!success) {
+            result.output = {};
+            session.workspace_.release();
+        }
+        session.executing_ = false;
+        sync_residency(session);
+        return result;
+    }
+}

src/core/execution_engine.h

@@ -0,0 +1,107 @@
+#ifndef __SD_CORE_EXECUTION_ENGINE_H__
+#define __SD_CORE_EXECUTION_ENGINE_H__
+
+#include <cstdint>
+#include <functional>
+#include <map>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "core/compute_workspace.h"
+#include "core/ggml_graph_cut.h"
+#include "core/runner_cache.h"
+#include "core/tensor.hpp"
+#include "device_residency_manager.h"
+
+namespace sd {
+    struct ComputeTask {
+        std::string description;
+        ggml_cgraph* graph  = nullptr;
+        ggml_tensor* output = nullptr;
+        std::vector<ggml_tensor*> weights;
+        std::map<ggml_tensor*, const void*> inputs;
+        std::vector<std::pair<ggml_tensor*, std::string>> debug_tensors;
+        int n_threads  = 1;
+        bool no_return = false;
+        std::function<bool()> read_outputs;
+    };
+
+    struct ComputeResult {
+        ggml_status status = GGML_STATUS_FAILED;
+        Tensor<float> output;
+    };
+
+    class ExecutionSession {
+        friend class ExecutionEngine;
+        friend class GraphExecution;
+
+        ggml_backend_t backend_;
+        std::weak_ptr<DeviceResidencyManager> residency_manager_;
+        RunnerCache cache_;
+        GraphCutTensorCache cut_cache_;
+        ComputeWorkspace workspace_;
+        std::unordered_set<ggml_tensor*> params_;
+        std::vector<ggml_backend_t> extra_backends_;
+        size_t max_backend_bytes_ = 0;
+        bool layer_split_enabled_ = false;
+        std::vector<size_t> backend_limits_;
+        ggml_graph_cut::PlanCache plan_cache_;
+        std::unordered_map<const ggml_tensor*, ggml_backend_t> weight_assignments_;
+        std::unordered_map<const ggml_tensor*, ggml_backend_t> node_assignments_;
+        bool primary_notice_logged_ = false;
+        std::map<ggml_backend_t, size_t> logged_compute_bytes_;
+        size_t logged_segment_count_ = 0;
+        bool started_                = false;
+        bool executing_              = false;
+
+        size_t resident_bytes(ggml_backend_t backend) const;
+        uintptr_t owner_id() const { return reinterpret_cast<uintptr_t>(this); }
+
+    public:
+        ExecutionSession(ggml_backend_t backend, const std::shared_ptr<DeviceResidencyManager>& manager);
+        ExecutionSession(const ExecutionSession&)            = delete;
+        ExecutionSession& operator=(const ExecutionSession&) = delete;
+
+        RunnerCache& cache() { return cache_; }
+        bool started() const { return started_; }
+        size_t compute_bytes() const { return workspace_.bytes(backend_); }
+        void clear_plan_cache() { plan_cache_.graph_cut_plans.clear(); }
+        void set_max_backend_bytes(size_t bytes) { max_backend_bytes_ = bytes; }
+        void set_layer_split_enabled(bool enabled);
+        void set_backend_limits(const std::vector<size_t>& limits);
+        void set_runtime_backends(const std::vector<ggml_backend_t>& backends);
+    };
+
+    class ExecutionEngine {
+        friend class GraphExecution;
+
+        std::map<uintptr_t, std::unique_ptr<ExecutionSession>> sessions_;
+        DeviceResidencyManager* manager_;
+
+        void sync_residency(ExecutionSession& session);
+
+    public:
+        explicit ExecutionEngine(DeviceResidencyManager* manager = nullptr)
+            : manager_(manager) {}
+        ~ExecutionEngine();
+        ExecutionSession& create_session(ggml_backend_t backend, const std::shared_ptr<DeviceResidencyManager>& manager);
+        void release_session(ExecutionSession& session);
+        bool start(ExecutionSession& session);
+        void end(ExecutionSession& session);
+        void clear_cache(ExecutionSession& session);
+        void discard_graph(ExecutionSession& session);
+        bool ensure_capacity(DeviceResidencyManager& manager,
+                             const DeviceMemoryRequest& request,
+                             const std::vector<ggml_tensor*>& required_params,
+                             const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
+                             const std::vector<ggml_tensor*>& protected_params);
+        // Graphs, input data and output callbacks remain borrowed until execute returns.
+        ComputeResult execute(ExecutionSession& session, const ComputeTask& task);
+    };
+}
+
+#endif  // __SD_CORE_EXECUTION_ENGINE_H__

src/core/ggml_runner.cpp

@@ -1,15 +1,11 @@
 #include <algorithm>
 #include <exception>
-#include <map>
 #include <utility>
 
 #include "core/ggml_extend.h"
 #include "core/ggml_extend_backend.h"
 #include "core/ggml_runner.h"
 #include "core/ggml_tensor_utils.h"
-#include "core/layer_split_partition.h"
-#include "core/segment_graph_bindings.h"
-#include "core/segment_weight_pipeline.h"
 
 using namespace sd;
 
@@ -101,16 +97,6 @@ std::vector<ggml_tensor*> GGMLRunner::collect_used_param_tensors(ggml_cgraph* gf
     return used_params;
 }
 
-void GGMLRunner::evict_compute_backend_param_tensors(const std::vector<ggml_tensor*>& tensors) {
-    if (tensors.empty()) {
-        return;
-    }
-    auto manager = residency_manager.lock();
-    if (manager != nullptr) {
-        manager->evict_compute_backend_params(tensors);
-    }
-}
-
 void GGMLRunner::prepare_build_in_tensor_before() {
     one_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_F32, 1);
     ggml_set_name(one_tensor, "ggml_runner_build_in_tensor:one");
@@ -148,7 +134,7 @@ ggml_cgraph* GGMLRunner::get_compute_graph(get_graph_cb_t get_graph) {
             ggml_build_forward_expand(gf, entry.first);
         }
     }
-    for (const auto& entry : cache_.outputs()) {
+    for (const auto& entry : execution_session_.cache().outputs()) {
         if (entry.second != nullptr) {
             ggml_build_forward_expand(gf, entry.second);
         }
@@ -172,365 +158,22 @@ bool GGMLRunner::prepare_compute_graph(get_graph_cb_t get_graph,
     return true;
 }
 
-ggml_backend_t GGMLRunner::backend_for_weight(const ggml_tensor* tensor) const {
-    if (tensor == nullptr || tensor->buffer == nullptr) {
-        return nullptr;
-    }
-    if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS ||
-        ggml_backend_buffer_is_host(tensor->buffer)) {
-        return nullptr;
-    }
-    ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer));
-    if (dev == nullptr) {
-        return nullptr;
-    }
-    if (ggml_backend_get_device(runtime_backend) == dev) {
-        return runtime_backend;
-    }
-    for (ggml_backend_t backend : extra_runtime_backends) {
-        if (ggml_backend_get_device(backend) == dev) {
-            return backend;
-        }
-    }
-    return nullptr;
-}
-
-static bool backend_can_allocate_op(ggml_backend_t backend, const ggml_tensor* node) {
-    // Vulkan reports a preferred allocation chunk size, not a hard tensor limit.
-    if (sd_backend_is(backend, "Vulkan")) {
-        return true;
-    }
-    auto buft        = ggml_backend_get_default_buffer_type(backend);
-    size_t alignment = ggml_backend_buft_get_alignment(buft);
-    size_t max_size  = ggml_backend_buft_get_max_size(buft);
-    max_size -= max_size % alignment;
-    if (ggml_backend_buft_get_alloc_size(buft, node) > max_size) {
-        return false;
-    }
-    // Cross-backend copies must also fit the destination's allocation limit.
-    for (auto src : node->src) {
-        if (src != nullptr && ggml_backend_buft_get_alloc_size(buft, src) > max_size) {
-            return false;
-        }
-    }
-    return true;
-}
-
-void GGMLRunner::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph) {
-    if (sched == nullptr || gf == nullptr) {
-        return;
-    }
-    ggml_backend_t current = 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);
-        auto node_assignment = graph_cut_layer_split_node_assignments_.find(original_graph == nullptr ? node : ggml_graph_node(original_graph, i));
-        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;
-        }
-        if (!backend_can_allocate_op(current, node)) {
-            auto cpu = workspace_.cpu_backend();
-            if (cpu != nullptr && ggml_backend_supports_op(cpu, node)) {
-                ggml_backend_sched_set_tensor_backend(sched, node, cpu);
-            }
-        } else if (ggml_backend_supports_op(current, node)) {
-            ggml_backend_sched_set_tensor_backend(sched, node, current);
-        }
-    }
-}
-
-size_t GGMLRunner::retained_runtime_buffer_bytes(ggml_backend_t backend) const {
-    backend      = backend == nullptr ? runtime_backend : backend;
-    size_t bytes = workspace_.bytes(backend);
-    if (backend == runtime_backend) {
-        const size_t cache_bytes = cache_.resident_bytes(ggml_backend_get_device(backend));
-        bytes                    = cache_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cache_bytes;
-        const size_t cut_bytes   = cut_cache_.resident_bytes(ggml_backend_get_device(backend));
-        bytes                    = cut_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cut_bytes;
-    }
-    return bytes;
-}
-
-void GGMLRunner::sync_runtime_residency() {
-    if (auto manager = residency_manager.lock()) {
-        manager->update_runtime_residency(reinterpret_cast<uintptr_t>(this),
-                                          runtime_backend, retained_runtime_buffer_bytes());
-        for (auto backend : extra_runtime_backends) {
-            manager->update_runtime_residency(reinterpret_cast<uintptr_t>(this),
-                                              backend, retained_runtime_buffer_bytes(backend));
-        }
-    }
-}
-
-std::optional<sd::Tensor<float>> GGMLRunner::read_graph_tensor(ggml_tensor* tensor, const char* label) {
-    if (tensor == nullptr) {
-        LOG_ERROR("%s %s tensor is null", get_desc().c_str(), label);
-        return std::nullopt;
-    }
-    if (tensor->type != GGML_TYPE_F32) {
-        LOG_ERROR("%s %s tensor type mismatch: got %s",
-                  get_desc().c_str(),
-                  label,
-                  ggml_type_name(tensor->type));
-        return std::nullopt;
-    }
-    ggml_backend_buffer_t buf = sd::ggml_graph_cut::tensor_buffer(tensor);
-    if (buf == nullptr) {
-        LOG_ERROR("%s %s tensor buffer missing: name=%s op=%s buffer=%p view_src=%p view_src_buffer=%p data=%p",
-                  get_desc().c_str(),
-                  label,
-                  tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
-                  ggml_op_name(tensor->op),
-                  tensor->buffer,
-                  tensor->view_src,
-                  tensor->view_src ? tensor->view_src->buffer : nullptr,
-                  tensor->data);
-        return std::nullopt;
-    }
-
-    return sd::make_sd_tensor_from_ggml<float>(tensor);
-}
-
-void GGMLRunner::copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy) {
-    GGML_ASSERT(gf != nullptr);
-    std::unordered_set<const ggml_tensor*> graph_tensor_set;
-    const int n_leafs = sd::ggml_graph_cut::leaf_count(gf);
-    const int n_nodes = ggml_graph_n_nodes(gf);
-    graph_tensor_set.reserve(static_cast<size_t>(n_leafs + n_nodes));
-    for (int i = 0; i < n_leafs; ++i) {
-        graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i));
-    }
-    for (int i = 0; i < n_nodes; ++i) {
-        graph_tensor_set.insert(ggml_graph_node(gf, i));
-    }
-
-    for (auto& kv : backend_tensor_data_map) {
-        auto tensor = kv.first;
-        auto data   = kv.second;
-        if (tensor == nullptr || data == nullptr) {
-            continue;
-        }
-        const char* name = ggml_get_name(tensor);
-        if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) {
-            continue;
-        }
-        if (tensor->buffer == nullptr) {
-            LOG_WARN("%s skip backend tensor copy: tensor buffer not set, name='%s', ne=[%lld,%lld,%lld,%lld], type=%s",
-                     get_desc().c_str(),
-                     name != nullptr ? name : "",
-                     (long long)tensor->ne[0],
-                     (long long)tensor->ne[1],
-                     (long long)tensor->ne[2],
-                     (long long)tensor->ne[3],
-                     ggml_type_name(tensor->type));
-            continue;
-        }
-
-        ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
-        if (buf == nullptr) {
-            LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p",
-                     get_desc().c_str(),
-                     tensor && tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
-                     tensor ? ggml_op_name(tensor->op) : "<null>",
-                     data,
-                     tensor ? tensor->view_src : nullptr,
-                     (tensor && tensor->view_src) ? tensor->view_src->buffer : nullptr);
-            continue;
-        }
-
-        ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor));
-    }
-
-    if (clear_after_copy) {
-        backend_tensor_data_map.clear();
-    }
-}
-
-const GGMLRunner::GraphCutPlan& GGMLRunner::resolve_graph_cut_plan(ggml_cgraph* gf) {
-    GGML_ASSERT(gf != nullptr);
-    return sd::ggml_graph_cut::resolve_plan(runtime_backend,
-                                            gf,
-                                            &graph_cut_plan_cache_,
-                                            params_tensor_set_,
-                                            get_desc().c_str());
-}
-
-const GGMLRunner::GraphCutPlan& GGMLRunner::resolve_graph_cut_layer_split_plan(ggml_cgraph* gf) {
-    return resolve_graph_cut_plan(gf);
-}
-
-bool GGMLRunner::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;
-    }
-
-    const auto& plan = resolve_graph_cut_layer_split_plan(gf);
-    if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) {
-        auto manager = residency_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_VERBOSE("%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 = residency_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;
-}
-
 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()) {
-        manager->set_workspace_reclaimer(reinterpret_cast<uintptr_t>(this), [this]() {
-            if (!workspace_.release()) {
-                return false;
-            }
-            sync_runtime_residency();
-            return true;
-        });
-    }
-    runner_started_ = true;
-    return true;
+    return execution_engine_->start(execution_session_);
 }
 
 void GGMLRunner::runner_end() {
     GGML_ASSERT(!graph_active_);
-    if (!runner_started_) {
-        return;
-    }
-    workspace_.release();
+    execution_engine_->end(execution_session_);
     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()) {
-        manager->clear_prefetched_params(reinterpret_cast<uintptr_t>(this));
-        std::vector<ggml_tensor*> tensors;
-        for (auto tensor : params_tensor_set_) {
-            auto* parameter = manager->resolve_param_tensor(const_cast<ggml_tensor*>(tensor));
-            if (parameter != nullptr)
-                tensors.push_back(parameter);
-        }
-        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;
 }
 
 GGMLRunner::GGMLRunner(ggml_backend_t backend,
                        std::shared_ptr<DeviceResidencyManager> manager)
-    : runtime_backend(backend),
-      cache_(backend),
-      cut_cache_(backend),
-      workspace_(backend),
+    : execution_engine_(manager != nullptr ? manager->execution_engine() : std::make_shared<sd::ExecutionEngine>()),
+      execution_session_(execution_engine_->create_session(backend, manager)),
+      runtime_backend(backend),
       residency_manager(manager) {
     GGML_ASSERT(runtime_backend != nullptr);
     alloc_params_ctx();
@@ -538,6 +181,7 @@ GGMLRunner::GGMLRunner(ggml_backend_t backend,
 
 GGMLRunner::~GGMLRunner() {
     runner_end();
+    execution_engine_->release_session(execution_session_);
     free_compute_ctx();
     free_params_ctx();
 }
@@ -575,8 +219,7 @@ void GGMLRunner::reset_compute_ctx() {
 }
 
 void GGMLRunner::free_cache_ctx_and_buffer() {
-    cache_.clear();
-    sync_runtime_residency();
+    execution_engine_->clear_cache(execution_session_);
 }
 
 void GGMLRunner::set_backend_tensor_data(ggml_tensor* tensor, const void* data) {
@@ -610,7 +253,7 @@ void GGMLRunner::cache(const std::string name, ggml_tensor* tensor, ggml_cgraph*
     if (tensor != nullptr) {
         ggml_set_output(tensor);
     }
-    cache_.stage(name, tensor);
+    execution_session_.cache().stage(name, tensor);
     if (graph != nullptr && tensor != nullptr) {
         // Schedule the cache output here so its source can be reused before graph end.
         ggml_build_forward_expand(graph, tensor);
@@ -646,439 +289,54 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
         GGMLRunner& runner;
         const bool& success;
         ~GraphEndGuard() {
-            if (!runner.workspace_.segment_end()) {
-                runner.last_compute_status_ = GGML_STATUS_FAILED;
+            if (!success) {
+                runner.execution_engine_->discard_graph(runner.execution_session_);
             }
-            runner.cache_.graph_end(false);
-            runner.cut_cache_.clear();
             runner.free_compute_ctx();
             runner.graph_active_ = false;
-            if (!success) {
-                runner.workspace_.release();
-            }
-            runner.sync_runtime_residency();
         }
     } graph_guard{*this, success};
 
+    try {
         ggml_cgraph* graph = nullptr;
         if (!prepare_compute_graph(get_graph, &graph)) {
             return std::nullopt;
         }
         params_tensor_set_dirty_ = true;
         rebuild_params_tensor_set();
-    if (auto manager = residency_manager.lock()) {
-        for (int i = 0; i < sd::ggml_graph_cut::leaf_count(graph); ++i) {
-            auto* parameter = manager->resolve_param_tensor(sd::ggml_graph_cut::leaf_tensor(graph, i));
-            if (parameter != nullptr)
-                params_tensor_set_.insert(parameter);
+        sd::ComputeTask task;
+        task.description     = get_desc();
+        task.graph           = graph;
+        task.output          = no_return ? nullptr : ggml_get_tensor(compute_ctx, final_result_name.c_str());
+        task.weights         = collect_used_param_tensors(graph);
+        task.inputs          = backend_tensor_data_map;
+        task.debug_tensors   = debug_tensors;
+        task.n_threads       = n_threads;
+        task.no_return       = no_return;
+        task.read_outputs    = read_outputs;
+        auto result          = execution_engine_->execute(execution_session_, task);
+        last_compute_status_ = result.status;
+        success              = result.status == GGML_STATUS_SUCCESS;
+        if (success) {
+            return std::move(result.output);
         }
-    }
-    std::optional<sd::Tensor<float>> output;
-    try {
-        output = execute_graph(graph, n_threads, no_return, read_outputs);
     } catch (const std::bad_alloc&) {
         last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
         LOG_ERROR("%s graph allocation failed", get_desc().c_str());
-        return std::nullopt;
     } catch (const std::exception& error) {
-        last_compute_status_ = GGML_STATUS_FAILED;
-        LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(),
-                  ggml_backend_name(runtime_backend), error.what());
+        LOG_ERROR("%s graph submission failed: %s", get_desc().c_str(), error.what());
+    }
     return std::nullopt;
 }
-    success = output.has_value();
-    if (success) {
-        cache_.graph_end(true);
-        last_compute_status_ = GGML_STATUS_SUCCESS;
-    }
-    return output;
-}
 
 void GGMLRunner::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;
-    }
+    execution_session_.set_layer_split_enabled(enabled);
 }
 
 void GGMLRunner::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;
+    execution_session_.set_backend_limits(limits);
 }
 
 void GGMLRunner::set_runtime_backends(const std::vector<ggml_backend_t>& backends) {
-    extra_runtime_backends.clear();
-    for (ggml_backend_t backend : backends) {
-        if (backend == nullptr || backend == runtime_backend) {
-            continue;
-        }
-        if (std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), backend) ==
-            extra_runtime_backends.end()) {
-            extra_runtime_backends.push_back(backend);
-        }
-    }
-    workspace_.set_extra_backends(extra_runtime_backends);
-    graph_cut_layer_split_assignments_.clear();
-    graph_cut_layer_split_node_assignments_.clear();
-    graph_cut_layer_split_primary_notice_logged_ = false;
-}
-
-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, 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;
-        }
-        auto placement = graph_cut_layer_split_assignments_.find(tensor);
-        return placement == graph_cut_layer_split_assignments_.end() ? runtime_backend : placement->second;
-    };
-    auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) {
-        pin_multi_device_nodes(scheduler, copy, graph);
-    };
-    return workspace_.measure(graph, external_backend, assign_nodes, !is_multi_device());
-}
-
-std::vector<DeviceMemoryRequest> GGMLRunner::memory_requests(
-    const std::vector<BackendBufferSize>& sizes,
-    size_t pending_cache_bytes) const {
-    std::vector<DeviceMemoryRequest> requests;
-    for (const auto& size : sizes) {
-        const size_t retained    = retained_runtime_buffer_bytes(size.backend);
-        const size_t reusable    = workspace_.bytes(size.backend);
-        const size_t cache_bytes = size.backend == runtime_backend ? pending_cache_bytes : 0;
-        const size_t pending     = add_bytes(size.bytes > reusable ? size.bytes - reusable : 0, cache_bytes);
-        size_t limit             = max_graph_vram_bytes;
-        if (is_multi_device()) {
-            size_t index = 0;
-            if (size.backend != runtime_backend) {
-                auto position = std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), size.backend);
-                index         = static_cast<size_t>(position - extra_runtime_backends.begin()) + 1;
-            }
-            if (index < graph_cut_layer_split_backend_vram_limits_.size()) {
-                limit = graph_cut_layer_split_backend_vram_limits_[index];
-            }
-        }
-        requests.push_back({size.backend, reinterpret_cast<uintptr_t>(this), pending,
-                            retained, limit});
-    }
-    return requests;
-}
-
-bool GGMLRunner::fits(const std::vector<DeviceMemoryRequest>& requests,
-                      const std::vector<ggml_tensor*>& params) const {
-    auto manager = residency_manager.lock();
-    if (manager == nullptr) {
-        return params.empty();
-    }
-    for (const auto& request : requests) {
-        if (!manager->fits_compute_backend_capacity(request, params)) {
-            return false;
-        }
-    }
-    return true;
-}
-
-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);
-    }
-    if (workspace_.cpu_backend() != nullptr) {
-        sd_backend_cpu_set_n_threads(workspace_.cpu_backend(), 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());
-#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;
-        LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status));
-        return false;
-    }
-    const std::string description = get_desc();
-    if (!debug_tensors.empty()) {
-        std::unordered_set<const ggml_tensor*> graph_tensors;
-        const int leaf_count = ggml_graph_cut::leaf_count(graph);
-        const int node_count = ggml_graph_n_nodes(graph);
-        graph_tensors.reserve(static_cast<size_t>(leaf_count + node_count));
-        for (int index = 0; index < leaf_count; ++index) {
-            graph_tensors.insert(ggml_graph_cut::leaf_tensor(graph, index));
-        }
-        for (int index = 0; index < node_count; ++index) {
-            graph_tensors.insert(ggml_graph_node(graph, index));
-        }
-
-        for (const auto& entry : debug_tensors) {
-            ggml_tensor* tensor = entry.first;
-            if (tensor == nullptr || graph_tensors.find(tensor) == graph_tensors.end()) {
-                continue;
-            }
-            ggml_backend_buffer_t buffer =
-                tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer;
-            if (buffer == nullptr) {
-                LOG_WARN("%s skip debug tensor '%s': tensor buffer not set",
-                         description.c_str(),
-                         entry.second.c_str());
-                continue;
-            }
-            if (tensor->type != GGML_TYPE_F32) {
-                LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s",
-                         description.c_str(),
-                         entry.second.c_str(),
-                         ggml_type_name(tensor->type));
-                continue;
-            }
-            auto debug_tensor = make_sd_tensor_from_ggml<float>(tensor);
-            print_sd_tensor(debug_tensor, false, entry.second.c_str());
-        }
-    }
-
-    return true;
-}
-
-std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs) {
-    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);
-    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;
-        }
-        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();
-        auto& segment               = monolithic_plan.segments.back();
-        segment.group_name          = "graph";
-        segment.compute_buffer_size = cached_plan.compute_buffer_size;
-        segment.internal_node_indices.reserve(ggml_graph_n_nodes(graph));
-        segment.input_refs.reserve(ggml_graph_cut::leaf_count(graph));
-        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
-            segment.internal_node_indices.push_back(i);
-        }
-        for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) {
-            auto tensor = ggml_graph_cut::leaf_tensor(graph, i);
-            ggml_graph_cut::Segment::InputRef input;
-            input.leaf_index = i;
-            input.type       = canonical_param_tensor(tensor) != nullptr
-                                   ? ggml_graph_cut::Segment::INPUT_PARAM
-                                   : ggml_graph_cut::Segment::INPUT_EXTERNAL;
-            segment.input_refs.push_back(input);
-        }
-    }
-    const auto& plan            = segmented ? cached_plan : monolithic_plan;
-    const bool segments_changed = plan.segments.size() != logged_segment_count_;
-    if (segments_changed && (segmented || logged_segment_count_ > 1)) {
-        LOG_VERBOSE("%s using %zu segment%s", get_desc().c_str(),
-                    plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
-    }
-    SegmentGraphBindings bindings(cut_cache_, plan, graph);
-    SegmentWeightPipeline weights(manager, runtime_backend, reinterpret_cast<uintptr_t>(this),
-                                  graph, plan, params_tensor_set_,
-                                  segmented && manager != nullptr && manager->prefetch_enabled());
-
-    std::map<ggml_backend_t, size_t> peak_compute_bytes;
-    auto track_compute_buffer = [&](ggml_backend_t backend) {
-        if (backend != nullptr) {
-            auto& peak = peak_compute_bytes[backend];
-            peak       = std::max(peak, workspace_.bytes(backend));
-        }
-    };
-    std::optional<Tensor<float>> output = Tensor<float>();
-    for (size_t index = 0; index < plan.segments.size(); ++index) {
-        const auto& segment = plan.segments[index];
-        const bool last     = index + 1 == plan.segments.size();
-        auto fail_segment   = [&](const char* phase) {
-            LOG_ERROR("%s segment %zu/%zu (%s) failed during %s", get_desc().c_str(),
-                      index + 1, plan.segments.size(), segment.group_name.c_str(), phase);
-            return std::nullopt;
-        };
-        cut_cache_.prune(segment.live_cut_names);
-        bindings.reset(segment);
-        if (!bindings.bind_cached_inputs(segment, get_desc().c_str())) {
-            return fail_segment("input binding");
-        }
-        ggml_context* segment_context = nullptr;
-        auto segment_graph            = segmented
-                                            ? ggml_graph_cut::build_segment_graph(graph, segment, &segment_context)
-                                            : graph;
-        struct SegmentCleanup {
-            GGMLRunner& runner;
-            SegmentWeightPipeline& weights;
-            SegmentGraphBindings& bindings;
-            ggml_context* context;
-            ~SegmentCleanup() {
-                if (!runner.workspace_.segment_end()) {
-                    runner.last_compute_status_ = GGML_STATUS_FAILED;
-                }
-                bindings.restore();
-                weights.segment_end();
-                ggml_free(context);
-                runner.sync_runtime_residency();
-            }
-        } segment_cleanup{*this, weights, bindings, segment_context};
-
-        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);
-        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))) {
-                return true;
-            }
-            if (workspace_.release_excess(measurement)) {
-                sync_runtime_residency();
-                requests = memory_requests(measurement.buffers, new_cache_bytes);
-            }
-            const bool ready = weights.ensure_segment_capacity(index, requests);
-            if (!ready && manager != nullptr) {
-                last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
-            }
-            return ready;
-        };
-        if (!weights.segment_start(index, ensure_capacity)) {
-            return fail_segment("weight preparation");
-        }
-        // Preparing weights can execute LoRA graphs and reclaim an idle workspace.
-        if (!workspace_.prepare(measurement)) {
-            last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
-            return fail_segment("workspace preparation");
-        }
-        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);
-            })) {
-            last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
-            return fail_segment("workspace allocation");
-        }
-        for (const auto& size : measurement.buffers) {
-            track_compute_buffer(size.backend);
-        }
-        if (workspace_.scheduler() != nullptr) {
-            track_compute_buffer(workspace_.cpu_backend());
-        }
-        if (!ensure_capacity()) {
-            return fail_segment("allocated capacity check");
-        }
-        copy_data_to_backend_tensor(segment_graph, false);
-        auto prefetch_requests = memory_requests(measurement.buffers, new_cache_bytes);
-        if (!prefetch_requests.empty()) {
-            weights.enqueue_next(index, prefetch_requests.front());
-        }
-        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, [&]() { weights.submit_next(); })) {
-            return fail_segment("execution");
-        }
-        auto cache_status = cache_.capture(segment_graph);
-        if (cache_status == GGML_STATUS_SUCCESS) {
-            cache_status = cut_cache_.capture(graph, segment, get_desc().c_str());
-        }
-        if (cache_status != GGML_STATUS_SUCCESS) {
-            last_compute_status_ = cache_status;
-            return fail_segment("output caching");
-        }
-        sync_runtime_residency();
-        if (last) {
-            if (read_outputs && !read_outputs()) {
-                return fail_segment("output finalization");
-            }
-            if (!no_return) {
-                auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str());
-                output      = read_graph_tensor(result, "output");
-                if (!output.has_value()) {
-                    return fail_segment("output readback");
-                }
-            }
-        }
-        if (!workspace_.segment_end()) {
-            last_compute_status_ = GGML_STATUS_FAILED;
-            return fail_segment("workspace synchronization");
-        }
-        // Final outputs and their callbacks may still be views of consumed cuts.
-        cut_cache_.prune(segment.future_cut_names);
-    }
-    if (segments_changed || peak_compute_bytes != logged_compute_bytes_) {
-        for (const auto& entry : peak_compute_bytes) {
-            LOG_VERBOSE("%s compute buffer size: %.2f MB(%s) on %s (peak across %zu segment%s)",
-                        get_desc().c_str(), entry.second / (1024.0 * 1024.0),
-                        sd_backend_is_cpu(entry.first) ? "RAM" : "VRAM", ggml_backend_name(entry.first),
-                        plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
-        }
-        logged_compute_bytes_ = std::move(peak_compute_bytes);
-        logged_segment_count_ = plan.segments.size();
-    }
-    return output;
+    execution_session_.set_runtime_backends(backends);
 }

src/core/ggml_runner.h

@@ -12,9 +12,7 @@
 #include <utility>
 #include <vector>
 
-#include "core/compute_workspace.h"
-#include "core/ggml_graph_cut.h"
-#include "core/runner_cache.h"
+#include "core/execution_engine.h"
 #include "core/tensor_ggml.hpp"
 #include "core/util.h"
 #include "device_residency_manager.h"
@@ -132,44 +130,21 @@ struct GGMLRunnerContext {
 
 struct GGMLRunner {
 private:
-    std::map<ggml_backend_t, size_t> logged_compute_bytes_;
-    size_t logged_segment_count_     = 0;
+    std::shared_ptr<sd::ExecutionEngine> execution_engine_;
+    sd::ExecutionSession& execution_session_;
     ggml_status last_compute_status_ = GGML_STATUS_SUCCESS;
 
-    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, 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:
     typedef std::function<ggml_cgraph*()> get_graph_cb_t;
-    using GraphCutPlan = sd::ggml_graph_cut::Plan;
 
     ggml_backend_t runtime_backend = nullptr;
 
     ggml_context* params_ctx = nullptr;
 
-    sd::RunnerCache cache_;
-    sd::GraphCutTensorCache cut_cache_;
-    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;
 
-    size_t max_graph_vram_bytes        = 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)
-
     std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
     std::weak_ptr<DeviceResidencyManager> residency_manager;
     bool params_tensor_set_dirty_ = true;
@@ -193,11 +168,7 @@ protected:
     bool circular_x_enabled    = false;
     bool circular_y_enabled    = false;
 
-    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) {
@@ -239,8 +210,6 @@ protected:
 
     std::vector<ggml_tensor*> collect_used_param_tensors(ggml_cgraph* gf);
 
-    void evict_compute_backend_param_tensors(const std::vector<ggml_tensor*>& tensors);
-
     void prepare_build_in_tensor_before();
 
     void prepare_build_in_tensor_after(ggml_cgraph* gf);
@@ -252,39 +221,18 @@ protected:
     bool prepare_compute_graph(get_graph_cb_t get_graph,
                                ggml_cgraph** gf_out);
 
-    ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const;
-
-    // Weightless ops have no scheduler anchor, so pin them to the most recent
-    // weight device. Views must stay unpinned or cross-device copies can be
-    // skipped for their consumers.
-    void pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph = nullptr);
-
-    bool is_multi_device() const {
-        return !extra_runtime_backends.empty();
-    }
-
     size_t reusable_compute_buffer_bytes() const {
-        return workspace_.bytes(runtime_backend);
+        return execution_session_.compute_bytes();
     }
 
-    size_t retained_runtime_buffer_bytes(ggml_backend_t backend = nullptr) const;
-
-    void sync_runtime_residency();
-
-    std::optional<sd::Tensor<float>> read_graph_tensor(ggml_tensor* tensor, const char* label);
-
-    void copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy = true);
-
-    const GraphCutPlan& resolve_graph_cut_plan(ggml_cgraph* gf);
-
-    const GraphCutPlan& resolve_graph_cut_layer_split_plan(ggml_cgraph* gf);
-
-    bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf);
+    bool has_cache_tensors() const {
+        return !execution_session_.cache().empty();
+    }
 
 public:
     bool runner_start();
 
-    bool runner_started() const { return runner_started_; }
+    bool runner_started() const { return execution_session_.started(); }
 
     void runner_end();
 
@@ -335,7 +283,7 @@ public:
     void cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph = nullptr);
 
     ggml_tensor* get_cache_tensor_by_name(const std::string& name) {
-        return cache_.get(name);
+        return execution_session_.cache().get(name);
     }
 
     std::optional<sd::Tensor<float>> compute(get_graph_cb_t get_graph,
@@ -353,7 +301,7 @@ public:
     void set_sage_attention_enabled(bool enabled) {
         if (sage_attn_enabled != enabled) {
             free_cache_ctx_and_buffer();
-            graph_cut_plan_cache_.graph_cut_plans.clear();
+            execution_session_.clear_plan_cache();
             sage_attn_enabled = enabled;
         }
     }
@@ -381,7 +329,7 @@ public:
     }
 
     void set_max_graph_vram_bytes(size_t max_vram_bytes) {
-        max_graph_vram_bytes = max_vram_bytes;
+        execution_session_.set_max_backend_bytes(max_vram_bytes);
     }
 
     void set_graph_cut_layer_split_enabled(bool enabled);

src/device_residency_manager.h

@@ -3,12 +3,17 @@
 
 #include <cstdint>
 #include <functional>
+#include <memory>
 #include <vector>
 
 #include "ggml-backend.h"
 
 struct ggml_tensor;
 
+namespace sd {
+    class ExecutionEngine;
+}
+
 enum class WeightPrefetchResult {
     Scheduled,
     AlreadyResident,
@@ -36,12 +41,44 @@ struct DeviceMemoryRequest {
     }
 };
 
+struct DeviceMemoryCapacity {
+    size_t required_device_bytes  = 0;
+    size_t required_budget_bytes  = 0;
+    size_t available_device_bytes = SIZE_MAX;
+    size_t available_budget_bytes = SIZE_MAX;
+    bool requires_reclamation     = false;
+    bool valid                    = true;
+    bool wrapped_device_report    = false;
+
+    bool fits(bool after_reclamation = false) const {
+        return valid && !requires_reclamation &&
+               (after_reclamation || !wrapped_device_report) &&
+               required_device_bytes <= available_device_bytes &&
+               required_budget_bytes <= available_budget_bytes;
+    }
+};
+
+struct WeightResidencyBlock {
+    ggml_backend_buffer_t buffer = nullptr;
+    ggml_backend_t backend       = nullptr;
+    std::vector<ggml_tensor*> tensors;
+    size_t bytes            = 0;
+    uint64_t last_use_epoch = 0;
+    bool requires_demotion  = false;
+};
+
 struct DeviceResidencyManager {
+private:
+    std::shared_ptr<sd::ExecutionEngine> execution_engine_;
+
+public:
     virtual ~DeviceResidencyManager() = default;
 
+    std::shared_ptr<sd::ExecutionEngine> execution_engine();
+
     virtual bool segmented_compute_enabled() const          = 0;
     virtual bool prefetch_enabled() const                   = 0;
-    virtual void set_workspace_reclaimer(uintptr_t owner_id, std::function<bool()> reclaim) = 0;
+    virtual void register_runtime_owner(uintptr_t owner_id) = 0;
     virtual void remove_runtime_owner(uintptr_t owner_id)   = 0;
     // Capacity requests select their backend's weights; protection spans all backends.
     virtual bool fits_compute_backend_capacity(const DeviceMemoryRequest& request,
@@ -62,6 +99,18 @@ struct DeviceResidencyManager {
         const std::vector<ggml_tensor*>& required_params,
         const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
         const std::vector<ggml_tensor*>& protected_params) = 0;
+    virtual DeviceMemoryCapacity inspect_compute_backend_capacity(
+        const DeviceMemoryRequest& request,
+        const std::vector<ggml_tensor*>& params,
+        bool log_details = false) const {
+        DeviceMemoryCapacity result;
+        result.valid = false;
+        return result;
+    }
+    virtual std::vector<WeightResidencyBlock> evictable_compute_backend_blocks(
+        ggml_backend_t backend,
+        const std::vector<ggml_tensor*>& protected_params) const { return {}; }
+    virtual bool evict_compute_backend_block(const WeightResidencyBlock& block) { return false; }
     virtual WeightPrefetchResult prefetch_params(
         uintptr_t owner_id,
         const std::vector<ggml_tensor*>& tensors) = 0;

src/model/diffusion/qwen_image_2_1.hpp

@@ -576,7 +576,7 @@ namespace Qwen {
                 result = run(cache);
             }
             if (result.empty() && last_compute_status() == GGML_STATUS_ALLOC_FAILED &&
-                (cache.mode != QwenImage21PrefixCache::Mode::NONE || !cache_.empty())) {
+                (cache.mode != QwenImage21PrefixCache::Mode::NONE || has_cache_tensors())) {
                 // The failed graph has ended before persistent inputs are released.
                 free_cache_ctx_and_buffer();
                 prefix_cache_disabled = true;

src/model_manager.cpp

@@ -1,11 +1,9 @@
 #include "model_manager.h"
 
 #include <algorithm>
-#include <chrono>
 #include <cstdint>
 #include <iterator>
 #include <mutex>
-#include <thread>
 #include <tuple>
 #include <unordered_set>
 
@@ -150,7 +148,7 @@ bool ModelManager::set_loras(std::vector<LoraSpec> loras, SDVersion version) {
     }
     if (lora_version_ == version && lora_specs_equal(loras_, loras))
         return true;
-    if (!workspace_reclaimers_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) {
+    if (!runtime_owners_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) {
             return state->pin_count != 0;
         })) {
         LOG_ERROR("cannot change LoRA configuration during execution");
@@ -166,7 +164,7 @@ bool ModelManager::set_loras(std::vector<LoraSpec> loras, SDVersion version) {
 std::set<std::string> ModelManager::tensor_names() const {
     std::set<std::string> names;
     for (const auto& state : tensor_states_) {
-        if (state != nullptr && state->component != ModelComponent::LoRA) {
+        if (state != nullptr && state->component != ModelComponent::LoRA && model_sources_.count(state->source_file) == 0) {
             names.insert(state->name);
         }
     }
@@ -234,7 +232,7 @@ bool ModelManager::register_param_tensors(ModelComponent component,
         *registered_tensor_size += estimate_tensors_size(tensors);
     }
 
-    const auto scoped_sources = source_file != 0 ? model_loader_.file_tensors(source_file, source_version) : String2TensorStorage{};
+    const auto scoped_sources = source_file != 0 ? source_loader(source_file).file_tensors(source_file, source_version) : String2TensorStorage{};
     const auto& sources       = source_file != 0 ? scoped_sources : model_loader_.get_tensor_storage_map();
     std::unordered_set<ggml_tensor*> new_tensors;
     std::vector<std::unique_ptr<TensorState>> new_states;
@@ -424,6 +422,20 @@ bool ModelManager::validate_registered_tensors() {
     return ok;
 }
 
+bool ModelManager::validate_param_tensors(const std::vector<ggml_tensor*>& tensors) {
+    std::vector<TensorState*> states;
+    if (!resolve_required_tensor_states(tensors, states)) {
+        return false;
+    }
+    for (auto* state : states) {
+        if (!validate_tensor(*state)) {
+            return false;
+        }
+        state->metadata_validated = true;
+    }
+    return true;
+}
+
 bool ModelManager::load_tensors_to_params_backend(const std::vector<TensorState*>& states) {
     std::vector<TensorState*> need_load;
     need_load.reserve(states.size());
@@ -500,7 +512,7 @@ bool ModelManager::load_tensors_to_params_backend(const std::vector<TensorState*
 }
 
 bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorState*>& states) {
-    std::map<std::pair<ggml_backend_t, ggml_backend_buffer_type_t>, std::vector<TensorState*>> states_by_staging_target;
+    std::map<std::tuple<ggml_backend_t, ggml_backend_buffer_type_t, ModelLoader::FileId>, std::vector<TensorState*>> states_by_staging_target;
     for (TensorState* state : states) {
         if (state == nullptr || should_ignore(*state) || is_optional_missing_tensor(state->name)) {
             continue;
@@ -526,12 +538,12 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorStat
         if (staging_buft == nullptr) {
             staging_buft = ggml_backend_get_default_buffer_type(state->compute_backend);
         }
-        states_by_staging_target[{state->compute_backend, staging_buft}].push_back(state);
+        states_by_staging_target[{state->compute_backend, staging_buft, state->source_file}].push_back(state);
     }
 
     for (const auto& pair : states_by_staging_target) {
-        ggml_backend_t compute_backend                 = pair.first.first;
-        ggml_backend_buffer_type_t staging_buft        = pair.first.second;
+        ggml_backend_t compute_backend                 = std::get<0>(pair.first);
+        ggml_backend_buffer_type_t staging_buft        = std::get<1>(pair.first);
         const std::vector<TensorState*>& target_states = pair.second;
         if (target_states.empty()) {
             continue;
@@ -646,7 +658,7 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
     std::map<ApplyTarget, LoraApplyGroup> groups;
     for (TensorState* state : states) {
         if (state == nullptr || state->tensor == nullptr || state->component == ModelComponent::LoRA ||
-            state->component == ModelComponent::LatentUpsampler || should_ignore(*state) || is_optional_missing_tensor(state->name)) {
+            state->component == ModelComponent::LatentUpsampler || model_sources_.count(state->source_file) != 0 || should_ignore(*state) || is_optional_missing_tensor(state->name)) {
             continue;
         }
         if (state->applied_lora_epoch == current_lora_epoch_) {
@@ -743,6 +755,9 @@ void ModelManager::reset_lora_applied_params() {
 }
 
 bool ModelManager::should_ignore(const TensorState& state) const {
+    if (model_sources_.count(state.source_file) != 0) {
+        return false;
+    }
     for (const auto& ignore_prefix : common_ignore_tensors_) {
         if (starts_with(state.name, ignore_prefix)) {
             return true;
@@ -869,7 +884,7 @@ bool ModelManager::can_mmap_storage(const TensorState& state) const {
 
 bool ModelManager::alloc_params_buffers(const std::vector<TensorState*>& states,
                                         std::vector<ParamsStorageBlock*>& created_storage_blocks) {
-    std::map<std::pair<ggml_backend_buffer_type_t, int>, std::vector<TensorState*>> states_by_buffer_type;
+    std::map<std::tuple<ggml_backend_buffer_type_t, int, ModelLoader::FileId>, std::vector<TensorState*>> states_by_buffer_type;
     for (TensorState* state : states) {
         if (state == nullptr || state->tensor == nullptr) {
             continue;
@@ -878,11 +893,11 @@ bool ModelManager::alloc_params_buffers(const std::vector<TensorState*>& states,
         if (params_buft == nullptr) {
             return false;
         }
-        states_by_buffer_type[{params_buft, static_cast<int>(state->residency_mode)}].push_back(state);
+        states_by_buffer_type[{params_buft, static_cast<int>(state->residency_mode), state->source_file}].push_back(state);
     }
 
     for (const auto& pair : states_by_buffer_type) {
-        ggml_backend_buffer_type_t params_buft  = pair.first.first;
+        ggml_backend_buffer_type_t params_buft  = std::get<0>(pair.first);
         const std::vector<TensorState*>& states = pair.second;
         size_t alignment                        = ggml_backend_buft_get_alignment(params_buft);
         size_t max_size                         = ggml_backend_buft_get_max_size(params_buft);
@@ -1010,7 +1025,7 @@ bool ModelManager::load_tensor_parts(TensorState& state) {
     };
     const bool success = state.source_file == 0
                              ? model_loader_.load_tensors(callback, enable_mmap_, &names, false)
-                             : model_loader_.load_file_tensors(state.source_file, state.source_version, callback, names, enable_mmap_);
+                             : source_loader(state.source_file).load_file_tensors(state.source_file, state.source_version, callback, names, enable_mmap_);
     if (!success || loaded != names) {
         return false;
     }
@@ -1098,7 +1113,7 @@ bool ModelManager::load_tensors(const std::vector<TensorState*>& states) {
             };
             const auto file = group.first.first;
             bool success    = file == 0 ? model_loader_.load_tensors(callback, enable_mmap_, &names)
-                                        : model_loader_.load_file_tensors(file, group.first.second, callback, names, enable_mmap_);
+                                        : source_loader(file).load_file_tensors(file, group.first.second, callback, names, enable_mmap_);
             if (!success || loaded != names)
                 return false;
             for (auto& entry : batch) {
@@ -1316,7 +1331,7 @@ void ModelManager::erase_params_storage_block(ParamsStorageBlock* block) {
 void ModelManager::release_all() {
     clear_all_prefetched_params();
     runtime_residencies_.clear();
-    workspace_reclaimers_.clear();
+    runtime_owners_.clear();
     for (auto& state : tensor_states_) {
         state->pin_count          = 0;
         state->applied_lora_epoch = UINT64_MAX;
@@ -1643,11 +1658,17 @@ void ModelManager::evict_compute_backend_params(const std::vector<ggml_tensor*>&
     if (tensors.empty()) {
         return;
     }
-    std::vector<TensorState*> required_states;
-    if (!resolve_required_tensor_states(tensors, required_states)) {
+    std::unordered_set<TensorState*> target_states;
+    // Session cleanup may follow unregistration, after tensor metadata has been freed.
+    for (auto* tensor : tensors) {
+        const auto found = tensor_states_by_tensor_.find(tensor);
+        if (found != tensor_states_by_tensor_.end()) {
+            target_states.insert(found->second);
+        }
+    }
+    if (target_states.empty()) {
         return;
     }
-    std::unordered_set<TensorState*> target_states(required_states.begin(), required_states.end());
 
     for (const auto& block : compute_staging_blocks_) {
         const bool intersects = std::any_of(
@@ -1743,12 +1764,12 @@ WeightResidencyInfo ModelManager::inspect_compute_backend_params(
     return info;
 }
 
-void ModelManager::set_workspace_reclaimer(uintptr_t owner_id, std::function<bool()> reclaim) {
-    workspace_reclaimers_[owner_id] = std::move(reclaim);
+void ModelManager::register_runtime_owner(uintptr_t owner_id) {
+    runtime_owners_.insert(owner_id);
 }
 
 void ModelManager::remove_runtime_owner(uintptr_t owner_id) {
-    workspace_reclaimers_.erase(owner_id);
+    runtime_owners_.erase(owner_id);
     for (auto it = runtime_residencies_.begin(); it != runtime_residencies_.end();) {
         if (it->first.first == owner_id) {
             it = runtime_residencies_.erase(it);
@@ -1807,7 +1828,7 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
             LOG_VERBOSE("ignoring impossible free memory report on %s (%.2f MB free / %.2f MB total), using tracked residency",
                         ggml_backend_name(backend), free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0));
             free_bytes                   = total_bytes;
-            result.requires_reclamation = true;
+            result.wrapped_device_report = true;
         }
         if (total_bytes > 0) {
             free_bytes = std::min(free_bytes, resident < total_bytes ? total_bytes - resident : 0);
@@ -1818,6 +1839,7 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
     if (request.max_backend_bytes > 0) {
         const size_t resident         = add(compute_backend_resident_bytes(request.compute_backend),
                                             other_runtime_resident_bytes(request.owner_id, request.compute_backend));
+        result.requires_reclamation   = result.requires_reclamation || resident > request.max_backend_bytes;
         result.available_budget_bytes = resident < request.max_backend_bytes
                                             ? request.max_backend_bytes - resident
                                             : 0;
@@ -1842,6 +1864,7 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
         if (entry.second > 0) {
             const size_t resident         = add(compute_backend_resident_bytes(entry.first),
                                                 other_runtime_resident_bytes(request.owner_id, entry.first));
+            result.requires_reclamation   = result.requires_reclamation || resident > entry.second;
             const size_t available        = resident < entry.second ? entry.second - resident : 0;
             result.available_budget_bytes = std::min(result.available_budget_bytes, available);
         }
@@ -1849,6 +1872,139 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
     return result;
 }
 
+DeviceMemoryCapacity ModelManager::inspect_compute_backend_capacity(
+    const DeviceMemoryRequest& request,
+    const std::vector<ggml_tensor*>& params,
+    bool log_details) const {
+    std::vector<TensorState*> states;
+    if (request.compute_backend == nullptr || !resolve_required_tensor_states(params, states, request.compute_backend)) {
+        DeviceMemoryCapacity result;
+        result.valid = false;
+        return result;
+    }
+    return check_capacity(request, states, log_details);
+}
+
+std::vector<WeightResidencyBlock> ModelManager::evictable_compute_backend_blocks(
+    ggml_backend_t backend,
+    const std::vector<ggml_tensor*>& protected_params) const {
+    std::vector<WeightResidencyBlock> result;
+    if (backend == nullptr) {
+        return result;
+    }
+    std::vector<TensorState*> protected_list;
+    if (!resolve_required_tensor_states(protected_params, protected_list)) {
+        return result;
+    }
+    std::unordered_set<TensorState*> protected_states(protected_list.begin(), protected_list.end());
+    for (const auto& entry : prefetch_blocks_) {
+        protected_states.insert(entry.second->states.begin(), entry.second->states.end());
+    }
+    const auto device = ggml_backend_get_device(backend);
+    const bool cpu    = sd_backend_is_cpu(backend);
+    auto on_device    = [&](ggml_backend_buffer_t buffer) {
+        if (buffer == nullptr) {
+            return false;
+        }
+        if (cpu) {
+            return ggml_backend_buffer_is_host(buffer);
+        }
+        if (ggml_backend_buffer_is_host(buffer)) {
+            return false;
+        }
+        const auto type  = ggml_backend_buffer_get_type(buffer);
+        const auto split = split_buffer_devices_.find(type);
+        if (split != split_buffer_devices_.end()) {
+            return std::any_of(split->second.begin(), split->second.end(), [&](const auto& entry) {
+                return ggml_backend_get_device(entry.first) == device;
+            });
+        }
+        return ggml_backend_buft_get_device(type) == device;
+    };
+    auto append = [&](ggml_backend_buffer_t buffer, const std::vector<TensorState*>& states, bool storage) {
+        if (!on_device(buffer) || states.empty()) {
+            return;
+        }
+        WeightResidencyBlock candidate;
+        candidate.buffer  = buffer;
+        candidate.backend = backend;
+        candidate.bytes   = ggml_backend_buffer_get_size(buffer);
+        for (auto* state : states) {
+            if (state == nullptr) {
+                continue;
+            }
+            if (state->pin_count != 0 || protected_states.count(state) != 0) {
+                return;
+            }
+            if (storage) {
+                if (state->staged_to_compute_backend) {
+                    return;
+                }
+                if (state->residency_mode != ResidencyMode::Disk) {
+                    if (cpu || state->params_backend == nullptr || sd_backend_is_cpu(state->params_backend) ||
+                        state->compute_backend != state->params_backend ||
+                        state->sources.empty() || state->applied_lora_epoch != UINT64_MAX ||
+                        state->split_buffer_type != nullptr || state->tensor->view_src != nullptr) {
+                        return;
+                    }
+                    candidate.requires_demotion = true;
+                }
+            }
+            candidate.tensors.push_back(state->tensor);
+            candidate.last_use_epoch = std::max(candidate.last_use_epoch, state->last_use_epoch);
+        }
+        result.push_back(std::move(candidate));
+    };
+    for (const auto& block : compute_staging_blocks_) {
+        std::vector<TensorState*> states;
+        states.reserve(block->staged_tensors.size());
+        for (const auto& entry : block->staged_tensors) {
+            states.push_back(entry.first);
+        }
+        append(block->buffer, states, false);
+    }
+    for (const auto& block : params_storage_blocks_) {
+        // Loader-retained mappings cannot be reclaimed by freeing a storage block.
+        if (block->mmap_tensor_stores.empty()) {
+            append(block->buffer, block->states, true);
+        }
+    }
+    return result;
+}
+
+bool ModelManager::evict_compute_backend_block(const WeightResidencyBlock& block) {
+    const auto candidates = evictable_compute_backend_blocks(block.backend, {});
+    const auto found      = std::find_if(candidates.begin(), candidates.end(), [&](const auto& candidate) {
+        return candidate.buffer == block.buffer && candidate.tensors == block.tensors &&
+               candidate.requires_demotion == block.requires_demotion;
+    });
+    if (found == candidates.end()) {
+        return false;
+    }
+    for (auto it = compute_staging_blocks_.begin(); it != compute_staging_blocks_.end(); ++it) {
+        if ((*it)->buffer == block.buffer) {
+            free_compute_staging_block(**it);
+            compute_staging_blocks_.erase(it);
+            return true;
+        }
+    }
+    for (auto it = params_storage_blocks_.begin(); it != params_storage_blocks_.end(); ++it) {
+        if ((*it)->buffer == block.buffer) {
+            if (block.requires_demotion) {
+                for (auto* state : (*it)->states) {
+                    if (state != nullptr) {
+                        state->residency_mode = ResidencyMode::Disk;
+                    }
+                }
+            }
+            free_params_storage_block(**it);
+            params_storage_blocks_.erase(it);
+            return true;
+        }
+    }
+    return false;
+}
+
 bool ModelManager::fits_compute_backend_capacity(
     const DeviceMemoryRequest& request,
     const std::vector<ggml_tensor*>& required_params) const {
@@ -1865,184 +2021,5 @@ bool ModelManager::ensure_compute_backend_capacity(
     const std::vector<ggml_tensor*>& required_params,
     const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
     const std::vector<ggml_tensor*>& protected_params) {
-    std::vector<TensorState*> required_states;
-    if (!resolve_required_tensor_states(required_params, required_states, request.compute_backend)) {
-        return false;
-    }
-
-    ggml_backend_t compute_backend = request.compute_backend;
-    if (compute_backend == nullptr) {
-        LOG_ERROR("model manager cannot reclaim memory for a null compute backend");
-        return false;
-    }
-    if (sd_backend_is_cpu(compute_backend)) {
-        return true;
-    }
-
-    auto fits = [&]() {
-        const auto capacity = check_capacity(request, required_states);
-        return !capacity.requires_reclamation && capacity.fits();
-    };
-    if (fits()) {
-        return true;
-    }
-    for (const auto& entry : workspace_reclaimers_) {
-        if (entry.first != request.owner_id) {
-            entry.second();
-            if (fits()) {
-                return true;
-            }
-        }
-    }
-
-    std::unordered_set<TensorState*> protected_states;
-    std::vector<TensorState*> resolved_protected;
-    if (!resolve_required_tensor_states(protected_params, resolved_protected)) {
-        return false;
-    }
-    protected_states.insert(resolved_protected.begin(), resolved_protected.end());
-    for (const auto& entry : prefetch_blocks_) {
-        if (entry.second != nullptr) {
-            protected_states.insert(entry.second->states.begin(), entry.second->states.end());
-        }
-    }
-
-    std::unordered_set<TensorState*> eviction_states;
-    auto add_evictable_state = [&](TensorState* state) {
-        if (state == nullptr || state->compute_backend != compute_backend ||
-            state->pin_count > 0 || protected_states.find(state) != protected_states.end()) {
-            return;
-        }
-        const bool reloadable = state->residency_mode == ResidencyMode::Disk ||
-                                state->compute_backend != state->params_backend;
-        const bool resident   = state->compute_backend == state->params_backend
-                                    ? state->loaded_to_params_backend
-                                    : state->staged_to_compute_backend;
-        if (reloadable && resident) {
-            eviction_states.insert(state);
-        }
-    };
-    size_t last_eviction_count   = 0;
-    auto release_eviction_states = [&]() {
-        if (eviction_states.size() != last_eviction_count) {
-            last_eviction_count = eviction_states.size();
-            release_compute_staging_blocks(false, &eviction_states);
-            release_params_storage_blocks(false, &eviction_states);
-        }
-        // Device memory reports can lag behind releases even when no new candidates are added.
-        return fits();
-    };
-
-    for (const auto& candidate_params : preferred_eviction_order) {
-        std::vector<TensorState*> candidate_states;
-        if (!resolve_required_tensor_states(candidate_params, candidate_states)) {
-            return false;
-        }
-        for (TensorState* state : candidate_states) {
-            add_evictable_state(state);
-        }
-        if (release_eviction_states()) {
-            return true;
-        }
-    }
-
-    std::vector<TensorState*> global_candidates;
-    global_candidates.reserve(tensor_states_.size());
-    for (const auto& state : tensor_states_) {
-        if (state != nullptr && eviction_states.find(state.get()) == eviction_states.end()) {
-            global_candidates.push_back(state.get());
-        }
-    }
-    std::stable_sort(global_candidates.begin(),
-                     global_candidates.end(),
-                     [](const TensorState* lhs, const TensorState* rhs) {
-                         return lhs->last_use_epoch < rhs->last_use_epoch;
-                     });
-    for (TensorState* state : global_candidates) {
-        add_evictable_state(state);
-        if (release_eviction_states()) {
-            return true;
-        }
-    }
-
-    // The loader retains mmap buffers after storage blocks are released, so
-    // demoting mapped states would only force duplicate allocations on reload.
-    std::unordered_set<TensorState*> mapped_states;
-    for (const auto& block : params_storage_blocks_) {
-        if (!block->mmap_tensor_stores.empty()) {
-            mapped_states.insert(block->states.begin(), block->states.end());
-        }
-    }
-    std::unordered_set<TensorState*> required_set(required_states.begin(), required_states.end());
-    std::vector<TensorState*> demoted_states;
-    auto revert_demotions = [&]() {
-        for (TensorState* state : demoted_states) {
-            if (state->loaded_to_params_backend) {
-                state->residency_mode = ResidencyMode::ParamBackend;
-            }
-        }
-    };
-    auto can_demote_to_disk = [&](TensorState* state) {
-        return state != nullptr && state->tensor != nullptr && state->tensor->view_src == nullptr &&
-               state->residency_mode == ResidencyMode::ParamBackend &&
-               state->params_backend != nullptr && state->params_backend == state->compute_backend &&
-               state->compute_backend == compute_backend && !sd_backend_is_cpu(state->params_backend) &&
-               state->pin_count == 0 && !state->staged_to_compute_backend && state->loaded_to_params_backend &&
-               protected_states.find(state) == protected_states.end() &&
-               required_set.find(state) == required_set.end() &&
-               mapped_states.find(state) == mapped_states.end() &&
-               !state->sources.empty() && state->applied_lora_epoch == UINT64_MAX &&
-               state->split_buffer_type == nullptr;
-    };
-    for (TensorState* state : global_candidates) {
-        if (!can_demote_to_disk(state)) {
-            continue;
-        }
-        state->residency_mode = ResidencyMode::Disk;
-        demoted_states.push_back(state);
-        add_evictable_state(state);
-        if (release_eviction_states()) {
-            revert_demotions();
-            return true;
-        }
-    }
-    revert_demotions();
-
-    auto capacity = check_capacity(request, required_states);
-    if (!capacity.requires_reclamation && capacity.fits()) {
-        return true;
-    }
-    if (sd_backend_is(compute_backend, "Vulkan") &&
-        capacity.required_budget_bytes <= capacity.available_budget_bytes) {
-        ggml_backend_synchronize(compute_backend);
-        sd_backend_clear_buffer_cache(compute_backend);
-        // Vulkan heap usage can lag behind buffer releases. Fast memory queries
-        // must not turn that reporting delay into an allocation failure.
-        const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(100);
-        for (;;) {
-            capacity = check_capacity(request, required_states);
-            if (!capacity.requires_reclamation && capacity.fits()) {
-                return true;
-            }
-            if (capacity.required_budget_bytes > capacity.available_budget_bytes ||
-                std::chrono::steady_clock::now() >= deadline) {
-                break;
-            }
-            std::this_thread::sleep_for(std::chrono::milliseconds(10));
-        }
-    }
-    capacity                           = check_capacity(request, required_states, true);
-    const std::string available_device = capacity.available_device_bytes == SIZE_MAX
-                                             ? "unknown"
-                                             : sd_format("%.2f MB", capacity.available_device_bytes / (1024.0 * 1024.0));
-    const std::string available_budget = capacity.available_budget_bytes == SIZE_MAX
-                                             ? "unlimited"
-                                             : sd_format("%.2f MB", capacity.available_budget_bytes / (1024.0 * 1024.0));
-    LOG_WARN("model manager cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %s device / %s budget",
-             ggml_backend_name(compute_backend),
-             capacity.required_device_bytes / (1024.0 * 1024.0),
-             capacity.required_budget_bytes / (1024.0 * 1024.0),
-             available_device.c_str(),
-             available_budget.c_str());
-    return false;
+    return execution_engine()->ensure_capacity(*this, request, required_params, preferred_eviction_order, protected_params);
 }

src/model_manager.h

@@ -91,6 +91,8 @@ private:
     };
 
     ModelLoader model_loader_;
+    // Auxiliary models retain their own catalogs and conversion settings.
+    std::map<ModelLoader::FileId, ModelLoader> model_sources_;
     std::vector<std::unique_ptr<TensorState>> tensor_states_;
     std::map<const ggml_tensor*, TensorState*> tensor_states_by_tensor_;
     mutable std::list<ResolvedTensorStates> resolved_tensor_states_;
@@ -101,7 +103,7 @@ private:
     std::map<uintptr_t, std::unique_ptr<PrefetchBlock>> prefetch_blocks_;
     std::map<ggml_backend_t, ggml_backend_t> prefetch_backends_;
     std::map<std::pair<uintptr_t, ggml_backend_t>, RuntimeResidency> runtime_residencies_;
-    std::map<uintptr_t, std::function<bool()>> workspace_reclaimers_;
+    std::unordered_set<uintptr_t> runtime_owners_;
     bool warned_split_lora_skip_ = false;
     std::set<std::string> common_ignore_tensors_;
     std::vector<LoraSpec> loras_;
@@ -131,6 +133,7 @@ private:
                                         std::vector<TensorState*>& required_states,
                                         ggml_backend_t compute_backend = nullptr) const;
     bool should_ignore(const TensorState& state) const;
+    const ModelLoader& source_loader(ModelLoader::FileId file) const;
     bool is_optional_missing_tensor(const std::string& name) const;
     bool validate_tensor(const TensorState& state) const;
 
@@ -148,17 +151,7 @@ private:
     size_t compute_backend_alloc_size(const std::vector<TensorState*>& states,
                                       bool missing_only) const;
     size_t compute_backend_resident_bytes(ggml_backend_t compute_backend) const;
-    struct CapacityCheck {
-        size_t required_device_bytes  = 0;
-        size_t required_budget_bytes  = 0;
-        size_t available_device_bytes = SIZE_MAX;
-        size_t available_budget_bytes = SIZE_MAX;
-        bool requires_reclamation     = false;
-        bool fits() const {
-            return required_device_bytes <= available_device_bytes &&
-                   required_budget_bytes <= available_budget_bytes;
-        }
-    };
+    using CapacityCheck = DeviceMemoryCapacity;
     CapacityCheck check_capacity(const DeviceMemoryRequest& request,
                                  const std::vector<TensorState*>& states,
                                  bool log_details = false) const;
@@ -186,6 +179,8 @@ public:
     bool add_file(const std::string& path, const std::string& prefix = "", ModelLoader::FileId* id = nullptr, bool force = false);
     bool del_file(ModelLoader::FileId id);
     bool refresh_files();
+    ModelLoader::FileId register_model_source(ModelLoader loader);
+    bool unregister_model_source(ModelLoader::FileId file);
     ModelLoader::FileVersions source_versions(const std::set<ModelComponent>& components, const ModelLoader& loader) const;
     size_t registered_params_size(const std::set<ModelComponent>& components) const;
 
@@ -267,12 +262,13 @@ public:
     }
 
     bool validate_registered_tensors();
+    bool validate_param_tensors(const std::vector<ggml_tensor*>& 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 set_workspace_reclaimer(uintptr_t owner_id, std::function<bool()> reclaim) override;
+    void register_runtime_owner(uintptr_t owner_id) override;
     void remove_runtime_owner(uintptr_t owner_id) override;
     bool fits_compute_backend_capacity(const DeviceMemoryRequest& request,
                                        const std::vector<ggml_tensor*>& required_params) const override;
@@ -290,6 +286,14 @@ public:
         const std::vector<ggml_tensor*>& required_params,
         const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
         const std::vector<ggml_tensor*>& protected_params) override;
+    DeviceMemoryCapacity inspect_compute_backend_capacity(
+        const DeviceMemoryRequest& request,
+        const std::vector<ggml_tensor*>& params,
+        bool log_details = false) const override;
+    std::vector<WeightResidencyBlock> evictable_compute_backend_blocks(
+        ggml_backend_t backend,
+        const std::vector<ggml_tensor*>& protected_params) const override;
+    bool evict_compute_backend_block(const WeightResidencyBlock& block) override;
     WeightPrefetchResult prefetch_params(
         uintptr_t owner_id,
         const std::vector<ggml_tensor*>& tensors) override;

src/model_manager_files.cpp

@@ -74,7 +74,7 @@ void ModelManager::invalidate_sources(const std::unordered_set<TensorState*>& st
 }
 
 bool ModelManager::set_loader(ModelLoader loader) {
-    if (!workspace_reclaimers_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) {
+    if (!runtime_owners_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) {
             return state->pin_count != 0;
         })) {
         LOG_ERROR("cannot update model sources during execution");
@@ -87,7 +87,10 @@ bool ModelManager::set_loader(ModelLoader loader) {
         auto key   = std::make_pair(state.source_file, state.source_version);
         auto found = scoped.find(key);
         if (found == scoped.end())
-            found = scoped.emplace(key, loader.file_tensors(key.first, key.second)).first;
+            found = scoped.emplace(key, model_sources_.count(key.first) != 0
+                                            ? source_loader(key.first).file_tensors(key.first, key.second)
+                                            : loader.file_tensors(key.first, key.second))
+                        .first;
         return found->second;
     };
     bool lora_changed = false;
@@ -147,7 +150,9 @@ ModelLoader::FileVersions ModelManager::source_versions(const std::set<ModelComp
             continue;
         }
         if (state->source_file != 0) {
-            versions[state->source_file] = loader.file_revision(state->source_file);
+            versions[state->source_file] = model_sources_.count(state->source_file) != 0
+                                               ? source_loader(state->source_file).file_revision(state->source_file)
+                                               : loader.file_revision(state->source_file);
             continue;
         }
         for (const auto& source : find_tensor_sources(*state, sources)) {
@@ -167,3 +172,34 @@ size_t ModelManager::registered_params_size(const std::set<ModelComponent>& comp
     }
     return bytes;
 }
+
+const ModelLoader& ModelManager::source_loader(ModelLoader::FileId file) const {
+    auto found = model_sources_.find(file);
+    return found == model_sources_.end() ? model_loader_ : found->second;
+}
+
+ModelLoader::FileId ModelManager::register_model_source(ModelLoader loader) {
+    const auto versions = loader.file_versions();
+    if (versions.size() != 1) {
+        LOG_ERROR("a model source must contain exactly one root model file");
+        return 0;
+    }
+    const auto file = versions.begin()->first;
+    if (model_loader_.file_revision(file) != 0 || model_sources_.count(file) != 0) {
+        LOG_ERROR("model source is already registered");
+        return 0;
+    }
+    loader.set_n_threads(n_threads_);
+    model_sources_.emplace(file, std::move(loader));
+    return file;
+}
+
+bool ModelManager::unregister_model_source(ModelLoader::FileId file) {
+    if (std::any_of(tensor_states_.begin(), tensor_states_.end(), [file](const auto& state) {
+            return state->source_file == file;
+        })) {
+        LOG_ERROR("cannot unregister a model source while its weights are registered");
+        return false;
+    }
+    return model_sources_.erase(file) != 0;
+}

src/pipeline/image.cpp

@@ -9,10 +9,10 @@
 #include "core/rng.hpp"
 #include "diffusion_engine.h"
 #include "model/vae/vae.hpp"
+#include "pipeline/image_upscale.h"
 #include "request.h"
 #include "runtime/denoiser.hpp"
 #include "runtime/image_preprocess.h"
-#include "upscaler.h"
 
 namespace sd::pipeline {
 
@@ -751,7 +751,7 @@ namespace sd::pipeline {
     static sd::Tensor<float> upscale_hires_latent(StableDiffusionGGML* sd,
                                                   const sd::Tensor<float>& latent,
                                                   const GenerationRequest& request,
-                                                  UpscalerGGML* upscaler) {
+                                                  ImageUpscaleModel* upscaler) {
         if (sd->get_cancel_flag() == SD_CANCEL_ALL) {
             LOG_ERROR("cancelling hires latent upscale");
             return {};
@@ -837,7 +837,7 @@ namespace sd::pipeline {
             sd->generation_progress.report("upscaling", "hires");
             sd::Tensor<float> upscaled_tensor;
             if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) {
-                upscaled_tensor = upscaler->upscale_tensor(decoded);
+                upscaled_tensor = upscale_image_tensor(*upscaler, decoded, sd->n_threads, request.hires.upscale_tile_size);
                 if (upscaled_tensor.empty()) {
                     LOG_ERROR("hires model upscale failed");
                     return {};
@@ -1034,22 +1034,18 @@ namespace sd::pipeline {
             LOG_INFO("hires fix: upscaling to %dx%d", request.hires.target_width, request.hires.target_height);
 
             sd->generation_progress.report("upscaling", "hires");
-            std::unique_ptr<UpscalerGGML> hires_upscaler;
+            std::unique_ptr<ImageUpscaleModel> hires_upscaler;
             if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) {
                 if (sd->get_cancel_flag() == SD_CANCEL_ALL) {
                     LOG_ERROR("cancelling generation before hires model load");
                     return false;
                 }
                 LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path);
-                hires_upscaler                    = std::make_unique<UpscalerGGML>(sd->n_threads,
-                                                                                   false,
-                                                                                   request.hires.upscale_tile_size,
-                                                                                   sd->backend_spec,
-                                                                                   sd->params_backend_spec);
-                const size_t max_graph_vram_bytes = sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER);
-                hires_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes);
-                if (!hires_upscaler->load_from_file(request.hires.model_path,
-                                                    sd->n_threads)) {
+                hires_upscaler = load_image_upscale_model(sd->model_manager,
+                                                          sd->backend_manager,
+                                                          request.hires.model_path,
+                                                          sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER));
+                if (hires_upscaler == nullptr) {
                     LOG_ERROR("load hires model upscaler failed");
                     return false;
                 }

src/pipeline/image_upscale.cpp

@@ -0,0 +1,138 @@
+#include "pipeline/image_upscale.h"
+
+#include "core/ggml_extend_backend.h"
+#include "core/util.h"
+#include "model/upscaler/esrgan.hpp"
+#include "model/upscaler/mosr.hpp"
+#include "model_manager.h"
+#include "runtime/tiling.h"
+
+namespace sd::pipeline {
+    ImageUpscaleModel::ImageUpscaleModel(const std::shared_ptr<ModelManager>& manager)
+        : manager_(manager) {}
+
+    ImageUpscaleModel::~ImageUpscaleModel() {
+        if (auto* value = runner()) {
+            value->runner_end();
+        }
+        if (auto manager = manager_.lock()) {
+            GGML_ASSERT(manager->unregister_param_tensors(weights_));
+            if (source_file_ != 0) {
+                GGML_ASSERT(manager->unregister_model_source(source_file_));
+            }
+        }
+    }
+
+    GGMLRunner* ImageUpscaleModel::runner() const {
+        return mosr_ != nullptr ? static_cast<GGMLRunner*>(mosr_.get()) : esrgan_.get();
+    }
+
+    int ImageUpscaleModel::scale() const {
+        if (mosr_ != nullptr) {
+            return mosr_->config.scale;
+        }
+        return esrgan_ != nullptr ? esrgan_->config.scale : 1;
+    }
+
+    std::unique_ptr<ImageUpscaleModel> load_image_upscale_model(
+        const std::shared_ptr<ModelManager>& manager,
+        SDBackendManager& backends,
+        const std::string& path,
+        size_t max_vram_bytes,
+        bool conv_direct,
+        ggml_type weight_type) {
+        if (manager == nullptr) {
+            return nullptr;
+        }
+        const auto backend        = backends.runtime_backend(SDBackendModule::UPSCALER);
+        const auto params_backend = backends.params_backend(SDBackendModule::UPSCALER);
+        if (backend == nullptr || params_backend == nullptr) {
+            LOG_ERROR("failed to initialize upscaler backends");
+            return nullptr;
+        }
+        ModelLoader loader;
+        if (!loader.init_from_file_and_convert_name(path, "", VERSION_ESRGAN)) {
+            LOG_ERROR("init upscaler model loader failed: '%s'", path.c_str());
+            return nullptr;
+        }
+        loader.set_wtype_override(weight_type);
+        const auto tensors  = loader.get_tensor_storage_map();
+        auto model          = std::unique_ptr<ImageUpscaleModel>(new ImageUpscaleModel(manager));
+        model->source_file_ = manager->register_model_source(std::move(loader));
+        if (model->source_file_ == 0) {
+            return nullptr;
+        }
+        std::map<std::string, ggml_tensor*> weights;
+        if (tensors.count("gblocks.0.weight") != 0) {
+            model->mosr_ = std::make_unique<MoSR>(backend, tensors, manager);
+            if (model->mosr_->mosr_net == nullptr) {
+                LOG_ERROR("unsupported MoSR weights: '%s'", path.c_str());
+                return nullptr;
+            }
+            model->mosr_->get_param_tensors(weights);
+        } else {
+            model->esrgan_ = std::make_unique<ESRGAN>(backend, tensors, manager);
+            if (model->esrgan_->rrdb_net == nullptr) {
+                LOG_ERROR("unsupported ESRGAN weights: '%s'", path.c_str());
+                return nullptr;
+            }
+            model->esrgan_->get_param_tensors(weights);
+        }
+        for (const auto& entry : weights) {
+            model->weights_.push_back(entry.second);
+        }
+        const auto mode = backends.params_backend_is_disk(SDBackendModule::UPSCALER)
+                              ? ModelManager::ResidencyMode::Disk
+                              : ModelManager::ResidencyMode::ParamBackend;
+        if (!manager->register_param_tensors(ModelComponent::Upscaler, std::move(weights), mode,
+                                             backend, params_backend, nullptr, false, false, nullptr,
+                                             model->source_file_, VERSION_ESRGAN) ||
+            !manager->validate_param_tensors(model->weights_)) {
+            LOG_ERROR("register upscaler weights failed: '%s'", path.c_str());
+            return nullptr;
+        }
+        model->runner()->set_max_graph_vram_bytes(max_vram_bytes);
+        model->runner()->set_conv2d_direct_enabled(conv_direct);
+        LOG_INFO("Upscaler architecture: %s, weight type: %s", model->runner()->get_desc().c_str(), ggml_type_name(weight_type));
+        return model;
+    }
+
+    Tensor<float> upscale_image_tensor(ImageUpscaleModel& model, const Tensor<float>& input, int n_threads, int tile_size) {
+        if (input.empty() || input.dim() != 4 || (input.shape()[2] != 3 && input.shape()[2] != 4)) {
+            LOG_ERROR("upscaler expects a 4D RGB or RGBA image tensor");
+            return {};
+        }
+        struct RunnerEnd {
+            GGMLRunner* runner;
+            ~RunnerEnd() { runner->runner_end(); }
+        } cleanup{model.runner()};
+        const bool has_alpha = input.shape()[2] == 4;
+        auto rgb             = has_alpha ? sd::ops::slice(input, 2, 0, 3) : input;
+        const auto compute   = [&](const Tensor<float>& tensor) {
+            return model.mosr_ != nullptr ? model.mosr_->compute(n_threads, tensor) : model.esrgan_->compute(n_threads, tensor);
+        };
+        const int scale = model.scale();
+        Tensor<float> output;
+        if (tile_size <= 0 || (input.shape()[0] <= tile_size && input.shape()[1] <= tile_size)) {
+            output = compute(rgb);
+        } else {
+            output = process_tiles_2d(rgb,
+                                      static_cast<int>(input.shape()[0] * scale),
+                                      static_cast<int>(input.shape()[1] * scale),
+                                      scale, tile_size, tile_size, 0.25f, false, false, compute);
+        }
+        if (output.empty()) {
+            LOG_ERROR("upscaler compute failed");
+            return {};
+        }
+        if (has_alpha) {
+            auto alpha = sd::ops::slice(input, 2, 3, 4);
+            auto shape = alpha.shape();
+            shape[0]   = output.shape()[0];
+            shape[1]   = output.shape()[1];
+            alpha      = sd::ops::interpolate(alpha, shape, sd::ops::InterpolateMode::Bilinear);
+            output     = sd::ops::concat(output, alpha, 2);
+        }
+        return output;
+    }
+}

src/pipeline/image_upscale.h

@@ -0,0 +1,57 @@
+#ifndef __SD_PIPELINE_IMAGE_UPSCALE_H__
+#define __SD_PIPELINE_IMAGE_UPSCALE_H__
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "core/tensor.hpp"
+#include "ggml.h"
+
+class ModelManager;
+class SDBackendManager;
+struct GGMLRunner;
+struct ESRGAN;
+struct MoSR;
+
+namespace sd::pipeline {
+    class ImageUpscaleModel {
+        friend std::unique_ptr<ImageUpscaleModel> load_image_upscale_model(
+            const std::shared_ptr<ModelManager>&,
+            SDBackendManager&,
+            const std::string&,
+            size_t,
+            bool,
+            ggml_type);
+        friend Tensor<float> upscale_image_tensor(ImageUpscaleModel&, const Tensor<float>&, int, int);
+
+        std::weak_ptr<ModelManager> manager_;
+        std::unique_ptr<ESRGAN> esrgan_;
+        std::unique_ptr<MoSR> mosr_;
+        std::vector<ggml_tensor*> weights_;
+        uint64_t source_file_ = 0;
+
+        explicit ImageUpscaleModel(const std::shared_ptr<ModelManager>& manager);
+
+    public:
+        ~ImageUpscaleModel();
+        GGMLRunner* runner() const;
+        int scale() const;
+    };
+
+    std::unique_ptr<ImageUpscaleModel> load_image_upscale_model(
+        const std::shared_ptr<ModelManager>& manager,
+        SDBackendManager& backends,
+        const std::string& path,
+        size_t max_vram_bytes = 0,
+        bool conv_direct      = false,
+        ggml_type weight_type = GGML_TYPE_F16);
+
+    Tensor<float> upscale_image_tensor(ImageUpscaleModel& model,
+                                       const Tensor<float>& input,
+                                       int n_threads,
+                                       int tile_size);
+}
+
+#endif  // __SD_PIPELINE_IMAGE_UPSCALE_H__

src/upscaler.cpp

@@ -2,8 +2,9 @@
 #include "core/ggml_extend_backend.h"
 #include "core/util.h"
 #include "generation_info.h"
+#include "model/upscaler/esrgan.hpp"
+#include "model/upscaler/mosr.hpp"
 #include "model_loader.h"
-#include "runtime/tiling.h"
 #include "stable-diffusion.h"
 
 #include <cstdlib>
@@ -19,188 +20,44 @@ UpscalerGGML::UpscalerGGML(int n_threads,
       direct(direct),
       tile_size(tile_size),
       backend_spec(std::move(backend_spec)),
-      params_backend_spec(std::move(params_backend_spec)) {
-}
+      params_backend_spec(std::move(params_backend_spec)) {}
 
-UpscalerGGML::~UpscalerGGML() {
-    // ModelManager holds raw ggml tensor pointers owned by the runner context.
-    model_manager.reset();
-    esrgan_upscaler.reset();
-    mosr_upscaler.reset();
-}
+UpscalerGGML::~UpscalerGGML() = default;
 
 void UpscalerGGML::set_max_graph_vram_bytes(size_t max_vram_bytes) {
     max_graph_vram_bytes = max_vram_bytes;
-    if (runner() != nullptr) {
-        runner()->set_max_graph_vram_bytes(max_vram_bytes);
+    if (model != nullptr) {
+        model->runner()->set_max_graph_vram_bytes(max_vram_bytes);
     }
 }
 
-GGMLRunner* UpscalerGGML::runner() {
-    if (mosr_upscaler) {
-        return mosr_upscaler.get();
-    }
-    return esrgan_upscaler.get();
-}
-
 int UpscalerGGML::get_scale() {
-    if (mosr_upscaler) {
-        return mosr_upscaler->config.scale;
-    }
-    if (esrgan_upscaler) {
-        return esrgan_upscaler->config.scale;
-    }
-    return 1;
+    return model != nullptr ? model->scale() : 1;
 }
 
-sd::Tensor<float> UpscalerGGML::compute(const sd::Tensor<float>& input_tensor) {
-    if (mosr_upscaler) {
-        return mosr_upscaler->compute(n_threads, input_tensor);
-    }
-    return esrgan_upscaler->compute(n_threads, input_tensor);
-}
-
-bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
-                                  int n_threads) {
+bool UpscalerGGML::load_from_file(const std::string& path, int n_threads) {
     sd::ParallelScope tensor_scope(&tensor_executor);
     ggml_log_set(sd_ggml_log_callback, nullptr);
-
+    model.reset();
     std::string error;
-    if (!backend_manager.init(backend_spec.c_str(),
-                              params_backend_spec.c_str(),
-                              /*split_mode_spec=*/nullptr,
-                              &error)) {
+    if (!backend_manager.init(backend_spec.c_str(), params_backend_spec.c_str(), nullptr, &error)) {
         LOG_ERROR("upscaler backend config failed: %s", error.c_str());
         return false;
     }
-    auto backend_for = [&](SDBackendModule module) {
-        ggml_backend_t module_backend = backend_manager.runtime_backend(module);
-        if (module_backend == nullptr) {
-            LOG_ERROR("failed to initialize %s backend", sd_backend_module_name(module));
-        }
-        return module_backend;
-    };
-    auto params_backend_for = [&](SDBackendModule module) {
-        ggml_backend_t module_backend = backend_manager.params_backend(module);
-        if (module_backend == nullptr) {
-            LOG_ERROR("failed to initialize %s params backend", sd_backend_module_name(module));
-        }
-        return module_backend;
-    };
-    auto ensure_backend_pair = [&](SDBackendModule module) {
-        if (backend_for(module) == nullptr) {
-            return false;
-        }
-        return params_backend_for(module) != nullptr;
-    };
-    if (!ensure_backend_pair(SDBackendModule::UPSCALER)) {
-        return false;
-    }
-
     model_manager = std::make_shared<ModelManager>();
     model_manager->set_n_threads(n_threads);
     model_manager->set_enable_mmap(false);
-
-    ModelLoader model_loader;
-    if (!model_loader.init_from_file_and_convert_name(esrgan_path, "", VERSION_ESRGAN)) {
-        LOG_ERROR("init model loader from file failed: '%s'", esrgan_path.c_str());
-        return false;
-    }
-    model_loader.set_wtype_override(model_data_type);
-    LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type));
-    const auto& tensor_storage_map = model_loader.get_tensor_storage_map();
-    std::map<std::string, ggml_tensor*> tensors;
-    if (tensor_storage_map.count("gblocks.0.weight") > 0) {
-        mosr_upscaler = std::make_shared<MoSR>(backend_for(SDBackendModule::UPSCALER),
-                                               tensor_storage_map,
-                                               model_manager);
-        if (mosr_upscaler->mosr_net == nullptr) {
-            LOG_ERROR("init mosr model from metadata failed: '%s'", esrgan_path.c_str());
-            return false;
-        }
-        mosr_upscaler->get_param_tensors(tensors);
-    } else {
-        esrgan_upscaler = std::make_shared<ESRGAN>(backend_for(SDBackendModule::UPSCALER),
-                                                   tensor_storage_map,
-                                                   model_manager);
-        if (esrgan_upscaler->rrdb_net == nullptr) {
-            LOG_ERROR("init esrgan model from metadata failed: '%s'", esrgan_path.c_str());
-            return false;
-        }
-        esrgan_upscaler->get_param_tensors(tensors);
-    }
-    LOG_INFO("Upscaler architecture: %s", runner()->get_desc().c_str());
-    runner()->set_max_graph_vram_bytes(max_graph_vram_bytes);
-    if (direct) {
-        runner()->set_conv2d_direct_enabled(true);
-    }
-
-    if (!model_manager->set_loader(model_loader) ||
-        !model_manager->register_param_tensors(ModelComponent::Upscaler,
-                                               std::move(tensors),
-                                               backend_manager.params_backend_is_disk(SDBackendModule::UPSCALER) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend,
-                                               backend_for(SDBackendModule::UPSCALER),
-                                               params_backend_for(SDBackendModule::UPSCALER)) ||
-        !model_manager->validate_registered_tensors()) {
-        LOG_ERROR("register upscaler tensors with model manager failed");
-        return false;
-    }
-    return true;
+    model = sd::pipeline::load_image_upscale_model(model_manager, backend_manager, path,
+                                                   max_graph_vram_bytes, direct, model_data_type);
+    return model != nullptr;
 }
 
 sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_tensor) {
     sd::ParallelScope tensor_scope(&tensor_executor);
-    if (input_tensor.empty() || input_tensor.dim() != 4 ||
-        (input_tensor.shape()[2] != 3 && input_tensor.shape()[2] != 4)) {
-        LOG_ERROR("upscaler expects a 4D RGB or RGBA image tensor");
+    if (model == nullptr) {
         return {};
     }
-
-    const bool has_alpha = input_tensor.shape()[2] == 4;
-    sd::Tensor<float> rgb;
-    if (has_alpha) {
-        rgb = sd::ops::slice(input_tensor, 2, 0, 3);
-    }
-    const sd::Tensor<float>& model_input = has_alpha ? rgb : input_tensor;
-    sd::Tensor<float> upscaled;
-    const int scale = get_scale();
-    if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) {
-        upscaled = compute(model_input);
-    } else {
-        auto on_processing = [&](const sd::Tensor<float>& input_tile) -> sd::Tensor<float> {
-            auto output_tile = compute(input_tile);
-            if (output_tile.empty()) {
-                LOG_ERROR("upscaler compute failed while processing a tile");
-                return {};
-            }
-            return output_tile;
-        };
-
-        upscaled = process_tiles_2d(model_input,
-                                    static_cast<int>(input_tensor.shape()[0] * scale),
-                                    static_cast<int>(input_tensor.shape()[1] * scale),
-                                    scale,
-                                    tile_size,
-                                    tile_size,
-                                    0.25f,
-                                    false,
-                                    false,
-                                    on_processing);
-    }
-    runner()->runner_end();
-    if (upscaled.empty()) {
-        LOG_ERROR("upscaler compute failed");
-        return {};
-    }
-    if (has_alpha) {
-        auto alpha       = sd::ops::slice(input_tensor, 2, 3, 4);
-        auto alpha_shape = alpha.shape();
-        alpha_shape[0]   = upscaled.shape()[0];
-        alpha_shape[1]   = upscaled.shape()[1];
-        alpha            = sd::ops::interpolate(alpha, alpha_shape, sd::ops::InterpolateMode::Bilinear);
-        upscaled         = sd::ops::concat(upscaled, alpha, 2);
-    }
-    return upscaled;
+    return sd::pipeline::upscale_image_tensor(*model, input_tensor, n_threads, tile_size);
 }
 
 sd_image_t UpscalerGGML::upscale(sd_image_t input_image, uint32_t upscale_factor) {

src/upscaler.h

@@ -3,9 +3,8 @@
 
 #include "core/ggml_extend_backend.h"
 #include "core/tensor.hpp"
-#include "model/upscaler/esrgan.hpp"
-#include "model/upscaler/mosr.hpp"
 #include "model_manager.h"
+#include "pipeline/image_upscale.h"
 #include "stable-diffusion.h"
 
 #include <memory>
@@ -15,9 +14,7 @@ struct UpscalerGGML {
     SDBackendManager backend_manager;
     std::shared_ptr<ModelManager> model_manager;
     ggml_type model_data_type = GGML_TYPE_F16;
-    std::shared_ptr<ESRGAN> esrgan_upscaler;
-    std::shared_ptr<MoSR> mosr_upscaler;
-    std::string esrgan_path;
+    std::unique_ptr<sd::pipeline::ImageUpscaleModel> model;
     int n_threads;
     sd::ParallelExecutor tensor_executor;
     bool direct                 = false;
@@ -36,9 +33,7 @@ struct UpscalerGGML {
     bool load_from_file(const std::string& esrgan_path,
                         int n_threads);
     void set_max_graph_vram_bytes(size_t max_vram_bytes);
-    GGMLRunner* runner();
     int get_scale();
-    sd::Tensor<float> compute(const sd::Tensor<float>& input_tensor);
     sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor);
     sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor);
 };