diff --git a/CMakeLists.txt b/CMakeLists.txt index 3c2b7b4a8..e64ccb717 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -290,6 +290,10 @@ else() add_library(${SD_LIB} STATIC ${SD_LIB_SOURCES}) endif() +if(MSVC) + target_compile_options(${SD_LIB} PRIVATE $<$:/bigobj>) +endif() + if(APPLE) sd_set_macos_rpaths(${SD_LIB}) endif() diff --git a/docs/performance.md b/docs/performance.md index ccfe4778a..f5d39fa83 100644 --- a/docs/performance.md +++ b/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 ` 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": diff --git a/examples/common/common.cpp b/examples/common/common.cpp index d46be91a4..2053b6540 100644 --- a/examples/common/common.cpp +++ b/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(); diff --git a/examples/common/common.h b/examples/common/common.h index 34b4a013b..0e726fab8 100644 --- a/examples/common/common.h +++ b/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; diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index d5cda6d9a..9c85114f7 100644 --- a/include/stable-diffusion.h +++ b/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; diff --git a/src/core/ggml_extend.hpp b/src/core/ggml_extend.hpp index c262969d3..01d51644e 100644 --- a/src/core/ggml_extend.hpp +++ b/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 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* cache_keep_names = nullptr) { + const std::unordered_set* cache_keep_names = nullptr, + const std::function& before_compute = {}) { std::vector graph_param_tensors; std::vector 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(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 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(segment_graph, + std::function before_compute; + if (prefetch.enabled()) { + before_compute = [&, seg_idx]() { + if (!prefetch.enqueue_next(seg_idx)) { + disable_prefetch(); + } + }; + } + auto segment_output = execute_graph(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) { diff --git a/src/core/ggml_graph_cut.cpp b/src/core/ggml_graph_cut.cpp index 542b7fe16..dec747166 100644 --- a/src/core/ggml_graph_cut.cpp +++ b/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; diff --git a/src/core/ggml_graph_cut.h b/src/core/ggml_graph_cut.h index d955049d8..c2594ceec 100644 --- a/src/core/ggml_graph_cut.h +++ b/src/core/ggml_graph_cut.h @@ -122,8 +122,10 @@ namespace sd::ggml_graph_cut { const std::unordered_set& 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__ diff --git a/src/core/layer_stream_prefetch.cpp b/src/core/layer_stream_prefetch.cpp new file mode 100644 index 000000000..5b737be72 --- /dev/null +++ b/src/core/layer_stream_prefetch.cpp @@ -0,0 +1,124 @@ +#include "core/layer_stream_prefetch.h" + +#include + +#include "core/ggml_graph_cut.h" +#include "weight_manager.h" + +namespace sd { + static ggml_tensor* canonical_param( + ggml_tensor* tensor, + const std::unordered_set& 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& manager, + uintptr_t owner_id, + ggml_cgraph* graph, + const ggml_graph_cut::Plan& plan, + const std::unordered_set& 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 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 active_params( + segment_params_[segment_index].begin(), + segment_params_[segment_index].end()); + std::vector 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; + } +} diff --git a/src/core/layer_stream_prefetch.h b/src/core/layer_stream_prefetch.h new file mode 100644 index 000000000..2457a1c6d --- /dev/null +++ b/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 +#include +#include +#include +#include + +struct ggml_cgraph; +struct ggml_tensor; +struct RunnerWeightManager; + +namespace sd::ggml_graph_cut { + struct Plan; +} + +namespace sd { + class LayerStreamPrefetch { + private: + std::weak_ptr manager_; + uintptr_t owner_id_ = 0; + std::vector> segment_params_; + std::vector 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& manager, + uintptr_t owner_id, + ggml_cgraph* graph, + const ggml_graph_cut::Plan& plan, + const std::unordered_set& 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__ diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 8db303dda..72e813749 100644 --- a/src/model_manager.cpp +++ b/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 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 storage_blocks_to_release; @@ -608,6 +610,7 @@ bool ModelManager::apply_loras_to_params(const std::vector& 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& tenso return false; } + clear_all_prefetched_params(); for (TensorState* state : required_states) { if (state == nullptr || state->tensor == nullptr) { continue; diff --git a/src/model_manager.h b/src/model_manager.h index 3dc633864..d6910d1de 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -61,12 +61,24 @@ class ModelManager : public RunnerWeightManager { std::vector> staged_tensors; }; + struct PrefetchBlock { + std::vector 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> staged_tensors; + }; + ModelLoader model_loader_; std::vector> tensor_states_; std::map tensor_states_by_name_; std::vector> params_storage_blocks_; std::vector> compute_staging_blocks_; std::map split_buffer_types_; + std::map> prefetch_blocks_; + std::map prefetch_backends_; bool warned_split_lora_skip_ = false; std::set common_ignore_tensors_; std::vector loras_; @@ -79,6 +91,13 @@ class ModelManager : public RunnerWeightManager { void finish_compute_backend_usage(const std::vector& 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& tensors, std::vector& required_states) const; bool should_ignore(const TensorState& state) const; @@ -182,6 +201,11 @@ class ModelManager : public RunnerWeightManager { bool prepare_params(const std::vector& tensors) override; void release_compute_backend_params(const std::vector& tensors) override; void release_params_backend_params(const std::vector& tensors) override; + bool prefetch_params(uintptr_t owner_id, + const std::vector& tensors) override; + bool activate_prefetched_params(uintptr_t owner_id, + const std::vector& tensors) override; + void clear_prefetched_params(uintptr_t owner_id) override; }; #endif // __MODEL_MANAGER_H__ diff --git a/src/model_manager_prefetch.cpp b/src/model_manager_prefetch.cpp new file mode 100644 index 000000000..3d39a9982 --- /dev/null +++ b/src/model_manager_prefetch.cpp @@ -0,0 +1,271 @@ +#include "model_manager.h" + +#include +#include + +#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& tensors) { + clear_prefetched_params(owner_id); + if (tensors.empty()) { + return true; + } + + std::vector required_states; + if (!resolve_required_tensor_states(tensors, required_states) || + !load_tensors_to_params_backend(required_states)) { + return false; + } + + std::vector 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(); + 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& tensors) { + std::vector 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 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(); + 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 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(); +} diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index d2193e14c..217212cc0 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -248,8 +248,9 @@ class StableDiffusionGGML { 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 @@ class StableDiffusionGGML { 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 @@ class StableDiffusionGGML { 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 @@ class StableDiffusionGGML { 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), diff --git a/src/weight_manager.h b/src/weight_manager.h index 82f6d03e4..c54e33c45 100644 --- a/src/weight_manager.h +++ b/src/weight_manager.h @@ -1,6 +1,7 @@ #ifndef __WEIGHT_MANAGER_H__ #define __WEIGHT_MANAGER_H__ +#include #include #include "ggml-backend.h" @@ -14,6 +15,11 @@ struct RunnerWeightManager { virtual bool prepare_params(const std::vector& tensors) = 0; virtual void release_compute_backend_params(const std::vector& tensors) = 0; virtual void release_params_backend_params(const std::vector& tensors) = 0; + virtual bool prefetch_params(uintptr_t owner_id, + const std::vector& tensors) = 0; + virtual bool activate_prefetched_params(uintptr_t owner_id, + const std::vector& tensors) = 0; + virtual void clear_prefetched_params(uintptr_t owner_id) = 0; }; #endif // __WEIGHT_MANAGER_H__