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

feat: drive layer split from graph-cut segments (#1762)

Browse Source
This commit is contained in:
leejet
2026-07-07 23:16:52 +08:00
committed by GitHub
parent bb84971129
commit 9ef6e7398f
8 changed files with 618 additions and 213 deletions
Download Patch File Download Diff File Expand all files Collapse all files

126
src/conditioning/conditioner.hpp
View File

@@ -118,6 +118,8 @@ public:
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 set_graph_cut_layer_split_enabled(bool enabled) {}
virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {}
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) {}
@@ -181,6 +183,27 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
text_model->set_runtime_backends(backends);
if (sd_version_is_sdxl(version)) {
text_model2->set_runtime_backends(backends);
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
text_model->set_graph_cut_layer_split_enabled(enabled);
if (sd_version_is_sdxl(version)) {
text_model2->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
text_model->set_graph_cut_layer_split_backend_vram_limits(limits);
if (sd_version_is_sdxl(version)) {
text_model2->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
void set_flash_attention_enabled(bool enabled) override {
text_model->set_flash_attention_enabled(enabled);
if (sd_version_is_sdxl(version)) {
@@ -639,11 +662,41 @@ struct SD3CLIPEmbedder : public Conditioner {
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (clip_l) {
clip_l->set_runtime_backends(backends);
}
if (clip_g) {
clip_g->set_runtime_backends(backends);
}
if (t5) {
t5->set_runtime_backends(backends);
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_graph_cut_layer_split_enabled(enabled);
}
if (clip_g) {
clip_g->set_graph_cut_layer_split_enabled(enabled);
}
if (t5) {
t5->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
if (clip_l) {
clip_l->set_graph_cut_layer_split_backend_vram_limits(limits);
}
if (clip_g) {
clip_g->set_graph_cut_layer_split_backend_vram_limits(limits);
}
if (t5) {
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
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");
@@ -1010,11 +1063,32 @@ struct FluxCLIPEmbedder : public Conditioner {
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (clip_l) {
clip_l->set_runtime_backends(backends);
}
if (t5) {
t5->set_runtime_backends(backends);
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_graph_cut_layer_split_enabled(enabled);
}
if (t5) {
t5->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
if (clip_l) {
clip_l->set_graph_cut_layer_split_backend_vram_limits(limits);
}
if (t5) {
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
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");
@@ -1278,6 +1352,18 @@ struct T5CLIPEmbedder : public Conditioner {
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
if (t5) {
t5->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
if (t5) {
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
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");
@@ -1482,6 +1568,18 @@ struct MiniT2IConditioner : public Conditioner {
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
if (t5) {
t5->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
if (t5) {
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
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");
@@ -1576,6 +1674,14 @@ struct AnimaConditioner : public Conditioner {
llm->set_runtime_backends(backends);
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
llm->set_graph_cut_layer_split_enabled(enabled);
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm");
}
@@ -1729,6 +1835,18 @@ struct LLMEmbedder : public Conditioner {
llm->set_runtime_backends(backends);
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
if (llm) {
llm->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
if (llm) {
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm");
}
@@ -2406,6 +2524,14 @@ struct LTXAVEmbedder : public Conditioner {
llm->set_runtime_backends(backends);
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
llm->set_graph_cut_layer_split_enabled(enabled);
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm");
}

178
src/core/ggml_extend.hpp
View File

@@ -21,10 +21,12 @@
#include <sstream>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include "core/ggml_extend_backend.h"
#include "core/ggml_graph_cut.h"
#include "core/layer_split_partition.h"
#include "ggml-alloc.h"
#include "ggml-backend.h"
#include "ggml.h"
@@ -1745,6 +1747,8 @@ protected:
size_t max_graph_vram_bytes = 0;
bool stream_layers_enabled = false;
size_t observed_max_effective_budget_ = 0;
bool graph_cut_layer_split_enabled = false;
std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
std::vector<ggml_backend_t> extra_runtime_backends; // borrowed (SDBackendManager-owned)
ggml_backend_sched_t sched = nullptr; // owned, multi-device only
@@ -1776,6 +1780,9 @@ protected:
sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_;
std::unordered_set<const ggml_tensor*> params_tensor_set_;
std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_assignments_;
std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_node_assignments_;
bool graph_cut_layer_split_primary_notice_logged_ = false;
template <typename T>
static sd::Tensor<T> take_or_empty(std::optional<sd::Tensor<T>> tensor) {
@@ -1874,6 +1881,20 @@ protected:
params_tensor_set_dirty_ = false;
}
ggml_tensor* canonical_param_tensor(ggml_tensor* tensor) {
if (tensor == nullptr) {
return nullptr;
}
if (params_tensor_set_.find(tensor) != params_tensor_set_.end()) {
return tensor;
}
if (tensor->view_src != nullptr &&
params_tensor_set_.find(tensor->view_src) != params_tensor_set_.end()) {
return tensor->view_src;
}
return nullptr;
}
std::vector<ggml_tensor*> collect_used_param_tensors(ggml_cgraph* gf) {
std::vector<ggml_tensor*> used_params;
rebuild_params_tensor_set();
@@ -1886,12 +1907,8 @@ protected:
seen_params.reserve(static_cast<size_t>(n_leafs));
for (int i = 0; i < n_leafs; ++i) {
ggml_tensor* leaf = sd::ggml_graph_cut::leaf_tensor(gf, i);
ggml_tensor* param_leaf = leaf;
if (param_leaf != nullptr && params_tensor_set_.find(param_leaf) == params_tensor_set_.end()) {
param_leaf = param_leaf->view_src;
}
ggml_tensor* param_leaf = canonical_param_tensor(leaf);
if (param_leaf != nullptr &&
params_tensor_set_.find(param_leaf) != params_tensor_set_.end() &&
seen_params.insert(param_leaf).second) {
used_params.push_back(param_leaf);
}
@@ -2101,11 +2118,17 @@ protected:
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);
ggml_tensor* node = ggml_graph_node(gf, i);
auto node_assignment = graph_cut_layer_split_node_assignments_.find(node);
if (node_assignment != graph_cut_layer_split_node_assignments_.end()) {
current = node_assignment->second;
}
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_assignment == graph_cut_layer_split_node_assignments_.end()) {
current = weight_backend;
}
}
}
if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
@@ -2435,6 +2458,123 @@ protected:
return true;
}
bool resolve_graph_cut_layer_split_plan(ggml_cgraph* gf,
GraphCutPlan* plan_out) {
GGML_ASSERT(plan_out != nullptr);
GGML_ASSERT(gf != nullptr);
*plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend,
gf,
&graph_cut_plan_cache_,
0,
params_tensor_set_,
get_desc().c_str());
return true;
}
bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
graph_cut_layer_split_node_assignments_.clear();
if (!graph_cut_layer_split_enabled) {
return true;
}
if (!is_multi_device()) {
LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", get_desc().c_str());
return false;
}
GraphCutPlan plan;
if (!resolve_graph_cut_layer_split_plan(gf, &plan)) {
return false;
}
if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) {
auto manager = weight_manager.lock();
if (manager == nullptr) {
LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
return false;
}
std::vector<ggml_tensor*> graph_params = collect_used_param_tensors(gf);
if (!graph_params.empty() &&
!manager->assign_compute_backend(graph_params, runtime_backend)) {
LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s",
get_desc().c_str(),
sd::layer_split_backend_device_display_name(runtime_backend).c_str());
return false;
}
for (ggml_tensor* param : graph_params) {
if (param != nullptr) {
graph_cut_layer_split_assignments_[param] = 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);
if (node != nullptr) {
graph_cut_layer_split_node_assignments_[node] = runtime_backend;
}
}
if (!graph_cut_layer_split_primary_notice_logged_) {
LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
get_desc().c_str(),
sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
graph_params.size());
graph_cut_layer_split_primary_notice_logged_ = true;
} else {
LOG_DEBUG("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
get_desc().c_str(),
sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
graph_params.size());
}
return true;
}
std::vector<ggml_backend_t> split_backends;
split_backends.reserve(extra_runtime_backends.size() + 1);
split_backends.push_back(runtime_backend);
for (ggml_backend_t backend : extra_runtime_backends) {
if (backend != nullptr) {
split_backends.push_back(backend);
}
}
auto manager = weight_manager.lock();
if (manager == nullptr) {
LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
return false;
}
sd::GraphCutLayerSplitAssignment assignment;
auto canonicalize_param = [this](ggml_tensor* tensor) {
return canonical_param_tensor(tensor);
};
if (!sd::partition_graph_cut_layer_split(get_desc().c_str(),
gf,
plan,
split_backends,
graph_cut_layer_split_backend_vram_limits_,
max_graph_vram_bytes,
graph_cut_layer_split_assignments_,
canonicalize_param,
&assignment)) {
return false;
}
for (size_t i = 0; i < split_backends.size(); i++) {
if (assignment.tensors_by_backend[i].empty()) {
continue;
}
if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) {
LOG_ERROR("%s graph-cut layer split failed to assign params to %s",
get_desc().c_str(),
sd::layer_split_backend_device_display_name(split_backends[i]).c_str());
return false;
}
}
graph_cut_layer_split_node_assignments_ = std::move(assignment.node_assignments);
sd::log_graph_cut_layer_split_assignment(get_desc().c_str(), split_backends, assignment);
return true;
}
struct PersistentExternalBinding {
ggml_backend_buffer_t buffer = nullptr;
void* data = nullptr;
@@ -2972,6 +3112,11 @@ public:
GGML_ASSERT(gf != nullptr);
rebuild_params_tensor_set();
if (!assign_graph_cut_layer_split_backends(gf)) {
free_compute_ctx();
return std::nullopt;
}
if (can_attempt_graph_cut_segmented_compute()) {
GraphCutPlan plan;
if (!resolve_graph_cut_plan(gf, &plan)) {
@@ -3025,6 +3170,22 @@ public:
stream_layers_enabled = enabled;
}
void set_graph_cut_layer_split_enabled(bool enabled) {
graph_cut_layer_split_enabled = enabled;
if (!enabled) {
graph_cut_layer_split_assignments_.clear();
graph_cut_layer_split_node_assignments_.clear();
graph_cut_layer_split_primary_notice_logged_ = false;
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {
graph_cut_layer_split_backend_vram_limits_ = limits;
graph_cut_layer_split_assignments_.clear();
graph_cut_layer_split_node_assignments_.clear();
graph_cut_layer_split_primary_notice_logged_ = false;
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {
extra_runtime_backends.clear();
for (ggml_backend_t backend : backends) {
@@ -3036,6 +3197,9 @@ public:
extra_runtime_backends.push_back(backend);
}
}
graph_cut_layer_split_assignments_.clear();
graph_cut_layer_split_node_assignments_.clear();
graph_cut_layer_split_primary_notice_logged_ = false;
if (is_multi_device() && stream_layers_enabled) {
LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring",
get_desc().c_str());

314
src/core/layer_split_partition.cpp
View File

@@ -1,9 +1,11 @@
#include "core/layer_split_partition.h"
#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <limits>
#include <unordered_set>
#include <utility>
#include "core/util.h"
@@ -62,160 +64,194 @@ namespace sd {
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 graph_cut_layer_split_backend_vram_limit(const std::vector<size_t>& backend_vram_limits,
size_t backend_index,
size_t primary_backend_vram_limit) {
if (backend_index < backend_vram_limits.size()) {
return backend_vram_limits[backend_index];
}
return backend_index == 0 ? primary_backend_vram_limit : 0;
}
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 ? layer_split_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.
static std::vector<int64_t> graph_cut_layer_split_backend_capacities(const std::vector<ggml_backend_t>& backends,
const std::vector<size_t>& backend_vram_limits,
size_t primary_backend_vram_limit) {
std::vector<int64_t> capacities(backends.size(), std::numeric_limits<int64_t>::max() / 4);
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;
if (free_bytes > 0) {
capacities[i] = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes, 0);
}
size_t limit_bytes = graph_cut_layer_split_backend_vram_limit(backend_vram_limits,
i,
primary_backend_vram_limit);
if (limit_bytes > 0) {
capacities[i] = std::min<int64_t>(capacities[i], (int64_t)limit_bytes);
}
}
return capacities;
}
bool partition_graph_cut_layer_split(const char* desc,
ggml_cgraph* gf,
const sd::ggml_graph_cut::Plan& plan,
const std::vector<ggml_backend_t>& split_backends,
const std::vector<size_t>& backend_vram_limits,
size_t primary_backend_vram_limit,
std::unordered_map<const ggml_tensor*, ggml_backend_t>& param_assignments,
const std::function<ggml_tensor*(ggml_tensor*)>& canonical_param_tensor,
GraphCutLayerSplitAssignment* assignment_out) {
GGML_ASSERT(gf != nullptr);
GGML_ASSERT(assignment_out != nullptr);
GGML_ASSERT(canonical_param_tensor != nullptr);
GGML_ASSERT(!split_backends.empty());
GraphCutLayerSplitAssignment assignment;
assignment.segment_count = plan.segments.size();
assignment.tensors_by_backend.resize(split_backends.size());
assignment.bytes_by_backend.resize(split_backends.size(), 0);
assignment.first_segment_by_backend.resize(split_backends.size(), plan.segments.size());
assignment.last_segment_by_backend.resize(split_backends.size(), 0);
std::vector<std::vector<ggml_tensor*>> segment_params(plan.segments.size());
std::vector<int64_t> segment_param_bytes(plan.segments.size(), 0);
std::unordered_set<ggml_tensor*> seen_params;
for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
std::vector<ggml_tensor*> params = sd::ggml_graph_cut::param_tensors(gf, plan.segments[seg_idx]);
for (ggml_tensor* raw_param : params) {
ggml_tensor* param = canonical_param_tensor(raw_param);
if (param == nullptr || !seen_params.insert(param).second) {
continue;
}
segment_params[seg_idx].push_back(param);
segment_param_bytes[seg_idx] += (int64_t)ggml_nbytes(param);
}
used += bytes;
}
for (const auto& kv : tensors) {
size_t target = 0;
int idx = split_tensors.count(kv.first) != 0 ? layer_split_tensor_block_index(kv.first) : -1;
if (idx >= 0) {
while (target < boundaries.size() && idx >= boundaries[target]) {
target++;
int64_t total_param_bytes = 0;
for (int64_t bytes : segment_param_bytes) {
total_param_bytes += bytes;
}
if (total_param_bytes <= 0) {
LOG_ERROR("%s graph-cut layer split found no graph params to assign", desc);
return false;
}
std::vector<int64_t> backend_capacities = graph_cut_layer_split_backend_capacities(split_backends,
backend_vram_limits,
primary_backend_vram_limit);
std::vector<ggml_backend_t> backend_by_segment(plan.segments.size(), split_backends[0]);
size_t current_backend = 0;
int64_t current_used = 0;
for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
int64_t bytes = segment_param_bytes[seg_idx];
while (current_backend + 1 < split_backends.size() &&
bytes > 0 &&
current_used + bytes > backend_capacities[current_backend]) {
current_backend++;
current_used = 0;
}
if (bytes > 0 && current_used + bytes > backend_capacities[current_backend]) {
LOG_ERROR("%s graph-cut layer split: segment %zu needs %.1f MB on %s, but only %.1f MB is available under current VRAM limits",
desc,
seg_idx,
(current_used + bytes) / (1024.0 * 1024.0),
layer_split_backend_device_display_name(split_backends[current_backend]).c_str(),
backend_capacities[current_backend] / (1024.0 * 1024.0));
return false;
}
current_used += bytes;
backend_by_segment[seg_idx] = split_backends[current_backend];
for (ggml_tensor* param : segment_params[seg_idx]) {
ggml_backend_t target_backend = split_backends[current_backend];
auto assigned_it = param_assignments.find(param);
if (assigned_it == param_assignments.end()) {
param_assignments[param] = target_backend;
assignment.has_new_param_assignment = true;
} else {
target_backend = assigned_it->second;
}
target = std::min(target, backends.size() - 1);
target = layer_split_supported_target(desc, kv.first, kv.second, backends, target);
auto backend_it = std::find(split_backends.begin(), split_backends.end(), target_backend);
if (backend_it == split_backends.end()) {
LOG_ERROR("%s graph-cut layer split tensor '%s' is assigned to an unavailable backend",
desc,
ggml_get_name(param));
return false;
}
size_t backend_idx = (size_t)std::distance(split_backends.begin(), backend_it);
assignment.first_segment_by_backend[backend_idx] = std::min(assignment.first_segment_by_backend[backend_idx], seg_idx);
assignment.last_segment_by_backend[backend_idx] = std::max(assignment.last_segment_by_backend[backend_idx], seg_idx + 1);
assignment.tensors_by_backend[backend_idx].push_back(param);
assignment.bytes_by_backend[backend_idx] += (int64_t)ggml_nbytes(param);
}
}
const int n_nodes = ggml_graph_n_nodes(gf);
for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
ggml_backend_t backend = backend_by_segment[seg_idx];
const auto& segment = plan.segments[seg_idx];
for (int node_index : segment.internal_node_indices) {
if (node_index < 0 || node_index >= n_nodes) {
continue;
}
ggml_tensor* node = ggml_graph_node(gf, node_index);
if (node != nullptr) {
assignment.node_assignments[node] = backend;
}
}
for (int node_index : segment.output_node_indices) {
if (node_index < 0 || node_index >= n_nodes) {
continue;
}
ggml_tensor* node = ggml_graph_node(gf, node_index);
if (node != nullptr) {
assignment.node_assignments[node] = backend;
}
}
}
*assignment_out = std::move(assignment);
return true;
}
void log_graph_cut_layer_split_assignment(const char* desc,
const std::vector<ggml_backend_t>& split_backends,
const GraphCutLayerSplitAssignment& assignment) {
for (size_t i = 0; i < split_backends.size(); i++) {
if (i >= assignment.tensors_by_backend.size() ||
assignment.tensors_by_backend[i].empty()) {
continue;
}
size_t first_segment = assignment.first_segment_by_backend[i] == assignment.segment_count
? 0
: assignment.first_segment_by_backend[i];
size_t last_segment = assignment.last_segment_by_backend[i];
if (assignment.has_new_param_assignment) {
LOG_INFO("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
desc,
layer_split_backend_device_display_name(split_backends[i]).c_str(),
first_segment,
last_segment,
assignment.tensors_by_backend[i].size(),
assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
} else {
auto target_it = non_block_targets.find(kv.first);
if (target_it != non_block_targets.end()) {
target = target_it->second;
}
LOG_DEBUG("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
desc,
layer_split_backend_device_display_name(split_backends[i]).c_str(),
first_segment,
last_segment,
assignment.tensors_by_backend[i].size(),
assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
}
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

34
src/core/layer_split_partition.h
View File

@@ -1,23 +1,43 @@
#ifndef __SD_CORE_LAYER_SPLIT_PARTITION_H__
#define __SD_CORE_LAYER_SPLIT_PARTITION_H__
#include <map>
#include <cstdint>
#include <functional>
#include <string>
#include <unordered_map>
#include <vector>
#include "ggml-backend.h"
#include "ggml.h"
#include "core/ggml_graph_cut.h"
namespace sd {
struct GraphCutLayerSplitAssignment {
std::vector<std::vector<ggml_tensor*>> tensors_by_backend;
std::vector<int64_t> bytes_by_backend;
std::vector<size_t> first_segment_by_backend;
std::vector<size_t> last_segment_by_backend;
std::unordered_map<const ggml_tensor*, ggml_backend_t> node_assignments;
size_t segment_count = 0;
bool has_new_param_assignment = false;
};
std::string layer_split_backend_device_display_name(ggml_backend_t backend);
int layer_split_tensor_block_index(const std::string& name);
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);
bool partition_graph_cut_layer_split(const char* desc,
ggml_cgraph* gf,
const sd::ggml_graph_cut::Plan& plan,
const std::vector<ggml_backend_t>& split_backends,
const std::vector<size_t>& backend_vram_limits,
size_t primary_backend_vram_limit,
std::unordered_map<const ggml_tensor*, ggml_backend_t>& param_assignments,
const std::function<ggml_tensor*(ggml_tensor*)>& canonical_param_tensor,
GraphCutLayerSplitAssignment* assignment_out);
void log_graph_cut_layer_split_assignment(const char* desc,
const std::vector<ggml_backend_t>& split_backends,
const GraphCutLayerSplitAssignment& assignment);
} // namespace sd

68
src/model_manager.cpp
View File

@@ -134,7 +134,8 @@ bool ModelManager::register_param_tensors(const std::string& desc,
ggml_backend_t compute_backend,
ggml_backend_t params_backend,
size_t* registered_tensor_size,
bool allow_split_buffer) {
bool allow_split_buffer,
bool params_follow_compute_backend) {
if (desc.empty()) {
LOG_ERROR("model manager tensor desc is empty");
return false;
@@ -158,14 +159,15 @@ bool ModelManager::register_param_tensors(const std::string& desc,
}
ggml_set_name(tensor, name.c_str());
auto state = std::make_unique<TensorState>();
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;
auto state = std::make_unique<TensorState>();
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;
state->params_follow_compute_backend = params_follow_compute_backend;
new_states.push_back(std::move(state));
}
@@ -919,6 +921,54 @@ bool ModelManager::resolve_required_tensor_states(const std::vector<ggml_tensor*
return true;
}
bool ModelManager::assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
ggml_backend_t compute_backend) {
if (tensors.empty()) {
return true;
}
if (compute_backend == nullptr) {
LOG_ERROR("model manager cannot assign tensors to a null compute backend");
return false;
}
std::vector<TensorState*> required_states;
if (!resolve_required_tensor_states(tensors, required_states)) {
return false;
}
for (TensorState* state : required_states) {
if (state == nullptr || state->tensor == nullptr) {
continue;
}
const bool params_follow_compute = state->params_follow_compute_backend ||
state->residency_mode == ResidencyMode::Disk;
const bool compute_changes = state->compute_backend != compute_backend;
const bool params_changes = params_follow_compute && state->params_backend != compute_backend;
if (!compute_changes && !params_changes) {
continue;
}
if (state->active_prepare_count > 0 || state->staged_to_compute_backend) {
LOG_ERROR("model manager cannot move active tensor '%s' to another compute backend",
state->name.c_str());
return false;
}
if (params_changes && state->loaded_to_params_backend) {
LOG_ERROR("model manager cannot move loaded tensor '%s' to another params backend",
state->name.c_str());
return false;
}
state->compute_backend = compute_backend;
if (params_follow_compute) {
state->params_backend = compute_backend;
}
}
return true;
}
bool ModelManager::prepare_params(const std::vector<ggml_tensor*>& tensors) {
if (tensors.empty()) {
return true;

18
src/model_manager.h
View File

@@ -33,11 +33,12 @@ private:
ggml_tensor* tensor = nullptr;
std::string desc;
ResidencyMode residency_mode = ResidencyMode::ParamBackend;
ggml_backend_t compute_backend = nullptr;
ggml_backend_t params_backend = nullptr;
bool allow_split_buffer = false;
bool metadata_validated = false;
ResidencyMode residency_mode = ResidencyMode::ParamBackend;
ggml_backend_t compute_backend = nullptr;
ggml_backend_t params_backend = nullptr;
bool allow_split_buffer = false;
bool params_follow_compute_backend = false;
bool metadata_validated = false;
int active_prepare_count = 0;
@@ -129,8 +130,9 @@ public:
ResidencyMode residency_mode,
ggml_backend_t compute_backend,
ggml_backend_t params_backend,
size_t* registered_tensor_size = nullptr,
bool allow_split_buffer = false);
size_t* registered_tensor_size = nullptr,
bool allow_split_buffer = false,
bool params_follow_compute_backend = false);
template <typename Runner>
bool register_runner_params(const std::string& desc,
@@ -170,6 +172,8 @@ public:
bool validate_registered_tensors();
bool load_all_params_eagerly();
bool assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
ggml_backend_t compute_backend) override;
bool prepare_params(const std::vector<ggml_tensor*>& tensors) override;
void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) override;
void release_params_backend_params(const std::vector<ggml_tensor*>& tensors) override;

89
src/stable-diffusion.cpp
View File

@@ -268,6 +268,15 @@ public:
return max_vram_assignment.bytes_for_backend(backend_for(module));
}
std::vector<size_t> layer_split_vram_limits_for_backends(const std::vector<ggml_backend_t>& backends) {
std::vector<size_t> limits;
limits.reserve(backends.size());
for (ggml_backend_t backend : backends) {
limits.push_back(max_vram_assignment.bytes_for_backend(backend));
}
return limits;
}
bool ensure_backend_pair(SDBackendModule module) {
if (backend_for(module) == nullptr) {
return false;
@@ -427,8 +436,9 @@ public:
params_mem_size);
}
// Register each layer-split partition with its compute backend; the
// ModelManager handles allocation, staging, and LoRA by backend.
// Register graph-cut layer-split tensors on the primary backend first.
// The first real graph assigns each param tensor to a runtime backend
// before weights are loaded or staged.
template <typename T>
bool register_layer_split_runner_params(const std::string& desc,
const std::shared_ptr<T>& model,
@@ -459,52 +469,29 @@ public:
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);
model->set_graph_cut_layer_split_backend_vram_limits(layer_split_vram_limits_for_backends(module_backends));
model->set_graph_cut_layer_split_enabled(true);
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;
}
ggml_backend_t initial_params_backend = params_follow_runtime ? module_backends[0] : params_backend_for(module);
if (initial_params_backend == nullptr) {
return false;
}
return true;
LOG_INFO("%s graph-cut layer split: deferring %zu tensors across %zu runtime backends until first graph",
desc.c_str(),
group_tensors.size(),
module_backends.size());
return model_manager->register_param_tensors(desc,
std::move(group_tensors),
residency_mode,
module_backends[0],
initial_params_backend,
params_mem_size,
false,
params_follow_runtime);
}
bool init_backend() {
@@ -529,6 +516,16 @@ public:
return false;
}
bool graph_cut_layer_split_active() {
for (SDBackendModule module : {SDBackendModule::DIFFUSION, SDBackendModule::TE}) {
if (backend_manager.split_mode(module) == SDSplitMode::LAYER &&
backend_manager.runtime_backends(module).size() > 1) {
return true;
}
}
return false;
}
std::shared_ptr<RNG> get_rng(rng_type_t rng_type) {
if (rng_type == STD_DEFAULT_RNG) {
return std::make_shared<STDDefaultRNG>();
@@ -785,6 +782,10 @@ public:
LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring");
stream_layers = false;
}
if (eager_load && graph_cut_layer_split_active()) {
LOG_WARN("--eager-load is not supported with graph-cut layer split; weights will be prepared lazily");
eager_load = false;
}
std::map<ggml_type, uint32_t> wtype_stat = model_loader.get_wtype_stat();
std::map<ggml_type, uint32_t> conditioner_wtype_stat = model_loader.get_conditioner_wtype_stat();

4
src/weight_manager.h
View File

@@ -3,10 +3,14 @@
#include <vector>
#include "ggml-backend.h"
struct ggml_tensor;
struct RunnerWeightManager {
virtual ~RunnerWeightManager() = default;
virtual bool assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
ggml_backend_t compute_backend) = 0;
virtual bool prepare_params(const std::vector<ggml_tensor*>& tensors) = 0;
virtual void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;
virtual void release_params_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;