feat: prefetch streamed layers during compute (#1905)

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

CMakeLists.txt

@@ -290,6 +290,10 @@ else()
     add_library(${SD_LIB} STATIC ${SD_LIB_SOURCES})
 endif()
 
+if(MSVC)
+    target_compile_options(${SD_LIB} PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/bigobj>)
+endif()
+
 if(APPLE)
     sd_set_macos_rpaths(${SD_LIB})
 endif()

docs/performance.md

@@ -58,7 +58,8 @@ See [backend selection](./backend.md) for full syntax.
 `--offload-to-cpu` alone keeps every parameter in system RAM and stages it to the runtime backend on first use, then leaves it resident there. If the diffusion model is larger than the runtime backend's free memory (e.g. Flux dev at bf16 on an 8 GiB GPU), that residency stops fitting during the sampling loop and generation fails. Two additional flags make it fit by trading a small amount of speed for room:
 
 - `--max-vram <GiB>` sets a VRAM budget the graph-cut segmenter respects. It cuts each forward pass into segments sized to fit the budget, running them in sequence and freeing intermediate activations between them. Negative values auto-detect free VRAM and spare the given amount (`--max-vram -1` uses most of the free VRAM and keeps ~1 GiB headroom), a positive value caps the budget, `0` disables segmentation.
-- `--stream-layers` streams the diffusion model's transformer blocks one at a time. Each block's parameters are copied from the CPU to the runtime backend just before it runs and evicted when the residency budget is reached. Prefetching hides most of the copy latency behind compute. This flag only takes effect when the diffusion params backend is CPU, so it must be combined with `--offload-to-cpu` (or an explicit `--params-backend diffusion=cpu`); a warning is logged and the flag is ignored otherwise.
+- `--stream-layers` streams the diffusion model's transformer blocks one at a time. While one block computes, the next block's parameters are prefetched automatically from the CPU on a separate transfer queue; parameters are evicted when the residency budget is reached. This flag only takes effect when the diffusion params backend is CPU, so it must be combined with `--offload-to-cpu` (or an explicit `--params-backend diffusion=cpu`); a warning is logged and the flag is ignored otherwise.
+- `--disable-prefetch` disables the asynchronous next-block prefetch while retaining synchronous `--stream-layers` execution. This is mainly useful for debugging or backends where transfer and compute do not overlap effectively.
 
 The three flags stack. The recommended shape for "biggest model my card can host":
 

examples/common/common.cpp

@@ -520,6 +520,10 @@ ArgOptions SDContextParams::get_options() {
          "--stream-layers",
          "enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram; defaults to false)",
          true, &stream_layers},
+        {"",
+         "--disable-prefetch",
+         "disable asynchronous layer prefetch while keeping synchronous --stream-layers behavior (defaults to false)",
+         true, &disable_prefetch},
         {"",
          "--eager-load",
          "load all params into the params backend at model-load time instead of lazily on first use (defaults to false)",
@@ -832,6 +836,7 @@ std::string SDContextParams::to_string() const {
         << "  offload_params_to_cpu: " << (offload_params_to_cpu ? "true" : "false") << ",\n"
         << "  max_vram: \"" << max_vram << "\",\n"
         << "  stream_layers: " << (stream_layers ? "true" : "false") << ",\n"
+        << "  disable_prefetch: " << (disable_prefetch ? "true" : "false") << ",\n"
         << "  eager_load: " << (eager_load ? "true" : "false") << ",\n"
         << "  backend: \"" << backend << "\",\n"
         << "  params_backend: \"" << params_backend << "\",\n"
@@ -904,6 +909,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
     sd_ctx_params.vae_format                      = str_to_vae_format(vae_format);
     sd_ctx_params.max_vram                        = max_vram.c_str();
     sd_ctx_params.stream_layers                   = stream_layers;
+    sd_ctx_params.disable_prefetch                = disable_prefetch;
     sd_ctx_params.eager_load                      = eager_load;
     sd_ctx_params.backend                         = effective_backend.c_str();
     sd_ctx_params.params_backend                  = effective_params_backend.c_str();

examples/common/common.h

@@ -151,6 +151,7 @@ struct SDContextParams {
     bool offload_params_to_cpu  = false;
     std::string max_vram        = "0";
     bool stream_layers          = false;
+    bool disable_prefetch       = false;
     bool eager_load             = false;
     std::string backend;
     std::string params_backend;

include/stable-diffusion.h

@@ -231,6 +231,7 @@ typedef struct {
     enum sd_vae_format_t vae_format;
     const char* max_vram;  // GiB budget or backend assignment spec for graph-cut segmented param offload (0 = disabled, -1 = auto)
     bool stream_layers;  // Enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram)
+    bool disable_prefetch;  // Disable asynchronous layer prefetch while retaining synchronous stream_layers behavior
     bool eager_load;  // Load all params into the params backend at model-load time instead of lazily on first use
     const char* backend;
     const char* params_backend;

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.
+    // Mark leading segments resident after reserving streamed execution headroom.
-    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);
 }  // 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 stream_layers    = false;
-    bool eager_load    = 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__