upstream-archive/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-07-31 15:20:41 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat: add multi-device layer split (--backend "diffusion=cuda0&cuda1") (#1734)

Browse Source
This commit is contained in:
Piotr Wilkin (ilintar)
2026-07-04 10:06:33 +02:00
committed by GitHub
parent 2574f59365
commit 7bcd189639
11 changed files with 748 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

74
src/conditioning/conditioner.hpp
View File

@@ -116,6 +116,8 @@ public:
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
virtual void set_stream_layers_enabled(bool enabled) {}
virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {}
virtual void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {}
virtual void set_flash_attention_enabled(bool enabled) = 0;
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
virtual void runner_done() {}
@@ -635,6 +637,18 @@ struct SD3CLIPEmbedder : public Conditioner {
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (t5) {
t5->set_runtime_backends(backends);
}
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& 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 +1008,18 @@ struct FluxCLIPEmbedder : public Conditioner {
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (t5) {
t5->set_runtime_backends(backends);
}
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& 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 +1252,18 @@ struct T5CLIPEmbedder : public Conditioner {
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (t5) {
t5->set_runtime_backends(backends);
}
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& 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 +1456,18 @@ struct MiniT2IConditioner : public Conditioner {
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (t5) {
t5->set_runtime_backends(backends);
}
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& 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 +1552,14 @@ struct AnimaConditioner : public Conditioner {
llm->set_stream_layers_enabled(enabled);
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
llm->set_runtime_backends(backends);
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& 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 +1705,14 @@ struct LLMEmbedder : public Conditioner {
llm->set_stream_layers_enabled(enabled);
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
llm->set_runtime_backends(backends);
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& 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 +2382,14 @@ struct LTXAVEmbedder : public Conditioner {
projector->set_max_graph_vram_bytes(max_vram_bytes);
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
llm->set_runtime_backends(backends);
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm");
}
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
llm->set_weight_adapter(adapter);
projector->set_weight_adapter(adapter);

195
src/core/ggml_extend.hpp
View File

@@ -1746,6 +1746,11 @@ protected:
bool stream_layers_enabled = false;
size_t observed_max_effective_budget_ = 0;
std::vector<ggml_backend_t> 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<WeightAdapter> weight_adapter = nullptr;
std::weak_ptr<RunnerWeightManager> weight_manager;
std::unordered_set<const ggml_tensor*> kept_compute_param_tensor_set;
@@ -2013,7 +2018,121 @@ protected:
return true;
}
// Pass explicit buffer types: synthesized defaults can make CUDA devices
// report supporting each other's buffers and skip a required copy.
bool ensure_sched(ggml_cgraph* gf) {
if (sched != nullptr) {
return true;
}
std::vector<ggml_backend_t> 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<ggml_backend_buffer_type_t> 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<size_t>(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;
}
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;
}
// Weightless ops have no scheduler anchor, so pin them to the most recent
// weight device. Views must stay unpinned or cross-device copies can be
// skipped for their 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 +2348,14 @@ protected:
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 +2611,14 @@ protected:
};
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 +2627,27 @@ protected:
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,
gf,
sd_get_backend_eval_callback(),
sd_get_backend_eval_callback_data());
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 +2824,10 @@ public:
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 +2868,20 @@ public:
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 +3017,31 @@ public:
}
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;
}
void set_runtime_backends(const std::vector<ggml_backend_t>& 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 {

79
src/core/ggml_extend_backend.cpp
View File

@@ -669,8 +669,47 @@ void SDBackendManager::reset() {
params_assignment_ = {};
}
static std::vector<std::string> split_device_list(const std::string& value) {
std::vector<std::string> 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<std::string> 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<ggml_backend_t> SDBackendManager::runtime_backends(SDBackendModule module) {
std::vector<ggml_backend_t> 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 +735,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 +772,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 +790,41 @@ bool SDBackendManager::validate(std::string* error) const {
}
return false;
};
auto validate_runtime_name = [&](const std::string& name) -> bool {
if (name.find('&') == std::string::npos) {
return validate_single_runtime_name(name);
}
std::vector<std::string> 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) ||

4
src/core/ggml_extend_backend.h
View File

@@ -6,6 +6,7 @@
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include "ggml-backend.h"
#include "ggml.h"
@@ -58,9 +59,12 @@ public:
ggml_backend_t runtime_backend(SDBackendModule module);
ggml_backend_t params_backend(SDBackendModule module);
std::vector<ggml_backend_t> 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;
bool params_backend_follows_runtime(SDBackendModule module) const;
bool runtime_backend_supports_host_buffer(SDBackendModule module);
private:

194
src/core/layer_split_partition.cpp Normal file
View File

@@ -0,0 +1,194 @@
#include "core/layer_split_partition.h"
#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include "core/util.h"
namespace sd {
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 += std::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 std::atoi(name.substr(pos, end - pos).c_str());
}
}
return -1;
}
std::string layer_split_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";
}
static bool layer_split_backend_supports_tensor(ggml_backend_t backend, const ggml_tensor* tensor) {
return backend != nullptr && tensor != nullptr && ggml_backend_supports_op(backend, tensor);
}
static size_t layer_split_supported_target(const std::string& desc,
const std::string& tensor_name,
const ggml_tensor* tensor,
const std::vector<ggml_backend_t>& backends,
size_t preferred) {
if (tensor == nullptr || backends.empty()) {
return preferred;
}
size_t preferred_safe = std::min(preferred, backends.size() - 1);
if (layer_split_backend_supports_tensor(backends[preferred_safe], tensor)) {
return preferred_safe;
}
for (size_t i = 0; i < backends.size(); i++) {
if (layer_split_backend_supports_tensor(backends[i], tensor)) {
LOG_WARN("%s layer split: moving tensor '%s' from %s to %s because the preferred backend cannot run op=%s type=%s nbytes=%.2f MB",
desc.c_str(),
tensor_name.c_str(),
layer_split_backend_device_display_name(backends[preferred_safe]).c_str(),
layer_split_backend_device_display_name(backends[i]).c_str(),
ggml_op_name(tensor->op),
ggml_type_name(tensor->type),
ggml_nbytes(tensor) / (1024.0 * 1024.0));
return i;
}
}
LOG_WARN("%s layer split: tensor '%s' is not supported by any split backend: op=%s type=%s nbytes=%.2f MB",
desc.c_str(),
tensor_name.c_str(),
ggml_op_name(tensor->op),
ggml_type_name(tensor->type),
ggml_nbytes(tensor) / (1024.0 * 1024.0));
return preferred_safe;
}
std::vector<std::map<std::string, ggml_tensor*>> partition_layer_split_tensors(
const std::string& desc,
const std::map<std::string, ggml_tensor*>& tensors,
const std::map<std::string, ggml_tensor*>& split_tensors,
const std::vector<ggml_backend_t>& backends) {
std::vector<std::map<std::string, ggml_tensor*>> partitions(backends.size());
if (backends.empty()) {
LOG_WARN("%s: no backend available for a layer split", desc.c_str());
return partitions;
}
std::map<int, int64_t> block_bytes;
std::map<std::string, size_t> non_block_targets;
std::vector<int64_t> other_bytes_by_backend(backends.size(), 0);
int64_t total_block_bytes = 0;
int64_t total_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 {
size_t target = layer_split_supported_target(desc, kv.first, kv.second, backends, 0);
non_block_targets[kv.first] = target;
other_bytes_by_backend[target] += bytes;
total_other_bytes += bytes;
}
}
if (n_blocks == 0) {
LOG_WARN("%s: no transformer blocks found for a layer split; keeping tensors on compatible backends starting from %s",
desc.c_str(),
layer_split_backend_device_display_name(backends[0]).c_str());
for (const auto& kv : tensors) {
size_t target = 0;
auto target_it = non_block_targets.find(kv.first);
if (target_it != non_block_targets.end()) {
target = target_it->second;
}
partitions[target][kv.first] = kv.second;
}
return partitions;
}
// Reserve compute headroom and subtract each device's actual non-block
// bytes from its block budget.
constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024;
std::vector<double> 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>((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<int64_t> block_budgets(backends.size(), 0);
const int64_t total_bytes = total_block_bytes + total_other_bytes;
for (size_t i = 0; i < backends.size(); i++) {
int64_t budget = (int64_t)((double)total_bytes * device_weights[i] / weight_sum);
budget = std::max<int64_t>(budget - other_bytes_by_backend[i], 0);
block_budgets[i] = budget;
}
std::vector<int> 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);
target = layer_split_supported_target(desc, kv.first, kv.second, backends, target);
} else {
auto target_it = non_block_targets.find(kv.first);
if (target_it != non_block_targets.end()) {
target = target_it->second;
}
}
partitions[target][kv.first] = kv.second;
}
int range_start = 0;
for (size_t i = 0; i < backends.size(); i++) {
int range_end = boundaries[i];
const char* non_block_suffix = other_bytes_by_backend[i] > 0 ? " + non-block tensors" : "";
LOG_INFO("%s layer split: %s <- blocks [%d, %d)%s",
desc.c_str(),
layer_split_backend_device_display_name(backends[i]).c_str(),
range_start,
range_end,
non_block_suffix);
range_start = range_end;
}
return partitions;
}
} // namespace sd

23
src/core/layer_split_partition.h Normal file
View File

@@ -0,0 +1,23 @@
#ifndef __SD_CORE_LAYER_SPLIT_PARTITION_H__
#define __SD_CORE_LAYER_SPLIT_PARTITION_H__
#include <map>
#include <string>
#include <vector>
#include "ggml-backend.h"
#include "ggml.h"
namespace sd {
std::string layer_split_backend_device_display_name(ggml_backend_t backend);
std::vector<std::map<std::string, ggml_tensor*>> partition_layer_split_tensors(
const std::string& desc,
const std::map<std::string, ggml_tensor*>& tensors,
const std::map<std::string, ggml_tensor*>& split_tensors,
const std::vector<ggml_backend_t>& backends);
} // namespace sd
#endif // __SD_CORE_LAYER_SPLIT_PARTITION_H__

26
src/core/util.cpp
View File

@@ -4,6 +4,8 @@
#include <cmath>
#include <codecvt>
#include <cstdarg>
#include <cstdlib>
#include <cstring>
#include <exception>
#include <fstream>
#include <locale>
@@ -25,6 +27,7 @@
#include <unistd.h>
#endif
#include "ggml-backend.h"
#include "ggml.h"
#include "stable-diffusion.h"
@@ -997,3 +1000,26 @@ std::vector<std::pair<std::string, float>> split_quotation_attention(
}
return result;
}
size_t sd_list_devices(char* buffer, size_t buffer_size) {
if (ggml_backend_dev_count() == 0) {
// dynamic-backend builds discover their backend modules at runtime
ggml_backend_load_all();
}
std::ostringstream oss;
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);
oss << (name ? name : "") << '\t' << (desc ? desc : "") << '\n';
}
std::string devices = oss.str();
if (buffer != nullptr && buffer_size > 0) {
size_t copy_size = std::min(devices.size(), buffer_size - 1);
memcpy(buffer, devices.data(), copy_size);
buffer[copy_size] = '\0';
}
return devices.size();
}

112
src/stable-diffusion.cpp
View File

@@ -2,11 +2,14 @@
#include <cmath>
#include <cstdlib>
#include <set>
#include <type_traits>
#include <unordered_set>
#include <utility>
#include <vector>
#include "core/ggml_extend.hpp"
#include "core/ggml_graph_cut.h"
#include "core/layer_split_partition.h"
#include "core/rng.hpp"
#include "core/rng_mt19937.hpp"
@@ -170,6 +173,13 @@ static float get_cache_reuse_threshold(const sd_cache_params_t& params) {
/*=============================================== StableDiffusionGGML ================================================*/
template <typename T, typename = void>
struct has_set_runtime_backends : std::false_type {};
template <typename T>
struct has_set_runtime_backends<T,
std::void_t<decltype(std::declval<T&>().set_runtime_backends(
std::declval<const std::vector<ggml_backend_t>&>()))>> : std::true_type {};
static_assert(std::atomic<sd_cancel_mode_t>::is_always_lock_free,
"sd_cancel_mode_t must be lock-free");
@@ -275,14 +285,114 @@ public:
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<ggml_backend_t> module_backends = backend_manager.runtime_backends(module);
if (module_backends.size() > 1) {
if constexpr (has_set_runtime_backends<T>::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),
sd::layer_split_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);
}
// Register each layer-split partition with its compute backend; the
// ModelManager handles allocation, staging, and LoRA by backend.
template <typename T>
bool register_layer_split_runner_params(const std::string& desc,
const std::shared_ptr<T>& model,
SDBackendModule module,
const std::vector<ggml_backend_t>& module_backends,
std::map<std::string, ggml_tensor*> 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 <module>=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(),
sd::layer_split_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<std::string, ggml_tensor*> split_tensors;
if constexpr (std::is_base_of_v<Conditioner, T>) {
model->get_layer_split_param_tensors(split_tensors);
} else {
split_tensors = group_tensors;
}
auto partitions = sd::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(),