From 538d8cdd53da6a89f40b06295c9e8f07dfedb8b5 Mon Sep 17 00:00:00 2001 From: Piotr Wilkin Date: Fri, 3 Jul 2026 16:21:02 +0200 Subject: [PATCH 1/8] feat: add multi-device layer split (--backend "diffusion=cuda0&cuda1") A --backend module assignment can now list several devices separated by '&'. The module's transformer blocks are partitioned into contiguous ranges sized proportionally to each device's free memory (minus a fixed compute headroom) and registered with the ModelManager with per-tensor compute backends; the existing allocation/staging/LoRA/residency machinery handles the weights unchanged. The module's graphs execute on a ggml_backend_sched spanning the devices, pinning each node to the device of the most recently consumed weight (view ops are never pinned) and splitting each graph exactly once. Supported for the diffusion and te modules; for te the dominant encoder (t5xxl or the LLM) splits while small sub-runners stay on the main device. Graph-cut segmentation and --stream-layers are disabled for split modules. Adds --list-devices to print the ggml device names accepted by the backend specs. Manual placement only; row/tensor split and auto-fit are follow-ups. Co-Authored-By: Claude Fable 5 --- docs/backend.md | 34 ++++ examples/common/common.cpp | 9 ++ include/stable-diffusion.h | 5 + src/conditioning/conditioner.hpp | 83 ++++++++++ src/core/ggml_extend.hpp | 218 +++++++++++++++++++++++++- src/core/ggml_extend_backend.cpp | 82 +++++++++- src/core/ggml_extend_backend.h | 10 ++ src/core/util.cpp | 14 ++ src/stable-diffusion.cpp | 259 ++++++++++++++++++++++++++++++- 9 files changed, 706 insertions(+), 8 deletions(-) diff --git a/docs/backend.md b/docs/backend.md index 29ac8031b..805143a3a 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -51,6 +51,40 @@ Module names are case-insensitive. Hyphens and underscores in module names are i sd-cli -m model.safetensors -p "a cat" --backend all=cuda0,te=cpu ``` +## Multiple devices per module (layer split) + +A `--backend` module assignment can list several devices separated by `&`: + +```shell +sd-cli -m model.safetensors -p "a cat" --backend "diffusion=cuda0&cuda1" +``` + +The module's transformer blocks are then distributed across the listed devices +in contiguous ranges sized proportionally to each device's free memory (minus a +compute-buffer headroom of about 2 GiB per device), and the +module's graphs are executed with a `ggml_backend_sched` that runs each block +on the device holding its weights, copying the residual stream at the range +boundaries. The first device in the list is the module's main device: it also +holds the non-block tensors (embeddings, final norms, small sub-runners such as +CLIP models or projectors) and the graph inputs/outputs. + +Layer split is supported for the `diffusion` and `te` modules. For `te` it +applies to the dominant text encoder (`t5xxl` or the LLM); other modules accept +only a single device. If the module has no recognizable transformer blocks, the +assignment falls back to the first listed device. + +`--params-backend` accepts no device lists. If the module has no explicit +params assignment, each block range's parameters are loaded directly to (and, +with `--params-backend diffusion=disk`, released directly from) its own device; +an explicit assignment such as `te=cpu` keeps the parameters on that backend +and stages each range to its device on demand. + +Layer split cannot be combined with `--max-vram` graph-cut segmentation or +`--stream-layers` for the split module; those are single-device mechanisms and +are disabled for it. + +Use `--list-devices` to see the device names available on the system. + ## Modules | Module | Purpose | Accepted names | diff --git a/examples/common/common.cpp b/examples/common/common.cpp index ac1cf32f6..acea9171e 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -663,6 +663,15 @@ ArgOptions SDContextParams::get_options() { "but it usually offers faster inference speed and, in some cases, lower memory usage. " "The at_runtime mode, on the other hand, is exactly the opposite.", on_lora_apply_mode_arg}, + {"", + "--list-devices", + "list available ggml backend devices (one 'namedescription' per line) and exit; " + "the names are the device names accepted by --backend and --params-backend", + [](int /*argc*/, const char** /*argv*/, int /*index*/) { + sd_list_devices(); + std::exit(0); + return 0; + }}, }; return options; diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 26bb61ef8..f04e8cddb 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -534,6 +534,11 @@ SD_API void disable_imatrix_collection(void); SD_API const char* sd_commit(void); SD_API const char* sd_version(void); +// List available ggml backend devices to stdout, one `namedescription` +// per line. The names are the device names accepted by the --backend / +// --params-backend assignment specs. +SD_API void sd_list_devices(void); + // for C API, caller needs to call free_sd_images to free the memory after use // This helps avoid CRT problems on Windows when memory is allocated in the library but freed in the caller, which may use a different CRT. SD_API void free_sd_images(sd_image_t* result_images, int num_images); diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index d63303a82..3627a166c 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -116,6 +116,15 @@ struct Conditioner { virtual void get_param_tensors(std::map& tensors) = 0; virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {} virtual void set_stream_layers_enabled(bool enabled) {} + // Multi-device (layer split) hooks. When the TE module is assigned more + // than one runtime backend, set_runtime_backends hands the backend list to + // the conditioner's dominant runner (t5 / llm), and + // get_layer_split_param_tensors returns that runner's tensors — the ones + // whose transformer blocks may be distributed across the devices. The + // defaults mean "no layer split support"; smaller sub-runners (clip_l, + // projectors, ...) always stay on the main backend. + virtual void set_runtime_backends(const std::vector& backends) {} + virtual void get_layer_split_param_tensors(std::map& tensors) {} virtual void set_flash_attention_enabled(bool enabled) = 0; virtual void set_weight_adapter(const std::shared_ptr& adapter) {} virtual void runner_done() {} @@ -635,6 +644,18 @@ struct SD3CLIPEmbedder : public Conditioner { } } + void set_runtime_backends(const std::vector& backends) override { + if (t5) { + t5->set_runtime_backends(backends); + } + } + + void get_layer_split_param_tensors(std::map& tensors) override { + if (t5) { + t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + } + } + void set_flash_attention_enabled(bool enabled) override { if (clip_l) { clip_l->set_flash_attention_enabled(enabled); @@ -994,6 +1015,18 @@ struct FluxCLIPEmbedder : public Conditioner { } } + void set_runtime_backends(const std::vector& backends) override { + if (t5) { + t5->set_runtime_backends(backends); + } + } + + void get_layer_split_param_tensors(std::map& tensors) override { + if (t5) { + t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + } + } + void set_flash_attention_enabled(bool enabled) override { if (clip_l) { clip_l->set_flash_attention_enabled(enabled); @@ -1226,6 +1259,18 @@ struct T5CLIPEmbedder : public Conditioner { } } + void set_runtime_backends(const std::vector& backends) override { + if (t5) { + t5->set_runtime_backends(backends); + } + } + + void get_layer_split_param_tensors(std::map& tensors) override { + if (t5) { + t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + } + } + void set_flash_attention_enabled(bool enabled) override { if (t5) { t5->set_flash_attention_enabled(enabled); @@ -1418,6 +1463,18 @@ struct MiniT2IConditioner : public Conditioner { } } + void set_runtime_backends(const std::vector& backends) override { + if (t5) { + t5->set_runtime_backends(backends); + } + } + + void get_layer_split_param_tensors(std::map& tensors) override { + if (t5) { + t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + } + } + void set_flash_attention_enabled(bool enabled) override { if (t5) { t5->set_flash_attention_enabled(enabled); @@ -1502,6 +1559,14 @@ struct AnimaConditioner : public Conditioner { llm->set_stream_layers_enabled(enabled); } + void set_runtime_backends(const std::vector& backends) override { + llm->set_runtime_backends(backends); + } + + void get_layer_split_param_tensors(std::map& tensors) override { + llm->get_param_tensors(tensors, "text_encoders.llm"); + } + void set_flash_attention_enabled(bool enabled) override { llm->set_flash_attention_enabled(enabled); } @@ -1647,6 +1712,14 @@ struct LLMEmbedder : public Conditioner { llm->set_stream_layers_enabled(enabled); } + void set_runtime_backends(const std::vector& backends) override { + llm->set_runtime_backends(backends); + } + + void get_layer_split_param_tensors(std::map& tensors) override { + llm->get_param_tensors(tensors, "text_encoders.llm"); + } + void set_flash_attention_enabled(bool enabled) override { llm->set_flash_attention_enabled(enabled); } @@ -2316,6 +2389,16 @@ struct LTXAVEmbedder : public Conditioner { projector->set_max_graph_vram_bytes(max_vram_bytes); } + // Layer split applies to the heavy LLM only; the small projector stays on + // the main backend. + void set_runtime_backends(const std::vector& backends) override { + llm->set_runtime_backends(backends); + } + + void get_layer_split_param_tensors(std::map& tensors) override { + llm->get_param_tensors(tensors, "text_encoders.llm"); + } + void set_weight_adapter(const std::shared_ptr& adapter) override { llm->set_weight_adapter(adapter); projector->set_weight_adapter(adapter); diff --git a/src/core/ggml_extend.hpp b/src/core/ggml_extend.hpp index c4138a235..b80e059f3 100644 --- a/src/core/ggml_extend.hpp +++ b/src/core/ggml_extend.hpp @@ -1746,6 +1746,16 @@ struct GGMLRunner { bool stream_layers_enabled = false; size_t observed_max_effective_budget_ = 0; + // Multi-device execution (layer split): when the module is assigned more + // than one runtime backend (--backend "diffusion=cuda0&cuda1"), the + // runner's weights are registered with per-tensor compute backends and the + // graph is executed with a ggml_backend_sched that routes each op to the + // device holding its weights. runtime_backend stays the main device. + std::vector extra_runtime_backends; // borrowed (SDBackendManager-owned) + ggml_backend_sched_t sched = nullptr; // owned, multi-device only + ggml_backend_t cpu_fallback_backend = nullptr; // owned, sched requires a trailing CPU backend + bool multi_device_eval_callback_warned = false; + std::shared_ptr weight_adapter = nullptr; std::weak_ptr weight_manager; std::unordered_set kept_compute_param_tensor_set; @@ -2013,7 +2023,140 @@ struct GGMLRunner { return true; } + // Build the multi-device scheduler on first use. Backend order: the main + // runtime backend, the extra device backends, then a CPU backend last + // (ggml_backend_sched_new requires the final backend to be CPU). An + // explicit buffer-type array is passed instead of nullptr: the sched uses + // these in buffer_supported() to decide whether a cross-backend src needs + // a copy, and with the synthesized defaults CUDA devices can spuriously + // report supporting each other's buffers, skipping a required copy. The + // CPU slot uses the main device's host (pinned) buffer type when + // available, as llama.cpp does. + bool ensure_sched(ggml_cgraph* gf) { + if (sched != nullptr) { + return true; + } + std::vector backends; + backends.reserve(extra_runtime_backends.size() + 2); + backends.push_back(runtime_backend); + for (ggml_backend_t backend : extra_runtime_backends) { + backends.push_back(backend); + } + if (cpu_fallback_backend == nullptr && !sd_backend_is_cpu(runtime_backend)) { + cpu_fallback_backend = sd_backend_cpu_init(); + } + if (cpu_fallback_backend != nullptr) { + backends.push_back(cpu_fallback_backend); + } + + std::vector bufts; + bufts.reserve(backends.size()); + ggml_backend_dev_t main_dev = ggml_backend_get_device(runtime_backend); + for (ggml_backend_t backend : backends) { + ggml_backend_buffer_type_t buft = nullptr; + if (backend == cpu_fallback_backend && main_dev != nullptr) { + buft = ggml_backend_dev_host_buffer_type(main_dev); + } + if (buft == nullptr) { + buft = ggml_backend_get_default_buffer_type(backend); + } + bufts.push_back(buft); + } + + size_t graph_size = MAX_GRAPH_SIZE; + if (gf != nullptr) { + graph_size = std::max(graph_size, (size_t)ggml_graph_n_nodes(gf)); + } + sched = ggml_backend_sched_new(backends.data(), + bufts.data(), + (int)backends.size(), + graph_size, + /*parallel=*/false, + /*op_offload=*/false); + if (sched == nullptr) { + LOG_ERROR("%s: failed to create backend sched", get_desc().c_str()); + return false; + } + return true; + } + + // Map a weight tensor to the runner backend whose device holds its data, + // or nullptr for non-weight tensors and weights outside the sched devices + // (e.g. mmap'd or pinned host buffers). + ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const { + if (tensor == nullptr || tensor->buffer == nullptr) { + return nullptr; + } + if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS || + ggml_backend_buffer_is_host(tensor->buffer)) { + return nullptr; + } + ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer)); + if (dev == nullptr) { + return nullptr; + } + if (ggml_backend_get_device(runtime_backend) == dev) { + return runtime_backend; + } + for (ggml_backend_t backend : extra_runtime_backends) { + if (ggml_backend_get_device(backend) == dev) { + return backend; + } + } + return nullptr; + } + + // Pin compute nodes to their layer's device. The sched anchors + // weight-bearing ops (matmuls) to the weight's device, but weightless ops + // (norm, residual add, cont) have no anchor and its placement heuristic + // can land them on the wrong device, which is then read without a + // cross-device copy. llama.cpp pins each layer-boundary norm to the + // layer's device for the same reason (llama_context::graph_compute). This + // generalizes that: walk the graph in execution order, track the device of + // the most recently consumed weight (= the current layer's device), and + // pin every node to it, so the sched only copies the residual stream at + // layer boundaries. View ops must never be pinned: a view assigned to a + // different backend than its view_src's data makes the sched skip the + // cross-device copy for consumers. + void pin_multi_device_nodes(ggml_cgraph* gf) { + if (sched == nullptr || gf == nullptr) { + return; + } + ggml_backend_t current = runtime_backend; + const int n_nodes = ggml_graph_n_nodes(gf); + for (int i = 0; i < n_nodes; i++) { + ggml_tensor* node = ggml_graph_node(gf, i); + for (int s = 0; s < GGML_MAX_SRC; s++) { + ggml_backend_t weight_backend = backend_for_weight(node->src[s]); + if (weight_backend != nullptr) { + current = weight_backend; + } + } + if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || + node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) { + continue; + } + if (ggml_backend_supports_op(current, node)) { + ggml_backend_sched_set_tensor_backend(sched, node, current); + } + } + } + + bool is_multi_device() const { + return !extra_runtime_backends.empty(); + } + bool alloc_compute_buffer(ggml_cgraph* gf) { + if (is_multi_device()) { + // The sched replaces the gallocr. Do NOT ggml_backend_sched_reserve + // the graph here: reserve runs split_graph, which rewires the + // graph's src pointers to sched-internal copy tensors, and the + // later ggml_backend_sched_alloc_graph would split the already + // rewired graph, silently corrupting every cross-backend input. A + // graph must be split at most once; the alloc in execute_graph + // performs the real allocation. + return ensure_sched(gf); + } if (compute_allocr != nullptr) { return true; } @@ -2229,12 +2372,14 @@ struct GGMLRunner { plan.valid && max_graph_vram_bytes > 0 && plan.segments.size() > 1 && - !sd_backend_is_cpu(runtime_backend); + !sd_backend_is_cpu(runtime_backend) && + !is_multi_device(); } bool can_attempt_graph_cut_segmented_compute() const { return max_graph_vram_bytes > 0 && - !sd_backend_is_cpu(runtime_backend); + !sd_backend_is_cpu(runtime_backend) && + !is_multi_device(); } bool resolve_graph_cut_plan(ggml_cgraph* gf, @@ -2490,7 +2635,14 @@ struct GGMLRunner { }; ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer); - if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) { + if (is_multi_device()) { + ggml_backend_sched_reset(sched); + pin_multi_device_nodes(gf); // reset clears the pins; re-apply before alloc + if (!ggml_backend_sched_alloc_graph(sched, gf)) { + LOG_ERROR("%s sched alloc compute graph failed", get_desc().c_str()); + return std::nullopt; + } + } else if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) { LOG_ERROR("%s alloc compute graph failed", get_desc().c_str()); return std::nullopt; } @@ -2499,11 +2651,27 @@ struct GGMLRunner { if (sd_backend_is_cpu(runtime_backend)) { sd_backend_cpu_set_n_threads(runtime_backend, n_threads); } + if (cpu_fallback_backend != nullptr) { + sd_backend_cpu_set_n_threads(cpu_fallback_backend, n_threads); + } - ggml_status status = sd_backend_graph_compute_with_eval_callback(runtime_backend, + ggml_status status; + if (is_multi_device()) { + if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) { + LOG_WARN("%s: eval callback is not supported with multiple runtime backends; ignoring", + get_desc().c_str()); + multi_device_eval_callback_warned = true; + } + status = ggml_backend_sched_graph_compute(sched, gf); + if (status == GGML_STATUS_SUCCESS) { + ggml_backend_sched_synchronize(sched); + } + } else { + status = sd_backend_graph_compute_with_eval_callback(runtime_backend, gf, sd_get_backend_eval_callback(), sd_get_backend_eval_callback_data()); + } if (status != GGML_STATUS_SUCCESS) { LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status)); return std::nullopt; @@ -2680,6 +2848,10 @@ struct GGMLRunner { free_params_ctx(); free_compute_ctx(); free_cache_ctx_and_buffer(); + if (cpu_fallback_backend != nullptr) { + ggml_backend_free(cpu_fallback_backend); + cpu_fallback_backend = nullptr; + } } virtual GGMLRunnerContext get_context() { @@ -2720,10 +2892,20 @@ struct GGMLRunner { ggml_gallocr_free(compute_allocr); compute_allocr = nullptr; } + if (sched != nullptr) { + ggml_backend_sched_free(sched); + sched = nullptr; + } } // do copy after alloc graph void set_backend_tensor_data(ggml_tensor* tensor, const void* data) { + if (is_multi_device()) { + // The sched only assigns a backend (and thus a buffer) to tensors + // that participate in the graph; flag standalone data tensors as + // inputs so they get one. + ggml_set_input(tensor); + } backend_tensor_data_map[tensor] = data; } @@ -2859,8 +3041,36 @@ struct GGMLRunner { } void set_stream_layers_enabled(bool enabled) { + if (enabled && is_multi_device()) { + LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring", + get_desc().c_str()); + return; + } stream_layers_enabled = enabled; } + + // Hand the runner its module's runtime backend list (layer split). The + // list is the SDBackendManager assignment order with backends[0] = the + // runner's main runtime backend; only the extra devices are kept. Layer + // streaming works through graph-cut segmentation, which is a + // single-backend mechanism, so it is turned off for a split runner. + void set_runtime_backends(const std::vector& backends) { + extra_runtime_backends.clear(); + for (ggml_backend_t backend : backends) { + if (backend == nullptr || backend == runtime_backend) { + continue; + } + if (std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), backend) == + extra_runtime_backends.end()) { + extra_runtime_backends.push_back(backend); + } + } + if (is_multi_device() && stream_layers_enabled) { + LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring", + get_desc().c_str()); + stream_layers_enabled = false; + } + } }; class GGMLBlock { diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index f29bdb696..aaf01da42 100644 --- a/src/core/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -669,8 +669,49 @@ void SDBackendManager::reset() { params_assignment_ = {}; } +// Split an '&'-separated device list ("cuda0&cuda1") into its entries. +// A plain single name yields one entry. +static std::vector split_device_list(const std::string& value) { + std::vector names; + for (const std::string& raw : split_copy(value, '&')) { + const std::string name = trim_copy(raw); + if (!name.empty()) { + names.push_back(name); + } + } + return names; +} + +static std::string primary_device_name(const std::string& value) { + std::vector names = split_device_list(value); + return names.empty() ? std::string() : names.front(); +} + ggml_backend_t SDBackendManager::runtime_backend(SDBackendModule module) { - return init_cached_backend(runtime_assignment_.get(module)); + return init_cached_backend(primary_device_name(runtime_assignment_.get(module))); +} + +std::vector SDBackendManager::runtime_backends(SDBackendModule module) { + std::vector backends; + for (const std::string& name : split_device_list(runtime_assignment_.get(module))) { + ggml_backend_t backend = init_cached_backend(name); + if (backend == nullptr) { + LOG_ERROR("failed to initialize backend '%s' for module %s", + name.c_str(), + sd_backend_module_name(module)); + continue; + } + if (std::find(backends.begin(), backends.end(), backend) == backends.end()) { + backends.push_back(backend); + } + } + if (backends.empty()) { + ggml_backend_t backend = runtime_backend(module); + if (backend != nullptr) { + backends.push_back(backend); + } + } + return backends; } ggml_backend_t SDBackendManager::params_backend(SDBackendModule module) { @@ -696,6 +737,10 @@ bool SDBackendManager::params_backend_is_disk(SDBackendModule module) const { return is_disk_backend_token(params_assignment_.get(module)); } +bool SDBackendManager::params_backend_follows_runtime(SDBackendModule module) const { + return params_assignment_.get(module).empty(); +} + bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule module) { ggml_backend_t backend = runtime_backend(module); if (backend == nullptr) { @@ -729,7 +774,7 @@ bool SDBackendManager::init(const char* backend_spec, } bool SDBackendManager::validate(std::string* error) const { - auto validate_runtime_name = [&](const std::string& name) -> bool { + auto validate_single_runtime_name = [&](const std::string& name) -> bool { if (is_default_backend_token(name)) { return true; } @@ -747,11 +792,42 @@ bool SDBackendManager::validate(std::string* error) const { } return false; }; + auto validate_runtime_name = [&](const std::string& name) -> bool { + // A runtime assignment may be an '&'-separated device list. + if (name.find('&') == std::string::npos) { + return validate_single_runtime_name(name); + } + std::vector names = split_device_list(name); + if (names.empty()) { + if (error != nullptr) { + *error = "invalid backend device list '" + name + "'"; + } + return false; + } + for (const std::string& entry : names) { + if (is_default_backend_token(entry)) { + if (error != nullptr) { + *error = "default backend token is not allowed in a device list '" + name + "'"; + } + return false; + } + if (!validate_single_runtime_name(entry)) { + return false; + } + } + return true; + }; auto validate_params_name = [&](const std::string& name) -> bool { if (is_disk_backend_token(name)) { return true; } - return validate_runtime_name(name); + if (name.find('&') != std::string::npos) { + if (error != nullptr) { + *error = "params_backend does not accept device lists ('" + name + "')"; + } + return false; + } + return validate_single_runtime_name(name); }; if (!validate_runtime_name(runtime_assignment_.default_name) || diff --git a/src/core/ggml_extend_backend.h b/src/core/ggml_extend_backend.h index 19b71d432..b5f8ae00c 100644 --- a/src/core/ggml_extend_backend.h +++ b/src/core/ggml_extend_backend.h @@ -6,6 +6,7 @@ #include #include #include +#include #include "ggml-backend.h" #include "ggml.h" @@ -58,9 +59,18 @@ class SDBackendManager { ggml_backend_t runtime_backend(SDBackendModule module); ggml_backend_t params_backend(SDBackendModule module); + // All runtime backends assigned to a module, in assignment order. A module + // is assigned more than one backend with an '&'-separated device list, + // e.g. --backend "diffusion=cuda0&cuda1". The first entry is the main + // backend (the same one runtime_backend() returns); duplicates are folded. + std::vector runtime_backends(SDBackendModule module); + bool runtime_backend_is_cpu(SDBackendModule module); bool params_backend_is_cpu(SDBackendModule module); bool params_backend_is_disk(SDBackendModule module) const; + // True when the module has no explicit params assignment, so params + // placement follows the runtime backend (per device for layer splits). + bool params_backend_follows_runtime(SDBackendModule module) const; bool runtime_backend_supports_host_buffer(SDBackendModule module); private: diff --git a/src/core/util.cpp b/src/core/util.cpp index 6d2479f9f..362a20529 100644 --- a/src/core/util.cpp +++ b/src/core/util.cpp @@ -25,6 +25,7 @@ #include #endif +#include "ggml-backend.h" #include "ggml.h" #include "stable-diffusion.h" @@ -997,3 +998,16 @@ std::vector> split_quotation_attention( } return result; } + +void sd_list_devices(void) { + if (ggml_backend_dev_count() == 0) { + // dynamic-backend builds discover their backend modules at runtime + ggml_backend_load_all(); + } + for (size_t i = 0; i < ggml_backend_dev_count(); i++) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + const char* name = ggml_backend_dev_name(dev); + const char* desc = ggml_backend_dev_description(dev); + printf("%s\t%s\n", name ? name : "", desc ? desc : ""); + } +} diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index eb592f414..6dd812e87 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -2,7 +2,9 @@ #include #include #include +#include #include +#include #include #include "core/ggml_extend.hpp" @@ -170,6 +172,155 @@ static float get_cache_reuse_threshold(const sd_cache_params_t& params) { /*=============================================== StableDiffusionGGML ================================================*/ +// Parse a transformer block index out of a weight name, or -1 if the tensor +// does not belong to a block ("model.diffusion_model.transformer_blocks.12.*" +// -> 12, "text_encoders.llm.model.layers.30.*" -> 30). +static int tensor_block_index(const std::string& name) { + static const char* block_keywords[] = {"transformer_blocks.", "joint_blocks.", "double_blocks.", + "single_blocks.", "blocks.", "block.", "layers."}; + for (const char* keyword : block_keywords) { + size_t pos = name.find(keyword); + if (pos == std::string::npos) { + continue; + } + pos += strlen(keyword); + size_t end = pos; + while (end < name.size() && name[end] >= '0' && name[end] <= '9') { + end++; + } + if (end > pos && (end == name.size() || name[end] == '.')) { + return atoi(name.substr(pos, end - pos).c_str()); + } + } + return -1; +} + +static std::string backend_device_display_name(ggml_backend_t backend) { + ggml_backend_dev_t dev = ggml_backend_get_device(backend); + const char* name = dev != nullptr ? ggml_backend_dev_name(dev) : ggml_backend_name(backend); + return name != nullptr ? name : "unknown"; +} + +// Distribute a module's param tensors across its runtime backends for a layer +// split. The tensors in split_tensors (the module's dominant transformer) are +// grouped into contiguous block ranges sized proportionally to each device's +// free memory; every other tensor (embeddings, final norms, small sub-runners) +// stays on the main backend, which pays for them with a smaller block range. +// Returns one name->tensor map per backend, in backend order. When no blocks +// are found, everything lands in the first partition (no split). +static std::vector> partition_layer_split_tensors( + const std::string& desc, + const std::map& tensors, + const std::map& split_tensors, + const std::vector& backends) { + std::vector> partitions(backends.size()); + + std::map block_bytes; + int64_t total_block_bytes = 0; + int64_t other_bytes = 0; + int n_blocks = 0; + for (const auto& kv : tensors) { + int64_t bytes = (int64_t)ggml_nbytes(kv.second); + int idx = split_tensors.count(kv.first) != 0 ? tensor_block_index(kv.first) : -1; + if (idx >= 0) { + block_bytes[idx] += bytes; + total_block_bytes += bytes; + n_blocks = std::max(n_blocks, idx + 1); + } else { + other_bytes += bytes; + } + } + if (n_blocks == 0) { + LOG_WARN("%s: no transformer blocks found for a layer split; keeping the module on %s", + desc.c_str(), + backend_device_display_name(backends[0]).c_str()); + partitions[0] = tensors; + return partitions; + } + + // Weight each device by its free memory minus a fixed compute headroom: + // every device participating in a layer split also hosts a share of the + // scheduler's compute buffers (the activations of its block range), which + // for large models runs into gigabytes; without the headroom the weight + // share fills the device exactly and the compute allocation OOMs. The main + // backend additionally holds the non-block tensors, so those are + // subtracted from its block budget below. + constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024; + std::vector device_weights(backends.size(), 1.0); + double weight_sum = 0.0; + for (size_t i = 0; i < backends.size(); i++) { + ggml_backend_dev_t dev = ggml_backend_get_device(backends[i]); + size_t free_bytes = 0, total_bytes = 0; + if (dev != nullptr) { + ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); + } + // Keep a small share even for tight devices instead of dropping them. + int64_t usable_bytes = std::max((int64_t)free_bytes - compute_headroom_bytes, + (int64_t)free_bytes / 8); + device_weights[i] = usable_bytes > 0 ? (double)usable_bytes : 1.0; + weight_sum += device_weights[i]; + } + + std::vector block_budgets(backends.size(), 0); + for (size_t i = 0; i < backends.size(); i++) { + int64_t budget = (int64_t)((double)(total_block_bytes + other_bytes) * device_weights[i] / weight_sum); + if (i == 0) { + budget = std::max(budget - other_bytes, 0); + } + block_budgets[i] = budget; + } + + // Assign contiguous block ranges: boundaries[i] is the first block index + // NOT owned by backend i. Every backend keeps at least one block while + // blocks remain, and the last backend absorbs the remainder. + std::vector boundaries(backends.size(), n_blocks); + size_t current = 0; + int64_t used = 0; + for (int b = 0; b < n_blocks; b++) { + int64_t bytes = block_bytes.count(b) != 0 ? block_bytes[b] : 0; + if (current + 1 < backends.size() && used > 0 && used + bytes > block_budgets[current]) { + boundaries[current] = b; + current++; + used = 0; + } + used += bytes; + } + + for (const auto& kv : tensors) { + size_t target = 0; + int idx = split_tensors.count(kv.first) != 0 ? tensor_block_index(kv.first) : -1; + if (idx >= 0) { + while (target < boundaries.size() && idx >= boundaries[target]) { + target++; + } + target = std::min(target, backends.size() - 1); + } + partitions[target][kv.first] = kv.second; + } + + int range_start = 0; + for (size_t i = 0; i < backends.size(); i++) { + int range_end = boundaries[i]; + LOG_INFO("%s layer split: %s <- blocks [%d, %d)%s", + desc.c_str(), + backend_device_display_name(backends[i]).c_str(), + range_start, + range_end, + i == 0 ? " + non-block tensors" : ""); + range_start = range_end; + } + return partitions; +} + +// Detects the multi-device hook shared by GGMLRunner and Conditioner; runner +// types without it (e.g. generation extensions) never layer-split. +template +struct has_set_runtime_backends : std::false_type {}; +template +struct has_set_runtime_backends().set_runtime_backends( + std::declval&>()))>> : std::true_type {}; + static_assert(std::atomic::is_always_lock_free, "sd_cancel_mode_t must be lock-free"); @@ -275,14 +426,120 @@ class StableDiffusionGGML { if (model_manager == nullptr) { return true; } + ModelManager::ResidencyMode residency_mode = + backend_manager.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend; + + std::vector module_backends = backend_manager.runtime_backends(module); + if (module_backends.size() > 1) { + if constexpr (has_set_runtime_backends::value) { + if (module == SDBackendModule::DIFFUSION || module == SDBackendModule::TE) { + return register_layer_split_runner_params(desc, + model, + module, + module_backends, + std::move(group_tensors), + residency_mode, + params_mem_size); + } + } + LOG_WARN("%s module does not support multiple runtime backends; using %s", + sd_backend_module_name(module), + backend_device_display_name(module_backends[0]).c_str()); + } return model_manager->register_param_tensors(desc, std::move(group_tensors), - backend_manager.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend, + residency_mode, backend_for(module), params_backend_for(module), params_mem_size); } + // Layer split registration: partition the module's tensors into contiguous + // transformer-block ranges (one per runtime backend) and register each + // range with that backend as its per-tensor compute backend — the + // ModelManager's existing allocation/staging/LoRA paths already group by + // backend and buffer type, so no special weight handling is needed. When + // the module has no explicit params assignment (or uses disk residency), + // each range's params follow its own device so weights load straight to + // (and release straight from) the device that computes with them. + template + bool register_layer_split_runner_params(const std::string& desc, + const std::shared_ptr& model, + SDBackendModule module, + const std::vector& module_backends, + std::map group_tensors, + ModelManager::ResidencyMode residency_mode, + size_t* params_mem_size) { + bool has_cpu_device = false; + for (ggml_backend_t backend : module_backends) { + has_cpu_device = has_cpu_device || sd_backend_is_cpu(backend); + } + if (has_cpu_device) { + // The scheduler reserves the CPU slot for its fallback backend, and + // CPU weight participation is what --params-backend =cpu is + // for; a CPU device in a split list is almost certainly a mistake. + LOG_WARN( + "%s: layer split across a CPU device is not supported; using %s " + "(use --params-backend %s=cpu to keep weights in RAM)", + desc.c_str(), + backend_device_display_name(module_backends[0]).c_str(), + sd_backend_module_name(module)); + return model_manager->register_param_tensors(desc, + std::move(group_tensors), + residency_mode, + module_backends[0], + params_backend_for(module), + params_mem_size); + } + + std::map split_tensors; + if constexpr (std::is_base_of_v) { + model->get_layer_split_param_tensors(split_tensors); + } else { + split_tensors = group_tensors; + } + + auto partitions = partition_layer_split_tensors(desc, group_tensors, split_tensors, module_backends); + bool is_split = false; + for (size_t i = 1; i < partitions.size(); i++) { + if (!partitions[i].empty()) { + is_split = true; + break; + } + } + if (!is_split) { + return model_manager->register_param_tensors(desc, + std::move(group_tensors), + residency_mode, + module_backends[0], + params_backend_for(module), + params_mem_size); + } + + model->set_runtime_backends(module_backends); + const bool params_follow_runtime = backend_manager.params_backend_follows_runtime(module) || + backend_manager.params_backend_is_disk(module); + for (size_t i = 0; i < module_backends.size(); i++) { + if (partitions[i].empty()) { + continue; + } + ggml_backend_t partition_params_backend = + params_follow_runtime ? module_backends[i] : params_backend_for(module); + if (partition_params_backend == nullptr) { + return false; + } + if (!model_manager->register_param_tensors(desc, + std::move(partitions[i]), + residency_mode, + module_backends[i], + partition_params_backend, + params_mem_size)) { + return false; + } + } + return true; + } + bool init_backend() { std::string error; if (!backend_manager.init(backend_spec.c_str(), From a9f261008ad2a128e282c8320ea6c7549339bd72 Mon Sep 17 00:00:00 2001 From: Piotr Wilkin Date: Fri, 3 Jul 2026 16:43:17 +0200 Subject: [PATCH 2/8] feat: add row split support via backend split buffer types (--split-mode row) --split-mode selects how a module assigned multiple runtime devices distributes its weights: layer (default) or row. In row mode the module keeps executing on its main device while its transformer-block matmul weights are allocated in the backend's row-split buffer type (resolved through the "ggml_backend_split_buffer_type" proc, CUDA only for now), which slices each weight's rows across the devices in proportion to free memory and runs the matmuls multi-GPU internally. The ModelManager owns the split buffer types (set_split_buffer_type): params_buffer_type_for returns the split type for eligible tensors when params live on the compute backend, and the staging path groups by (backend, buffer type) and allocates with ggml_backend_alloc_ctx_tensors_from_buft so cpu/disk params residency stages straight into split buffers. Eligibility is limited to contiguous 2D weights of at least 256x256 inside transformer blocks: anything else is consumed by non-matmul ops or sliced into views, which split buffers do not support. Direct LoRA application skips row-split tensors; the automatic LoRA mode selects at_runtime when row split is active. Falls back to a layer split when the backend has no split buffer type or the devices span registries. Co-Authored-By: Claude Fable 5 --- docs/backend.md | 30 +++++++ examples/common/common.cpp | 9 ++ examples/common/common.h | 1 + include/stable-diffusion.h | 1 + src/core/ggml_extend_backend.cpp | 64 ++++++++++++- src/core/ggml_extend_backend.h | 24 +++++ src/model_manager.cpp | 86 +++++++++++++++--- src/model_manager.h | 22 ++++- src/stable-diffusion.cpp | 148 ++++++++++++++++++++++++++++++- src/upscaler.cpp | 1 + 10 files changed, 369 insertions(+), 17 deletions(-) diff --git a/docs/backend.md b/docs/backend.md index 805143a3a..52a9acf20 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -85,6 +85,36 @@ are disabled for it. Use `--list-devices` to see the device names available on the system. +### Row split (`--split-mode row`) + +`--split-mode` selects how a multi-device module distributes its weights: +`layer` (the default, described above) or `row`. It accepts a single mode or +per-module assignments: + +```shell +sd-cli -m model.safetensors -p "a cat" --backend "diffusion=cuda0&cuda1" --split-mode row +sd-cli -m model.safetensors -p "a cat" --backend "diffusion=cuda0&cuda1,te=cuda0&cuda1" --split-mode diffusion=row,te=layer +``` + +In row mode the module keeps executing on its main (first listed) device, but +its transformer-block matmul weights are allocated in the backend's row-split +buffer type, which slices each weight's rows across the listed devices in +proportion to free memory and runs those matmuls on all devices in parallel. +Compared to a layer split this uses all GPUs within every layer (instead of +sequentially device by device) at the cost of a cross-device reduction per +matmul — usually the faster option when the devices have fast interconnect. + +Row split requires backend support for split buffers and is currently +available on CUDA only; on other backends (or when the listed devices belong +to different backend registries) the module falls back to a layer split. +Embeddings, normalization weights, biases and other non-block tensors stay in +regular buffers on the main device. + +Direct ("immediately") LoRA application cannot patch row-split tensors; with +`--split-mode row` the automatic LoRA mode selects runtime application, and an +explicit `--lora-apply-mode immediately` skips the split tensors with a +warning. + ## Modules | Module | Purpose | Accepted names | diff --git a/examples/common/common.cpp b/examples/common/common.cpp index acea9171e..2c2c1278d 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -468,6 +468,13 @@ ArgOptions SDContextParams::get_options() { "parameter backend assignment, e.g. disk, cpu, or diffusion=disk,clip=cpu", (int)',', ¶ms_backend}, + {"", + "--split-mode", + "weight distribution for modules assigned multiple devices (--backend \"diffusion=cuda0&cuda1\"): " + "layer (whole transformer blocks per device, default) or row (matmul rows split across devices, CUDA only). " + "Accepts a single mode or per-module assignments, e.g. row or diffusion=row,te=layer", + (int)',', + &split_mode}, {"", "--rpc-servers", "comma-separated list of RPC servers to connect to for offloading, in the format host:port, e.g. localhost:50052,192.168.1.3:50052", @@ -827,6 +834,7 @@ std::string SDContextParams::to_string() const { << " eager_load: " << (eager_load ? "true" : "false") << ",\n" << " backend: \"" << backend << "\",\n" << " params_backend: \"" << params_backend << "\",\n" + << " split_mode: \"" << split_mode << "\",\n" << " enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n" << " control_net_cpu: " << (control_net_cpu ? "true" : "false") << ",\n" << " clip_on_cpu: " << (clip_on_cpu ? "true" : "false") << ",\n" @@ -907,6 +915,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { 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(); + sd_ctx_params.split_mode = split_mode.c_str(); sd_ctx_params.rpc_servers = rpc_servers.c_str(); return sd_ctx_params; } diff --git a/examples/common/common.h b/examples/common/common.h index 941fa3317..daa15e72a 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -151,6 +151,7 @@ struct SDContextParams { bool eager_load = false; std::string backend; std::string params_backend; + std::string split_mode; std::string rpc_servers; std::string effective_backend; std::string effective_params_backend; diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index f04e8cddb..bcd0cc025 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -227,6 +227,7 @@ typedef struct { 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; + const char* split_mode; // weight distribution for multi-device modules: layer (default) or row, or per-module assignments e.g. "diffusion=row" const char* rpc_servers; } sd_ctx_params_t; diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index aaf01da42..66e90f570 100644 --- a/src/core/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -760,6 +760,7 @@ bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule modu bool SDBackendManager::init(const char* backend_spec, const char* params_backend_spec, + const char* split_mode_spec, std::string* error) { reset(); @@ -769,10 +770,54 @@ bool SDBackendManager::init(const char* backend_spec, if (!sd_parse_backend_assignment(SAFE_STR(params_backend_spec), ¶ms_assignment_, error)) { return false; } + if (!sd_parse_backend_assignment(SAFE_STR(split_mode_spec), &split_mode_assignment_, error)) { + return false; + } return validate(error); } +SDSplitMode SDBackendManager::split_mode(SDBackendModule module) const { + return lower_copy(trim_copy(split_mode_assignment_.get(module))) == "row" ? SDSplitMode::ROW + : SDSplitMode::LAYER; +} + +ggml_backend_buffer_type_t SDBackendManager::split_buffer_type(ggml_backend_t backend, + const std::vector& tensor_split) { + if (backend == nullptr) { + return nullptr; + } + ggml_backend_dev_t dev = ggml_backend_get_device(backend); + if (dev == nullptr) { + return nullptr; + } + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev); + if (reg == nullptr) { + return nullptr; + } + auto fn = (ggml_backend_split_buffer_type_t)ggml_backend_reg_get_proc_address(reg, "ggml_backend_split_buffer_type"); + if (fn == nullptr) { + return nullptr; // backend has no row-split support + } + // main_device is the backend's device index WITHIN its registry. + int main_device = -1; + const size_t dev_count = ggml_backend_reg_dev_count(reg); + for (size_t i = 0; i < dev_count; ++i) { + if (ggml_backend_reg_dev_get(reg, i) == dev) { + main_device = (int)i; + break; + } + } + if (main_device < 0) { + return nullptr; + } + // The backend reads a fixed-size proportion array (e.g. CUDA scans + // GGML_CUDA_MAX_DEVICES entries); pad generously to stay in bounds. + std::vector padded_split(std::max(tensor_split.size(), 64), 0.0f); + std::copy(tensor_split.begin(), tensor_split.end(), padded_split.begin()); + return fn(main_device, padded_split.data()); +} + bool SDBackendManager::validate(std::string* error) const { auto validate_single_runtime_name = [&](const std::string& name) -> bool { if (is_default_backend_token(name)) { @@ -830,8 +875,20 @@ bool SDBackendManager::validate(std::string* error) const { return validate_single_runtime_name(name); }; + auto validate_split_mode_name = [&](const std::string& name) -> bool { + const std::string lower = lower_copy(trim_copy(name)); + if (lower.empty() || lower == "layer" || lower == "row") { + return true; + } + if (error != nullptr) { + *error = "invalid split mode '" + name + "' (expected layer or row)"; + } + return false; + }; + if (!validate_runtime_name(runtime_assignment_.default_name) || - !validate_params_name(params_assignment_.default_name)) { + !validate_params_name(params_assignment_.default_name) || + !validate_split_mode_name(split_mode_assignment_.default_name)) { return false; } for (const auto& kv : runtime_assignment_.module_names) { @@ -844,6 +901,11 @@ bool SDBackendManager::validate(std::string* error) const { return false; } } + for (const auto& kv : split_mode_assignment_.module_names) { + if (!validate_split_mode_name(kv.second)) { + return false; + } + } return true; } diff --git a/src/core/ggml_extend_backend.h b/src/core/ggml_extend_backend.h index b5f8ae00c..b3878a7ff 100644 --- a/src/core/ggml_extend_backend.h +++ b/src/core/ggml_extend_backend.h @@ -38,10 +38,20 @@ struct SDBackendHandleDeleter { using SDBackendHandle = std::unique_ptr; +// How a module with a multi-device runtime assignment distributes its weights: +// LAYER assigns whole transformer blocks to devices (backend-generic), ROW +// splits each matmul weight row-wise via the backend's split buffer type +// (currently CUDA only). +enum class SDSplitMode { + LAYER, + ROW, +}; + class SDBackendManager { private: SDBackendAssignment runtime_assignment_; SDBackendAssignment params_assignment_; + SDBackendAssignment split_mode_assignment_; std::unordered_map backends_; public: @@ -53,6 +63,7 @@ class SDBackendManager { bool init(const char* backend_spec, const char* params_backend_spec, + const char* split_mode_spec, std::string* error); void reset(); @@ -65,6 +76,19 @@ class SDBackendManager { // backend (the same one runtime_backend() returns); duplicates are folded. std::vector runtime_backends(SDBackendModule module); + // The module's --split-mode assignment ("layer" when unset). + SDSplitMode split_mode(SDBackendModule module) const; + + // Row (tensor) split buffer type for the backend, with tensor_split + // holding per-registry-device weight proportions (normalized by the + // backend) and the backend's own device as the main device. Returns + // nullptr when the backend does not publish the + // "ggml_backend_split_buffer_type" proc (currently CUDA only; callers + // fall back to a layer split). The buffer type is cached by the backend + // and must not be freed. + ggml_backend_buffer_type_t split_buffer_type(ggml_backend_t backend, + const std::vector& tensor_split); + bool runtime_backend_is_cpu(SDBackendModule module); bool params_backend_is_cpu(SDBackendModule module); bool params_backend_is_disk(SDBackendModule module) const; diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 7095ec6a9..2eaa13460 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -100,12 +100,46 @@ size_t estimate_tensors_size(const std::map& tensors) return size; } +void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft) { + if (compute_backend == nullptr) { + return; + } + if (split_buft == nullptr) { + split_buffer_types_.erase(compute_backend); + return; + } + split_buffer_types_[compute_backend] = split_buft; +} + +bool ModelManager::tensor_shape_supports_split_buffer(const ggml_tensor* tensor) { + // Split buffers assert on views and non-contiguous tensors, split along + // rows (ne[1]) and are only usable as matmul weights. The size floor keeps + // small 2D block tensors (modulation/scale-shift tables) out: those are + // sliced into views inside the graphs, and views of split tensors are not + // supported. + return tensor != nullptr && + tensor->view_src == nullptr && + ggml_is_contiguous(tensor) && + ggml_n_dims(tensor) == 2 && + tensor->ne[0] >= 256 && + tensor->ne[1] >= 256; +} + +ggml_backend_buffer_type_t ModelManager::split_buffer_type_for(const TensorState& state) const { + if (!state.allow_split_buffer || !tensor_shape_supports_split_buffer(state.tensor)) { + return nullptr; + } + auto it = split_buffer_types_.find(state.compute_backend); + return it != split_buffer_types_.end() ? it->second : nullptr; +} + bool ModelManager::register_param_tensors(const std::string& desc, std::map tensors, ResidencyMode residency_mode, ggml_backend_t compute_backend, ggml_backend_t params_backend, - size_t* registered_tensor_size) { + size_t* registered_tensor_size, + bool allow_split_buffer) { if (desc.empty()) { LOG_ERROR("model manager tensor desc is empty"); return false; @@ -129,13 +163,14 @@ bool ModelManager::register_param_tensors(const std::string& desc, } ggml_set_name(tensor, name.c_str()); - auto state = std::make_unique(); - state->name = name; - state->tensor = tensor; - state->desc = desc; - state->residency_mode = residency_mode; - state->compute_backend = compute_backend; - state->params_backend = params_backend; + auto state = std::make_unique(); + state->name = name; + state->tensor = tensor; + state->desc = desc; + state->residency_mode = residency_mode; + state->compute_backend = compute_backend; + state->params_backend = params_backend; + state->allow_split_buffer = allow_split_buffer; new_states.push_back(std::move(state)); } @@ -237,7 +272,10 @@ bool ModelManager::load_tensors_to_params_backend(const std::vector& states) { - std::map> states_by_compute_backend; + // Group by (compute backend, staging buffer type): split-eligible tensors + // stage into the backend's row-split buffer type, the rest into its + // default buffer type. + std::map, std::vector> states_by_staging_target; for (TensorState* state : states) { if (state == nullptr || should_ignore(*state) || is_optional_missing_tensor(state->name)) { continue; @@ -257,11 +295,16 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vectorname.c_str()); return false; } - states_by_compute_backend[state->compute_backend].push_back(state); + ggml_backend_buffer_type_t staging_buft = split_buffer_type_for(*state); + if (staging_buft == nullptr) { + staging_buft = ggml_backend_get_default_buffer_type(state->compute_backend); + } + states_by_staging_target[{state->compute_backend, staging_buft}].push_back(state); } - for (const auto& pair : states_by_compute_backend) { - ggml_backend_t compute_backend = pair.first; + for (const auto& pair : states_by_staging_target) { + ggml_backend_t compute_backend = pair.first.first; + ggml_backend_buffer_type_t staging_buft = pair.first.second; const std::vector& states = pair.second; if (states.empty()) { continue; @@ -285,7 +328,7 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector& states LOG_ERROR("model manager compute backend is null for lora target tensor '%s'", state->name.c_str()); return false; } + if (state->tensor->buffer != nullptr && + ggml_backend_buffer_get_type(state->tensor->buffer) == split_buffer_type_for(*state)) { + // Direct LoRA application builds add/cpy graphs over the model + // tensors, and row-split tensors only support matmul; skip them + // here (use the at_runtime LoRA apply mode with row split). + if (!warned_split_lora_skip_) { + LOG_WARN("model manager skipping direct lora application to row-split tensors " + "(use --lora-apply-mode at_runtime with row split)"); + warned_split_lora_skip_ = true; + } + state->applied_lora_epoch = current_lora_epoch_; + continue; + } if (state->tensor->data == nullptr) { LOG_ERROR("model manager lora target tensor '%s' is not prepared", state->name.c_str()); return false; @@ -694,6 +750,10 @@ ggml_backend_buffer_type_t ModelManager::params_buffer_type_for(const TensorStat if (compute_dev != nullptr) { params_buft = ggml_backend_dev_host_buffer_type(compute_dev); } + } else if (state.params_backend == state.compute_backend) { + // When params live on the compute backend, split-eligible tensors go + // straight into the row-split buffer type. + params_buft = split_buffer_type_for(state); } if (params_buft == nullptr) { params_buft = ggml_backend_get_default_buffer_type(state.params_backend); diff --git a/src/model_manager.h b/src/model_manager.h index 9225e3ea6..69ad4c9c3 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -36,7 +36,12 @@ class ModelManager : public RunnerWeightManager { ResidencyMode residency_mode = ResidencyMode::ParamBackend; ggml_backend_t compute_backend = nullptr; ggml_backend_t params_backend = nullptr; - bool metadata_validated = false; + // Set at registration for tensors that may live in the compute + // backend's row-split buffer type (weights only ever consumed by + // matmuls). Only takes effect when a split buffer type was registered + // for the compute backend (set_split_buffer_type). + bool allow_split_buffer = false; + bool metadata_validated = false; int active_prepare_count = 0; @@ -63,6 +68,9 @@ class ModelManager : public RunnerWeightManager { std::map tensor_states_by_name_; std::vector> params_storage_blocks_; std::vector> compute_staging_blocks_; + // Row-split buffer type per compute backend (backend-cached, not owned). + std::map split_buffer_types_; + bool warned_split_lora_skip_ = false; std::set common_ignore_tensors_; std::vector loras_; SDVersion lora_version_ = VERSION_COUNT; @@ -91,6 +99,7 @@ class ModelManager : public RunnerWeightManager { bool stage_tensors_to_compute_backend(const std::vector& states); ggml_backend_buffer_type_t params_buffer_type_for(const TensorState& state) const; + ggml_backend_buffer_type_t split_buffer_type_for(const TensorState& state) const; void release_compute_staging_blocks(bool force = false, const std::unordered_set* target_states = nullptr); void release_params_storage_blocks(bool force = false, @@ -114,6 +123,14 @@ class ModelManager : public RunnerWeightManager { void set_writable_mmap(bool writable_mmap) { writable_mmap_ = writable_mmap; } void set_common_ignore_tensors(std::set ignore_tensors); void set_loras(std::vector loras, SDVersion version); + // Register the row-split buffer type used for split-eligible tensors whose + // compute backend is `compute_backend`. The buffer type is owned by the + // backend (never freed here). + void set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft); + + // Shape constraints of the backend split buffer types: contiguous, + // non-view, rank-2 (split along rows, matmul weights only). + static bool tensor_shape_supports_split_buffer(const ggml_tensor* tensor); std::set tensor_names() const; @@ -122,7 +139,8 @@ class ModelManager : public RunnerWeightManager { ResidencyMode residency_mode, ggml_backend_t compute_backend, ggml_backend_t params_backend, - size_t* registered_tensor_size = nullptr); + size_t* registered_tensor_size = nullptr, + bool allow_split_buffer = false); template bool register_runner_params(const std::string& desc, diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 6dd812e87..cca0ad1b8 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -359,6 +359,7 @@ class StableDiffusionGGML { bool eager_load = false; std::string backend_spec; std::string params_backend_spec; + std::string split_mode_spec; bool is_using_v_parameterization = false; bool is_using_edm_v_parameterization = false; @@ -433,6 +434,15 @@ class StableDiffusionGGML { if (module_backends.size() > 1) { if constexpr (has_set_runtime_backends::value) { if (module == SDBackendModule::DIFFUSION || module == SDBackendModule::TE) { + if (backend_manager.split_mode(module) == SDSplitMode::ROW) { + return register_row_split_runner_params(desc, + model, + module, + module_backends, + std::move(group_tensors), + residency_mode, + params_mem_size); + } return register_layer_split_runner_params(desc, model, module, @@ -454,6 +464,122 @@ class StableDiffusionGGML { params_mem_size); } + // Row split registration (--split-mode row): the module keeps its main + // runtime backend and executes there, but its transformer-block matmul + // weights are allocated in the backend's row-split buffer type, which + // distributes each weight's rows across the listed devices (the backend + // runs those matmuls multi-GPU internally). Everything else — embeddings, + // norms, biases, non-block tensors — stays in regular buffers on the main + // device. Falls back to a layer split when the backend has no split buffer + // type (non-CUDA) or the devices span different registries. + template + bool register_row_split_runner_params(const std::string& desc, + const std::shared_ptr& model, + SDBackendModule module, + const std::vector& module_backends, + std::map group_tensors, + ModelManager::ResidencyMode residency_mode, + size_t* params_mem_size) { + ggml_backend_t main_backend = module_backends[0]; + + auto fall_back_to_layer_split = [&](const char* reason) { + LOG_WARN("%s: row split unavailable (%s); falling back to layer split", desc.c_str(), reason); + return register_layer_split_runner_params(desc, + model, + module, + module_backends, + std::move(group_tensors), + residency_mode, + params_mem_size); + }; + + // All devices must belong to the main backend's registry, and the + // tensor_split proportions are indexed by registry device index. + ggml_backend_dev_t main_dev = ggml_backend_get_device(main_backend); + ggml_backend_reg_t reg = main_dev != nullptr ? ggml_backend_dev_backend_reg(main_dev) : nullptr; + if (reg == nullptr) { + return fall_back_to_layer_split("no backend registry"); + } + const size_t reg_dev_count = ggml_backend_reg_dev_count(reg); + std::vector tensor_split(reg_dev_count, 0.0f); + constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024; + for (ggml_backend_t backend : module_backends) { + ggml_backend_dev_t dev = ggml_backend_get_device(backend); + int reg_index = -1; + for (size_t i = 0; i < reg_dev_count; i++) { + if (ggml_backend_reg_dev_get(reg, i) == dev) { + reg_index = (int)i; + break; + } + } + if (reg_index < 0) { + return fall_back_to_layer_split("devices span different backend registries"); + } + size_t free_bytes = 0, total_bytes = 0; + ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); + int64_t usable_bytes = std::max((int64_t)free_bytes - compute_headroom_bytes, + (int64_t)free_bytes / 8); + tensor_split[reg_index] = usable_bytes > 0 ? (float)((double)usable_bytes / (1024.0 * 1024.0)) : 1.0f; + } + + ggml_backend_buffer_type_t split_buft = backend_manager.split_buffer_type(main_backend, tensor_split); + if (split_buft == nullptr) { + return fall_back_to_layer_split("backend has no split buffer type"); + } + model_manager->set_split_buffer_type(main_backend, split_buft); + + std::map split_tensors; + if constexpr (std::is_base_of_v) { + model->get_layer_split_param_tensors(split_tensors); + } else { + split_tensors = group_tensors; + } + + // Row-split candidates: transformer-block 2D weights of the module's + // dominant transformer. Non-block tensors (embeddings, output heads) + // may be consumed by non-matmul ops, which split buffers do not + // support, so they stay in regular buffers. + std::map row_split_map; + std::map regular_map; + size_t row_split_bytes = 0; + for (const auto& kv : group_tensors) { + if (split_tensors.count(kv.first) != 0 && + tensor_block_index(kv.first) >= 0 && + ModelManager::tensor_shape_supports_split_buffer(kv.second)) { + row_split_map[kv.first] = kv.second; + row_split_bytes += ggml_nbytes(kv.second); + } else { + regular_map[kv.first] = kv.second; + } + } + if (row_split_map.empty()) { + return fall_back_to_layer_split("no row-splittable transformer block weights found"); + } + + LOG_INFO("%s row split: %zu tensors (%.1f MB) split across %zu devices (main %s)", + desc.c_str(), + row_split_map.size(), + row_split_bytes / (1024.f * 1024.f), + module_backends.size(), + backend_device_display_name(main_backend).c_str()); + + if (!model_manager->register_param_tensors(desc, + std::move(row_split_map), + residency_mode, + main_backend, + params_backend_for(module), + params_mem_size, + /*allow_split_buffer=*/true)) { + return false; + } + return model_manager->register_param_tensors(desc, + std::move(regular_map), + residency_mode, + main_backend, + params_backend_for(module), + params_mem_size); + } + // Layer split registration: partition the module's tensors into contiguous // transformer-block ranges (one per runtime backend) and register each // range with that backend as its per-tensor compute backend — the @@ -544,6 +670,7 @@ class StableDiffusionGGML { std::string error; if (!backend_manager.init(backend_spec.c_str(), params_backend_spec.c_str(), + split_mode_spec.c_str(), &error)) { LOG_ERROR("backend config failed: %s", error.c_str()); return false; @@ -551,6 +678,17 @@ class StableDiffusionGGML { return ensure_backend_pair(SDBackendModule::DIFFUSION); } + // True when any split-capable module distributes row-split weights. + bool row_split_active() { + for (SDBackendModule module : {SDBackendModule::DIFFUSION, SDBackendModule::TE}) { + if (backend_manager.split_mode(module) == SDSplitMode::ROW && + backend_manager.runtime_backends(module).size() > 1) { + return true; + } + } + return false; + } + std::shared_ptr get_rng(rng_type_t rng_type) { if (rng_type == STD_DEFAULT_RNG) { return std::make_shared(); @@ -609,6 +747,7 @@ class StableDiffusionGGML { 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); + split_mode_spec = SAFE_STR(sd_ctx_params->split_mode); max_vram_assignment.reset(0.f); { std::string error; @@ -837,12 +976,16 @@ class StableDiffusionGGML { // Avoid full-model LoRA merge buffers on constrained setups. const bool params_offloaded = params_backend_for(SDBackendModule::DIFFUSION) != backend_for(SDBackendModule::DIFFUSION); const bool streaming_constrained = stream_layers || params_offloaded; - if (have_quantized_weight || streaming_constrained) { + if (have_quantized_weight || streaming_constrained || row_split_active()) { apply_lora_immediately = false; } else { apply_lora_immediately = true; } } else if (sd_ctx_params->lora_apply_mode == LORA_APPLY_IMMEDIATELY) { + if (row_split_active()) { + LOG_WARN("row-split tensors do not support the immediately LoRA apply mode; " + "LoRAs will not be applied to them (use --lora-apply-mode at_runtime)"); + } apply_lora_immediately = true; } else { apply_lora_immediately = false; @@ -3063,6 +3206,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO; sd_ctx_params->backend = nullptr; sd_ctx_params->params_backend = nullptr; + sd_ctx_params->split_mode = nullptr; sd_ctx_params->rpc_servers = nullptr; sd_ctx_params->pulid_weights_path = nullptr; } @@ -3102,6 +3246,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "eager_load: %s\n" "backend: %s\n" "params_backend: %s\n" + "split_mode: %s\n" "flash_attn: %s\n" "diffusion_flash_attn: %s\n" "circular_x: %s\n" @@ -3138,6 +3283,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { BOOL_STR(sd_ctx_params->eager_load), SAFE_STR(sd_ctx_params->backend), SAFE_STR(sd_ctx_params->params_backend), + SAFE_STR(sd_ctx_params->split_mode), BOOL_STR(sd_ctx_params->flash_attn), BOOL_STR(sd_ctx_params->diffusion_flash_attn), BOOL_STR(sd_ctx_params->circular_x), diff --git a/src/upscaler.cpp b/src/upscaler.cpp index 88a8a6336..dbb99af36 100644 --- a/src/upscaler.cpp +++ b/src/upscaler.cpp @@ -46,6 +46,7 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path, std::string error; if (!backend_manager.init(backend_spec.c_str(), params_backend_spec.c_str(), + /*split_mode_spec=*/nullptr, &error)) { LOG_ERROR("upscaler backend config failed: %s", error.c_str()); return false; From 44896d741e5feb32bdc550bf6c62db0799f45f8d Mon Sep 17 00:00:00 2001 From: Piotr Wilkin Date: Fri, 3 Jul 2026 16:54:47 +0200 Subject: [PATCH 3/8] feat: add automatic device placement (--auto-fit) --auto-fit derives the diffusion/te/vae placements from the model metadata and per-device memory budgets, then feeds them through the existing --backend / --params-backend assignment mechanism (the plan and the emitted specs are printed). Budgets reuse --max-vram: positive values cap a device, negative values mean free memory minus that many GiB, unset defaults to free minus a 512 MiB margin. When all components fit resident they are spread across the GPUs; otherwise the plan switches to time-share, giving the heavy components disk params residency (load per phase, free after) and splitting a component too large for any single device across all GPUs via the layer/row split mechanisms (--split-mode still selects which). Components that fit nowhere fall back to the CPU. A VAE decode that still runs out of memory retries once with tiling enabled (temporal for the LTX video VAE, spatial otherwise). Backend initialization now happens after model metadata is loaded so the planner can size components before any backend exists. auto_fit defaults to off. Co-Authored-By: Claude Fable 5 --- docs/backend.md | 27 ++++ examples/common/common.cpp | 8 + examples/common/common.h | 1 + include/stable-diffusion.h | 6 + src/backend_fit.hpp | 311 +++++++++++++++++++++++++++++++++++++ src/model_loader.h | 2 + src/stable-diffusion.cpp | 141 +++++++++++++++-- 7 files changed, 481 insertions(+), 15 deletions(-) create mode 100644 src/backend_fit.hpp diff --git a/docs/backend.md b/docs/backend.md index 52a9acf20..b88e44e6b 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -115,6 +115,33 @@ Direct ("immediately") LoRA application cannot patch row-split tensors; with explicit `--lora-apply-mode immediately` skips the split tensors with a warning. +## Automatic placement (`--auto-fit`) + +`--auto-fit` derives the `diffusion` / `te` / `vae` placements from the model +metadata and the per-device memory budgets, then feeds them into the same +backend assignment mechanism described above (the chosen specs are printed). +`--backend` and `--params-backend` are ignored while auto-fit is enabled. + +```shell +sd-cli -m model.safetensors -p "a cat" --auto-fit +sd-cli -m model.safetensors -p "a cat" --auto-fit --max-vram cuda0=8,cuda1=14 +sd-cli -m model.safetensors -p "a cat" --auto-fit --split-mode row +``` + +Budgets reuse `--max-vram`: a positive per-device value caps what auto-fit +plans with on that device, a negative value means "free memory minus that many +GiB", and with no budget set each device's free memory minus a 512 MiB margin +is used. (The same values still drive graph-cut segmented execution for +modules that end up on a single device.) + +When everything fits resident, components are simply spread across the +available GPUs. When it does not, auto-fit switches to time-share mode: the +heavy components get `disk` params residency (loaded for their phase, freed +after), and a component too large for any single device is split across all +GPUs with the layer/row split mechanism (`--split-mode` selects which, layer +by default). Components that fit nowhere fall back to the CPU. If a VAE decode +still runs out of memory, tiling is enabled and the decode retried once. + ## Modules | Module | Purpose | Accepted names | diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 2c2c1278d..b6df0ba9c 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -508,6 +508,12 @@ ArgOptions SDContextParams::get_options() { "--eager-load", "load all params into the params backend at model-load time instead of lazily on first use (defaults to false)", true, &eager_load}, + {"", + "--auto-fit", + "pick the diffusion/te/vae device placements automatically from the model size and the per-device " + "memory budgets (--max-vram; defaults to free memory minus a small margin). Overrides --backend and " + "--params-backend; may split modules across GPUs (--split-mode still selects layer or row)", + true, &auto_fit}, {"", "--force-sdxl-vae-conv-scale", "force use of conv scale on sdxl vae", @@ -835,6 +841,7 @@ std::string SDContextParams::to_string() const { << " backend: \"" << backend << "\",\n" << " params_backend: \"" << params_backend << "\",\n" << " split_mode: \"" << split_mode << "\",\n" + << " auto_fit: " << (auto_fit ? "true" : "false") << ",\n" << " enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n" << " control_net_cpu: " << (control_net_cpu ? "true" : "false") << ",\n" << " clip_on_cpu: " << (clip_on_cpu ? "true" : "false") << ",\n" @@ -916,6 +923,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { sd_ctx_params.backend = effective_backend.c_str(); sd_ctx_params.params_backend = effective_params_backend.c_str(); sd_ctx_params.split_mode = split_mode.c_str(); + sd_ctx_params.auto_fit = auto_fit; sd_ctx_params.rpc_servers = rpc_servers.c_str(); return sd_ctx_params; } diff --git a/examples/common/common.h b/examples/common/common.h index daa15e72a..654871359 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -152,6 +152,7 @@ struct SDContextParams { std::string backend; std::string params_backend; std::string split_mode; + bool auto_fit = false; std::string rpc_servers; std::string effective_backend; std::string effective_params_backend; diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index bcd0cc025..1968af268 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -228,6 +228,12 @@ typedef struct { const char* backend; const char* params_backend; const char* split_mode; // weight distribution for multi-device modules: layer (default) or row, or per-module assignments e.g. "diffusion=row" + // Pick the diffusion/te/vae device placements automatically from the model + // metadata and the per-device memory budgets (max_vram). When enabled, + // `backend` and `params_backend` are ignored and derived instead; the + // plan may use multi-device layer/row splits (`split_mode` still applies) + // and disk params residency so components time-share VRAM. + bool auto_fit; const char* rpc_servers; } sd_ctx_params_t; diff --git a/src/backend_fit.hpp b/src/backend_fit.hpp new file mode 100644 index 000000000..078794a5b --- /dev/null +++ b/src/backend_fit.hpp @@ -0,0 +1,311 @@ +#ifndef __SD_BACKEND_FIT_HPP__ +#define __SD_BACKEND_FIT_HPP__ + +#include +#include +#include +#include +#include + +#include "core/ggml_extend_backend.h" +#include "core/ggml_graph_cut.h" +#include "core/util.h" +#include "ggml-backend.h" +#include "model_loader.h" + +// Auto-fit (--auto-fit): derive --backend / --params-backend placements for +// the DiT / conditioner / VAE from the model metadata and per-device memory +// budgets. This is a policy layer only — it emits the same backend assignment +// specs a user would pass manually (including '&' device lists for the layer / +// row split mechanisms) and feeds them to SDBackendManager. +// +// Budgets come from --max-vram: a positive per-device value caps the memory +// auto-fit plans with on that device, a negative value means "free memory +// minus that many GiB", and an unset budget defaults to the device's free +// memory minus a 512 MiB margin. +namespace backend_fit { + +constexpr int64_t MiB = 1024ll * 1024; + +enum class ComponentKind { + DIT = 0, + VAE = 1, + CONDITIONER = 2, +}; + +struct Component { + ComponentKind kind; + const char* name; + int64_t params_bytes = 0; + // Estimated compute-buffer bytes the component needs alongside its params + // on every device it runs on. + int64_t reserve_bytes = 0; + bool splittable = false; +}; + +struct Device { + ggml_backend_dev_t dev = nullptr; + std::string name; + std::string description; + int64_t free_bytes = 0; + int64_t total_bytes = 0; + int64_t budget_bytes = 0; +}; + +struct Decision { + ComponentKind kind; + bool on_cpu = false; + std::vector device_idxs; // more than one entry = split device list +}; + +struct Plan { + bool valid = false; + // Components load lazily (params=disk) and time-share the budgets phase + // by phase instead of all being resident at once. + bool time_share = false; + std::vector decisions; +}; + +inline bool classify_tensor(const std::string& name, ComponentKind& out) { + auto contains = [&](const char* s) { return name.find(s) != std::string::npos; }; + + if (contains("model.diffusion_model.") || contains("unet.")) { + out = ComponentKind::DIT; + return true; + } + if (contains("first_stage_model.") || + name.rfind("vae.", 0) == 0 || + name.rfind("tae.", 0) == 0) { + out = ComponentKind::VAE; + return true; + } + if (contains("text_encoders") || + contains("cond_stage_model") || + contains("te.text_model.") || + contains("conditioner") || + name.rfind("text_encoder.", 0) == 0 || + // Connector / text projection layers that run on the conditioner + // backend (e.g. LTX-2's text_embedding_projection). + name.rfind("text_embedding_projection.", 0) == 0 || + contains(".aggregate_embed.")) { + out = ComponentKind::CONDITIONER; + return true; + } + return false; +} + +inline std::vector estimate_components(ModelLoader& loader, ggml_type override_wtype) { + const auto& storage = loader.get_tensor_storage_map(); + + int64_t bytes[3] = {0, 0, 0}; + for (const auto& [name, ts_const] : storage) { + TensorStorage ts = ts_const; + if (is_unused_tensor(ts.name)) { + continue; + } + ComponentKind kind; + if (!classify_tensor(ts.name, kind)) { + continue; + } + if (override_wtype != GGML_TYPE_COUNT && + loader.tensor_should_be_converted(ts, override_wtype)) { + ts.type = override_wtype; + } else if (ts.expected_type != GGML_TYPE_COUNT && ts.expected_type != ts.type) { + ts.type = ts.expected_type; + } + bytes[int(kind)] += (int64_t)ts.nbytes() + 64; + } + + // Built-in per-component compute reserves. These are per-device estimates + // for typical workloads; --max-vram can shrink the budgets when they turn + // out too small for a given resolution. + std::vector out; + out.push_back({ComponentKind::DIT, "DiT", bytes[int(ComponentKind::DIT)], 2048 * MiB, true}); + out.push_back({ComponentKind::VAE, "VAE", bytes[int(ComponentKind::VAE)], 1024 * MiB, false}); + out.push_back({ComponentKind::CONDITIONER, "Conditioner", bytes[int(ComponentKind::CONDITIONER)], 2048 * MiB, true}); + return out; +} + +inline std::vector enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) { + std::vector out; + for (size_t i = 0; i < ggml_backend_dev_count(); i++) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) { + continue; + } + Device d; + d.dev = dev; + d.name = ggml_backend_dev_name(dev); + d.description = ggml_backend_dev_description(dev); + size_t free_bytes = 0, total_bytes = 0; + ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); + d.free_bytes = (int64_t)free_bytes; + d.total_bytes = (int64_t)total_bytes; + + // MaxVramAssignment canonicalizes its keys to lowercase device names. + std::string budget_key = d.name; + std::transform(budget_key.begin(), budget_key.end(), budget_key.begin(), + [](unsigned char c) { return (char)std::tolower(c); }); + float gib = budgets.default_gib; + auto it = budgets.backend_gib.find(budget_key); + if (it != budgets.backend_gib.end()) { + gib = it->second; + } + if (gib > 0.f) { + d.budget_bytes = std::min((int64_t)(gib * 1024.0 * 1024.0 * 1024.0), d.free_bytes); + } else if (gib < 0.f) { + d.budget_bytes = d.free_bytes + (int64_t)(gib * 1024.0 * 1024.0 * 1024.0); + } else { + d.budget_bytes = d.free_bytes - 512 * MiB; + } + d.budget_bytes = std::max(d.budget_bytes, 0); + out.push_back(d); + } + return out; +} + +// Place each component, largest first. First try the eager regime (all +// components resident at once); when that overflows, fall back to the +// time-share regime, where params=disk residency loads each phase's weights +// on demand and frees them afterwards, so each component only has to fit on +// its own. A component that cannot fit a single device splits across all +// GPUs when its module supports it, and lands on the CPU otherwise. +inline Plan compute_plan(const std::vector& components, const std::vector& devices) { + Plan plan; + if (devices.empty()) { + return plan; + } + + std::vector order(components.size()); + for (size_t i = 0; i < order.size(); i++) { + order[i] = i; + } + std::sort(order.begin(), order.end(), [&](size_t a, size_t b) { + return components[a].params_bytes > components[b].params_bytes; + }); + + // Eager regime: components coexist; per device, the params sum plus the + // largest single reserve must fit the budget (compute buffers of the + // sequential phases do not coexist). + { + std::vector params_sum(devices.size(), 0); + std::vector max_reserve(devices.size(), 0); + std::vector decisions(components.size()); + bool ok = true; + for (size_t ci : order) { + const Component& comp = components[ci]; + decisions[ci].kind = comp.kind; + if (comp.params_bytes == 0) { + continue; + } + int best = -1; + for (size_t di = 0; di < devices.size(); di++) { + int64_t need = params_sum[di] + comp.params_bytes + std::max(max_reserve[di], comp.reserve_bytes); + if (need <= devices[di].budget_bytes && + (best < 0 || devices[di].budget_bytes - params_sum[di] > devices[best].budget_bytes - params_sum[best])) { + best = (int)di; + } + } + if (best < 0) { + ok = false; + break; + } + params_sum[best] += comp.params_bytes; + max_reserve[best] = std::max(max_reserve[best], comp.reserve_bytes); + decisions[ci].device_idxs.push_back((size_t)best); + } + if (ok) { + plan.valid = true; + plan.time_share = false; + plan.decisions = std::move(decisions); + return plan; + } + } + + // Time-share regime: each component fits on its own (params=disk). + plan.decisions.assign(components.size(), {}); + for (size_t ci : order) { + const Component& comp = components[ci]; + Decision& decision = plan.decisions[ci]; + decision.kind = comp.kind; + if (comp.params_bytes == 0) { + continue; + } + int best = -1; + for (size_t di = 0; di < devices.size(); di++) { + if (comp.params_bytes + comp.reserve_bytes <= devices[di].budget_bytes && + (best < 0 || devices[di].budget_bytes > devices[best].budget_bytes)) { + best = (int)di; + } + } + if (best >= 0) { + decision.device_idxs.push_back((size_t)best); + continue; + } + if (comp.splittable && devices.size() > 1) { + int64_t capacity = 0; + for (const Device& d : devices) { + capacity += std::max(d.budget_bytes - comp.reserve_bytes, 0); + } + if (comp.params_bytes <= capacity) { + // All GPUs, largest budget first (the first device is the + // module's main device). + std::vector idxs(devices.size()); + for (size_t i = 0; i < idxs.size(); i++) { + idxs[i] = i; + } + std::sort(idxs.begin(), idxs.end(), [&](size_t a, size_t b) { + return devices[a].budget_bytes > devices[b].budget_bytes; + }); + decision.device_idxs = std::move(idxs); + continue; + } + } + decision.on_cpu = true; + } + plan.valid = true; + plan.time_share = true; + return plan; +} + +inline void print_plan(const Plan& plan, + const std::vector& components, + const std::vector& devices) { + LOG_INFO("auto-fit plan%s:", plan.time_share ? " (time-share: params load per phase and free after)" : ""); + LOG_INFO(" devices:"); + for (const Device& d : devices) { + LOG_INFO(" %-12s %-32s free %6lld MiB, budget %6lld MiB", + d.name.c_str(), d.description.c_str(), + (long long)(d.free_bytes / MiB), (long long)(d.budget_bytes / MiB)); + } + LOG_INFO(" components:"); + for (size_t ci = 0; ci < components.size(); ci++) { + const Component& comp = components[ci]; + const Decision& decision = plan.decisions[ci]; + std::string target; + if (comp.params_bytes == 0) { + target = "(not present)"; + } else if (decision.on_cpu) { + target = "CPU"; + } else { + for (size_t k = 0; k < decision.device_idxs.size(); k++) { + if (k > 0) { + target += " & "; + } + target += devices[decision.device_idxs[k]].name; + } + if (decision.device_idxs.size() > 1) { + target += " (split)"; + } + } + LOG_INFO(" %-12s params %6lld MiB, compute reserve %5lld MiB -> %s", + comp.name, + (long long)(comp.params_bytes / MiB), + (long long)(comp.reserve_bytes / MiB), + target.c_str()); + } +} + +} // namespace backend_fit + +#endif // __SD_BACKEND_FIT_HPP__ diff --git a/src/model_loader.h b/src/model_loader.h index 4dc700f20..529f3e890 100644 --- a/src/model_loader.h +++ b/src/model_loader.h @@ -27,6 +27,8 @@ struct MmapTensorStore { std::shared_ptr mmbuffer; }; +bool is_unused_tensor(const std::string& name); + class ModelLoader { protected: SDVersion version_ = VERSION_COUNT; diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index cca0ad1b8..7b6109775 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -18,6 +18,7 @@ #include "model_manager.h" #include "stable-diffusion.h" +#include "backend_fit.hpp" #include "conditioning/conditioner.hpp" #include "extensions/generation_extension.h" #include "model/adapter/lora.hpp" @@ -360,6 +361,7 @@ class StableDiffusionGGML { std::string backend_spec; std::string params_backend_spec; std::string split_mode_spec; + bool auto_fit_enabled = false; bool is_using_v_parameterization = false; bool is_using_edm_v_parameterization = false; @@ -773,20 +775,9 @@ class StableDiffusionGGML { ggml_log_set(ggml_log_callback_default, nullptr); - if (!init_backend()) { - return false; - } - { - std::string error; - if (!max_vram_assignment.canonicalize_backend_keys(&error)) { - LOG_ERROR("%s", error.c_str()); - return false; - } - } - if (stream_layers && !backend_manager.params_backend_is_cpu(SDBackendModule::DIFFUSION)) { - LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring"); - stream_layers = false; - } + // Backend initialization happens after the model metadata is loaded, + // so --auto-fit can size the components and derive device placements + // before any backend is created. model_manager = std::make_shared(); model_manager->set_n_threads(n_threads); @@ -935,6 +926,92 @@ class StableDiffusionGGML { model_loader.set_wtype_override(wtype, tensor_type_rules); } + auto_fit_enabled = sd_ctx_params->auto_fit; + if (auto_fit_enabled) { + if (!backend_spec.empty() || !params_backend_spec.empty()) { + LOG_WARN("--auto-fit is enabled; ignoring --backend / --params-backend"); + } + // The budgets are keyed by device name; canonicalize against the + // registry (no backend needs to exist yet). + { + std::string error; + if (!max_vram_assignment.canonicalize_backend_keys(&error)) { + LOG_ERROR("%s", error.c_str()); + return false; + } + } + auto components = backend_fit::estimate_components(model_loader, wtype); + auto devices = backend_fit::enumerate_gpu_devices(max_vram_assignment); + auto plan = backend_fit::compute_plan(components, devices); + if (!plan.valid) { + LOG_WARN("auto-fit: no usable GPU devices; using the default backend"); + } else { + backend_fit::print_plan(plan, components, devices); + std::string runtime_spec; + std::string params_spec; + auto append_assignment = [](std::string& spec, const char* key, const std::string& value) { + if (!spec.empty()) { + spec += ","; + } + spec += key; + spec += "="; + spec += value; + }; + auto apply_decision = [&](backend_fit::ComponentKind kind, const char* module_key) { + for (size_t ci = 0; ci < components.size(); ci++) { + if (components[ci].kind != kind || components[ci].params_bytes == 0) { + continue; + } + const backend_fit::Decision& decision = plan.decisions[ci]; + if (decision.on_cpu) { + append_assignment(runtime_spec, module_key, "cpu"); + return; + } + if (decision.device_idxs.empty()) { + return; // default backend + } + std::string device_list; + for (size_t k = 0; k < decision.device_idxs.size(); k++) { + if (k > 0) { + device_list += "&"; + } + device_list += devices[decision.device_idxs[k]].name; + } + append_assignment(runtime_spec, module_key, device_list); + if (plan.time_share) { + append_assignment(params_spec, module_key, "disk"); + } + return; + } + }; + apply_decision(backend_fit::ComponentKind::DIT, "diffusion"); + apply_decision(backend_fit::ComponentKind::CONDITIONER, "te"); + apply_decision(backend_fit::ComponentKind::VAE, "vae"); + backend_spec = runtime_spec; + params_backend_spec = params_spec; + LOG_INFO("auto-fit: --backend \"%s\"%s%s%s", + backend_spec.empty() ? "(default)" : backend_spec.c_str(), + params_backend_spec.empty() ? "" : " --params-backend \"", + params_backend_spec.c_str(), + params_backend_spec.empty() ? "" : "\""); + } + } + + if (!init_backend()) { + return false; + } + { + std::string error; + if (!max_vram_assignment.canonicalize_backend_keys(&error)) { + LOG_ERROR("%s", error.c_str()); + return false; + } + } + if (stream_layers && !backend_manager.params_backend_is_cpu(SDBackendModule::DIFFUSION)) { + LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring"); + stream_layers = false; + } + std::map wtype_stat = model_loader.get_wtype_stat(); std::map conditioner_wtype_stat = model_loader.get_conditioner_wtype_stat(); std::map diffusion_model_wtype_stat = model_loader.get_diffusion_model_wtype_stat(); @@ -2850,7 +2927,38 @@ class StableDiffusionGGML { } auto latents = first_stage_model->diffusion_to_vae_latents(x); first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling); - return first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); + auto decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); + // Auto-fit picks the placement without knowing the output size; a + // full-frame decode can exceed the VAE compute reserve and fail + // gracefully with an empty result. Enable tiling and retry once + // (temporal for the LTX video VAE, spatial otherwise) instead of + // failing the whole generation. + if (decoded.empty() && auto_fit_enabled) { + bool changed = false; + if (decode_video && std::dynamic_pointer_cast(first_stage_model) != nullptr) { + if (!vae_tiling_params.temporal_tiling) { + vae_tiling_params.temporal_tiling = true; + changed = true; + } + } else if (!vae_tiling_params.enabled) { + vae_tiling_params.enabled = true; + if (vae_tiling_params.tile_size_x <= 0) { + vae_tiling_params.tile_size_x = 256; + } + if (vae_tiling_params.tile_size_y <= 0) { + vae_tiling_params.tile_size_y = 256; + } + changed = true; + } + if (changed) { + LOG_WARN("auto-fit: VAE decode failed (likely out of memory); retrying with %s tiling", + vae_tiling_params.temporal_tiling ? "temporal" : "spatial"); + first_stage_model->free_compute_buffer(); + first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling); + decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); + } + } + return decoded; } sd::Tensor normalize_ltx_video_latents(const sd::Tensor& x) { @@ -3207,6 +3315,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { sd_ctx_params->backend = nullptr; sd_ctx_params->params_backend = nullptr; sd_ctx_params->split_mode = nullptr; + sd_ctx_params->auto_fit = false; sd_ctx_params->rpc_servers = nullptr; sd_ctx_params->pulid_weights_path = nullptr; } @@ -3247,6 +3356,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "backend: %s\n" "params_backend: %s\n" "split_mode: %s\n" + "auto_fit: %s\n" "flash_attn: %s\n" "diffusion_flash_attn: %s\n" "circular_x: %s\n" @@ -3284,6 +3394,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { SAFE_STR(sd_ctx_params->backend), SAFE_STR(sd_ctx_params->params_backend), SAFE_STR(sd_ctx_params->split_mode), + BOOL_STR(sd_ctx_params->auto_fit), BOOL_STR(sd_ctx_params->flash_attn), BOOL_STR(sd_ctx_params->diffusion_flash_attn), BOOL_STR(sd_ctx_params->circular_x), From 0982b805c5a2ef9fe2af9035d2a77b1d4581b8cd Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 5 Jul 2026 00:22:11 +0800 Subject: [PATCH 4/8] refactor: move auto-fit planning into backend helper --- src/backend_fit.hpp | 117 +++++++++++++++++++++++++++++++++++++++ src/stable-diffusion.cpp | 93 +++---------------------------- 2 files changed, 125 insertions(+), 85 deletions(-) diff --git a/src/backend_fit.hpp b/src/backend_fit.hpp index 078794a5b..1c386b96e 100644 --- a/src/backend_fit.hpp +++ b/src/backend_fit.hpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include "core/ggml_extend_backend.h" @@ -12,6 +13,7 @@ #include "core/util.h" #include "ggml-backend.h" #include "model_loader.h" +#include "stable-diffusion.h" // Auto-fit (--auto-fit): derive --backend / --params-backend placements for // the DiT / conditioner / VAE from the model metadata and per-device memory @@ -306,6 +308,121 @@ inline void print_plan(const Plan& plan, } } +inline void append_assignment(std::string& spec, const char* key, const std::string& value) { + if (!spec.empty()) { + spec += ","; + } + spec += key; + spec += "="; + spec += value; +} + +inline void append_component_decision(const std::vector& components, + const std::vector& devices, + const Plan& plan, + ComponentKind kind, + const char* module_key, + std::string& runtime_spec, + std::string& params_spec) { + for (size_t ci = 0; ci < components.size(); ci++) { + if (components[ci].kind != kind || components[ci].params_bytes == 0) { + continue; + } + const Decision& decision = plan.decisions[ci]; + if (decision.on_cpu) { + append_assignment(runtime_spec, module_key, "cpu"); + return; + } + if (decision.device_idxs.empty()) { + return; + } + std::string device_list; + for (size_t k = 0; k < decision.device_idxs.size(); k++) { + if (k > 0) { + device_list += "&"; + } + device_list += devices[decision.device_idxs[k]].name; + } + append_assignment(runtime_spec, module_key, device_list); + if (plan.time_share) { + append_assignment(params_spec, module_key, "disk"); + } + return; + } +} + +inline bool derive_backend_specs(ModelLoader& loader, + ggml_type override_wtype, + sd::ggml_graph_cut::MaxVramAssignment& budgets, + std::string& runtime_spec, + std::string& params_spec) { + if (!runtime_spec.empty() || !params_spec.empty()) { + LOG_WARN("--auto-fit is enabled; ignoring --backend / --params-backend"); + } + + // The budgets are keyed by device name; canonicalize against the registry + // before SDBackendManager creates any backend instance. + { + std::string error; + if (!budgets.canonicalize_backend_keys(&error)) { + LOG_ERROR("%s", error.c_str()); + return false; + } + } + + auto components = estimate_components(loader, override_wtype); + auto devices = enumerate_gpu_devices(budgets); + auto plan = compute_plan(components, devices); + if (!plan.valid) { + LOG_WARN("auto-fit: no usable GPU devices; using the default backend"); + runtime_spec.clear(); + params_spec.clear(); + return true; + } + + print_plan(plan, components, devices); + + std::string derived_runtime_spec; + std::string derived_params_spec; + append_component_decision(components, devices, plan, ComponentKind::DIT, "diffusion", derived_runtime_spec, derived_params_spec); + append_component_decision(components, devices, plan, ComponentKind::CONDITIONER, "te", derived_runtime_spec, derived_params_spec); + append_component_decision(components, devices, plan, ComponentKind::VAE, "vae", derived_runtime_spec, derived_params_spec); + + runtime_spec = std::move(derived_runtime_spec); + params_spec = std::move(derived_params_spec); + + LOG_INFO("auto-fit: --backend \"%s\"%s%s%s", + runtime_spec.empty() ? "(default)" : runtime_spec.c_str(), + params_spec.empty() ? "" : " --params-backend \"", + params_spec.c_str(), + params_spec.empty() ? "" : "\""); + return true; +} + +inline bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) { + if (prefer_temporal_tiling) { + if (tiling_params.temporal_tiling) { + return false; + } + tiling_params.temporal_tiling = true; + } else { + if (tiling_params.enabled) { + return false; + } + tiling_params.enabled = true; + if (tiling_params.tile_size_x <= 0) { + tiling_params.tile_size_x = 256; + } + if (tiling_params.tile_size_y <= 0) { + tiling_params.tile_size_y = 256; + } + } + + LOG_WARN("auto-fit: VAE decode failed (likely out of memory); retrying with %s tiling", + tiling_params.temporal_tiling ? "temporal" : "spatial"); + return true; +} + } // namespace backend_fit #endif // __SD_BACKEND_FIT_HPP__ diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index b4d2d7350..690da69e8 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -766,72 +766,12 @@ class StableDiffusionGGML { } if (auto_fit_enabled) { - if (!backend_spec.empty() || !params_backend_spec.empty()) { - LOG_WARN("--auto-fit is enabled; ignoring --backend / --params-backend"); - } - // The budgets are keyed by device name; canonicalize against the - // registry (no backend needs to exist yet). - { - std::string error; - if (!max_vram_assignment.canonicalize_backend_keys(&error)) { - LOG_ERROR("%s", error.c_str()); - return false; - } - } - auto components = backend_fit::estimate_components(model_loader, wtype); - auto devices = backend_fit::enumerate_gpu_devices(max_vram_assignment); - auto plan = backend_fit::compute_plan(components, devices); - if (!plan.valid) { - LOG_WARN("auto-fit: no usable GPU devices; using the default backend"); - } else { - backend_fit::print_plan(plan, components, devices); - std::string runtime_spec; - std::string params_spec; - auto append_assignment = [](std::string& spec, const char* key, const std::string& value) { - if (!spec.empty()) { - spec += ","; - } - spec += key; - spec += "="; - spec += value; - }; - auto apply_decision = [&](backend_fit::ComponentKind kind, const char* module_key) { - for (size_t ci = 0; ci < components.size(); ci++) { - if (components[ci].kind != kind || components[ci].params_bytes == 0) { - continue; - } - const backend_fit::Decision& decision = plan.decisions[ci]; - if (decision.on_cpu) { - append_assignment(runtime_spec, module_key, "cpu"); - return; - } - if (decision.device_idxs.empty()) { - return; // default backend - } - std::string device_list; - for (size_t k = 0; k < decision.device_idxs.size(); k++) { - if (k > 0) { - device_list += "&"; - } - device_list += devices[decision.device_idxs[k]].name; - } - append_assignment(runtime_spec, module_key, device_list); - if (plan.time_share) { - append_assignment(params_spec, module_key, "disk"); - } - return; - } - }; - apply_decision(backend_fit::ComponentKind::DIT, "diffusion"); - apply_decision(backend_fit::ComponentKind::CONDITIONER, "te"); - apply_decision(backend_fit::ComponentKind::VAE, "vae"); - backend_spec = runtime_spec; - params_backend_spec = params_spec; - LOG_INFO("auto-fit: --backend \"%s\"%s%s%s", - backend_spec.empty() ? "(default)" : backend_spec.c_str(), - params_backend_spec.empty() ? "" : " --params-backend \"", - params_backend_spec.c_str(), - params_backend_spec.empty() ? "" : "\""); + if (!backend_fit::derive_backend_specs(model_loader, + wtype, + max_vram_assignment, + backend_spec, + params_backend_spec)) { + return false; } } @@ -2773,25 +2713,8 @@ class StableDiffusionGGML { // (temporal for the LTX video VAE, spatial otherwise) instead of // failing the whole generation. if (decoded.empty() && auto_fit_enabled) { - bool changed = false; - if (decode_video && std::dynamic_pointer_cast(first_stage_model) != nullptr) { - if (!vae_tiling_params.temporal_tiling) { - vae_tiling_params.temporal_tiling = true; - changed = true; - } - } else if (!vae_tiling_params.enabled) { - vae_tiling_params.enabled = true; - if (vae_tiling_params.tile_size_x <= 0) { - vae_tiling_params.tile_size_x = 256; - } - if (vae_tiling_params.tile_size_y <= 0) { - vae_tiling_params.tile_size_y = 256; - } - changed = true; - } - if (changed) { - LOG_WARN("auto-fit: VAE decode failed (likely out of memory); retrying with %s tiling", - vae_tiling_params.temporal_tiling ? "temporal" : "spatial"); + bool prefer_temporal_tiling = decode_video && std::dynamic_pointer_cast(first_stage_model) != nullptr; + if (backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) { first_stage_model->free_compute_buffer(); first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling); decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); From 05e799f6a2aba75bc61ba9846a551fbd5ce23336 Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 5 Jul 2026 00:31:29 +0800 Subject: [PATCH 5/8] refactor: split auto-fit helper into source file --- src/{backend_fit.hpp => backend_fit.cpp} | 78 ++++++++++-------------- src/backend_fit.h | 26 ++++++++ src/stable-diffusion.cpp | 14 ++--- 3 files changed, 65 insertions(+), 53 deletions(-) rename src/{backend_fit.hpp => backend_fit.cpp} (83%) create mode 100644 src/backend_fit.h diff --git a/src/backend_fit.hpp b/src/backend_fit.cpp similarity index 83% rename from src/backend_fit.hpp rename to src/backend_fit.cpp index 1c386b96e..a2c354b92 100644 --- a/src/backend_fit.hpp +++ b/src/backend_fit.cpp @@ -1,31 +1,17 @@ -#ifndef __SD_BACKEND_FIT_HPP__ -#define __SD_BACKEND_FIT_HPP__ +#include "backend_fit.h" #include #include #include -#include #include #include #include "core/ggml_extend_backend.h" -#include "core/ggml_graph_cut.h" #include "core/util.h" #include "ggml-backend.h" -#include "model_loader.h" -#include "stable-diffusion.h" - -// Auto-fit (--auto-fit): derive --backend / --params-backend placements for -// the DiT / conditioner / VAE from the model metadata and per-device memory -// budgets. This is a policy layer only — it emits the same backend assignment -// specs a user would pass manually (including '&' device lists for the layer / -// row split mechanisms) and feeds them to SDBackendManager. -// -// Budgets come from --max-vram: a positive per-device value caps the memory -// auto-fit plans with on that device, a negative value means "free memory -// minus that many GiB", and an unset budget defaults to the device's free -// memory minus a 512 MiB margin. -namespace backend_fit { + +namespace sd::backend_fit { +namespace { constexpr int64_t MiB = 1024ll * 1024; @@ -68,7 +54,7 @@ struct Plan { std::vector decisions; }; -inline bool classify_tensor(const std::string& name, ComponentKind& out) { +bool classify_tensor(const std::string& name, ComponentKind& out) { auto contains = [&](const char* s) { return name.find(s) != std::string::npos; }; if (contains("model.diffusion_model.") || contains("unet.")) { @@ -96,7 +82,7 @@ inline bool classify_tensor(const std::string& name, ComponentKind& out) { return false; } -inline std::vector estimate_components(ModelLoader& loader, ggml_type override_wtype) { +std::vector estimate_components(ModelLoader& loader, ggml_type override_wtype) { const auto& storage = loader.get_tensor_storage_map(); int64_t bytes[3] = {0, 0, 0}; @@ -128,7 +114,7 @@ inline std::vector estimate_components(ModelLoader& loader, ggml_type return out; } -inline std::vector enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) { +std::vector enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) { std::vector out; for (size_t i = 0; i < ggml_backend_dev_count(); i++) { ggml_backend_dev_t dev = ggml_backend_dev_get(i); @@ -172,7 +158,7 @@ inline std::vector enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVr // on demand and frees them afterwards, so each component only has to fit on // its own. A component that cannot fit a single device splits across all // GPUs when its module supports it, and lands on the CPU otherwise. -inline Plan compute_plan(const std::vector& components, const std::vector& devices) { +Plan compute_plan(const std::vector& components, const std::vector& devices) { Plan plan; if (devices.empty()) { return plan; @@ -227,9 +213,9 @@ inline Plan compute_plan(const std::vector& components, const std::ve // Time-share regime: each component fits on its own (params=disk). plan.decisions.assign(components.size(), {}); for (size_t ci : order) { - const Component& comp = components[ci]; - Decision& decision = plan.decisions[ci]; - decision.kind = comp.kind; + const Component& comp = components[ci]; + Decision& decision = plan.decisions[ci]; + decision.kind = comp.kind; if (comp.params_bytes == 0) { continue; } @@ -270,9 +256,9 @@ inline Plan compute_plan(const std::vector& components, const std::ve return plan; } -inline void print_plan(const Plan& plan, - const std::vector& components, - const std::vector& devices) { +void print_plan(const Plan& plan, + const std::vector& components, + const std::vector& devices) { LOG_INFO("auto-fit plan%s:", plan.time_share ? " (time-share: params load per phase and free after)" : ""); LOG_INFO(" devices:"); for (const Device& d : devices) { @@ -282,7 +268,7 @@ inline void print_plan(const Plan& plan, } LOG_INFO(" components:"); for (size_t ci = 0; ci < components.size(); ci++) { - const Component& comp = components[ci]; + const Component& comp = components[ci]; const Decision& decision = plan.decisions[ci]; std::string target; if (comp.params_bytes == 0) { @@ -308,7 +294,7 @@ inline void print_plan(const Plan& plan, } } -inline void append_assignment(std::string& spec, const char* key, const std::string& value) { +void append_assignment(std::string& spec, const char* key, const std::string& value) { if (!spec.empty()) { spec += ","; } @@ -317,13 +303,13 @@ inline void append_assignment(std::string& spec, const char* key, const std::str spec += value; } -inline void append_component_decision(const std::vector& components, - const std::vector& devices, - const Plan& plan, - ComponentKind kind, - const char* module_key, - std::string& runtime_spec, - std::string& params_spec) { +void append_component_decision(const std::vector& components, + const std::vector& devices, + const Plan& plan, + ComponentKind kind, + const char* module_key, + std::string& runtime_spec, + std::string& params_spec) { for (size_t ci = 0; ci < components.size(); ci++) { if (components[ci].kind != kind || components[ci].params_bytes == 0) { continue; @@ -351,11 +337,13 @@ inline void append_component_decision(const std::vector& components, } } -inline bool derive_backend_specs(ModelLoader& loader, - ggml_type override_wtype, - sd::ggml_graph_cut::MaxVramAssignment& budgets, - std::string& runtime_spec, - std::string& params_spec) { +} // namespace + +bool derive_backend_specs(ModelLoader& loader, + ggml_type override_wtype, + sd::ggml_graph_cut::MaxVramAssignment& budgets, + std::string& runtime_spec, + std::string& params_spec) { if (!runtime_spec.empty() || !params_spec.empty()) { LOG_WARN("--auto-fit is enabled; ignoring --backend / --params-backend"); } @@ -399,7 +387,7 @@ inline bool derive_backend_specs(ModelLoader& loader, return true; } -inline bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) { +bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) { if (prefer_temporal_tiling) { if (tiling_params.temporal_tiling) { return false; @@ -423,6 +411,4 @@ inline bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, b return true; } -} // namespace backend_fit - -#endif // __SD_BACKEND_FIT_HPP__ +} // namespace sd::backend_fit diff --git a/src/backend_fit.h b/src/backend_fit.h new file mode 100644 index 000000000..d9ab0eb99 --- /dev/null +++ b/src/backend_fit.h @@ -0,0 +1,26 @@ +#ifndef __SD_BACKEND_FIT_H__ +#define __SD_BACKEND_FIT_H__ + +#include + +#include "core/ggml_graph_cut.h" +#include "model_loader.h" +#include "stable-diffusion.h" + +// Auto-fit (--auto-fit): derive --backend / --params-backend placements for +// the DiT / conditioner / VAE from the model metadata and per-device memory +// budgets. +namespace sd::backend_fit { + +bool derive_backend_specs(ModelLoader& loader, + ggml_type override_wtype, + sd::ggml_graph_cut::MaxVramAssignment& budgets, + std::string& runtime_spec, + std::string& params_spec); + +bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, + bool prefer_temporal_tiling); + +} // namespace sd::backend_fit + +#endif // __SD_BACKEND_FIT_H__ diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 690da69e8..d8a7d71fb 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -19,7 +19,7 @@ #include "model_manager.h" #include "stable-diffusion.h" -#include "backend_fit.hpp" +#include "backend_fit.h" #include "conditioning/conditioner.hpp" #include "extensions/generation_extension.h" #include "model/adapter/lora.hpp" @@ -766,11 +766,11 @@ class StableDiffusionGGML { } if (auto_fit_enabled) { - if (!backend_fit::derive_backend_specs(model_loader, - wtype, - max_vram_assignment, - backend_spec, - params_backend_spec)) { + if (!sd::backend_fit::derive_backend_specs(model_loader, + wtype, + max_vram_assignment, + backend_spec, + params_backend_spec)) { return false; } } @@ -2714,7 +2714,7 @@ class StableDiffusionGGML { // failing the whole generation. if (decoded.empty() && auto_fit_enabled) { bool prefer_temporal_tiling = decode_video && std::dynamic_pointer_cast(first_stage_model) != nullptr; - if (backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) { + if (sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) { first_stage_model->free_compute_buffer(); first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling); decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); From 89ea2b6d3a23854bf872a74c9ea5bdeb86a0d19a Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 5 Jul 2026 00:34:28 +0800 Subject: [PATCH 6/8] clean code --- include/stable-diffusion.h | 5 - src/backend_fit.cpp | 702 ++++++++++++++++++------------------- src/backend_fit.h | 17 +- src/stable-diffusion.cpp | 9 - 4 files changed, 346 insertions(+), 387 deletions(-) diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 1c74ae7f8..c905698e5 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -228,11 +228,6 @@ typedef struct { const char* backend; const char* params_backend; const char* split_mode; // weight distribution for multi-device modules: layer (default) or row, or per-module assignments e.g. "diffusion=row" - // Pick the diffusion/te/vae device placements automatically from the model - // metadata and the per-device memory budgets (max_vram). When enabled, - // `backend` and `params_backend` are ignored and derived instead; the - // plan may use multi-device layer/row splits (`split_mode` still applies) - // and disk params residency so components time-share VRAM. bool auto_fit; const char* rpc_servers; } sd_ctx_params_t; diff --git a/src/backend_fit.cpp b/src/backend_fit.cpp index a2c354b92..460194740 100644 --- a/src/backend_fit.cpp +++ b/src/backend_fit.cpp @@ -11,404 +11,380 @@ #include "ggml-backend.h" namespace sd::backend_fit { -namespace { - -constexpr int64_t MiB = 1024ll * 1024; - -enum class ComponentKind { - DIT = 0, - VAE = 1, - CONDITIONER = 2, -}; - -struct Component { - ComponentKind kind; - const char* name; - int64_t params_bytes = 0; - // Estimated compute-buffer bytes the component needs alongside its params - // on every device it runs on. - int64_t reserve_bytes = 0; - bool splittable = false; -}; - -struct Device { - ggml_backend_dev_t dev = nullptr; - std::string name; - std::string description; - int64_t free_bytes = 0; - int64_t total_bytes = 0; - int64_t budget_bytes = 0; -}; - -struct Decision { - ComponentKind kind; - bool on_cpu = false; - std::vector device_idxs; // more than one entry = split device list -}; - -struct Plan { - bool valid = false; - // Components load lazily (params=disk) and time-share the budgets phase - // by phase instead of all being resident at once. - bool time_share = false; - std::vector decisions; -}; - -bool classify_tensor(const std::string& name, ComponentKind& out) { - auto contains = [&](const char* s) { return name.find(s) != std::string::npos; }; - - if (contains("model.diffusion_model.") || contains("unet.")) { - out = ComponentKind::DIT; - return true; - } - if (contains("first_stage_model.") || - name.rfind("vae.", 0) == 0 || - name.rfind("tae.", 0) == 0) { - out = ComponentKind::VAE; - return true; - } - if (contains("text_encoders") || - contains("cond_stage_model") || - contains("te.text_model.") || - contains("conditioner") || - name.rfind("text_encoder.", 0) == 0 || - // Connector / text projection layers that run on the conditioner - // backend (e.g. LTX-2's text_embedding_projection). - name.rfind("text_embedding_projection.", 0) == 0 || - contains(".aggregate_embed.")) { - out = ComponentKind::CONDITIONER; - return true; - } - return false; -} + namespace { + + constexpr int64_t MiB = 1024ll * 1024; + + enum class ComponentKind { + DIT = 0, + VAE = 1, + CONDITIONER = 2, + }; + + struct Component { + ComponentKind kind; + const char* name; + int64_t params_bytes = 0; + int64_t reserve_bytes = 0; + bool splittable = false; + }; + + struct Device { + ggml_backend_dev_t dev = nullptr; + std::string name; + std::string description; + int64_t free_bytes = 0; + int64_t total_bytes = 0; + int64_t budget_bytes = 0; + }; + + struct Decision { + ComponentKind kind; + bool on_cpu = false; + std::vector device_idxs; + }; + + struct Plan { + bool valid = false; + bool time_share = false; + std::vector decisions; + }; + + bool classify_tensor(const std::string& name, ComponentKind& out) { + auto contains = [&](const char* s) { return name.find(s) != std::string::npos; }; + + if (contains("model.diffusion_model.") || contains("unet.")) { + out = ComponentKind::DIT; + return true; + } + if (contains("first_stage_model.") || + name.rfind("vae.", 0) == 0 || + name.rfind("tae.", 0) == 0) { + out = ComponentKind::VAE; + return true; + } + if (contains("text_encoders") || + contains("cond_stage_model") || + contains("te.text_model.") || + contains("conditioner") || + name.rfind("text_encoder.", 0) == 0 || + name.rfind("text_embedding_projection.", 0) == 0 || + contains(".aggregate_embed.")) { + out = ComponentKind::CONDITIONER; + return true; + } + return false; + } -std::vector estimate_components(ModelLoader& loader, ggml_type override_wtype) { - const auto& storage = loader.get_tensor_storage_map(); + std::vector estimate_components(ModelLoader& loader, ggml_type override_wtype) { + const auto& storage = loader.get_tensor_storage_map(); - int64_t bytes[3] = {0, 0, 0}; - for (const auto& [name, ts_const] : storage) { - TensorStorage ts = ts_const; - if (is_unused_tensor(ts.name)) { - continue; - } - ComponentKind kind; - if (!classify_tensor(ts.name, kind)) { - continue; - } - if (override_wtype != GGML_TYPE_COUNT && - loader.tensor_should_be_converted(ts, override_wtype)) { - ts.type = override_wtype; - } else if (ts.expected_type != GGML_TYPE_COUNT && ts.expected_type != ts.type) { - ts.type = ts.expected_type; - } - bytes[int(kind)] += (int64_t)ts.nbytes() + 64; - } + int64_t bytes[3] = {0, 0, 0}; + for (const auto& [name, ts_const] : storage) { + TensorStorage ts = ts_const; + if (is_unused_tensor(ts.name)) { + continue; + } + ComponentKind kind; + if (!classify_tensor(ts.name, kind)) { + continue; + } + if (override_wtype != GGML_TYPE_COUNT && + loader.tensor_should_be_converted(ts, override_wtype)) { + ts.type = override_wtype; + } else if (ts.expected_type != GGML_TYPE_COUNT && ts.expected_type != ts.type) { + ts.type = ts.expected_type; + } + bytes[int(kind)] += (int64_t)ts.nbytes() + 64; + } - // Built-in per-component compute reserves. These are per-device estimates - // for typical workloads; --max-vram can shrink the budgets when they turn - // out too small for a given resolution. - std::vector out; - out.push_back({ComponentKind::DIT, "DiT", bytes[int(ComponentKind::DIT)], 2048 * MiB, true}); - out.push_back({ComponentKind::VAE, "VAE", bytes[int(ComponentKind::VAE)], 1024 * MiB, false}); - out.push_back({ComponentKind::CONDITIONER, "Conditioner", bytes[int(ComponentKind::CONDITIONER)], 2048 * MiB, true}); - return out; -} - -std::vector enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) { - std::vector out; - for (size_t i = 0; i < ggml_backend_dev_count(); i++) { - ggml_backend_dev_t dev = ggml_backend_dev_get(i); - if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) { - continue; + std::vector out; + out.push_back({ComponentKind::DIT, "DiT", bytes[int(ComponentKind::DIT)], 2048 * MiB, true}); + out.push_back({ComponentKind::VAE, "VAE", bytes[int(ComponentKind::VAE)], 1024 * MiB, false}); + out.push_back({ComponentKind::CONDITIONER, "Conditioner", bytes[int(ComponentKind::CONDITIONER)], 2048 * MiB, true}); + return out; } - Device d; - d.dev = dev; - d.name = ggml_backend_dev_name(dev); - d.description = ggml_backend_dev_description(dev); - size_t free_bytes = 0, total_bytes = 0; - ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); - d.free_bytes = (int64_t)free_bytes; - d.total_bytes = (int64_t)total_bytes; - - // MaxVramAssignment canonicalizes its keys to lowercase device names. - std::string budget_key = d.name; - std::transform(budget_key.begin(), budget_key.end(), budget_key.begin(), - [](unsigned char c) { return (char)std::tolower(c); }); - float gib = budgets.default_gib; - auto it = budgets.backend_gib.find(budget_key); - if (it != budgets.backend_gib.end()) { - gib = it->second; - } - if (gib > 0.f) { - d.budget_bytes = std::min((int64_t)(gib * 1024.0 * 1024.0 * 1024.0), d.free_bytes); - } else if (gib < 0.f) { - d.budget_bytes = d.free_bytes + (int64_t)(gib * 1024.0 * 1024.0 * 1024.0); - } else { - d.budget_bytes = d.free_bytes - 512 * MiB; + + std::vector enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) { + std::vector out; + for (size_t i = 0; i < ggml_backend_dev_count(); i++) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) { + continue; + } + Device d; + d.dev = dev; + d.name = ggml_backend_dev_name(dev); + d.description = ggml_backend_dev_description(dev); + size_t free_bytes = 0, total_bytes = 0; + ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); + d.free_bytes = (int64_t)free_bytes; + d.total_bytes = (int64_t)total_bytes; + + std::string budget_key = d.name; + std::transform(budget_key.begin(), budget_key.end(), budget_key.begin(), + [](unsigned char c) { return (char)std::tolower(c); }); + float gib = budgets.default_gib; + auto it = budgets.backend_gib.find(budget_key); + if (it != budgets.backend_gib.end()) { + gib = it->second; + } + if (gib > 0.f) { + d.budget_bytes = std::min((int64_t)(gib * 1024.0 * 1024.0 * 1024.0), d.free_bytes); + } else if (gib < 0.f) { + d.budget_bytes = d.free_bytes + (int64_t)(gib * 1024.0 * 1024.0 * 1024.0); + } else { + d.budget_bytes = d.free_bytes - 512 * MiB; + } + d.budget_bytes = std::max(d.budget_bytes, 0); + out.push_back(d); + } + return out; } - d.budget_bytes = std::max(d.budget_bytes, 0); - out.push_back(d); - } - return out; -} - -// Place each component, largest first. First try the eager regime (all -// components resident at once); when that overflows, fall back to the -// time-share regime, where params=disk residency loads each phase's weights -// on demand and frees them afterwards, so each component only has to fit on -// its own. A component that cannot fit a single device splits across all -// GPUs when its module supports it, and lands on the CPU otherwise. -Plan compute_plan(const std::vector& components, const std::vector& devices) { - Plan plan; - if (devices.empty()) { - return plan; - } - std::vector order(components.size()); - for (size_t i = 0; i < order.size(); i++) { - order[i] = i; - } - std::sort(order.begin(), order.end(), [&](size_t a, size_t b) { - return components[a].params_bytes > components[b].params_bytes; - }); - - // Eager regime: components coexist; per device, the params sum plus the - // largest single reserve must fit the budget (compute buffers of the - // sequential phases do not coexist). - { - std::vector params_sum(devices.size(), 0); - std::vector max_reserve(devices.size(), 0); - std::vector decisions(components.size()); - bool ok = true; - for (size_t ci : order) { - const Component& comp = components[ci]; - decisions[ci].kind = comp.kind; - if (comp.params_bytes == 0) { - continue; + Plan compute_plan(const std::vector& components, const std::vector& devices) { + Plan plan; + if (devices.empty()) { + return plan; } - int best = -1; - for (size_t di = 0; di < devices.size(); di++) { - int64_t need = params_sum[di] + comp.params_bytes + std::max(max_reserve[di], comp.reserve_bytes); - if (need <= devices[di].budget_bytes && - (best < 0 || devices[di].budget_bytes - params_sum[di] > devices[best].budget_bytes - params_sum[best])) { - best = (int)di; + + std::vector order(components.size()); + for (size_t i = 0; i < order.size(); i++) { + order[i] = i; + } + std::sort(order.begin(), order.end(), [&](size_t a, size_t b) { + return components[a].params_bytes > components[b].params_bytes; + }); + + { + std::vector params_sum(devices.size(), 0); + std::vector max_reserve(devices.size(), 0); + std::vector decisions(components.size()); + bool ok = true; + for (size_t ci : order) { + const Component& comp = components[ci]; + decisions[ci].kind = comp.kind; + if (comp.params_bytes == 0) { + continue; + } + int best = -1; + for (size_t di = 0; di < devices.size(); di++) { + int64_t need = params_sum[di] + comp.params_bytes + std::max(max_reserve[di], comp.reserve_bytes); + if (need <= devices[di].budget_bytes && + (best < 0 || devices[di].budget_bytes - params_sum[di] > devices[best].budget_bytes - params_sum[best])) { + best = (int)di; + } + } + if (best < 0) { + ok = false; + break; + } + params_sum[best] += comp.params_bytes; + max_reserve[best] = std::max(max_reserve[best], comp.reserve_bytes); + decisions[ci].device_idxs.push_back((size_t)best); + } + if (ok) { + plan.valid = true; + plan.time_share = false; + plan.decisions = std::move(decisions); + return plan; } } - if (best < 0) { - ok = false; - break; + + plan.decisions.assign(components.size(), {}); + for (size_t ci : order) { + const Component& comp = components[ci]; + Decision& decision = plan.decisions[ci]; + decision.kind = comp.kind; + if (comp.params_bytes == 0) { + continue; + } + int best = -1; + for (size_t di = 0; di < devices.size(); di++) { + if (comp.params_bytes + comp.reserve_bytes <= devices[di].budget_bytes && + (best < 0 || devices[di].budget_bytes > devices[best].budget_bytes)) { + best = (int)di; + } + } + if (best >= 0) { + decision.device_idxs.push_back((size_t)best); + continue; + } + if (comp.splittable && devices.size() > 1) { + int64_t capacity = 0; + for (const Device& d : devices) { + capacity += std::max(d.budget_bytes - comp.reserve_bytes, 0); + } + if (comp.params_bytes <= capacity) { + std::vector idxs(devices.size()); + for (size_t i = 0; i < idxs.size(); i++) { + idxs[i] = i; + } + std::sort(idxs.begin(), idxs.end(), [&](size_t a, size_t b) { + return devices[a].budget_bytes > devices[b].budget_bytes; + }); + decision.device_idxs = std::move(idxs); + continue; + } + } + decision.on_cpu = true; } - params_sum[best] += comp.params_bytes; - max_reserve[best] = std::max(max_reserve[best], comp.reserve_bytes); - decisions[ci].device_idxs.push_back((size_t)best); - } - if (ok) { plan.valid = true; - plan.time_share = false; - plan.decisions = std::move(decisions); + plan.time_share = true; return plan; } - } - // Time-share regime: each component fits on its own (params=disk). - plan.decisions.assign(components.size(), {}); - for (size_t ci : order) { - const Component& comp = components[ci]; - Decision& decision = plan.decisions[ci]; - decision.kind = comp.kind; - if (comp.params_bytes == 0) { - continue; - } - int best = -1; - for (size_t di = 0; di < devices.size(); di++) { - if (comp.params_bytes + comp.reserve_bytes <= devices[di].budget_bytes && - (best < 0 || devices[di].budget_bytes > devices[best].budget_bytes)) { - best = (int)di; - } - } - if (best >= 0) { - decision.device_idxs.push_back((size_t)best); - continue; - } - if (comp.splittable && devices.size() > 1) { - int64_t capacity = 0; + void print_plan(const Plan& plan, + const std::vector& components, + const std::vector& devices) { + LOG_INFO("auto-fit plan%s:", plan.time_share ? " (time-share: params load per phase and free after)" : ""); + LOG_INFO(" devices:"); for (const Device& d : devices) { - capacity += std::max(d.budget_bytes - comp.reserve_bytes, 0); + LOG_INFO(" %-12s %-32s free %6lld MiB, budget %6lld MiB", + d.name.c_str(), d.description.c_str(), + (long long)(d.free_bytes / MiB), (long long)(d.budget_bytes / MiB)); } - if (comp.params_bytes <= capacity) { - // All GPUs, largest budget first (the first device is the - // module's main device). - std::vector idxs(devices.size()); - for (size_t i = 0; i < idxs.size(); i++) { - idxs[i] = i; + LOG_INFO(" components:"); + for (size_t ci = 0; ci < components.size(); ci++) { + const Component& comp = components[ci]; + const Decision& decision = plan.decisions[ci]; + std::string target; + if (comp.params_bytes == 0) { + target = "(not present)"; + } else if (decision.on_cpu) { + target = "CPU"; + } else { + for (size_t k = 0; k < decision.device_idxs.size(); k++) { + if (k > 0) { + target += " & "; + } + target += devices[decision.device_idxs[k]].name; + } + if (decision.device_idxs.size() > 1) { + target += " (split)"; + } } - std::sort(idxs.begin(), idxs.end(), [&](size_t a, size_t b) { - return devices[a].budget_bytes > devices[b].budget_bytes; - }); - decision.device_idxs = std::move(idxs); - continue; + LOG_INFO(" %-12s params %6lld MiB, compute reserve %5lld MiB -> %s", + comp.name, + (long long)(comp.params_bytes / MiB), + (long long)(comp.reserve_bytes / MiB), + target.c_str()); } } - decision.on_cpu = true; - } - plan.valid = true; - plan.time_share = true; - return plan; -} - -void print_plan(const Plan& plan, - const std::vector& components, - const std::vector& devices) { - LOG_INFO("auto-fit plan%s:", plan.time_share ? " (time-share: params load per phase and free after)" : ""); - LOG_INFO(" devices:"); - for (const Device& d : devices) { - LOG_INFO(" %-12s %-32s free %6lld MiB, budget %6lld MiB", - d.name.c_str(), d.description.c_str(), - (long long)(d.free_bytes / MiB), (long long)(d.budget_bytes / MiB)); - } - LOG_INFO(" components:"); - for (size_t ci = 0; ci < components.size(); ci++) { - const Component& comp = components[ci]; - const Decision& decision = plan.decisions[ci]; - std::string target; - if (comp.params_bytes == 0) { - target = "(not present)"; - } else if (decision.on_cpu) { - target = "CPU"; - } else { - for (size_t k = 0; k < decision.device_idxs.size(); k++) { - if (k > 0) { - target += " & "; - } - target += devices[decision.device_idxs[k]].name; - } - if (decision.device_idxs.size() > 1) { - target += " (split)"; + + void append_assignment(std::string& spec, const char* key, const std::string& value) { + if (!spec.empty()) { + spec += ","; } + spec += key; + spec += "="; + spec += value; } - LOG_INFO(" %-12s params %6lld MiB, compute reserve %5lld MiB -> %s", - comp.name, - (long long)(comp.params_bytes / MiB), - (long long)(comp.reserve_bytes / MiB), - target.c_str()); - } -} -void append_assignment(std::string& spec, const char* key, const std::string& value) { - if (!spec.empty()) { - spec += ","; - } - spec += key; - spec += "="; - spec += value; -} - -void append_component_decision(const std::vector& components, - const std::vector& devices, - const Plan& plan, - ComponentKind kind, - const char* module_key, - std::string& runtime_spec, - std::string& params_spec) { - for (size_t ci = 0; ci < components.size(); ci++) { - if (components[ci].kind != kind || components[ci].params_bytes == 0) { - continue; - } - const Decision& decision = plan.decisions[ci]; - if (decision.on_cpu) { - append_assignment(runtime_spec, module_key, "cpu"); - return; - } - if (decision.device_idxs.empty()) { - return; - } - std::string device_list; - for (size_t k = 0; k < decision.device_idxs.size(); k++) { - if (k > 0) { - device_list += "&"; + void append_component_decision(const std::vector& components, + const std::vector& devices, + const Plan& plan, + ComponentKind kind, + const char* module_key, + std::string& runtime_spec, + std::string& params_spec) { + for (size_t ci = 0; ci < components.size(); ci++) { + if (components[ci].kind != kind || components[ci].params_bytes == 0) { + continue; + } + const Decision& decision = plan.decisions[ci]; + if (decision.on_cpu) { + append_assignment(runtime_spec, module_key, "cpu"); + return; + } + if (decision.device_idxs.empty()) { + return; + } + std::string device_list; + for (size_t k = 0; k < decision.device_idxs.size(); k++) { + if (k > 0) { + device_list += "&"; + } + device_list += devices[decision.device_idxs[k]].name; + } + append_assignment(runtime_spec, module_key, device_list); + if (plan.time_share) { + append_assignment(params_spec, module_key, "disk"); + } + return; } - device_list += devices[decision.device_idxs[k]].name; - } - append_assignment(runtime_spec, module_key, device_list); - if (plan.time_share) { - append_assignment(params_spec, module_key, "disk"); } - return; - } -} -} // namespace + } // namespace -bool derive_backend_specs(ModelLoader& loader, - ggml_type override_wtype, - sd::ggml_graph_cut::MaxVramAssignment& budgets, - std::string& runtime_spec, - std::string& params_spec) { - if (!runtime_spec.empty() || !params_spec.empty()) { - LOG_WARN("--auto-fit is enabled; ignoring --backend / --params-backend"); - } + bool derive_backend_specs(ModelLoader& loader, + ggml_type override_wtype, + sd::ggml_graph_cut::MaxVramAssignment& budgets, + std::string& runtime_spec, + std::string& params_spec) { + if (!runtime_spec.empty() || !params_spec.empty()) { + LOG_WARN("--auto-fit is enabled; ignoring --backend / --params-backend"); + } - // The budgets are keyed by device name; canonicalize against the registry - // before SDBackendManager creates any backend instance. - { - std::string error; - if (!budgets.canonicalize_backend_keys(&error)) { - LOG_ERROR("%s", error.c_str()); - return false; + { + std::string error; + if (!budgets.canonicalize_backend_keys(&error)) { + LOG_ERROR("%s", error.c_str()); + return false; + } } - } - auto components = estimate_components(loader, override_wtype); - auto devices = enumerate_gpu_devices(budgets); - auto plan = compute_plan(components, devices); - if (!plan.valid) { - LOG_WARN("auto-fit: no usable GPU devices; using the default backend"); - runtime_spec.clear(); - params_spec.clear(); - return true; - } + auto components = estimate_components(loader, override_wtype); + auto devices = enumerate_gpu_devices(budgets); + auto plan = compute_plan(components, devices); + if (!plan.valid) { + LOG_WARN("auto-fit: no usable GPU devices; using the default backend"); + runtime_spec.clear(); + params_spec.clear(); + return true; + } - print_plan(plan, components, devices); + print_plan(plan, components, devices); - std::string derived_runtime_spec; - std::string derived_params_spec; - append_component_decision(components, devices, plan, ComponentKind::DIT, "diffusion", derived_runtime_spec, derived_params_spec); - append_component_decision(components, devices, plan, ComponentKind::CONDITIONER, "te", derived_runtime_spec, derived_params_spec); - append_component_decision(components, devices, plan, ComponentKind::VAE, "vae", derived_runtime_spec, derived_params_spec); + std::string derived_runtime_spec; + std::string derived_params_spec; + append_component_decision(components, devices, plan, ComponentKind::DIT, "diffusion", derived_runtime_spec, derived_params_spec); + append_component_decision(components, devices, plan, ComponentKind::CONDITIONER, "te", derived_runtime_spec, derived_params_spec); + append_component_decision(components, devices, plan, ComponentKind::VAE, "vae", derived_runtime_spec, derived_params_spec); - runtime_spec = std::move(derived_runtime_spec); - params_spec = std::move(derived_params_spec); + runtime_spec = std::move(derived_runtime_spec); + params_spec = std::move(derived_params_spec); - LOG_INFO("auto-fit: --backend \"%s\"%s%s%s", - runtime_spec.empty() ? "(default)" : runtime_spec.c_str(), - params_spec.empty() ? "" : " --params-backend \"", - params_spec.c_str(), - params_spec.empty() ? "" : "\""); - return true; -} + LOG_INFO("auto-fit: --backend \"%s\"%s%s%s", + runtime_spec.empty() ? "(default)" : runtime_spec.c_str(), + params_spec.empty() ? "" : " --params-backend \"", + params_spec.c_str(), + params_spec.empty() ? "" : "\""); + return true; + } -bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) { - if (prefer_temporal_tiling) { - if (tiling_params.temporal_tiling) { - return false; - } - tiling_params.temporal_tiling = true; - } else { - if (tiling_params.enabled) { - return false; - } - tiling_params.enabled = true; - if (tiling_params.tile_size_x <= 0) { - tiling_params.tile_size_x = 256; - } - if (tiling_params.tile_size_y <= 0) { - tiling_params.tile_size_y = 256; + bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) { + if (prefer_temporal_tiling) { + if (tiling_params.temporal_tiling) { + return false; + } + tiling_params.temporal_tiling = true; + } else { + if (tiling_params.enabled) { + return false; + } + tiling_params.enabled = true; + if (tiling_params.tile_size_x <= 0) { + tiling_params.tile_size_x = 256; + } + if (tiling_params.tile_size_y <= 0) { + tiling_params.tile_size_y = 256; + } } - } - LOG_WARN("auto-fit: VAE decode failed (likely out of memory); retrying with %s tiling", - tiling_params.temporal_tiling ? "temporal" : "spatial"); - return true; -} + LOG_WARN("auto-fit: VAE decode failed (likely out of memory); retrying with %s tiling", + tiling_params.temporal_tiling ? "temporal" : "spatial"); + return true; + } } // namespace sd::backend_fit diff --git a/src/backend_fit.h b/src/backend_fit.h index d9ab0eb99..9ef298b3b 100644 --- a/src/backend_fit.h +++ b/src/backend_fit.h @@ -7,19 +7,16 @@ #include "model_loader.h" #include "stable-diffusion.h" -// Auto-fit (--auto-fit): derive --backend / --params-backend placements for -// the DiT / conditioner / VAE from the model metadata and per-device memory -// budgets. namespace sd::backend_fit { -bool derive_backend_specs(ModelLoader& loader, - ggml_type override_wtype, - sd::ggml_graph_cut::MaxVramAssignment& budgets, - std::string& runtime_spec, - std::string& params_spec); + bool derive_backend_specs(ModelLoader& loader, + ggml_type override_wtype, + sd::ggml_graph_cut::MaxVramAssignment& budgets, + std::string& runtime_spec, + std::string& params_spec); -bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, - bool prefer_temporal_tiling); + bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, + bool prefer_temporal_tiling); } // namespace sd::backend_fit diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index d8a7d71fb..34d2dc58b 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -614,10 +614,6 @@ class StableDiffusionGGML { ggml_log_set(ggml_log_callback_default, nullptr); - // Backend initialization happens after the model metadata is loaded, - // so --auto-fit can size the components and derive device placements - // before any backend is created. - model_manager = std::make_shared(); model_manager->set_n_threads(n_threads); model_manager->set_enable_mmap(enable_mmap); @@ -2707,11 +2703,6 @@ class StableDiffusionGGML { auto latents = first_stage_model->diffusion_to_vae_latents(x); first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling); auto decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); - // Auto-fit picks the placement without knowing the output size; a - // full-frame decode can exceed the VAE compute reserve and fail - // gracefully with an empty result. Enable tiling and retry once - // (temporal for the LTX video VAE, spatial otherwise) instead of - // failing the whole generation. if (decoded.empty() && auto_fit_enabled) { bool prefer_temporal_tiling = decode_video && std::dynamic_pointer_cast(first_stage_model) != nullptr; if (sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) { From e4e9f0afadfb5677219092d05bf2f97e9aa4d984 Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 5 Jul 2026 00:37:41 +0800 Subject: [PATCH 7/8] mv src/backend_fit* to src/core/ --- src/{ => core}/backend_fit.cpp | 0 src/{ => core}/backend_fit.h | 0 src/stable-diffusion.cpp | 2 +- 3 files changed, 1 insertion(+), 1 deletion(-) rename src/{ => core}/backend_fit.cpp (100%) rename src/{ => core}/backend_fit.h (100%) diff --git a/src/backend_fit.cpp b/src/core/backend_fit.cpp similarity index 100% rename from src/backend_fit.cpp rename to src/core/backend_fit.cpp diff --git a/src/backend_fit.h b/src/core/backend_fit.h similarity index 100% rename from src/backend_fit.h rename to src/core/backend_fit.h diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 34d2dc58b..dca2cb926 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -19,7 +19,7 @@ #include "model_manager.h" #include "stable-diffusion.h" -#include "backend_fit.h" +#include "core/backend_fit.h" #include "conditioning/conditioner.hpp" #include "extensions/generation_extension.h" #include "model/adapter/lora.hpp" From 7f77e35475109947fcaea164a354f2bebcd9623c Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 5 Jul 2026 00:38:09 +0800 Subject: [PATCH 8/8] format code --- src/stable-diffusion.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index dca2cb926..91c25ae79 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -19,8 +19,8 @@ #include "model_manager.h" #include "stable-diffusion.h" -#include "core/backend_fit.h" #include "conditioning/conditioner.hpp" +#include "core/backend_fit.h" #include "extensions/generation_extension.h" #include "model/adapter/lora.hpp" #include "model/diffusion/anima.hpp"