Compare commits

...

3 Commits

Author SHA1 Message Date
leejet
6314af404a docs: add shared agent instructions 2026-07-08 00:07:37 +08:00
leejet
885f01a555 chore: close inactive issues automatically 2026-07-07 23:37:59 +08:00
leejet
9ef6e7398f
feat: drive layer split from graph-cut segments (#1762) 2026-07-07 23:16:52 +08:00
12 changed files with 863 additions and 213 deletions

View File

@ -0,0 +1,47 @@
name: Close inactive issues
on:
schedule:
# Run daily. GitHub cron schedules use UTC.
- cron: "30 1 * * *"
workflow_dispatch:
inputs:
debug_only:
description: "Dry run: log intended actions without changing issues"
required: false
default: false
type: boolean
permissions:
issues: write
concurrency:
group: ${{ github.workflow }}
cancel-in-progress: false
jobs:
close-inactive-issues:
runs-on: ubuntu-latest
steps:
- name: Comment and close inactive issues
uses: actions/stale@v10
with:
days-before-issue-stale: 365
days-before-issue-close: 0
days-before-pr-stale: -1
days-before-pr-close: -1
stale-issue-label: issue:inactive
close-issue-label: issue:auto-closed
stale-issue-message: ""
close-issue-message: >
This issue has had no activity for one year. The latest version of
the code may already have fixed the problem.
If the issue still exists in the latest version, you can reopen
this issue at any time with updated reproduction details.
remove-issue-stale-when-updated: true
operations-per-run: 1000
debug-only: ${{ github.event_name == 'workflow_dispatch' && inputs.debug_only || false }}

4
.gitignore vendored
View File

@ -13,3 +13,7 @@ output*.png
models* models*
*.log *.log
preview.png preview.png
.claude/
CLAUDE.local.md
.agents/
.codex/

183
AGENTS.md Normal file
View File

@ -0,0 +1,183 @@
# Instructions for stable-diffusion.cpp
This document is for AI coding agents working in this repository. It should
describe agent-specific workflow, repository routing, editing boundaries, and
project-specific pitfalls.
For general contribution rules, including PR scope, commit conventions, code
style, dependency updates, security hygiene, and AI-assisted contribution policy,
see `CONTRIBUTING.md`.
---
## Agent Operating Rules
Before analyzing or modifying the repository:
1. Read this file.
2. Use `rg` / `rg --files` or directory listing commands to confirm the current
tree before relying on a path.
3. Start from `src/` and relevant `docs/` for runtime behavior.
4. Read the relevant code before editing.
5. Prefer the smallest change that fits the existing architecture.
6. Report focused verification and mention any tests not run.
Agents must not:
* Run `git push`, create PRs, or submit issue/PR comments on the user's behalf.
* Create commits unless the user explicitly requests that specific commit.
* Modify `ggml/`, `thirdparty/`, or `examples/server/frontend/` unless
explicitly requested and necessary.
* Read large local model files or tokenizer vocabulary files.
* Rewrite unrelated code for style-only reasons.
* Add secrets, model weights, generated binaries, local absolute paths, or
machine-specific output.
When a change is large, architectural, or likely to affect public behavior,
pause and present a short plan before editing.
---
## Repository Map and Editing Boundaries
This is a routing map for agents, not a full architecture document. The layout
can change, so verify paths before using them. Do not inspect excluded
large-data directories while checking the tree.
### Primary Project Code
Core implementation lives under `src/`.
Current source layout includes:
* `src/core/` - shared tensor, ggml integration, backend, graph, RNG, and utility
code.
* `src/model/` - model families and model components.
* `src/model_io/` - model file loading, GGUF, safetensors, pickle, and related
serialization helpers.
* `src/runtime/` - sampling, denoising, guidance, caching, preprocessing, and
runtime execution helpers.
* `src/tokenizers/` - tokenizer implementations.
* `src/conditioning/` - conditioning and prompt-related implementation.
* `src/extensions/` - optional feature extensions.
* top-level `src/*.cpp` and `src/*.h` files - public implementation entry
points, model loading, conversion, versioning, and shared managers.
`src/tokenizers/vocab/` contains large tokenizer vocabulary data. Do not read or
parse files in this directory; reference the path only when necessary.
### Public API
`include/` contains the C API exposed by the project. Currently the primary
public header is `include/stable-diffusion.h`.
Treat public headers as stable API. Avoid breaking compatibility unless the user
explicitly requests it. If public behavior changes, update relevant examples or
documentation.
### Examples
`examples/` contains programs demonstrating library usage.
* `examples/cli/` - command line program for running models, testing features,
and debugging.
* `examples/common/` - shared example support code.
* `examples/server/` - server application built on top of the library.
* `examples/server/frontend/` - git submodule containing independent frontend
code. Avoid modifying it unless explicitly requested.
### Documentation and Tooling
* `docs/` - documentation for supported models, build options, behavior, and
workflows.
* `scripts/` - development, model processing, build automation, formatting, and
tooling scripts.
* `cmake/` - CMake support modules.
* `docker/` - Docker-related project files.
* `assets/` - documentation assets; not runtime code.
### External, Local, and Generated State
* `ggml/` - git submodule for the ggml dependency.
* `thirdparty/` - vendored third-party dependencies.
* `models/` - local model storage. Ignore this directory and do not read model
files.
* `test/` - local testing scripts. Use only when relevant to the task.
* `build/`, `build_*`, and similar directories - generated build outputs.
Inspect them only when debugging a build result.
---
## Agent Workflow for Code Changes
1. Identify the relevant modules under `src/`.
2. Check whether the change touches the public API in `include/`.
3. Consult relevant `docs/` and examples before changing user-facing behavior.
4. Follow existing local patterns before adding new abstractions.
5. Keep edits scoped to the requested behavior.
6. Run the narrowest useful build, test, or inspection command available.
Follow `CONTRIBUTING.md` for formatting, naming, PR expectations, dependency
update policy, and security rules.
---
## Code Comments
Keep comments rare and useful.
Do not add comments that only describe what the code does. Add comments only
when the code cannot fully express the logic, the logic is unusually complex, or
there are historical reasons, invariants, constraints, compatibility concerns,
or known pitfalls that future maintainers need to understand.
Do not add task-specific comments that will be meaningless after review.
Examples from the current codebase:
```cpp
// GOOD: explains a safety constraint that is not obvious from the assignment.
// From src/model_io/pickle_io.cpp.
// Non-tensor checkpoint metadata can use REDUCE for arbitrary
// Python objects. Do not execute it; keep stack shape only.
stack.push_back(make_none_value());
// BAD: describes only what the next line does.
// Set the token count to zero.
token_count = 0;
```
---
## Text File Encoding
When reading or editing repository text files:
* Prefer UTF-8 with LF for Markdown, frontend source, JSON, and other text-first
project files unless the file already clearly uses a different encoding.
* Do not assume terminal output encoding matches file encoding on Windows.
* A file that looks garbled in PowerShell output may still be valid UTF-8.
* When inspecting UTF-8 files in PowerShell, prefer explicit UTF-8 reads such as:
* `[Console]::OutputEncoding = [System.Text.Encoding]::UTF8`
* `Get-Content -Encoding utf8 <path>`
* Avoid rewriting a file purely because console output looked garbled; verify
the actual file encoding first.
---
## Tensor and Layout Notes
Additional tensor/layout rules for this codebase:
* `sd::Tensor` shape order is not PyTorch/NumPy-style. `shape()[0]` is the
lowest and most contiguous dimension, and higher indices are higher
dimensions.
* Broadcasting for `sd::Tensor` must align dimensions from low to high dimension
indices. If one tensor has fewer dimensions, append implicit `1`s at the
higher-dimension end.
* `ggml_n_dims` / `ggml_n_dims(tensor)` can drop trailing singleton high
dimensions. Do not assume a logical trailing dimension of `1` will still be
counted in ggml metadata.
* Internal tensor-returning interfaces use an empty `sd::Tensor` to represent
null, absent, or failure states. Do not add `std::optional<sd::Tensor<...>>`
for internal APIs unless a distinct semantic state is truly required.

11
CLAUDE.md Normal file
View File

@ -0,0 +1,11 @@
@AGENTS.md
## Claude Code
Follow `AGENTS.md` as the shared repository instructions.
Do not duplicate contribution, style, PR, dependency, or security policy here;
use `CONTRIBUTING.md` as the canonical source for those rules.
Keep Claude-specific project notes in this file only when they do not apply to
other coding agents.

View File

@ -118,6 +118,8 @@ public:
virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {} virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
virtual void set_stream_layers_enabled(bool enabled) {} virtual void set_stream_layers_enabled(bool enabled) {}
virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {} 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 get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {}
virtual void set_flash_attention_enabled(bool enabled) = 0; virtual void set_flash_attention_enabled(bool enabled) = 0;
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {} 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 { void set_flash_attention_enabled(bool enabled) override {
text_model->set_flash_attention_enabled(enabled); text_model->set_flash_attention_enabled(enabled);
if (sd_version_is_sdxl(version)) { 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 { 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) { if (t5) {
t5->set_runtime_backends(backends); 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 { void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (t5) { if (t5) {
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); 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 { void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (clip_l) {
clip_l->set_runtime_backends(backends);
}
if (t5) { if (t5) {
t5->set_runtime_backends(backends); 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 { void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (t5) { if (t5) {
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); 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 { void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (t5) { if (t5) {
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); 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 { void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (t5) { if (t5) {
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
@ -1576,6 +1674,14 @@ struct AnimaConditioner : public Conditioner {
llm->set_runtime_backends(backends); 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 { void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm"); llm->get_param_tensors(tensors, "text_encoders.llm");
} }
@ -1729,6 +1835,18 @@ struct LLMEmbedder : public Conditioner {
llm->set_runtime_backends(backends); 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 { void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm"); llm->get_param_tensors(tensors, "text_encoders.llm");
} }
@ -2406,6 +2524,14 @@ struct LTXAVEmbedder : public Conditioner {
llm->set_runtime_backends(backends); 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 { void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm"); llm->get_param_tensors(tensors, "text_encoders.llm");
} }

View File

@ -21,10 +21,12 @@
#include <sstream> #include <sstream>
#include <string> #include <string>
#include <unordered_map> #include <unordered_map>
#include <unordered_set>
#include <vector> #include <vector>
#include "core/ggml_extend_backend.h" #include "core/ggml_extend_backend.h"
#include "core/ggml_graph_cut.h" #include "core/ggml_graph_cut.h"
#include "core/layer_split_partition.h"
#include "ggml-alloc.h" #include "ggml-alloc.h"
#include "ggml-backend.h" #include "ggml-backend.h"
#include "ggml.h" #include "ggml.h"
@ -1745,6 +1747,8 @@ protected:
size_t max_graph_vram_bytes = 0; size_t max_graph_vram_bytes = 0;
bool stream_layers_enabled = false; bool stream_layers_enabled = false;
size_t observed_max_effective_budget_ = 0; 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) std::vector<ggml_backend_t> extra_runtime_backends; // borrowed (SDBackendManager-owned)
ggml_backend_sched_t sched = nullptr; // owned, multi-device only ggml_backend_sched_t sched = nullptr; // owned, multi-device only
@ -1776,6 +1780,9 @@ protected:
sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_; sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_;
std::unordered_set<const ggml_tensor*> params_tensor_set_; 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> template <typename T>
static sd::Tensor<T> take_or_empty(std::optional<sd::Tensor<T>> tensor) { static sd::Tensor<T> take_or_empty(std::optional<sd::Tensor<T>> tensor) {
@ -1874,6 +1881,20 @@ protected:
params_tensor_set_dirty_ = false; 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*> collect_used_param_tensors(ggml_cgraph* gf) {
std::vector<ggml_tensor*> used_params; std::vector<ggml_tensor*> used_params;
rebuild_params_tensor_set(); rebuild_params_tensor_set();
@ -1886,12 +1907,8 @@ protected:
seen_params.reserve(static_cast<size_t>(n_leafs)); seen_params.reserve(static_cast<size_t>(n_leafs));
for (int i = 0; i < n_leafs; ++i) { for (int i = 0; i < n_leafs; ++i) {
ggml_tensor* leaf = sd::ggml_graph_cut::leaf_tensor(gf, i); ggml_tensor* leaf = sd::ggml_graph_cut::leaf_tensor(gf, i);
ggml_tensor* param_leaf = leaf; ggml_tensor* param_leaf = canonical_param_tensor(leaf);
if (param_leaf != nullptr && params_tensor_set_.find(param_leaf) == params_tensor_set_.end()) {
param_leaf = param_leaf->view_src;
}
if (param_leaf != nullptr && if (param_leaf != nullptr &&
params_tensor_set_.find(param_leaf) != params_tensor_set_.end() &&
seen_params.insert(param_leaf).second) { seen_params.insert(param_leaf).second) {
used_params.push_back(param_leaf); used_params.push_back(param_leaf);
} }
@ -2101,11 +2118,17 @@ protected:
ggml_backend_t current = runtime_backend; ggml_backend_t current = runtime_backend;
const int n_nodes = ggml_graph_n_nodes(gf); const int n_nodes = ggml_graph_n_nodes(gf);
for (int i = 0; i < n_nodes; i++) { 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++) { for (int s = 0; s < GGML_MAX_SRC; s++) {
ggml_backend_t weight_backend = backend_for_weight(node->src[s]); ggml_backend_t weight_backend = backend_for_weight(node->src[s]);
if (weight_backend != nullptr) { 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 || if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
@ -2435,6 +2458,123 @@ protected:
return true; 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 { struct PersistentExternalBinding {
ggml_backend_buffer_t buffer = nullptr; ggml_backend_buffer_t buffer = nullptr;
void* data = nullptr; void* data = nullptr;
@ -2972,6 +3112,11 @@ public:
GGML_ASSERT(gf != nullptr); GGML_ASSERT(gf != nullptr);
rebuild_params_tensor_set(); 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()) { if (can_attempt_graph_cut_segmented_compute()) {
GraphCutPlan plan; GraphCutPlan plan;
if (!resolve_graph_cut_plan(gf, &plan)) { if (!resolve_graph_cut_plan(gf, &plan)) {
@ -3025,6 +3170,22 @@ public:
stream_layers_enabled = enabled; 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) { void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {
extra_runtime_backends.clear(); extra_runtime_backends.clear();
for (ggml_backend_t backend : backends) { for (ggml_backend_t backend : backends) {
@ -3036,6 +3197,9 @@ public:
extra_runtime_backends.push_back(backend); 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) { if (is_multi_device() && stream_layers_enabled) {
LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring", LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring",
get_desc().c_str()); get_desc().c_str());

View File

@ -1,9 +1,11 @@
#include "core/layer_split_partition.h" #include "core/layer_split_partition.h"
#include <algorithm> #include <algorithm>
#include <cstdint>
#include <cstdlib> #include <cstdlib>
#include <cstring> #include <cstring>
#include <limits>
#include <unordered_set>
#include <utility>
#include "core/util.h" #include "core/util.h"
@ -62,160 +64,194 @@ namespace sd {
return name != nullptr ? name : "unknown"; return name != nullptr ? name : "unknown";
} }
static bool layer_split_backend_supports_tensor(ggml_backend_t backend, const ggml_tensor* tensor) { static size_t graph_cut_layer_split_backend_vram_limit(const std::vector<size_t>& backend_vram_limits,
return backend != nullptr && tensor != nullptr && ggml_backend_supports_op(backend, tensor); 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, static std::vector<int64_t> graph_cut_layer_split_backend_capacities(const std::vector<ggml_backend_t>& backends,
const std::string& tensor_name, const std::vector<size_t>& backend_vram_limits,
const ggml_tensor* tensor, size_t primary_backend_vram_limit) {
const std::vector<ggml_backend_t>& backends, std::vector<int64_t> capacities(backends.size(), std::numeric_limits<int64_t>::max() / 4);
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.
constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024; 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++) { for (size_t i = 0; i < backends.size(); i++) {
ggml_backend_dev_t dev = ggml_backend_get_device(backends[i]); ggml_backend_dev_t dev = ggml_backend_get_device(backends[i]);
size_t free_bytes = 0, total_bytes = 0; size_t free_bytes = 0, total_bytes = 0;
if (dev != nullptr) { if (dev != nullptr) {
ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
} }
// Keep a small share even for tight devices instead of dropping them. if (free_bytes > 0) {
int64_t usable_bytes = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes, capacities[i] = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes, 0);
(int64_t)free_bytes / 8); }
device_weights[i] = usable_bytes > 0 ? (double)usable_bytes : 1.0; size_t limit_bytes = graph_cut_layer_split_backend_vram_limit(backend_vram_limits,
weight_sum += device_weights[i]; i,
} primary_backend_vram_limit);
if (limit_bytes > 0) {
std::vector<int64_t> block_budgets(backends.size(), 0); capacities[i] = std::min<int64_t>(capacities[i], (int64_t)limit_bytes);
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); return capacities;
budget = std::max<int64_t>(budget - other_bytes_by_backend[i], 0); }
block_budgets[i] = budget;
} bool partition_graph_cut_layer_split(const char* desc,
ggml_cgraph* gf,
std::vector<int> boundaries(backends.size(), n_blocks); const sd::ggml_graph_cut::Plan& plan,
size_t current = 0; const std::vector<ggml_backend_t>& split_backends,
int64_t used = 0; const std::vector<size_t>& backend_vram_limits,
for (int b = 0; b < n_blocks; b++) { size_t primary_backend_vram_limit,
int64_t bytes = block_bytes.count(b) != 0 ? block_bytes[b] : 0; std::unordered_map<const ggml_tensor*, ggml_backend_t>& param_assignments,
if (current + 1 < backends.size() && used > 0 && used + bytes > block_budgets[current]) { const std::function<ggml_tensor*(ggml_tensor*)>& canonical_param_tensor,
boundaries[current] = b; GraphCutLayerSplitAssignment* assignment_out) {
current++; GGML_ASSERT(gf != nullptr);
used = 0; 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) { int64_t total_param_bytes = 0;
size_t target = 0; for (int64_t bytes : segment_param_bytes) {
int idx = split_tensors.count(kv.first) != 0 ? layer_split_tensor_block_index(kv.first) : -1; total_param_bytes += bytes;
if (idx >= 0) { }
while (target < boundaries.size() && idx >= boundaries[target]) { if (total_param_bytes <= 0) {
target++; 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 { } else {
auto target_it = non_block_targets.find(kv.first); LOG_DEBUG("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
if (target_it != non_block_targets.end()) { desc,
target = target_it->second; 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 } // namespace sd

View File

@ -1,23 +1,43 @@
#ifndef __SD_CORE_LAYER_SPLIT_PARTITION_H__ #ifndef __SD_CORE_LAYER_SPLIT_PARTITION_H__
#define __SD_CORE_LAYER_SPLIT_PARTITION_H__ #define __SD_CORE_LAYER_SPLIT_PARTITION_H__
#include <map> #include <cstdint>
#include <functional>
#include <string> #include <string>
#include <unordered_map>
#include <vector> #include <vector>
#include "ggml-backend.h" #include "ggml-backend.h"
#include "ggml.h" #include "ggml.h"
#include "core/ggml_graph_cut.h"
namespace sd { 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); std::string layer_split_backend_device_display_name(ggml_backend_t backend);
int layer_split_tensor_block_index(const std::string& name); int layer_split_tensor_block_index(const std::string& name);
bool partition_graph_cut_layer_split(const char* desc,
std::vector<std::map<std::string, ggml_tensor*>> partition_layer_split_tensors( ggml_cgraph* gf,
const std::string& desc, const sd::ggml_graph_cut::Plan& plan,
const std::map<std::string, ggml_tensor*>& tensors, const std::vector<ggml_backend_t>& split_backends,
const std::map<std::string, ggml_tensor*>& split_tensors, const std::vector<size_t>& backend_vram_limits,
const std::vector<ggml_backend_t>& backends); 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 } // namespace sd

View File

@ -134,7 +134,8 @@ bool ModelManager::register_param_tensors(const std::string& desc,
ggml_backend_t compute_backend, ggml_backend_t compute_backend,
ggml_backend_t params_backend, ggml_backend_t params_backend,
size_t* registered_tensor_size, size_t* registered_tensor_size,
bool allow_split_buffer) { bool allow_split_buffer,
bool params_follow_compute_backend) {
if (desc.empty()) { if (desc.empty()) {
LOG_ERROR("model manager tensor desc is empty"); LOG_ERROR("model manager tensor desc is empty");
return false; return false;
@ -158,14 +159,15 @@ bool ModelManager::register_param_tensors(const std::string& desc,
} }
ggml_set_name(tensor, name.c_str()); ggml_set_name(tensor, name.c_str());
auto state = std::make_unique<TensorState>(); auto state = std::make_unique<TensorState>();
state->name = name; state->name = name;
state->tensor = tensor; state->tensor = tensor;
state->desc = desc; state->desc = desc;
state->residency_mode = residency_mode; state->residency_mode = residency_mode;
state->compute_backend = compute_backend; state->compute_backend = compute_backend;
state->params_backend = params_backend; state->params_backend = params_backend;
state->allow_split_buffer = allow_split_buffer; state->allow_split_buffer = allow_split_buffer;
state->params_follow_compute_backend = params_follow_compute_backend;
new_states.push_back(std::move(state)); new_states.push_back(std::move(state));
} }
@ -919,6 +921,54 @@ bool ModelManager::resolve_required_tensor_states(const std::vector<ggml_tensor*
return true; 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) { bool ModelManager::prepare_params(const std::vector<ggml_tensor*>& tensors) {
if (tensors.empty()) { if (tensors.empty()) {
return true; return true;

View File

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

View File

@ -268,6 +268,15 @@ public:
return max_vram_assignment.bytes_for_backend(backend_for(module)); 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) { bool ensure_backend_pair(SDBackendModule module) {
if (backend_for(module) == nullptr) { if (backend_for(module) == nullptr) {
return false; return false;
@ -427,8 +436,9 @@ public:
params_mem_size); params_mem_size);
} }
// Register each layer-split partition with its compute backend; the // Register graph-cut layer-split tensors on the primary backend first.
// ModelManager handles allocation, staging, and LoRA by backend. // The first real graph assigns each param tensor to a runtime backend
// before weights are loaded or staged.
template <typename T> template <typename T>
bool register_layer_split_runner_params(const std::string& desc, bool register_layer_split_runner_params(const std::string& desc,
const std::shared_ptr<T>& model, const std::shared_ptr<T>& model,
@ -459,52 +469,29 @@ public:
params_mem_size); 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_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) || const bool params_follow_runtime = backend_manager.params_backend_follows_runtime(module) ||
backend_manager.params_backend_is_disk(module); backend_manager.params_backend_is_disk(module);
for (size_t i = 0; i < module_backends.size(); i++) { ggml_backend_t initial_params_backend = params_follow_runtime ? module_backends[0] : params_backend_for(module);
if (partitions[i].empty()) { if (initial_params_backend == nullptr) {
continue; return false;
}
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;
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() { bool init_backend() {
@ -529,6 +516,16 @@ public:
return false; 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) { std::shared_ptr<RNG> get_rng(rng_type_t rng_type) {
if (rng_type == STD_DEFAULT_RNG) { if (rng_type == STD_DEFAULT_RNG) {
return std::make_shared<STDDefaultRNG>(); return std::make_shared<STDDefaultRNG>();
@ -785,6 +782,10 @@ public:
LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring"); LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring");
stream_layers = false; 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> wtype_stat = model_loader.get_wtype_stat();
std::map<ggml_type, uint32_t> conditioner_wtype_stat = model_loader.get_conditioner_wtype_stat(); std::map<ggml_type, uint32_t> conditioner_wtype_stat = model_loader.get_conditioner_wtype_stat();

View File

@ -3,10 +3,14 @@
#include <vector> #include <vector>
#include "ggml-backend.h"
struct ggml_tensor; struct ggml_tensor;
struct RunnerWeightManager { struct RunnerWeightManager {
virtual ~RunnerWeightManager() = default; 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 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_compute_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;
virtual void release_params_backend_params(const std::vector<ggml_tensor*>& tensors) = 0; virtual void release_params_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;