feat: prefetch streamed layers during compute (#1905)

Co-authored-by: leejet <leejet714@gmail.com>

src/core/ggml_extend.hpp

@@ -27,6 +27,7 @@
 #include "core/ggml_extend_backend.h"
 #include "core/ggml_graph_cut.h"
 #include "core/layer_split_partition.h"
+#include "core/layer_stream_prefetch.h"
 #include "ggml-alloc.h"
 #include "ggml-backend.h"
 #include "ggml.h"
@@ -1788,6 +1789,7 @@ protected:
 
     size_t max_graph_vram_bytes           = 0;
     bool stream_layers_enabled            = false;
+    bool layer_prefetch_enabled           = true;
     size_t observed_max_effective_budget_ = 0;
     bool graph_cut_layer_split_enabled    = false;
     std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
@@ -2513,7 +2515,9 @@ protected:
                                                      params_tensor_set_,
                                                      get_desc().c_str());
         if (stream_layers_enabled) {
-            sd::ggml_graph_cut::annotate_residency(*plan_out, effective_budget);
+            sd::ggml_graph_cut::annotate_residency(*plan_out,
+                                                   effective_budget,
+                                                   layer_prefetch_enabled);
         }
         if (stream_layers_enabled) {
             if (budget_increased) {
@@ -2771,7 +2775,8 @@ protected:
                                                bool free_compute_params,
                                                bool preserve_backend_tensor_data_map,
                                                bool no_return                                          = false,
-                                               const std::unordered_set<std::string>* cache_keep_names = nullptr) {
+                                               const std::unordered_set<std::string>* cache_keep_names = nullptr,
+                                               const std::function<void()>& before_compute             = {}) {
         std::vector<ggml_tensor*> graph_param_tensors;
         std::vector<ggml_tensor*> params_to_prepare;
         if (!prepare_execute_graph_weights(gf, graph_param_tensors, params_to_prepare, !free_compute_params)) {
@@ -2835,6 +2840,9 @@ protected:
         }
 
         copy_data_to_backend_tensor(gf, !preserve_backend_tensor_data_map);
+        if (before_compute) {
+            before_compute();
+        }
         if (sd_backend_is_cpu(runtime_backend)) {
             sd_backend_cpu_set_n_threads(runtime_backend, n_threads);
         }
@@ -2940,6 +2948,20 @@ protected:
         free_compute_buffer();
         free_cache_ctx_and_buffer();
 
+        sd::LayerStreamPrefetch prefetch(weight_manager.lock(),
+                                         reinterpret_cast<uintptr_t>(this),
+                                         gf,
+                                         plan,
+                                         params_tensor_set_,
+                                         stream_layers_enabled && layer_prefetch_enabled);
+        auto disable_prefetch = [&]() {
+            if (layer_prefetch_enabled) {
+                LOG_WARN("%s layer prefetch failed; continuing with synchronous streaming",
+                         get_desc().c_str());
+            }
+            layer_prefetch_enabled = false;
+        };
+
         std::unordered_map<ggml_tensor*, PersistentExternalBinding> persistent_externals;
         snapshot_persistent_externals(plan, gf, persistent_externals);
 
@@ -2948,6 +2970,9 @@ protected:
             const auto& segment   = plan.segments[seg_idx];
             const bool is_last    = seg_idx + 1 == plan.segments.size();
             auto future_cut_names = sd::ggml_graph_cut::collect_future_input_names(gf, plan, seg_idx);
+            if (!prefetch.activate(seg_idx)) {
+                disable_prefetch();
+            }
             if (log_residency) {
                 LOG_DEBUG("%s graph cut executing segment %zu/%zu: %s (residency=%s)",
                           get_desc().c_str(),
@@ -2985,13 +3010,22 @@ protected:
             ggml_context* segment_graph_ctx = nullptr;
             ggml_cgraph* segment_graph      = sd::ggml_graph_cut::build_segment_graph(gf, segment, &segment_graph_ctx);
             const bool keep_segment_params  = segment.residency == sd::ggml_graph_cut::SegmentResidency::RESIDENT;
-            auto segment_output             = execute_graph<T>(segment_graph,
+            std::function<void()> before_compute;
+            if (prefetch.enabled()) {
+                before_compute = [&, seg_idx]() {
+                    if (!prefetch.enqueue_next(seg_idx)) {
+                        disable_prefetch();
+                    }
+                };
+            }
+            auto segment_output = execute_graph<T>(segment_graph,
                                                    n_threads,
                                                    true,
                                                    !keep_segment_params,
                                                    true,
                                                    !is_last || no_return,
-                                                   &future_cut_names);
+                                                   &future_cut_names,
+                                                   before_compute);
             ggml_free(segment_graph_ctx);
             if (!segment_output.has_value()) {
                 free_cache_ctx_and_buffer();
@@ -3239,6 +3273,10 @@ public:
         stream_layers_enabled = enabled;
     }
 
+    void set_layer_prefetch_enabled(bool enabled) {
+        layer_prefetch_enabled = enabled;
+    }
+
     void set_graph_cut_layer_split_enabled(bool enabled) {
         graph_cut_layer_split_enabled = enabled;
         if (!enabled) {

src/core/ggml_graph_cut.cpp

@@ -995,7 +995,9 @@ namespace sd::ggml_graph_cut {
         return resolved_plan;
     }
 
-    void annotate_residency(Plan& plan, size_t max_graph_vram_bytes) {
+    void annotate_residency(Plan& plan,
+                            size_t max_graph_vram_bytes,
+                            bool prefetch_enabled) {
         // Cached plans may be reused with a smaller live budget.
         for (auto& seg : plan.segments) {
             seg.residency = SegmentResidency::STREAMED;
@@ -1017,6 +1019,7 @@ namespace sd::ggml_graph_cut {
 
         // Leave room for the largest active streamed segment.
         size_t worst_streamed_footprint = 0;
+        size_t prefetch_headroom        = 0;
         for (const auto& seg : plan.segments) {
             const size_t seg_footprint = seg.input_param_bytes +
                                          seg.compute_buffer_size +
@@ -1026,9 +1029,19 @@ namespace sd::ggml_graph_cut {
             if (seg_footprint > worst_streamed_footprint) {
                 worst_streamed_footprint = seg_footprint;
             }
+            prefetch_headroom = std::max(prefetch_headroom, seg.input_param_bytes);
         }
         constexpr size_t safety = 512ull * 1024 * 1024;
-        const size_t reserved   = safety + worst_streamed_footprint;
+        if (worst_streamed_footprint > SIZE_MAX - safety) {
+            return;
+        }
+        size_t reserved = safety + worst_streamed_footprint;
+        if (prefetch_enabled) {
+            if (prefetch_headroom > SIZE_MAX - reserved) {
+                return;
+            }
+            reserved += prefetch_headroom;
+        }
 
         if (max_graph_vram_bytes <= reserved) {
             return;

src/core/ggml_graph_cut.h

@@ -122,8 +122,10 @@ namespace sd::ggml_graph_cut {
                       const std::unordered_set<const ggml_tensor*>& params_tensor_set,
                       const char* log_desc);
 
-    // Mark leading segments resident when they fit after streamed-segment headroom.
-    void annotate_residency(Plan& plan, size_t max_graph_vram_bytes);
+    // Mark leading segments resident after reserving streamed execution headroom.
+    void annotate_residency(Plan& plan,
+                            size_t max_graph_vram_bytes,
+                            bool prefetch_enabled);
 }  // namespace sd::ggml_graph_cut
 
 #endif  // __SD_CORE_GGML_GRAPH_CUT_H__

src/core/layer_stream_prefetch.cpp

@@ -0,0 +1,124 @@
+#include "core/layer_stream_prefetch.h"
+
+#include <utility>
+
+#include "core/ggml_graph_cut.h"
+#include "weight_manager.h"
+
+namespace sd {
+    static ggml_tensor* canonical_param(
+        ggml_tensor* tensor,
+        const std::unordered_set<const ggml_tensor*>& params) {
+        for (ggml_tensor* current = tensor; current != nullptr; current = current->view_src) {
+            if (params.find(current) != params.end()) {
+                return current;
+            }
+        }
+        return nullptr;
+    }
+
+    LayerStreamPrefetch::LayerStreamPrefetch(
+        const std::shared_ptr<RunnerWeightManager>& manager,
+        uintptr_t owner_id,
+        ggml_cgraph* graph,
+        const ggml_graph_cut::Plan& plan,
+        const std::unordered_set<const ggml_tensor*>& params,
+        bool enabled)
+        : manager_(manager),
+          owner_id_(owner_id),
+          enabled_(enabled && manager != nullptr) {
+        segment_params_.resize(plan.segments.size());
+        for (size_t segment_index = 0; segment_index < plan.segments.size(); ++segment_index) {
+            std::unordered_set<ggml_tensor*> seen;
+            for (ggml_tensor* tensor :
+                 ggml_graph_cut::param_tensors(graph, plan.segments[segment_index])) {
+                ggml_tensor* param = canonical_param(tensor, params);
+                if (param != nullptr && seen.insert(param).second) {
+                    segment_params_[segment_index].push_back(param);
+                }
+            }
+        }
+    }
+
+    LayerStreamPrefetch::~LayerStreamPrefetch() {
+        clear();
+    }
+
+    size_t LayerStreamPrefetch::next_parameter_segment(size_t segment_index) const {
+        for (size_t next = segment_index + 1; next < segment_params_.size(); ++next) {
+            if (!segment_params_[next].empty()) {
+                return next;
+            }
+        }
+        return SIZE_MAX;
+    }
+
+    void LayerStreamPrefetch::disable() {
+        clear();
+        enabled_ = false;
+    }
+
+    bool LayerStreamPrefetch::activate(size_t segment_index) {
+        if (!enabled_ || queued_segment_ == SIZE_MAX) {
+            return true;
+        }
+        if (queued_segment_ != segment_index) {
+            return true;
+        }
+
+        auto manager = manager_.lock();
+        if (manager == nullptr ||
+            !manager->activate_prefetched_params(owner_id_, queued_params_)) {
+            disable();
+            return false;
+        }
+        queued_params_.clear();
+        queued_segment_ = SIZE_MAX;
+        return true;
+    }
+
+    bool LayerStreamPrefetch::enqueue_next(size_t segment_index) {
+        if (!enabled_) {
+            return true;
+        }
+        if (queued_segment_ != SIZE_MAX) {
+            return true;
+        }
+
+        const size_t next_segment = next_parameter_segment(segment_index);
+        if (next_segment == SIZE_MAX) {
+            return true;
+        }
+
+        std::unordered_set<ggml_tensor*> active_params(
+            segment_params_[segment_index].begin(),
+            segment_params_[segment_index].end());
+        std::vector<ggml_tensor*> params;
+        params.reserve(segment_params_[next_segment].size());
+        for (ggml_tensor* param : segment_params_[next_segment]) {
+            if (active_params.find(param) == active_params.end()) {
+                params.push_back(param);
+            }
+        }
+        if (params.empty()) {
+            return true;
+        }
+
+        auto manager = manager_.lock();
+        if (manager == nullptr || !manager->prefetch_params(owner_id_, params)) {
+            disable();
+            return false;
+        }
+        queued_params_  = std::move(params);
+        queued_segment_ = next_segment;
+        return true;
+    }
+
+    void LayerStreamPrefetch::clear() {
+        if (auto manager = manager_.lock()) {
+            manager->clear_prefetched_params(owner_id_);
+        }
+        queued_params_.clear();
+        queued_segment_ = SIZE_MAX;
+    }
+}

src/core/layer_stream_prefetch.h

@@ -0,0 +1,48 @@
+#ifndef __SD_CORE_LAYER_STREAM_PREFETCH_H__
+#define __SD_CORE_LAYER_STREAM_PREFETCH_H__
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <unordered_set>
+#include <vector>
+
+struct ggml_cgraph;
+struct ggml_tensor;
+struct RunnerWeightManager;
+
+namespace sd::ggml_graph_cut {
+    struct Plan;
+}
+
+namespace sd {
+    class LayerStreamPrefetch {
+    private:
+        std::weak_ptr<RunnerWeightManager> manager_;
+        uintptr_t owner_id_ = 0;
+        std::vector<std::vector<ggml_tensor*>> segment_params_;
+        std::vector<ggml_tensor*> queued_params_;
+        size_t queued_segment_ = SIZE_MAX;
+        bool enabled_          = true;
+
+        size_t next_parameter_segment(size_t segment_index) const;
+        void disable();
+
+    public:
+        LayerStreamPrefetch(
+            const std::shared_ptr<RunnerWeightManager>& manager,
+            uintptr_t owner_id,
+            ggml_cgraph* graph,
+            const ggml_graph_cut::Plan& plan,
+            const std::unordered_set<const ggml_tensor*>& params,
+            bool enabled = true);
+        ~LayerStreamPrefetch();
+
+        bool enabled() const { return enabled_; }
+        bool activate(size_t segment_index);
+        bool enqueue_next(size_t segment_index);
+        void clear();
+    };
+}
+
+#endif  // __SD_CORE_LAYER_STREAM_PREFETCH_H__

src/model_manager.cpp

@@ -97,6 +97,7 @@ static bool device_supports_param_op(ggml_backend_dev_t device,
 
 ModelManager::~ModelManager() {
     release_all();
+    release_prefetch();
 }
 
 void ModelManager::set_common_ignore_tensors(std::set<std::string> ignore_tensors) {
@@ -255,6 +256,7 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg
         return true;
     }
 
+    clear_all_prefetched_params();
     release_compute_staging_blocks(false);
 
     std::vector<ParamsStorageBlock*> storage_blocks_to_release;
@@ -608,6 +610,7 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
 }
 
 void ModelManager::reset_lora_applied_params() {
+    clear_all_prefetched_params();
     release_compute_staging_blocks(true);
     release_params_storage_blocks(true);
     for (auto& state : tensor_states_) {
@@ -1044,6 +1047,7 @@ void ModelManager::erase_params_storage_block(ParamsStorageBlock* block) {
 }
 
 void ModelManager::release_all() {
+    clear_all_prefetched_params();
     for (auto& state : tensor_states_) {
         state->active_prepare_count = 0;
         state->applied_lora_epoch   = UINT64_MAX;
@@ -1097,6 +1101,7 @@ bool ModelManager::assign_compute_backend(const std::vector<ggml_tensor*>& tenso
         return false;
     }
 
+    clear_all_prefetched_params();
     for (TensorState* state : required_states) {
         if (state == nullptr || state->tensor == nullptr) {
             continue;

src/model_manager.h

@@ -61,12 +61,24 @@ private:
         std::vector<std::pair<TensorState*, ggml_tensor*>> staged_tensors;
     };
 
+    struct PrefetchBlock {
+        std::vector<TensorState*> states;
+        ggml_backend_t compute_backend  = nullptr;
+        ggml_backend_t transfer_backend = nullptr;
+        ggml_backend_event_t event      = nullptr;
+        ggml_context* staging_ctx       = nullptr;
+        ggml_backend_buffer_t buffer    = nullptr;
+        std::vector<std::pair<TensorState*, ggml_tensor*>> staged_tensors;
+    };
+
     ModelLoader model_loader_;
     std::vector<std::unique_ptr<TensorState>> tensor_states_;
     std::map<std::string, TensorState*> tensor_states_by_name_;
     std::vector<std::unique_ptr<ParamsStorageBlock>> params_storage_blocks_;
     std::vector<std::unique_ptr<ComputeStagingBlock>> compute_staging_blocks_;
     std::map<ggml_backend_t, ggml_backend_buffer_type_t> split_buffer_types_;
+    std::map<uintptr_t, std::unique_ptr<PrefetchBlock>> prefetch_blocks_;
+    std::map<ggml_backend_t, ggml_backend_t> prefetch_backends_;
     bool warned_split_lora_skip_ = false;
     std::set<std::string> common_ignore_tensors_;
     std::vector<LoraSpec> loras_;
@@ -79,6 +91,13 @@ private:
     void finish_compute_backend_usage(const std::vector<TensorState*>& states);
     void release_all();
 
+    ggml_backend_t prefetch_backend_for(ggml_backend_t compute_backend);
+    bool populate_prefetch_block(PrefetchBlock& block);
+    void synchronize_prefetch_block(PrefetchBlock& block);
+    void free_prefetch_block(PrefetchBlock& block);
+    void clear_all_prefetched_params();
+    void release_prefetch();
+
     bool resolve_required_tensor_states(const std::vector<ggml_tensor*>& tensors,
                                         std::vector<TensorState*>& required_states) const;
     bool should_ignore(const TensorState& state) const;
@@ -182,6 +201,11 @@ public:
     bool prepare_params(const std::vector<ggml_tensor*>& tensors) override;
     void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) override;
     void release_params_backend_params(const std::vector<ggml_tensor*>& tensors) override;
+    bool prefetch_params(uintptr_t owner_id,
+                         const std::vector<ggml_tensor*>& tensors) override;
+    bool activate_prefetched_params(uintptr_t owner_id,
+                                    const std::vector<ggml_tensor*>& tensors) override;
+    void clear_prefetched_params(uintptr_t owner_id) override;
 };
 
 #endif  // __MODEL_MANAGER_H__

src/model_manager_prefetch.cpp

@@ -0,0 +1,271 @@
+#include "model_manager.h"
+
+#include <algorithm>
+#include <utility>
+
+#include "core/ggml_extend_backend.h"
+#include "core/util.h"
+
+ggml_backend_t ModelManager::prefetch_backend_for(ggml_backend_t compute_backend) {
+    auto existing = prefetch_backends_.find(compute_backend);
+    if (existing != prefetch_backends_.end()) {
+        return existing->second;
+    }
+    if (compute_backend == nullptr) {
+        return nullptr;
+    }
+
+    ggml_backend_dev_t device = ggml_backend_get_device(compute_backend);
+    if (device == nullptr || ggml_backend_dev_type(device) == GGML_BACKEND_DEVICE_TYPE_CPU) {
+        return nullptr;
+    }
+    ggml_backend_t transfer_backend = ggml_backend_dev_init(device, nullptr);
+    if (transfer_backend == nullptr) {
+        LOG_WARN("model manager failed to create a prefetch backend for %s",
+                 ggml_backend_name(compute_backend));
+    }
+    prefetch_backends_[compute_backend] = transfer_backend;
+    return transfer_backend;
+}
+
+void ModelManager::synchronize_prefetch_block(PrefetchBlock& block) {
+    if (block.event != nullptr) {
+        ggml_backend_event_synchronize(block.event);
+        ggml_backend_event_free(block.event);
+        block.event = nullptr;
+    } else if (block.transfer_backend != nullptr) {
+        ggml_backend_synchronize(block.transfer_backend);
+    }
+    block.transfer_backend = nullptr;
+}
+
+void ModelManager::free_prefetch_block(PrefetchBlock& block) {
+    synchronize_prefetch_block(block);
+    block.staged_tensors.clear();
+    if (block.buffer != nullptr) {
+        ggml_backend_buffer_free(block.buffer);
+        block.buffer = nullptr;
+    }
+    if (block.staging_ctx != nullptr) {
+        ggml_free(block.staging_ctx);
+        block.staging_ctx = nullptr;
+    }
+}
+
+bool ModelManager::populate_prefetch_block(PrefetchBlock& block) {
+    if (block.states.empty() || block.compute_backend == nullptr) {
+        return false;
+    }
+
+    block.transfer_backend = prefetch_backend_for(block.compute_backend);
+    if (block.transfer_backend == nullptr) {
+        return false;
+    }
+
+    ggml_init_params init_params;
+    init_params.mem_size   = block.states.size() * ggml_tensor_overhead();
+    init_params.mem_buffer = nullptr;
+    init_params.no_alloc   = true;
+    block.staging_ctx      = ggml_init(init_params);
+    if (block.staging_ctx == nullptr) {
+        return false;
+    }
+
+    block.staged_tensors.reserve(block.states.size());
+    for (TensorState* state : block.states) {
+        if (state == nullptr || state->tensor == nullptr ||
+            state->tensor->buffer == nullptr || state->tensor->data == nullptr ||
+            state->params_backend == nullptr || state->staged_to_compute_backend ||
+            state->active_prepare_count > 0) {
+            return false;
+        }
+        ggml_tensor* staging_tensor = ggml_dup_tensor(block.staging_ctx, state->tensor);
+        ggml_set_name(staging_tensor, state->tensor->name);
+        block.staged_tensors.push_back({state, staging_tensor});
+    }
+
+    ggml_backend_buffer_type_t buffer_type =
+        ggml_backend_get_default_buffer_type(block.compute_backend);
+    if (buffer_type == nullptr) {
+        return false;
+    }
+    block.buffer = ggml_backend_alloc_ctx_tensors_from_buft(block.staging_ctx, buffer_type);
+    if (block.buffer == nullptr) {
+        return false;
+    }
+    ggml_backend_buffer_set_usage(block.buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+
+    for (const auto& pair : block.staged_tensors) {
+        TensorState* state          = pair.first;
+        ggml_tensor* staging_tensor = pair.second;
+        const bool host_source      = ggml_backend_buffer_is_host(state->tensor->buffer);
+        if (host_source &&
+            (!ggml_is_contiguous(state->tensor) || !ggml_is_contiguous(staging_tensor) ||
+             ggml_nbytes(state->tensor) != ggml_nbytes(staging_tensor))) {
+            return false;
+        }
+    }
+
+    for (const auto& pair : block.staged_tensors) {
+        TensorState* state          = pair.first;
+        ggml_tensor* staging_tensor = pair.second;
+        if (ggml_backend_buffer_is_host(state->tensor->buffer)) {
+            ggml_backend_tensor_set_async(block.transfer_backend,
+                                          staging_tensor,
+                                          state->tensor->data,
+                                          0,
+                                          ggml_nbytes(state->tensor));
+        } else {
+            ggml_backend_tensor_copy_async(state->params_backend,
+                                           block.transfer_backend,
+                                           state->tensor,
+                                           staging_tensor);
+        }
+    }
+
+    ggml_backend_dev_t device = ggml_backend_get_device(block.transfer_backend);
+    block.event               = ggml_backend_event_new(device);
+    if (block.event != nullptr) {
+        ggml_backend_event_record(block.event, block.transfer_backend);
+    }
+
+    LOG_DEBUG("model manager queued layer prefetch (%6.2f MB, %zu tensors) to %s",
+              ggml_backend_buffer_get_size(block.buffer) / (1024.f * 1024.f),
+              block.states.size(),
+              ggml_backend_name(block.compute_backend));
+    return true;
+}
+
+bool ModelManager::prefetch_params(uintptr_t owner_id,
+                                   const std::vector<ggml_tensor*>& tensors) {
+    clear_prefetched_params(owner_id);
+    if (tensors.empty()) {
+        return true;
+    }
+
+    std::vector<TensorState*> required_states;
+    if (!resolve_required_tensor_states(tensors, required_states) ||
+        !load_tensors_to_params_backend(required_states)) {
+        return false;
+    }
+
+    std::vector<TensorState*> states;
+    states.reserve(required_states.size());
+    ggml_backend_t compute_backend = nullptr;
+    for (TensorState* state : required_states) {
+        if (state == nullptr || should_ignore(*state) ||
+            is_optional_missing_tensor(state->name) ||
+            state->compute_backend == state->params_backend ||
+            state->staged_to_compute_backend || state->active_prepare_count > 0) {
+            continue;
+        }
+        if (compute_backend == nullptr) {
+            compute_backend = state->compute_backend;
+        } else if (compute_backend != state->compute_backend) {
+            return false;
+        }
+        states.push_back(state);
+    }
+    if (states.empty()) {
+        return true;
+    }
+    if (compute_backend == nullptr || sd_backend_is_cpu(compute_backend)) {
+        return false;
+    }
+
+    auto block             = std::make_unique<PrefetchBlock>();
+    block->states          = std::move(states);
+    block->compute_backend = compute_backend;
+    if (!populate_prefetch_block(*block)) {
+        free_prefetch_block(*block);
+        return false;
+    }
+    prefetch_blocks_[owner_id] = std::move(block);
+    return true;
+}
+
+bool ModelManager::activate_prefetched_params(
+    uintptr_t owner_id,
+    const std::vector<ggml_tensor*>& tensors) {
+    std::vector<TensorState*> required_states;
+    if (!resolve_required_tensor_states(tensors, required_states)) {
+        return false;
+    }
+
+    const bool already_staged = std::all_of(
+        required_states.begin(),
+        required_states.end(),
+        [&](TensorState* state) {
+            return state == nullptr || should_ignore(*state) ||
+                   is_optional_missing_tensor(state->name) ||
+                   state->compute_backend == state->params_backend ||
+                   state->staged_to_compute_backend;
+        });
+    if (already_staged) {
+        clear_prefetched_params(owner_id);
+        return true;
+    }
+
+    auto existing = prefetch_blocks_.find(owner_id);
+    if (existing == prefetch_blocks_.end()) {
+        return false;
+    }
+    std::unique_ptr<PrefetchBlock> block = std::move(existing->second);
+    prefetch_blocks_.erase(existing);
+    synchronize_prefetch_block(*block);
+
+    for (const auto& pair : block->staged_tensors) {
+        TensorState* state          = pair.first;
+        ggml_tensor* staging_tensor = pair.second;
+        if (state == nullptr || state->tensor == nullptr || staging_tensor == nullptr ||
+            state->staged_to_compute_backend || state->active_prepare_count > 0) {
+            free_prefetch_block(*block);
+            return false;
+        }
+    }
+    for (auto& pair : block->staged_tensors) {
+        TensorState* state          = pair.first;
+        ggml_tensor* staging_tensor = pair.second;
+        std::swap(state->tensor->buffer, staging_tensor->buffer);
+        std::swap(state->tensor->data, staging_tensor->data);
+        std::swap(state->tensor->extra, staging_tensor->extra);
+        state->staged_to_compute_backend = true;
+    }
+
+    auto staging_block             = std::make_unique<ComputeStagingBlock>();
+    staging_block->compute_backend = block->compute_backend;
+    staging_block->buffer          = block->buffer;
+    staging_block->staging_ctx     = block->staging_ctx;
+    staging_block->staged_tensors  = std::move(block->staged_tensors);
+    block->buffer                  = nullptr;
+    block->staging_ctx             = nullptr;
+    compute_staging_blocks_.push_back(std::move(staging_block));
+    return true;
+}
+
+void ModelManager::clear_prefetched_params(uintptr_t owner_id) {
+    auto existing = prefetch_blocks_.find(owner_id);
+    if (existing == prefetch_blocks_.end()) {
+        return;
+    }
+    std::unique_ptr<PrefetchBlock> block = std::move(existing->second);
+    prefetch_blocks_.erase(existing);
+    free_prefetch_block(*block);
+}
+
+void ModelManager::clear_all_prefetched_params() {
+    for (auto& entry : prefetch_blocks_) {
+        free_prefetch_block(*entry.second);
+    }
+    prefetch_blocks_.clear();
+}
+
+void ModelManager::release_prefetch() {
+    clear_all_prefetched_params();
+    for (auto& entry : prefetch_backends_) {
+        if (entry.second != nullptr) {
+            ggml_backend_free(entry.second);
+        }
+    }
+    prefetch_backends_.clear();
+}

src/stable-diffusion.cpp

@@ -248,8 +248,9 @@ public:
     sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr};
     bool enable_mmap                     = false;
     sd::ggml_graph_cut::MaxVramAssignment max_vram_assignment;
-    bool stream_layers = false;
-    bool eager_load    = false;
+    bool stream_layers    = false;
+    bool disable_prefetch = false;
+    bool eager_load       = false;
     std::string backend_spec;
     std::string params_backend_spec;
     std::string split_mode_spec;
@@ -862,6 +863,7 @@ public:
         n_threads           = sd_ctx_params->n_threads;
         enable_mmap         = sd_ctx_params->enable_mmap;
         stream_layers       = sd_ctx_params->stream_layers;
+        disable_prefetch    = sd_ctx_params->disable_prefetch;
         eager_load          = sd_ctx_params->eager_load;
         backend_spec        = SAFE_STR(sd_ctx_params->backend);
         params_backend_spec = SAFE_STR(sd_ctx_params->params_backend);
@@ -1356,6 +1358,7 @@ public:
 
             diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION));
             diffusion_model->set_stream_layers_enabled(stream_layers);
+            diffusion_model->set_layer_prefetch_enabled(!disable_prefetch);
             if (!register_runner_params("Diffusion model",
                                         diffusion_model,
                                         SDBackendModule::DIFFUSION,
@@ -1366,6 +1369,7 @@ public:
             if (high_noise_diffusion_model) {
                 high_noise_diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION));
                 high_noise_diffusion_model->set_stream_layers_enabled(stream_layers);
+                high_noise_diffusion_model->set_layer_prefetch_enabled(!disable_prefetch);
                 if (!register_runner_params("High noise diffusion model",
                                             high_noise_diffusion_model,
                                             SDBackendModule::DIFFUSION,
@@ -3572,6 +3576,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
     sd_ctx_params->lora_apply_mode      = LORA_APPLY_AUTO;
     sd_ctx_params->max_vram             = nullptr;
     sd_ctx_params->stream_layers        = false;
+    sd_ctx_params->disable_prefetch     = false;
     sd_ctx_params->eager_load           = false;
     sd_ctx_params->enable_mmap          = false;
     sd_ctx_params->diffusion_flash_attn = false;
@@ -3617,6 +3622,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              "prediction: %s\n"
              "max_vram: %s\n"
              "stream_layers: %s\n"
+             "disable_prefetch: %s\n"
              "eager_load: %s\n"
              "backend: %s\n"
              "params_backend: %s\n"
@@ -3651,6 +3657,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              sd_prediction_name(sd_ctx_params->prediction),
              SAFE_STR(sd_ctx_params->max_vram),
              BOOL_STR(sd_ctx_params->stream_layers),
+             BOOL_STR(sd_ctx_params->disable_prefetch),
              BOOL_STR(sd_ctx_params->eager_load),
              SAFE_STR(sd_ctx_params->backend),
              SAFE_STR(sd_ctx_params->params_backend),

src/weight_manager.h

@@ -1,6 +1,7 @@
 #ifndef __WEIGHT_MANAGER_H__
 #define __WEIGHT_MANAGER_H__
 
+#include <cstdint>
 #include <vector>
 
 #include "ggml-backend.h"
@@ -14,6 +15,11 @@ struct RunnerWeightManager {
     virtual bool prepare_params(const std::vector<ggml_tensor*>& tensors)                 = 0;
     virtual void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;
     virtual void release_params_backend_params(const std::vector<ggml_tensor*>& tensors)  = 0;
+    virtual bool prefetch_params(uintptr_t owner_id,
+                                 const std::vector<ggml_tensor*>& tensors)                = 0;
+    virtual bool activate_prefetched_params(uintptr_t owner_id,
+                                            const std::vector<ggml_tensor*>& tensors)     = 0;
+    virtual void clear_prefetched_params(uintptr_t owner_id)                              = 0;
 };
 
 #endif  // __WEIGHT_MANAGER_H__