perf: ratchet streaming budget so plan stops re-merging every step (#1611)

Signed-off-by: kei-g <km.8k6ce+github@gmail.com>

examples/cli/main.cpp

@@ -169,8 +169,9 @@ struct SDCliParams {
             return 1;
         };
 
-        auto on_help_arg = [&](int argc, const char** argv, int index) {
+        auto on_help_arg = [&](int argc, const char** argv, int index, bool& valid) {
             normal_exit = true;
+            valid       = true;
             return -1;
         };
 

examples/common/common.cpp

@@ -245,6 +245,7 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
         return false;
     };
 
+    bool valid = false;
     for (int i = 1; i < argc; i++) {
         arg            = argv[i];
         bool found_arg = false;
@@ -287,7 +288,7 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
                 break;
 
             if (match_and_apply(options.manual_options, [&](auto& option) {
-                    int ret = option.cb(argc, argv, i);
+                    int ret = option.cb(argc, argv, i, valid);
                     if (ret < 0) {
                         invalid_arg = true;
                         return;
@@ -299,7 +300,9 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
         }
 
         if (invalid_arg) {
-            LOG_ERROR("error: invalid parameter for argument: %s", arg.c_str());
+            if (!valid) {
+                LOG_ERROR("error: invalid parameter for argument: %s", arg.c_str());
+            }
             return false;
         }
         if (!found_arg) {

examples/common/common.h

@@ -56,11 +56,42 @@ struct BoolOption {
     bool* target;
 };
 
+struct ManualFunction {
+    std::function<int(int, const char**, int, bool&)> _func;
+
+    ManualFunction() = default;
+
+    ManualFunction(std::function<int(int argc, const char** argv, int index, bool& valid)> func)
+        : _func(std::move(func)) {
+    }
+
+    template <typename F>
+    ManualFunction(F func)
+        : _func(make_function(func)) {
+    }
+
+    int operator()(int argc, const char** argv, int index, bool& valid) const {
+        return _func(argc, argv, index, valid);
+    }
+
+private:
+    template <typename F>
+    static std::function<int(int, const char**, int, bool&)> make_function(F func) {
+        if constexpr (std::is_invocable_v<F, int, const char**, int, bool&>) {
+            return func;
+        } else {
+            return [func](int argc, const char** argv, int index, bool&) {
+                return func(argc, argv, index);
+            };
+        }
+    }
+};
+
 struct ManualOption {
     std::string short_name;
     std::string long_name;
     std::string desc;
-    std::function<int(int argc, const char** argv, int index)> cb;
+    ManualFunction cb;
 };
 
 struct ArgOptions {

examples/server/runtime.cpp

@@ -203,8 +203,9 @@ ArgOptions SDSvrParams::get_options() {
         {"", "--color", "colors the logging tags according to level", true, &color},
     };
 
-    auto on_help_arg = [&](int, const char**, int) {
+    auto on_help_arg = [&](int, const char**, int, bool& valid) {
         normal_exit = true;
+        valid       = true;
         return -1;
     };
 

src/ggml_extend.hpp

@@ -1706,8 +1706,9 @@ protected:
     std::unordered_set<ggml_tensor*> resident_param_set;
     uint64_t resident_state_token = 0;
 
-    size_t max_graph_vram_bytes = 0;
-    bool stream_layers_enabled  = false;
+    size_t max_graph_vram_bytes           = 0;
+    bool   stream_layers_enabled          = false;
+    size_t observed_max_effective_budget_ = 0;
 
     sd::layer_registry::LayerRegistry layer_registry_;
 
@@ -2446,15 +2447,25 @@ protected:
                 constexpr size_t safety_margin = 512ull * 1024 * 1024;
                 size_t free_clamp              = (free_vram > safety_margin) ? (free_vram - safety_margin) : 0;
                 if (free_clamp < effective_budget) {
-                    LOG_INFO("%s clamping streaming budget: actual free VRAM %.2f MB < user cap %.2f MB",
-                             get_desc().c_str(),
-                             free_clamp / (1024.0 * 1024.0),
-                             effective_budget / (1024.0 * 1024.0));
+                    LOG_DEBUG("%s clamping streaming budget: actual free VRAM %.2f MB < user cap %.2f MB",
+                              get_desc().c_str(),
+                              free_clamp / (1024.0 * 1024.0),
+                              effective_budget / (1024.0 * 1024.0));
                     effective_budget = free_clamp;
                 }
             }
         }
 
+        bool budget_increased = false;
+        if (stream_layers_enabled) {
+            if (effective_budget > observed_max_effective_budget_) {
+                observed_max_effective_budget_ = effective_budget;
+                budget_increased               = true;
+            } else {
+                effective_budget = observed_max_effective_budget_;
+            }
+        }
+
         if (effective_budget_out != nullptr) {
             *effective_budget_out = effective_budget;
         }
@@ -2466,9 +2477,15 @@ protected:
                                                      params_tensor_set_,
                                                      get_desc().c_str());
         if (stream_layers_enabled) {
-            LOG_INFO("%s streaming budget = %.2f MB",
-                     get_desc().c_str(),
-                     effective_budget / (1024.0 * 1024.0));
+            if (budget_increased) {
+                LOG_INFO("%s streaming budget = %.2f MB",
+                         get_desc().c_str(),
+                         effective_budget / (1024.0 * 1024.0));
+            } else {
+                LOG_DEBUG("%s streaming budget = %.2f MB",
+                          get_desc().c_str(),
+                          effective_budget / (1024.0 * 1024.0));
+            }
         }
         return true;
     }
@@ -3017,7 +3034,18 @@ public:
             LOG_DEBUG("%s skipping params allocation (no tensors)", get_desc().c_str());
             return true;
         }
-        params_buffer = ggml_backend_alloc_ctx_tensors(params_ctx, params_backend);
+        // Pinned host buffer when CPU-offloaded for DMA-direct H2D.
+        ggml_backend_buffer_type_t params_buft = nullptr;
+        if (params_backend != runtime_backend) {
+            ggml_backend_dev_t runtime_dev = ggml_backend_get_device(runtime_backend);
+            if (runtime_dev != nullptr) {
+                params_buft = ggml_backend_dev_host_buffer_type(runtime_dev);
+            }
+        }
+        if (params_buft == nullptr) {
+            params_buft = ggml_backend_get_default_buffer_type(params_backend);
+        }
+        params_buffer = ggml_backend_alloc_ctx_tensors_from_buft(params_ctx, params_buft);
         if (params_buffer == nullptr) {
             LOG_ERROR("%s alloc params backend buffer failed, num_tensors = %i",
                       get_desc().c_str(),
@@ -3042,6 +3070,7 @@ public:
             ggml_backend_buffer_free(params_buffer);
             params_buffer = nullptr;
         }
+        observed_max_effective_budget_ = 0;
     }
 
     size_t get_params_buffer_size() {

src/ggml_graph_cut.cpp

@@ -699,9 +699,9 @@ namespace sd::ggml_graph_cut {
         }
 
         if (log_desc != nullptr) {
-            LOG_INFO("%s graph cut max_vram budget merge took %lld ms",
-                     log_desc,
-                     ggml_time_ms() - t_budget_begin);
+            LOG_DEBUG("%s graph cut max_vram budget merge took %lld ms",
+                      log_desc,
+                      ggml_time_ms() - t_budget_begin);
         }
 
         return merged_plan;

src/stable-diffusion.cpp

@@ -1680,12 +1680,15 @@ public:
                                          const sd::Tensor<float>& init_latent,
                                          const sd::Tensor<float>& denoise_mask) {
         if (diffusion_model->get_desc() == "Wan2.2-TI2V-5B") {
-            auto new_timesteps = std::vector<float>(static_cast<size_t>(init_latent.shape()[2]), timesteps[0]);
+            int64_t frame_count = init_latent.shape()[2];
+            auto new_timesteps  = std::vector<float>(static_cast<size_t>(frame_count), timesteps[0]);
 
-            if (!denoise_mask.empty()) {
-                float value = denoise_mask.dim() == 5 ? denoise_mask.index(0, 0, 0, 0, 0) : denoise_mask.index(0, 0, 0, 0);
-                if (value == 0.f) {
-                    new_timesteps[0] = 0.f;
+            if (!denoise_mask.empty() && denoise_mask.dim() >= 4 && denoise_mask.shape()[2] == frame_count) {
+                for (int64_t frame = 0; frame < frame_count; ++frame) {
+                    float value = denoise_mask.dim() == 5 ? denoise_mask.index(0, 0, frame, 0, 0) : denoise_mask.index(0, 0, frame, 0);
+                    if (value == 0.f) {
+                        new_timesteps[static_cast<size_t>(frame)] = 0.f;
+                    }
                 }
             }
             return new_timesteps;