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__