From 538d8cdd53da6a89f40b06295c9e8f07dfedb8b5 Mon Sep 17 00:00:00 2001 From: Piotr Wilkin Date: Fri, 3 Jul 2026 16:21:02 +0200 Subject: [PATCH 1/3] 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/3] 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 bb80f5879210e48ff7dd8cfd0b29fcc458459109 Mon Sep 17 00:00:00 2001 From: leejet Date: Sat, 4 Jul 2026 16:30:33 +0800 Subject: [PATCH 3/3] clean code --- src/core/ggml_extend_backend.h | 1 - src/model_manager.cpp | 18 +++--------------- src/model_manager.h | 14 ++------------ src/stable-diffusion.cpp | 7 ++++--- 4 files changed, 9 insertions(+), 31 deletions(-) diff --git a/src/core/ggml_extend_backend.h b/src/core/ggml_extend_backend.h index f75cf1b85..1f3bc8b3a 100644 --- a/src/core/ggml_extend_backend.h +++ b/src/core/ggml_extend_backend.h @@ -38,7 +38,6 @@ struct SDBackendHandleDeleter { using SDBackendHandle = std::unique_ptr; -// How a module with a multi-device runtime assignment distributes its weights. enum class SDSplitMode { LAYER, ROW, diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 2eaa13460..c5bddcc9e 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -112,11 +112,6 @@ void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_ba } 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) && @@ -272,9 +267,6 @@ bool ModelManager::load_tensors_to_params_backend(const std::vector& states) { - // 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)) { @@ -395,12 +387,10 @@ bool ModelManager::apply_loras_to_params(const std::vector& states } 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)"); + 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_; @@ -751,8 +741,6 @@ ggml_backend_buffer_type_t ModelManager::params_buffer_type_for(const TensorStat 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) { diff --git a/src/model_manager.h b/src/model_manager.h index 69ad4c9c3..8ec7a7a17 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -36,12 +36,8 @@ class ModelManager : public RunnerWeightManager { ResidencyMode residency_mode = ResidencyMode::ParamBackend; ggml_backend_t compute_backend = nullptr; ggml_backend_t params_backend = nullptr; - // 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; + bool allow_split_buffer = false; + bool metadata_validated = false; int active_prepare_count = 0; @@ -68,7 +64,6 @@ 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_; @@ -123,13 +118,8 @@ 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; diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index fd407e6d4..bdc2e9253 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -367,7 +367,7 @@ class StableDiffusionGGML { 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); + (int64_t)free_bytes / 8); tensor_split[reg_index] = usable_bytes > 0 ? (float)((double)usable_bytes / (1024.0 * 1024.0)) : 1.0f; } @@ -821,8 +821,9 @@ class StableDiffusionGGML { } } 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)"); + 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 {