feat: prefetch streamed layers during compute

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 @@ struct GGMLRunner {
 
     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 @@ struct GGMLRunner {
                                                      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 @@ struct GGMLRunner {
                                                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 @@ struct GGMLRunner {
         }
 
         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 @@ struct GGMLRunner {
         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 @@ struct GGMLRunner {
             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 @@ struct GGMLRunner {
             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 @@ struct GGMLRunner {
         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;
+    }
+}