DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-11 19:36:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: unify runner lifecycles and weight residency (#1940)

Browse Source
This commit is contained in:
leejet authored and GitHub committed 2026-09-06 22:30:45 +08:00
1 parent 6c57cc3b38
commit 462d675018
78 files changed
+2963 -2051
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/conditioning/conditioner.hpp
+50 -142
View File
@@ -142,14 +142,13 @@ public:
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
virtual void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {}
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) {}
virtual void runner_done() {}
virtual void runner_end() {}
};
// ldm.modules.encoders.modules.FrozenCLIPEmbedder
@@ -202,13 +201,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
}
}
void set_stream_layers_enabled(bool enabled) override {
text_model->set_stream_layers_enabled(enabled);
if (sd_version_is_sdxl(version)) {
text_model2->set_stream_layers_enabled(enabled);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
text_model->set_runtime_backends(backends);
if (sd_version_is_sdxl(version)) {
@@ -244,10 +236,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
}
}
void runner_done() override {
text_model->runner_done();
void runner_end() override {
text_model->runner_end();
if (sd_version_is_sdxl(version)) {
text_model2->runner_done();
text_model2->runner_end();
}
}
@@ -461,9 +453,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
max_token_idx,
false,
clip_skip,
false,
true,
true);
false);
GGML_ASSERT(!chunk_hidden_states.empty());
if (sd_version_is_sdxl(version)) {
auto chunk_hidden_states2 = text_model2->compute(n_threads,
@@ -473,9 +463,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
max_token_idx,
false,
clip_skip,
false,
true,
true);
false);
GGML_ASSERT(!chunk_hidden_states2.empty());
chunk_hidden_states = sd::ops::concat(chunk_hidden_states, chunk_hidden_states2, 0);
@@ -487,9 +475,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
max_token_idx,
true,
clip_skip,
false,
true,
true);
false);
GGML_ASSERT(!pooled.empty());
}
}
@@ -608,7 +594,7 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(pixel_values, return_pooled, clip_skip);
};
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true));
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
}
};
@@ -675,18 +661,6 @@ struct SD3CLIPEmbedder : public Conditioner {
}
}
void set_stream_layers_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_stream_layers_enabled(enabled);
}
if (clip_g) {
clip_g->set_stream_layers_enabled(enabled);
}
if (t5) {
t5->set_stream_layers_enabled(enabled);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (clip_l) {
clip_l->set_runtime_backends(backends);
@@ -753,15 +727,15 @@ struct SD3CLIPEmbedder : public Conditioner {
}
}
void runner_done() override {
void runner_end() override {
if (clip_l) {
clip_l->runner_done();
clip_l->runner_end();
}
if (clip_g) {
clip_g->runner_done();
clip_g->runner_end();
}
if (t5) {
t5->runner_done();
t5->runner_end();
}
}
@@ -881,9 +855,7 @@ struct SD3CLIPEmbedder : public Conditioner {
max_token_idx,
false,
clip_skip,
false,
true,
true);
false);
GGML_ASSERT(!chunk_hidden_states_l.empty());
chunk_hidden_states_l = ::apply_token_weights(std::move(chunk_hidden_states_l), chunk_weights);
@@ -897,9 +869,7 @@ struct SD3CLIPEmbedder : public Conditioner {
max_token_idx,
true,
clip_skip,
false,
true,
true);
false);
GGML_ASSERT(!pooled_l.empty());
}
} else {
@@ -928,9 +898,7 @@ struct SD3CLIPEmbedder : public Conditioner {
max_token_idx,
false,
clip_skip,
false,
true,
true);
false);
GGML_ASSERT(!chunk_hidden_states_g.empty());
chunk_hidden_states_g = ::apply_token_weights(std::move(chunk_hidden_states_g), chunk_weights);
@@ -944,9 +912,7 @@ struct SD3CLIPEmbedder : public Conditioner {
max_token_idx,
true,
clip_skip,
false,
true,
true);
false);
GGML_ASSERT(!pooled_g.empty());
}
} else {
@@ -969,9 +935,7 @@ struct SD3CLIPEmbedder : public Conditioner {
chunk_hidden_states_t5 = t5->compute(n_threads,
input_ids,
sd::Tensor<float>(),
false,
true,
true);
false);
GGML_ASSERT(!chunk_hidden_states_t5.empty());
chunk_hidden_states_t5 = ::apply_token_weights(std::move(chunk_hidden_states_t5), chunk_weights);
} else {
@@ -1079,15 +1043,6 @@ struct FluxCLIPEmbedder : public Conditioner {
}
}
void set_stream_layers_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_stream_layers_enabled(enabled);
}
if (t5) {
t5->set_stream_layers_enabled(enabled);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (clip_l) {
clip_l->set_runtime_backends(backends);
@@ -1139,12 +1094,12 @@ struct FluxCLIPEmbedder : public Conditioner {
}
}
void runner_done() override {
void runner_end() override {
if (clip_l) {
clip_l->runner_done();
clip_l->runner_end();
}
if (t5) {
t5->runner_done();
t5->runner_end();
}
}
@@ -1247,9 +1202,7 @@ struct FluxCLIPEmbedder : public Conditioner {
max_token_idx,
true,
clip_skip,
false,
true,
true);
false);
GGML_ASSERT(!pooled.empty());
} else {
pooled = sd::Tensor<float>::zeros({768});
@@ -1268,9 +1221,7 @@ struct FluxCLIPEmbedder : public Conditioner {
chunk_hidden_states = t5->compute(n_threads,
input_ids,
sd::Tensor<float>(),
false,
true,
true);
false);
GGML_ASSERT(!chunk_hidden_states.empty());
chunk_hidden_states = ::apply_token_weights(std::move(chunk_hidden_states), chunk_weights);
if (zero_out_masked) {
@@ -1366,12 +1317,6 @@ struct T5CLIPEmbedder : public Conditioner {
}
}
void set_stream_layers_enabled(bool enabled) override {
if (t5) {
t5->set_stream_layers_enabled(enabled);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (t5) {
t5->set_runtime_backends(backends);
@@ -1408,9 +1353,9 @@ struct T5CLIPEmbedder : public Conditioner {
}
}
void runner_done() override {
void runner_end() override {
if (t5) {
t5->runner_done();
t5->runner_end();
}
}
@@ -1508,9 +1453,7 @@ struct T5CLIPEmbedder : public Conditioner {
auto chunk_hidden_states = t5->compute(n_threads,
input_ids,
t5_attn_mask_chunk,
false,
true,
true);
false);
GGML_ASSERT(!chunk_hidden_states.empty());
chunk_hidden_states = apply_token_weights(std::move(chunk_hidden_states), chunk_weights);
@@ -1582,12 +1525,6 @@ struct MiniT2IConditioner : public Conditioner {
}
}
void set_stream_layers_enabled(bool enabled) override {
if (t5) {
t5->set_stream_layers_enabled(enabled);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (t5) {
t5->set_runtime_backends(backends);
@@ -1624,9 +1561,9 @@ struct MiniT2IConditioner : public Conditioner {
}
}
void runner_done() override {
void runner_end() override {
if (t5) {
t5->runner_done();
t5->runner_end();
}
}
@@ -1657,9 +1594,7 @@ struct MiniT2IConditioner : public Conditioner {
sd::Tensor<float> hidden_states = t5->compute(n_threads,
input_ids,
sd::Tensor<float>::from_vector(t5_mask),
false,
true,
true);
false);
GGML_ASSERT(!hidden_states.empty());
result.c_crossattn = std::move(hidden_states);
result.c_vector = sd::Tensor<float>::from_vector(mask);
@@ -1696,10 +1631,6 @@ struct AnimaConditioner : public Conditioner {
llm->set_max_graph_vram_bytes(max_vram_bytes);
}
void set_stream_layers_enabled(bool enabled) override {
llm->set_stream_layers_enabled(enabled);
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
llm->set_runtime_backends(backends);
}
@@ -1724,8 +1655,8 @@ struct AnimaConditioner : public Conditioner {
llm->set_weight_adapter(adapter);
}
void runner_done() override {
llm->runner_done();
void runner_end() override {
llm->runner_end();
}
std::tuple<std::vector<int>, std::vector<float>, std::vector<int>, std::vector<float>> tokenize(std::string text) {
@@ -1787,9 +1718,7 @@ struct AnimaConditioner : public Conditioner {
{},
{},
false,
false,
true,
true);
false);
GGML_ASSERT(!hidden_states.empty());
hidden_states = apply_token_weights(std::move(hidden_states), qwen_weights);
auto t5_ids_tensor = sd::Tensor<int32_t>::from_vector(t5_tokens);
@@ -1887,13 +1816,6 @@ struct LLMEmbedder : public Conditioner {
}
}
void set_stream_layers_enabled(bool enabled) override {
llm->set_stream_layers_enabled(enabled);
if (byt5) {
byt5->set_stream_layers_enabled(enabled);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
llm->set_runtime_backends(backends);
if (byt5) {
@@ -1942,12 +1864,12 @@ struct LLMEmbedder : public Conditioner {
}
}
void runner_done() override {
void runner_end() override {
if (llm) {
llm->runner_done();
llm->runner_end();
}
if (byt5) {
byt5->runner_done();
byt5->runner_end();
}
}
@@ -2051,8 +1973,6 @@ struct LLMEmbedder : public Conditioner {
out_layers,
false,
false,
true,
true,
deepstack_image_embeds,
image_grids);
GGML_ASSERT(!hidden_states.empty());
@@ -2220,9 +2140,7 @@ struct LLMEmbedder : public Conditioner {
prompt += "<Picture " + std::to_string(++picture_index) + ">: ";
add_vision_outputs(llm->encode_image_outputs(n_threads,
resized,
false,
true,
true),
false),
static_cast<int>(resized.shape()[1]) / patch_size,
static_cast<int>(resized.shape()[0]) / patch_size);
continue;
@@ -2250,9 +2168,7 @@ struct LLMEmbedder : public Conditioner {
auto pair = sd::ops::concat(first.unsqueeze(2), second.unsqueeze(2), 2);
add_vision_outputs(llm->encode_video_block_outputs(n_threads,
pair,
false,
true,
true),
false),
static_cast<int>(first.shape()[1]) / patch_size,
static_cast<int>(first.shape()[0]) / patch_size);
}
@@ -2263,9 +2179,7 @@ struct LLMEmbedder : public Conditioner {
prompt += "<Picture " + std::to_string(i + 1) + ">: ";
add_vision_outputs(llm->encode_image_outputs(n_threads,
resized,
false,
true,
true),
false),
static_cast<int>(resized.shape()[1]) / patch_size,
static_cast<int>(resized.shape()[0]) / patch_size);
}
@@ -2352,7 +2266,7 @@ struct LLMEmbedder : public Conditioner {
LOG_DEBUG("resize LingBotVideo ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
std::string image_prefix = prompt + img_prompt + "<|vision_start|>";
@@ -2411,7 +2325,7 @@ struct LLMEmbedder : public Conditioner {
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
image_embeds.emplace_back(image_embed_idx, image_embed);
image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
@@ -2494,7 +2408,7 @@ struct LLMEmbedder : public Conditioner {
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
std::string image_prefix = prompt_prefix + img_prompt + "<|vision_start|>";
@@ -2562,7 +2476,7 @@ struct LLMEmbedder : public Conditioner {
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
std::string image_prefix = prompt + img_prompt + "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
@@ -2625,7 +2539,7 @@ struct LLMEmbedder : public Conditioner {
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
image_embeds.emplace_back(image_embed_idx, image_embed);
image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
@@ -2857,9 +2771,7 @@ struct LLMEmbedder : public Conditioner {
auto byt5_hidden_states = byt5->compute(n_threads,
input_ids,
sd::Tensor<float>(),
false,
true,
true);
false);
GGML_ASSERT(!byt5_hidden_states.empty());
extra_hidden_states_vec.push_back(std::move(byt5_hidden_states));
}
@@ -2960,13 +2872,11 @@ struct LTXAVTextProjectionRunner : public GGMLRunner {
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = true) {
bool auto_runner_end = true) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x);
};
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
}
};
@@ -3060,9 +2970,9 @@ struct LTXAVEmbedder : public Conditioner {
projector->set_weight_adapter(adapter);
}
void runner_done() override {
llm->runner_done();
projector->runner_done();
void runner_end() override {
llm->runner_end();
projector->runner_end();
}
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
@@ -3129,9 +3039,7 @@ struct LTXAVEmbedder : public Conditioner {
{},
{},
true,
false,
true,
true);
false);
GGML_ASSERT(!hidden_states.empty());
hidden_states = apply_token_weights(std::move(hidden_states), weights);
@@ -3190,7 +3098,7 @@ struct LTXAVEmbedder : public Conditioner {
}
hidden_states.reshape_({kNumStates * kHiddenSize, valid_tokens});
return projector->compute(n_threads, hidden_states, false, true, true);
return projector->compute(n_threads, hidden_states, false);
}
SDCondition get_learned_condition(int n_threads,
src/core/compute_workspace.cpp
+257
View File
@@ -0,0 +1,257 @@
#include "core/compute_workspace.h"
#include <algorithm>
#include <cstring>
#include <map>
#include <unordered_map>
#include <unordered_set>
#include "core/ggml_extend_backend.h"
#include "core/ggml_graph_cut.h"
#include "ggml-cpu.h"
#include "ggml/src/ggml-impl.h"
namespace sd {
ComputeWorkspace::~ComputeWorkspace() {
segment_end();
release();
ggml_backend_free(cpu_backend_);
}
void ComputeWorkspace::set_extra_backends(const std::vector<ggml_backend_t>& backends) {
if (extra_backends_ != backends) {
GGML_ASSERT(!active_);
release();
extra_backends_ = backends;
}
}
bool ComputeWorkspace::needs_scheduler(ggml_cgraph* graph) const {
if (!extra_backends_.empty()) {
return true;
}
for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
if (!ggml_backend_supports_op(backend_, ggml_graph_node(graph, i))) {
return true;
}
}
return false;
}
ggml_backend_sched_t ComputeWorkspace::make_scheduler(size_t graph_size) {
std::vector<ggml_backend_t> backends{backend_};
backends.insert(backends.end(), extra_backends_.begin(), extra_backends_.end());
if (!sd_backend_is_cpu(backend_)) {
if (cpu_backend_ == nullptr) {
cpu_backend_ = sd_backend_cpu_init();
}
if (cpu_backend_ == nullptr) {
return nullptr;
}
backends.push_back(cpu_backend_);
}
std::vector<ggml_backend_buffer_type_t> bufts;
for (auto backend : backends) {
auto buft = backend == cpu_backend_
? ggml_backend_dev_host_buffer_type(ggml_backend_get_device(backend_))
: nullptr;
bufts.push_back(buft != nullptr ? buft : ggml_backend_get_default_buffer_type(backend));
}
return ggml_backend_sched_new(backends.data(), bufts.data(), static_cast<int>(backends.size()),
graph_size, false, false);
}
bool ComputeWorkspace::measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const {
return measurement.scheduler == needs_scheduler(graph);
}
bool ComputeWorkspace::prepare(const Measurement& measurement) {
GGML_ASSERT(!active_);
if (measurement.buffers.empty()) {
return false;
}
const bool grows = std::any_of(measurement.buffers.begin(), measurement.buffers.end(),
[&](const BackendBufferSize& size) { return size.bytes > bytes(size.backend); });
if (measurement.scheduler != (scheduler_ != nullptr) || grows) {
release();
}
return true;
}
bool ComputeWorkspace::release_excess(const Measurement& measurement) {
return std::any_of(measurement.buffers.begin(), measurement.buffers.end(),
[&](const BackendBufferSize& size) { return bytes(size.backend) > size.bytes; }) &&
release();
}
bool ComputeWorkspace::allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes) {
GGML_ASSERT(!active_);
const bool use_scheduler = needs_scheduler(graph);
if (use_scheduler) {
if (allocator_ != nullptr) {
release();
}
const size_t capacity = static_cast<size_t>(graph->n_nodes + graph->n_leafs) + 8;
if (scheduler_ == nullptr || capacity > scheduler_capacity_) {
release();
scheduler_ = make_scheduler(capacity);
scheduler_capacity_ = capacity;
}
if (scheduler_ == nullptr) {
return false;
}
ggml_backend_sched_reset(scheduler_);
assign_nodes(scheduler_, graph);
// Scheduler allocation rewrites sources. Split the execution graph only once.
if (!ggml_backend_sched_alloc_graph(scheduler_, graph)) {
release();
return false;
}
} else {
if (scheduler_ != nullptr) {
release();
}
if (allocator_ == nullptr) {
allocator_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));
}
auto signature = ggml_graph_cut::graph_layout(graph, true);
if (signature != reservation_) {
if (!ggml_gallocr_reserve(allocator_, graph)) {
release();
return false;
}
reservation_ = std::move(signature);
++reservations_;
}
if (!ggml_gallocr_alloc_graph(allocator_, graph)) {
release();
return false;
}
}
active_ = true;
return true;
}
ComputeWorkspace::Measurement ComputeWorkspace::measure(
ggml_cgraph* graph,
size_t direct_bytes,
const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
const AssignNodes& assign_nodes) {
if (!needs_scheduler(graph)) {
return {{{backend_, direct_bytes}}, false};
}
std::vector<const ggml_tensor*> tensors;
std::unordered_set<const ggml_tensor*> seen;
auto visit = [&](const ggml_tensor* tensor) {
if (tensor != nullptr && seen.insert(tensor).second) {
tensors.push_back(tensor);
}
};
for (int i = 0; i < graph->n_nodes; ++i) {
visit(graph->nodes[i]);
}
for (int i = 0; i < graph->n_leafs; ++i) {
visit(graph->leafs[i]);
}
for (size_t i = 0; i < tensors.size(); ++i) {
visit(tensors[i]->view_src);
for (auto source : tensors[i]->src) {
visit(source);
}
}
const size_t graph_size = tensors.size() + 8;
auto context = ggml_init({tensors.size() * ggml_tensor_overhead() + ggml_graph_overhead_custom(graph_size, false), nullptr, true});
if (context == nullptr) {
return {};
}
std::unordered_map<const ggml_tensor*, ggml_tensor*> copies;
std::map<ggml_backend_t, ggml_backend_buffer_t> external_buffers;
for (auto tensor : tensors) {
auto copy = ggml_dup_tensor(context, tensor);
*copy = *tensor;
copies[tensor] = copy;
}
for (const auto& entry : copies) {
auto source = entry.first;
auto copy = entry.second;
copy->view_src = source->view_src == nullptr ? nullptr : copies.at(source->view_src);
for (int i = 0; i < GGML_MAX_SRC; ++i) {
copy->src[i] = source->src[i] == nullptr ? nullptr : copies.at(source->src[i]);
}
auto external = external_backend(source);
if (external != nullptr && source->view_src == nullptr) {
auto& buffer = external_buffers[external];
if (buffer == nullptr) {
buffer = ggml_backend_alloc_buffer(external, 0);
GGML_ASSERT(buffer != nullptr);
ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
}
copy->buffer = buffer;
copy->data = reinterpret_cast<void*>(static_cast<uintptr_t>(1));
copy->extra = nullptr;
}
}
auto copy_graph = ggml_new_graph_custom(context, graph_size, false);
copy_graph->n_nodes = graph->n_nodes;
copy_graph->n_leafs = graph->n_leafs;
for (int i = 0; i < graph->n_nodes; ++i) {
copy_graph->nodes[i] = copies.at(graph->nodes[i]);
}
for (int i = 0; i < graph->n_leafs; ++i) {
copy_graph->leafs[i] = copies.at(graph->leafs[i]);
}
Measurement result;
result.scheduler = true;
auto scheduler = make_scheduler(graph_size);
if (scheduler != nullptr) {
assign_nodes(scheduler, copy_graph);
std::vector<size_t> sizes(extra_backends_.size() + 2);
ggml_backend_sched_reserve_size(scheduler, copy_graph, sizes.data());
result.buffers.push_back({backend_, sizes[0]});
for (size_t i = 0; i < extra_backends_.size(); ++i) {
result.buffers.push_back({extra_backends_[i], sizes[i + 1]});
}
ggml_backend_sched_free(scheduler);
}
for (const auto& entry : external_buffers) {
ggml_backend_buffer_free(entry.second);
}
ggml_free(context);
return result;
}
void ComputeWorkspace::synchronize() const {
if (scheduler_ != nullptr) {
ggml_backend_sched_synchronize(scheduler_);
} else {
ggml_backend_synchronize(backend_);
}
}
void ComputeWorkspace::segment_end() {
if (active_) {
synchronize();
active_ = false;
}
}
bool ComputeWorkspace::release() {
if (active_) {
return false;
}
ggml_gallocr_free(allocator_);
allocator_ = nullptr;
ggml_backend_sched_free(scheduler_);
scheduler_ = nullptr;
scheduler_capacity_ = 0;
reservation_.clear();
return true;
}
size_t ComputeWorkspace::bytes(ggml_backend_t backend) const {
if (scheduler_ != nullptr) {
return ggml_backend_sched_get_buffer_size(scheduler_, backend);
}
return allocator_ != nullptr && backend == backend_ ? ggml_gallocr_get_buffer_size(allocator_, 0) : 0;
}
}
src/core/compute_workspace.h
+64
View File
@@ -0,0 +1,64 @@
#ifndef __SD_CORE_COMPUTE_WORKSPACE_H__
#define __SD_CORE_COMPUTE_WORKSPACE_H__
#include <functional>
#include <vector>
#include "ggml-alloc.h"
#include "ggml-backend.h"
namespace sd {
struct BackendBufferSize {
ggml_backend_t backend = nullptr;
size_t bytes = 0;
};
class ComputeWorkspace {
ggml_backend_t backend_;
std::vector<ggml_backend_t> extra_backends_;
ggml_backend_t cpu_backend_ = nullptr;
ggml_gallocr_t allocator_ = nullptr;
ggml_backend_sched_t scheduler_ = nullptr;
size_t scheduler_capacity_ = 0;
std::vector<uint64_t> reservation_;
bool active_ = false;
size_t reservations_ = 0;
ggml_backend_sched_t make_scheduler(size_t graph_size);
bool needs_scheduler(ggml_cgraph* graph) const;
public:
struct Measurement {
std::vector<BackendBufferSize> buffers;
bool scheduler = false;
};
using AssignNodes = std::function<void(ggml_backend_sched_t, ggml_cgraph*)>;
explicit ComputeWorkspace(ggml_backend_t backend)
: backend_(backend) {}
~ComputeWorkspace();
ComputeWorkspace(const ComputeWorkspace&) = delete;
ComputeWorkspace& operator=(const ComputeWorkspace&) = delete;
void set_extra_backends(const std::vector<ggml_backend_t>& backends);
bool measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const;
bool prepare(const Measurement& measurement);
bool release_excess(const Measurement& measurement);
bool allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes);
Measurement measure(
ggml_cgraph* graph,
size_t direct_bytes,
const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
const AssignNodes& assign_nodes);
void synchronize() const;
void segment_end();
bool release();
bool active() const { return active_; }
ggml_backend_sched_t scheduler() const { return scheduler_; }
ggml_backend_t cpu_backend() const { return cpu_backend_; }
size_t bytes(ggml_backend_t backend) const;
size_t reservation_count() const { return reservations_; }
};
}
#endif // __SD_CORE_COMPUTE_WORKSPACE_H__
src/core/ggml_extend.hpp
+150 -883
View File
File diff suppressed because it is too large. Load diff
src/core/ggml_graph_cut.cpp
+235 -339
View File
@@ -2,6 +2,7 @@
#include <algorithm>
#include <cctype>
#include <climits>
#include <cmath>
#include <cstring>
#include <map>
@@ -67,25 +68,6 @@ namespace sd::ggml_graph_cut {
return -1;
}
static Plan::InputShape input_shape(const ggml_tensor* tensor) {
Plan::InputShape shape;
if (tensor == nullptr) {
return shape;
}
shape.type = tensor->type;
for (int i = 0; i < GGML_MAX_DIMS; ++i) {
shape.ne[static_cast<size_t>(i)] = tensor->ne[i];
}
return shape;
}
static size_t graph_cut_segment_vram_bytes(const Segment& segment) {
return segment.compute_buffer_size +
segment.input_param_bytes +
segment.input_previous_cut_bytes +
segment.output_bytes;
}
static std::string lower_ascii_copy(std::string value) {
std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) {
return static_cast<char>(std::tolower(c));
@@ -291,55 +273,6 @@ namespace sd::ggml_graph_cut {
return max_vram_bytes_to_gib(resolve_auto_max_vram_bytes(-max_vram, backend));
}
static bool is_segment_output_needed_after(const Plan& plan,
size_t end_segment_index,
int output_node_index) {
if (end_segment_index + 1 >= plan.segments.size()) {
return false;
}
for (size_t seg_idx = end_segment_index + 1; seg_idx < plan.segments.size(); ++seg_idx) {
const auto& segment = plan.segments[seg_idx];
for (const auto& input_ref : segment.input_refs) {
if (input_ref.type == Segment::INPUT_PREVIOUS_CUT &&
input_ref.node_index == output_node_index) {
return true;
}
}
}
return false;
}
static Segment make_segment_seed(const Plan& plan,
size_t start_segment_index,
size_t end_segment_index) {
GGML_ASSERT(start_segment_index < plan.segments.size());
GGML_ASSERT(end_segment_index < plan.segments.size());
GGML_ASSERT(start_segment_index <= end_segment_index);
Segment seed;
const auto& start_segment = plan.segments[start_segment_index];
const auto& target_segment = plan.segments[end_segment_index];
std::unordered_set<int> seen_output_node_indices;
for (size_t seg_idx = start_segment_index; seg_idx <= end_segment_index; ++seg_idx) {
const bool is_boundary_segment = seg_idx == end_segment_index;
for (int output_node_index : plan.segments[seg_idx].output_node_indices) {
if ((is_boundary_segment ||
is_segment_output_needed_after(plan, end_segment_index, output_node_index)) &&
seen_output_node_indices.insert(output_node_index).second) {
seed.output_node_indices.push_back(output_node_index);
}
}
}
if (start_segment_index == end_segment_index) {
seed.group_name = target_segment.group_name;
} else {
seed.group_name = sd_format("%s..%s",
start_segment.group_name.c_str(),
target_segment.group_name.c_str());
}
return seed;
}
static void build_segment(ggml_cgraph* gf,
Plan& plan,
Segment& segment,
@@ -416,31 +349,7 @@ namespace sd::ggml_graph_cut {
}
return a.display_name < b.display_name;
});
segment.input_refs = input_refs;
for (const auto& input : input_refs) {
ggml_tensor* current_input = input_tensor(gf, input);
size_t tensor_bytes = current_input == nullptr
? 0
: (input.type == Segment::INPUT_PREVIOUS_CUT
? cache_tensor_bytes(current_input)
: ggml_nbytes(current_input));
switch (input.type) {
case Segment::INPUT_PREVIOUS_CUT:
segment.input_previous_cut_bytes += tensor_bytes;
break;
case Segment::INPUT_PARAM:
segment.input_param_bytes += tensor_bytes;
break;
case Segment::INPUT_EXTERNAL:
default:
segment.input_external_bytes += tensor_bytes;
break;
}
}
for (int output_node_index : segment.output_node_indices) {
ggml_tensor* output = ggml_graph_node(gf, output_node_index);
segment.output_bytes += cache_tensor_bytes(output);
}
segment.input_refs = input_refs;
segment.compute_buffer_size = measure_segment_compute_buffer(backend, gf, segment, log_desc);
for (int output_node_index : segment.output_node_indices) {
@@ -449,6 +358,70 @@ namespace sd::ggml_graph_cut {
plan.segments.push_back(std::move(segment));
}
static bool validate_plan(ggml_cgraph* gf,
const Plan& plan,
std::string* validation_error) {
auto fail = [&](const std::string& reason) {
if (validation_error != nullptr) {
*validation_error = reason;
}
return false;
};
if (!plan.has_cuts) {
return true;
}
if (plan.segments.size() <= 1) {
return fail("fewer than two segments");
}
const int n_nodes = ggml_graph_n_nodes(gf);
std::unordered_set<int> completed_outputs;
for (size_t segment_index = 0; segment_index < plan.segments.size(); ++segment_index) {
const Segment& segment = plan.segments[segment_index];
const std::string segment_label = "segment " + std::to_string(segment_index) +
" ('" + segment.group_name + "')";
if (segment.internal_node_indices.empty() || segment.output_node_indices.empty()) {
return fail(segment_label + " has no internal nodes or outputs");
}
for (const Segment::InputRef& input : segment.input_refs) {
if (input.type == Segment::INPUT_PREVIOUS_CUT) {
if (input.node_index < 0 || input.node_index >= n_nodes ||
completed_outputs.find(input.node_index) == completed_outputs.end()) {
return fail(segment_label + " references an unavailable cut node " +
std::to_string(input.node_index));
}
} else if (input.leaf_index < 0 || input.leaf_index >= gf->n_leafs) {
return fail(segment_label + " references an invalid leaf " +
std::to_string(input.leaf_index));
}
}
std::unordered_set<int> segment_nodes;
segment_nodes.reserve(segment.internal_node_indices.size());
for (int node_index : segment.internal_node_indices) {
if (node_index < 0 || node_index >= n_nodes) {
return fail(segment_label + " contains an invalid node " +
std::to_string(node_index));
}
if (!segment_nodes.insert(node_index).second) {
return fail(segment_label + " contains duplicate node " +
std::to_string(node_index));
}
}
for (int output_index : segment.output_node_indices) {
if (output_index < 0 || output_index >= n_nodes ||
segment_nodes.find(output_index) == segment_nodes.end()) {
return fail(segment_label + " has an output outside its node set: " +
std::to_string(output_index));
}
if (completed_outputs.find(output_index) != completed_outputs.end()) {
return fail(segment_label + " repeats output node " +
std::to_string(output_index));
}
completed_outputs.insert(output_index);
}
}
return true;
}
bool is_graph_cut_tensor(const ggml_tensor* tensor) {
if (tensor == nullptr || tensor->name[0] == '\0') {
return false;
@@ -509,26 +482,80 @@ namespace sd::ggml_graph_cut {
return ggml_nbytes(cache_src);
}
bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan) {
GGML_ASSERT(gf != nullptr);
if (ggml_graph_n_nodes(gf) != plan.n_nodes || gf->n_leafs != plan.n_leafs) {
return false;
}
for (const auto& input_shape_ref : plan.input_shapes) {
if (input_shape_ref.leaf_index < 0 || input_shape_ref.leaf_index >= gf->n_leafs) {
return false;
std::vector<uint64_t> graph_layout(ggml_cgraph* graph, bool include_bindings) {
std::vector<const ggml_tensor*> tensors;
std::unordered_map<const ggml_tensor*, size_t> indices;
auto add = [&](const ggml_tensor* tensor) {
if (tensor != nullptr && indices.emplace(tensor, tensors.size() + 1).second) {
tensors.push_back(tensor);
}
ggml_tensor* leaf = gf->leafs[input_shape_ref.leaf_index];
if (leaf == nullptr || input_shape_ref.type != leaf->type) {
return false;
};
for (int i = 0; i < graph->n_leafs; ++i) {
add(graph->leafs[i]);
}
for (int i = 0; i < graph->n_nodes; ++i) {
add(graph->nodes[i]);
}
for (size_t i = 0; i < tensors.size(); ++i) {
add(tensors[i]->view_src);
for (auto source : tensors[i]->src) {
add(source);
}
}
std::vector<uint64_t> signature;
signature.reserve(tensors.size() * 24);
signature.push_back(graph->n_nodes);
signature.push_back(graph->n_leafs);
for (int i = 0; i < graph->n_leafs; ++i) {
signature.push_back(indices.at(graph->leafs[i]));
}
for (int i = 0; i < graph->n_nodes; ++i) {
signature.push_back(indices.at(graph->nodes[i]));
}
for (auto tensor : tensors) {
signature.push_back(tensor->op);
signature.push_back(tensor->type);
signature.push_back(tensor->flags);
signature.push_back(tensor->view_offs);
if (include_bindings) {
signature.push_back(tensor->data != nullptr);
auto buffer = tensor_buffer(tensor);
signature.push_back(reinterpret_cast<uintptr_t>(buffer == nullptr ? nullptr : ggml_backend_buffer_get_type(buffer)));
}
for (int d = 0; d < GGML_MAX_DIMS; ++d) {
if (input_shape_ref.ne[static_cast<size_t>(d)] != leaf->ne[d]) {
return false;
}
signature.push_back(tensor->ne[d]);
signature.push_back(tensor->nb[d]);
}
signature.push_back(tensor->view_src == nullptr ? 0 : indices.at(tensor->view_src));
for (auto source : tensor->src) {
signature.push_back(source == nullptr ? 0 : indices.at(source));
}
for (int value : tensor->op_params) {
signature.push_back(static_cast<uint32_t>(value));
}
}
return true;
return signature;
}
bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan) {
GGML_ASSERT(gf != nullptr);
if (plan.leaf_names.size() != static_cast<size_t>(gf->n_leafs) ||
plan.layout != graph_layout(gf, false)) {
return false;
}
for (int i = 0; i < gf->n_leafs; ++i) {
if (plan.leaf_names[i] != gf->leafs[i]->name) {
return false;
}
}
std::vector<std::pair<int, std::string>> cut_markers;
for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) {
auto node = ggml_graph_node(gf, i);
if (is_graph_cut_tensor(node)) {
cut_markers.emplace_back(i, node->name);
}
}
return cut_markers == plan.cut_markers;
}
ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index) {
@@ -578,26 +605,6 @@ namespace sd::ggml_graph_cut {
return tensors;
}
std::unordered_set<std::string> collect_future_input_names(ggml_cgraph* gf,
const Plan& plan,
size_t current_segment_index) {
GGML_ASSERT(gf != nullptr);
std::unordered_set<std::string> future_input_names;
for (size_t seg_idx = current_segment_index + 1; seg_idx < plan.segments.size(); ++seg_idx) {
const auto& segment = plan.segments[seg_idx];
for (const auto& input_ref : segment.input_refs) {
if (input_ref.type != Segment::INPUT_PREVIOUS_CUT) {
continue;
}
ggml_tensor* current_input = input_tensor(gf, input_ref);
if (current_input != nullptr && current_input->name[0] != '\0') {
future_input_names.insert(current_input->name);
}
}
}
return future_input_names;
}
ggml_cgraph* build_segment_graph(ggml_cgraph* gf,
const Segment& segment,
ggml_context** graph_ctx_out) {
@@ -662,6 +669,10 @@ namespace sd::ggml_graph_cut {
continue;
}
ggml_set_output(output);
if (output->view_src != nullptr) {
// A consumed output view does not keep its storage alive in gallocr.
ggml_set_output(output->view_src);
}
}
for (int node_idx : segment.internal_node_indices) {
ggml_graph_add_node(segment_graph, ggml_graph_node(gf, node_idx));
@@ -716,6 +727,10 @@ namespace sd::ggml_graph_cut {
if (output != nullptr && saved_output_flags.find(output) == saved_output_flags.end()) {
saved_output_flags[output] = output->flags;
}
if (output != nullptr && output->view_src != nullptr &&
saved_output_flags.find(output->view_src) == saved_output_flags.end()) {
saved_output_flags[output->view_src] = output->view_src->flags;
}
}
ggml_context* graph_ctx = nullptr;
@@ -744,6 +759,46 @@ namespace sd::ggml_graph_cut {
return buffer_size;
}
static size_t measure_graph_compute_buffer(
ggml_backend_t backend,
ggml_cgraph* gf,
const std::unordered_set<const ggml_tensor*>& params_tensor_set) {
struct TensorRuntimeBinding {
ggml_backend_buffer_t buffer = nullptr;
void* data = nullptr;
void* extra = nullptr;
};
std::unordered_map<ggml_tensor*, TensorRuntimeBinding> saved_bindings;
auto mark_external = [&](ggml_tensor* tensor) {
if (tensor == nullptr || saved_bindings.find(tensor) != saved_bindings.end()) {
return;
}
saved_bindings[tensor] = {tensor->buffer, tensor->data, tensor->extra};
tensor->data = reinterpret_cast<void*>(static_cast<uintptr_t>(1));
};
for (int i = 0; i < leaf_count(gf); ++i) {
ggml_tensor* leaf = leaf_tensor(gf, i);
if (!is_params_tensor(params_tensor_set, leaf)) {
continue;
}
mark_external(leaf);
mark_external(leaf->view_src);
}
ggml_gallocr_t allocr = ggml_gallocr_new(
ggml_backend_get_default_buffer_type(backend));
size_t sizes[1] = {0};
ggml_gallocr_reserve_n_size(allocr, gf, nullptr, nullptr, sizes);
ggml_gallocr_free(allocr);
for (const auto& kv : saved_bindings) {
kv.first->buffer = kv.second.buffer;
kv.first->data = kv.second.data;
kv.first->extra = kv.second.extra;
}
return sizes[0];
}
Plan build_plan(ggml_backend_t backend,
ggml_cgraph* gf,
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
@@ -756,24 +811,22 @@ namespace sd::ggml_graph_cut {
if (n_nodes <= 0) {
return plan;
}
plan.n_nodes = n_nodes;
plan.n_leafs = gf->n_leafs;
plan.layout = graph_layout(gf, false);
for (int i = 0; i < gf->n_leafs; ++i) {
ggml_tensor* leaf = gf->leafs[i];
if (is_params_tensor(params_tensor_set, leaf)) {
continue;
}
auto shape = input_shape(leaf);
shape.leaf_index = i;
plan.input_shapes.push_back(shape);
plan.leaf_names.emplace_back(gf->leafs[i]->name);
}
plan.compute_buffer_size =
measure_graph_compute_buffer(backend, gf, params_tensor_set);
std::unordered_map<const ggml_tensor*, int> producer_index;
producer_index.reserve(static_cast<size_t>(n_nodes));
for (int i = 0; i < n_nodes; ++i) {
producer_index[ggml_graph_node(gf, i)] = i;
ggml_tensor* node = ggml_graph_node(gf, i);
producer_index[node] = i;
if (is_graph_cut_tensor(node)) {
plan.cut_markers.push_back({i, node->name});
}
}
std::vector<Segment> grouped_segments;
std::unordered_map<std::string, size_t> group_to_segment;
for (int i = 0; i < n_nodes; ++i) {
@@ -824,11 +877,24 @@ namespace sd::ggml_graph_cut {
if (final_output_index < 0) {
final_output_index = n_nodes - 1;
}
ggml_tensor* final_output = final_output_index >= 0 ? ggml_graph_node(gf, final_output_index) : nullptr;
if (final_output != nullptr && available_cut_output_node_indices.find(final_output_index) == available_cut_output_node_indices.end()) {
Segment final_segment;
final_segment.group_name = "ggml_runner.final";
Segment final_segment;
final_segment.group_name = "ggml_runner.final";
if (final_output_index >= 0 &&
available_cut_output_node_indices.find(final_output_index) ==
available_cut_output_node_indices.end()) {
final_segment.output_node_indices.push_back(final_output_index);
}
for (int i = 0; i < n_nodes; ++i) {
ggml_tensor* node = ggml_graph_node(gf, i);
if (i == final_output_index || node == nullptr ||
(node->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 ||
available_cut_output_node_indices.find(i) !=
available_cut_output_node_indices.end()) {
continue;
}
final_segment.output_node_indices.push_back(i);
}
if (!final_segment.output_node_indices.empty()) {
build_segment(gf,
plan,
final_segment,
@@ -839,223 +905,53 @@ namespace sd::ggml_graph_cut {
log_desc);
}
return plan;
}
Plan apply_max_vram_budget(ggml_cgraph* gf,
const Plan& base_plan,
size_t max_graph_vram_bytes,
ggml_backend_t backend,
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
const char* log_desc) {
GGML_ASSERT(backend != nullptr);
GGML_ASSERT(gf != nullptr);
int64_t t_budget_begin = ggml_time_ms();
if (max_graph_vram_bytes == 0 || !base_plan.has_cuts || base_plan.segments.size() <= 1) {
return base_plan;
}
const int n_nodes = ggml_graph_n_nodes(gf);
std::unordered_map<const ggml_tensor*, int> producer_index;
producer_index.reserve(static_cast<size_t>(n_nodes));
for (int i = 0; i < n_nodes; ++i) {
producer_index[ggml_graph_node(gf, i)] = i;
}
Plan merged_plan;
merged_plan.available = true;
merged_plan.has_cuts = base_plan.has_cuts;
merged_plan.valid = base_plan.valid;
merged_plan.n_nodes = base_plan.n_nodes;
merged_plan.n_leafs = base_plan.n_leafs;
std::unordered_set<int> available_cut_output_node_indices;
available_cut_output_node_indices.reserve(static_cast<size_t>(n_nodes));
size_t start_segment_index = 0;
while (start_segment_index < base_plan.segments.size()) {
Plan single_plan;
auto single_available_cut_output_node_indices = available_cut_output_node_indices;
auto single_seed = make_segment_seed(base_plan,
start_segment_index,
start_segment_index);
build_segment(gf,
single_plan,
single_seed,
producer_index,
single_available_cut_output_node_indices,
backend,
params_tensor_set,
log_desc);
GGML_ASSERT(!single_plan.segments.empty());
size_t best_end_segment_index = start_segment_index;
bool can_merge_next_segment = graph_cut_segment_vram_bytes(single_plan.segments.back()) <= max_graph_vram_bytes;
while (can_merge_next_segment && best_end_segment_index + 1 < base_plan.segments.size()) {
const size_t next_end_segment_index = best_end_segment_index + 1;
Plan candidate_plan;
auto candidate_available_cut_output_node_indices = available_cut_output_node_indices;
auto candidate_seed = make_segment_seed(base_plan,
start_segment_index,
next_end_segment_index);
build_segment(gf,
candidate_plan,
candidate_seed,
producer_index,
candidate_available_cut_output_node_indices,
backend,
params_tensor_set,
log_desc);
GGML_ASSERT(!candidate_plan.segments.empty());
const auto& candidate_segment = candidate_plan.segments.back();
const size_t candidate_bytes = graph_cut_segment_vram_bytes(candidate_segment);
if (candidate_bytes > max_graph_vram_bytes) {
break;
std::unordered_set<std::string> future_cut_names;
for (auto segment = plan.segments.rbegin(); segment != plan.segments.rend(); ++segment) {
segment->future_cut_names = future_cut_names;
segment->live_cut_names = future_cut_names;
for (const auto& input : segment->input_refs) {
if (input.type != Segment::INPUT_PREVIOUS_CUT) {
continue;
}
best_end_segment_index = next_end_segment_index;
segment->live_cut_names.insert(input.display_name);
future_cut_names.insert(input.display_name);
}
auto best_seed = make_segment_seed(base_plan,
start_segment_index,
best_end_segment_index);
build_segment(gf,
merged_plan,
best_seed,
producer_index,
available_cut_output_node_indices,
backend,
params_tensor_set,
log_desc);
start_segment_index = best_end_segment_index + 1;
}
if (log_desc != nullptr && merged_plan.segments.size() != base_plan.segments.size()) {
LOG_INFO("%s graph cut max_vram=%.2f MB merged %zu segments -> %zu segments",
std::string plan_validation_error;
plan.valid = validate_plan(gf, plan, &plan_validation_error);
if (!plan.valid && log_desc != nullptr) {
LOG_WARN("%s graph cut plan validation failed (%s); using monolithic execution",
log_desc,
max_graph_vram_bytes / 1024.0 / 1024.0,
base_plan.segments.size(),
merged_plan.segments.size());
plan_validation_error.c_str());
}
if (log_desc != nullptr) {
LOG_DEBUG("%s graph cut max_vram budget merge took %lld ms",
log_desc,
ggml_time_ms() - t_budget_begin);
}
return merged_plan;
return plan;
}
Plan resolve_plan(ggml_backend_t backend,
ggml_cgraph* gf,
PlanCache* cache,
size_t max_graph_vram_bytes,
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
const char* log_desc) {
GGML_ASSERT(backend != nullptr);
GGML_ASSERT(gf != nullptr);
GGML_ASSERT(cache != nullptr);
int64_t t_prepare_begin = ggml_time_ms();
Plan base_plan;
int64_t t_plan_begin = ggml_time_ms();
if (cache->graph_cut_plan.available && plan_matches_graph(gf, cache->graph_cut_plan)) {
base_plan = cache->graph_cut_plan;
} else {
base_plan = build_plan(backend, gf, params_tensor_set, log_desc);
cache->graph_cut_plan = base_plan;
cache->graph_cut_plan.available = true;
cache->budgeted_graph_cut_plan.available = false;
if (log_desc != nullptr) {
LOG_INFO("%s build cached graph cut plan done (taking %lld ms)", log_desc, ggml_time_ms() - t_plan_begin);
}
if (cache->graph_cut_plan.available &&
plan_matches_graph(gf, cache->graph_cut_plan)) {
return cache->graph_cut_plan;
}
Plan resolved_plan = base_plan;
if (max_graph_vram_bytes > 0 && base_plan.has_cuts) {
if (cache->budgeted_graph_cut_plan.available &&
cache->budgeted_graph_cut_plan_max_vram_bytes == max_graph_vram_bytes &&
plan_matches_graph(gf, cache->budgeted_graph_cut_plan)) {
resolved_plan = cache->budgeted_graph_cut_plan;
} else {
resolved_plan = apply_max_vram_budget(gf,
base_plan,
max_graph_vram_bytes,
backend,
params_tensor_set,
log_desc);
cache->budgeted_graph_cut_plan = resolved_plan;
cache->budgeted_graph_cut_plan.available = true;
cache->budgeted_graph_cut_plan_max_vram_bytes = max_graph_vram_bytes;
}
}
return resolved_plan;
}
void annotate_residency(Plan& plan,
size_t max_graph_vram_bytes,
bool prefetch_enabled) {
// Cached plans may be reused with a smaller live budget.
for (auto& seg : plan.segments) {
seg.residency = SegmentResidency::STREAMED;
}
if (max_graph_vram_bytes == 0 || plan.segments.size() < 2) {
return;
}
bool any_param_bearing = false;
for (const auto& seg : plan.segments) {
if (seg.input_param_bytes > 0) {
any_param_bearing = true;
break;
}
}
if (!any_param_bearing) {
return;
}
// Leave room for the largest active streamed segment.
size_t worst_streamed_footprint = 0;
size_t prefetch_headroom = 0;
for (const auto& seg : plan.segments) {
const size_t seg_footprint = seg.input_param_bytes +
seg.compute_buffer_size +
seg.output_bytes +
seg.input_previous_cut_bytes +
seg.input_external_bytes;
if (seg_footprint > worst_streamed_footprint) {
worst_streamed_footprint = seg_footprint;
}
prefetch_headroom = std::max(prefetch_headroom, seg.input_param_bytes);
}
constexpr size_t safety = 512ull * 1024 * 1024;
if (worst_streamed_footprint > SIZE_MAX - safety) {
return;
}
size_t reserved = safety + worst_streamed_footprint;
if (prefetch_enabled) {
if (prefetch_headroom > SIZE_MAX - reserved) {
return;
}
reserved += prefetch_headroom;
}
if (max_graph_vram_bytes <= reserved) {
return;
}
const size_t available = max_graph_vram_bytes - reserved;
size_t cumulative = 0;
for (auto& seg : plan.segments) {
if (cumulative + seg.input_param_bytes > available) {
break;
}
seg.residency = SegmentResidency::RESIDENT;
cumulative += seg.input_param_bytes;
int64_t t_plan_begin = ggml_time_ms();
Plan plan = build_plan(backend, gf, params_tensor_set, log_desc);
cache->graph_cut_plan = plan;
if (log_desc != nullptr) {
LOG_INFO("%s build cached graph cut plan done (taking %lld ms)",
log_desc,
ggml_time_ms() - t_plan_begin);
}
return plan;
}
} // namespace sd::ggml_graph_cut
src/core/ggml_graph_cut.h
+13 -41
View File
@@ -6,19 +6,13 @@
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#include "ggml-backend.h"
#include "ggml.h"
namespace sd::ggml_graph_cut {
// Streaming residency for a segment's params.
enum class SegmentResidency : uint8_t {
STREAMED = 0,
RESIDENT = 1,
};
struct Segment {
enum InputType {
INPUT_EXTERNAL = 0,
@@ -33,38 +27,28 @@ namespace sd::ggml_graph_cut {
int node_index = -1;
};
size_t compute_buffer_size = 0;
size_t output_bytes = 0;
size_t input_external_bytes = 0;
size_t input_previous_cut_bytes = 0;
size_t input_param_bytes = 0;
size_t compute_buffer_size = 0;
std::string group_name;
std::vector<int> internal_node_indices;
std::vector<int> output_node_indices;
std::vector<InputRef> input_refs;
SegmentResidency residency = SegmentResidency::STREAMED;
std::unordered_set<std::string> future_cut_names;
std::unordered_set<std::string> live_cut_names;
};
struct Plan {
struct InputShape {
int leaf_index = -1;
ggml_type type = GGML_TYPE_COUNT;
std::array<int64_t, GGML_MAX_DIMS> ne = {0, 0, 0, 0};
};
bool available = false;
bool has_cuts = false;
bool valid = true;
int n_nodes = 0;
int n_leafs = 0;
std::vector<InputShape> input_shapes;
bool available = false;
bool has_cuts = false;
bool valid = true;
size_t compute_buffer_size = 0;
std::vector<uint64_t> layout;
std::vector<std::string> leaf_names;
std::vector<std::pair<int, std::string>> cut_markers;
std::vector<Segment> segments;
};
struct PlanCache {
Plan graph_cut_plan;
Plan budgeted_graph_cut_plan;
size_t budgeted_graph_cut_plan_max_vram_bytes = 0;
};
static constexpr const char* GGML_RUNNER_CUT_PREFIX = "ggml_runner_cut:";
@@ -89,13 +73,12 @@ namespace sd::ggml_graph_cut {
ggml_backend_buffer_t tensor_buffer(const ggml_tensor* tensor);
ggml_tensor* cache_source_tensor(ggml_tensor* tensor);
size_t cache_tensor_bytes(const ggml_tensor* tensor);
// Plans ignore runtime bindings; allocator reservations must include them.
std::vector<uint64_t> graph_layout(ggml_cgraph* graph, bool include_bindings);
bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan);
ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index);
ggml_tensor* input_tensor(ggml_cgraph* gf, const Segment::InputRef& input_ref);
std::vector<ggml_tensor*> param_tensors(ggml_cgraph* gf, const Segment& segment);
std::unordered_set<std::string> collect_future_input_names(ggml_cgraph* gf,
const Plan& plan,
size_t current_segment_index);
ggml_cgraph* build_segment_graph(ggml_cgraph* gf,
const Segment& segment,
ggml_context** graph_ctx_out);
@@ -109,23 +92,12 @@ namespace sd::ggml_graph_cut {
ggml_cgraph* gf,
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
const char* log_desc);
Plan apply_max_vram_budget(ggml_cgraph* gf,
const Plan& base_plan,
size_t max_graph_vram_bytes,
ggml_backend_t backend,
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
const char* log_desc);
Plan resolve_plan(ggml_backend_t backend,
ggml_cgraph* gf,
PlanCache* cache,
size_t max_graph_vram_bytes,
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
const char* log_desc);
// Mark leading segments resident after reserving streamed execution headroom.
void annotate_residency(Plan& plan,
size_t max_graph_vram_bytes,
bool prefetch_enabled);
} // namespace sd::ggml_graph_cut
#endif // __SD_CORE_GGML_GRAPH_CUT_H__
src/core/ggml_runner.cpp
+296
View File
@@ -0,0 +1,296 @@
#include <algorithm>
#include <map>
#include <utility>
#include "core/ggml_extend.hpp"
#include "core/segment_graph_bindings.h"
#include "core/segment_weight_pipeline.h"
using namespace sd;
static size_t add_bytes(size_t a, size_t b) {
return b > SIZE_MAX - a ? SIZE_MAX : a + b;
}
ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, size_t direct_bytes) {
auto external_backend = [&](const ggml_tensor* tensor) -> ggml_backend_t {
if (!params_tensor_set_.count(tensor)) {
return nullptr;
}
auto placement = graph_cut_layer_split_assignments_.find(tensor);
return placement == graph_cut_layer_split_assignments_.end() ? runtime_backend : placement->second;
};
auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) {
pin_multi_device_nodes(scheduler, copy, graph);
};
return workspace_.measure(graph, direct_bytes, external_backend, assign_nodes);
}
std::vector<DeviceMemoryRequest> GGMLRunner::memory_requests(
const std::vector<BackendBufferSize>& sizes,
size_t pending_cache_bytes) const {
std::vector<DeviceMemoryRequest> requests;
for (const auto& size : sizes) {
const size_t retained = retained_runtime_buffer_bytes(size.backend);
const size_t reusable = workspace_.bytes(size.backend);
const size_t cache_bytes = size.backend == runtime_backend ? pending_cache_bytes : 0;
const size_t pending = add_bytes(size.bytes > reusable ? size.bytes - reusable : 0, cache_bytes);
size_t limit = max_graph_vram_bytes;
if (is_multi_device()) {
size_t index = 0;
if (size.backend != runtime_backend) {
auto position = std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), size.backend);
index = static_cast<size_t>(position - extra_runtime_backends.begin()) + 1;
}
if (index < graph_cut_layer_split_backend_vram_limits_.size()) {
limit = graph_cut_layer_split_backend_vram_limits_[index];
}
}
requests.push_back({size.backend, reinterpret_cast<uintptr_t>(this), pending,
retained, limit});
}
return requests;
}
bool GGMLRunner::fits(const std::vector<DeviceMemoryRequest>& requests,
const std::vector<ggml_tensor*>& params) const {
auto manager = residency_manager.lock();
if (manager == nullptr) {
return params.empty();
}
for (const auto& request : requests) {
if (!manager->fits_compute_backend_capacity(request, params)) {
return false;
}
}
return true;
}
bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) {
if (sd_backend_is_cpu(runtime_backend)) {
sd_backend_cpu_set_n_threads(runtime_backend, n_threads);
}
if (workspace_.cpu_backend() != nullptr) {
sd_backend_cpu_set_n_threads(workspace_.cpu_backend(), n_threads);
}
auto scheduler = workspace_.scheduler();
ggml_status status;
if (scheduler != nullptr) {
if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) {
LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring", get_desc().c_str());
multi_device_eval_callback_warned = true;
}
status = ggml_backend_sched_graph_compute(scheduler, graph);
} else {
status = sd_backend_graph_compute_with_eval_callback(runtime_backend, graph,
sd_get_backend_eval_callback(),
sd_get_backend_eval_callback_data());
}
workspace_.synchronize();
if (status != GGML_STATUS_SUCCESS) {
LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status));
return false;
}
const std::string description = get_desc();
if (!debug_tensors.empty()) {
std::unordered_set<const ggml_tensor*> graph_tensors;
const int leaf_count = ggml_graph_cut::leaf_count(graph);
const int node_count = ggml_graph_n_nodes(graph);
graph_tensors.reserve(static_cast<size_t>(leaf_count + node_count));
for (int index = 0; index < leaf_count; ++index) {
graph_tensors.insert(ggml_graph_cut::leaf_tensor(graph, index));
}
for (int index = 0; index < node_count; ++index) {
graph_tensors.insert(ggml_graph_node(graph, index));
}
for (const auto& entry : debug_tensors) {
ggml_tensor* tensor = entry.first;
if (tensor == nullptr || graph_tensors.find(tensor) == graph_tensors.end()) {
continue;
}
ggml_backend_buffer_t buffer =
tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer;
if (buffer == nullptr) {
LOG_WARN("%s skip debug tensor '%s': tensor buffer not set",
description.c_str(),
entry.second.c_str());
continue;
}
if (tensor->type != GGML_TYPE_F32) {
LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s",
description.c_str(),
entry.second.c_str(),
ggml_type_name(tensor->type));
continue;
}
auto debug_tensor = make_sd_tensor_from_ggml<float>(tensor);
print_sd_tensor(debug_tensor, false, entry.second.c_str());
}
}
return true;
}
std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs) {
if (!assign_graph_cut_layer_split_backends(graph)) {
return std::nullopt;
}
const auto params = collect_used_param_tensors(graph);
ggml_graph_cut::Plan plan;
if (!resolve_graph_cut_plan(graph, &plan)) {
return std::nullopt;
}
const auto full_measurement = measure(graph, plan.compute_buffer_size);
if (full_measurement.buffers.empty()) {
return std::nullopt;
}
auto manager = residency_manager.lock();
const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) &&
manager != nullptr && manager->segmented_compute_enabled() &&
plan.valid && plan.has_cuts && plan.segments.size() > 1 &&
!fits(memory_requests(full_measurement.buffers, cache_.pending_bytes(graph)), params);
if (!segmented) {
ggml_graph_cut::Segment segment;
segment.group_name = "graph";
segment.compute_buffer_size = plan.compute_buffer_size;
for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
segment.internal_node_indices.push_back(i);
}
for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) {
auto tensor = ggml_graph_cut::leaf_tensor(graph, i);
ggml_graph_cut::Segment::InputRef input;
input.leaf_index = i;
input.type = canonical_param_tensor(tensor) != nullptr
? ggml_graph_cut::Segment::INPUT_PARAM
: ggml_graph_cut::Segment::INPUT_EXTERNAL;
segment.input_refs.push_back(input);
}
plan.segments = {std::move(segment)};
}
const bool segments_changed = plan.segments.size() != logged_segment_count_;
if (segments_changed && (segmented || logged_segment_count_ > 1)) {
LOG_DEBUG("%s using %zu segment%s", get_desc().c_str(),
plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
}
SegmentGraphBindings bindings(cut_cache_, plan, graph);
SegmentWeightPipeline weights(manager, runtime_backend, reinterpret_cast<uintptr_t>(this),
graph, plan, params_tensor_set_,
segmented && manager != nullptr && manager->prefetch_enabled());
std::map<ggml_backend_t, size_t> peak_compute_bytes;
auto track_compute_buffer = [&](ggml_backend_t backend) {
if (backend != nullptr) {
auto& peak = peak_compute_bytes[backend];
peak = std::max(peak, workspace_.bytes(backend));
}
};
std::optional<Tensor<float>> output = Tensor<float>();
for (size_t index = 0; index < plan.segments.size(); ++index) {
const auto& segment = plan.segments[index];
const bool last = index + 1 == plan.segments.size();
auto fail_segment = [&](const char* phase) {
LOG_ERROR("%s segment %zu/%zu (%s) failed during %s", get_desc().c_str(),
index + 1, plan.segments.size(), segment.group_name.c_str(), phase);
return std::nullopt;
};
cut_cache_.prune(segment.live_cut_names);
bindings.reset(segment);
if (!bindings.bind_cached_inputs(segment, get_desc().c_str())) {
return fail_segment("input binding");
}
ggml_context* segment_context = nullptr;
auto segment_graph = segmented
? ggml_graph_cut::build_segment_graph(graph, segment, &segment_context)
: graph;
struct SegmentCleanup {
GGMLRunner& runner;
SegmentWeightPipeline& weights;
SegmentGraphBindings& bindings;
ggml_context* context;
~SegmentCleanup() {
runner.workspace_.segment_end();
bindings.restore();
weights.segment_end();
ggml_free(context);
runner.sync_runtime_residency();
}
} segment_cleanup{*this, weights, bindings, segment_context};
auto measurement = segmented ? measure(segment_graph, segment.compute_buffer_size) : full_measurement;
if (!workspace_.prepare(measurement)) {
return fail_segment("workspace preparation");
}
const size_t cut_bytes = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment);
const size_t new_cache_bytes = add_bytes(cut_bytes, cache_.pending_bytes(segment_graph));
auto ensure_capacity = [&]() {
sync_runtime_residency();
auto requests = memory_requests(measurement.buffers, new_cache_bytes);
if (!fits(requests, weights.params(index)) && workspace_.release_excess(measurement)) {
sync_runtime_residency();
requests = memory_requests(measurement.buffers, new_cache_bytes);
}
return weights.ensure_segment_capacity(index, requests);
};
if (!weights.segment_start(index, ensure_capacity)) {
return fail_segment("weight preparation");
}
// Preparing weights can execute LoRA graphs and reclaim an idle workspace.
if (!workspace_.measurement_matches(segment_graph, measurement)) {
measurement = measure(segment_graph, segment.compute_buffer_size);
}
if (!workspace_.prepare(measurement) || !ensure_capacity()) {
return fail_segment("workspace capacity check");
}
if (!workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) {
pin_multi_device_nodes(scheduler, current);
})) {
return fail_segment("workspace allocation");
}
for (const auto& size : measurement.buffers) {
track_compute_buffer(size.backend);
}
if (workspace_.scheduler() != nullptr) {
track_compute_buffer(workspace_.cpu_backend());
}
if (!ensure_capacity()) {
return fail_segment("allocated capacity check");
}
copy_data_to_backend_tensor(segment_graph, false);
auto prefetch_requests = memory_requests(measurement.buffers, new_cache_bytes);
if (!prefetch_requests.empty()) {
weights.enqueue_next(index, prefetch_requests.front());
}
if (!execute_segment(segment_graph, n_threads) ||
!cache_.capture(segment_graph) ||
!cut_cache_.capture(graph, segment, get_desc().c_str())) {
return fail_segment("execution or output caching");
}
sync_runtime_residency();
if (last) {
if (read_outputs && !read_outputs()) {
return fail_segment("output finalization");
}
if (!no_return) {
auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str());
output = read_graph_tensor<float>(result, "output");
if (!output.has_value()) {
return fail_segment("output readback");
}
}
}
// Final outputs and their callbacks may still be views of consumed cuts.
cut_cache_.prune(segment.future_cut_names);
}
if (segments_changed || peak_compute_bytes != logged_compute_bytes_) {
for (const auto& entry : peak_compute_bytes) {
LOG_DEBUG("%s compute buffer size: %.2f MB(%s) on %s (peak across %zu segment%s)",
get_desc().c_str(), entry.second / (1024.0 * 1024.0),
sd_backend_is_cpu(entry.first) ? "RAM" : "VRAM", ggml_backend_name(entry.first),
plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
}
logged_compute_bytes_ = std::move(peak_compute_bytes);
logged_segment_count_ = plan.segments.size();
}
return output;
}
src/core/layer_split_partition.cpp
+12 -4
View File
@@ -145,19 +145,24 @@ namespace sd {
std::vector<int64_t> backend_capacities = graph_cut_layer_split_backend_capacities(split_backends,
backend_vram_limits,
primary_backend_vram_limit);
// Existing placements may already occupy the reported free VRAM. Reuse
// them; execution checks missing weights and reclaims memory as needed.
const bool reuse_assignments = std::all_of(seen_params.begin(), seen_params.end(), [&](ggml_tensor* param) {
return param_assignments.count(param) != 0;
});
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() &&
while (!reuse_assignments && 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]) {
if (!reuse_assignments && 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,
@@ -167,7 +172,6 @@ namespace sd {
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];
@@ -186,12 +190,16 @@ namespace sd {
ggml_get_name(param));
return false;
}
size_t backend_idx = (size_t)std::distance(split_backends.begin(), backend_it);
size_t backend_idx = (size_t)std::distance(split_backends.begin(), backend_it);
if (reuse_assignments) {
current_backend = backend_idx;
}
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);
}
backend_by_segment[seg_idx] = split_backends[current_backend];
}
const int n_nodes = ggml_graph_n_nodes(gf);
src/core/layer_stream_prefetch.cpp
-124
View File
@@ -1,124 +0,0 @@
#include "core/layer_stream_prefetch.h"
#include <utility>
#include "core/ggml_graph_cut.h"
#include "weight_manager.h"
namespace sd {
static ggml_tensor* canonical_param(
ggml_tensor* tensor,
const std::unordered_set<const ggml_tensor*>& params) {
for (ggml_tensor* current = tensor; current != nullptr; current = current->view_src) {
if (params.find(current) != params.end()) {
return current;
}
}
return nullptr;
}
LayerStreamPrefetch::LayerStreamPrefetch(
const std::shared_ptr<RunnerWeightManager>& manager,
uintptr_t owner_id,
ggml_cgraph* graph,
const ggml_graph_cut::Plan& plan,
const std::unordered_set<const ggml_tensor*>& params,
bool enabled)
: manager_(manager),
owner_id_(owner_id),
enabled_(enabled && manager != nullptr) {
segment_params_.resize(plan.segments.size());
for (size_t segment_index = 0; segment_index < plan.segments.size(); ++segment_index) {
std::unordered_set<ggml_tensor*> seen;
for (ggml_tensor* tensor :
ggml_graph_cut::param_tensors(graph, plan.segments[segment_index])) {
ggml_tensor* param = canonical_param(tensor, params);
if (param != nullptr && seen.insert(param).second) {
segment_params_[segment_index].push_back(param);
}
}
}
}
LayerStreamPrefetch::~LayerStreamPrefetch() {
clear();
}
size_t LayerStreamPrefetch::next_parameter_segment(size_t segment_index) const {
for (size_t next = segment_index + 1; next < segment_params_.size(); ++next) {
if (!segment_params_[next].empty()) {
return next;
}
}
return SIZE_MAX;
}
void LayerStreamPrefetch::disable() {
clear();
enabled_ = false;
}
bool LayerStreamPrefetch::activate(size_t segment_index) {
if (!enabled_ || queued_segment_ == SIZE_MAX) {
return true;
}
if (queued_segment_ != segment_index) {
return true;
}
auto manager = manager_.lock();
if (manager == nullptr ||
!manager->activate_prefetched_params(owner_id_, queued_params_)) {
disable();
return false;
}
queued_params_.clear();
queued_segment_ = SIZE_MAX;
return true;
}
bool LayerStreamPrefetch::enqueue_next(size_t segment_index) {
if (!enabled_) {
return true;
}
if (queued_segment_ != SIZE_MAX) {
return true;
}
const size_t next_segment = next_parameter_segment(segment_index);
if (next_segment == SIZE_MAX) {
return true;
}
std::unordered_set<ggml_tensor*> active_params(
segment_params_[segment_index].begin(),
segment_params_[segment_index].end());
std::vector<ggml_tensor*> params;
params.reserve(segment_params_[next_segment].size());
for (ggml_tensor* param : segment_params_[next_segment]) {
if (active_params.find(param) == active_params.end()) {
params.push_back(param);
}
}
if (params.empty()) {
return true;
}
auto manager = manager_.lock();
if (manager == nullptr || !manager->prefetch_params(owner_id_, params)) {
disable();
return false;
}
queued_params_ = std::move(params);
queued_segment_ = next_segment;
return true;
}
void LayerStreamPrefetch::clear() {
if (auto manager = manager_.lock()) {
manager->clear_prefetched_params(owner_id_);
}
queued_params_.clear();
queued_segment_ = SIZE_MAX;
}
}
src/core/layer_stream_prefetch.h
-48
View File
@@ -1,48 +0,0 @@
#ifndef __SD_CORE_LAYER_STREAM_PREFETCH_H__
#define __SD_CORE_LAYER_STREAM_PREFETCH_H__
#include <cstddef>
#include <cstdint>
#include <memory>
#include <unordered_set>
#include <vector>
struct ggml_cgraph;
struct ggml_tensor;
struct RunnerWeightManager;
namespace sd::ggml_graph_cut {
struct Plan;
}
namespace sd {
class LayerStreamPrefetch {
private:
std::weak_ptr<RunnerWeightManager> manager_;
uintptr_t owner_id_ = 0;
std::vector<std::vector<ggml_tensor*>> segment_params_;
std::vector<ggml_tensor*> queued_params_;
size_t queued_segment_ = SIZE_MAX;
bool enabled_ = true;
size_t next_parameter_segment(size_t segment_index) const;
void disable();
public:
LayerStreamPrefetch(
const std::shared_ptr<RunnerWeightManager>& manager,
uintptr_t owner_id,
ggml_cgraph* graph,
const ggml_graph_cut::Plan& plan,
const std::unordered_set<const ggml_tensor*>& params,
bool enabled = true);
~LayerStreamPrefetch();
bool enabled() const { return enabled_; }
bool activate(size_t segment_index);
bool enqueue_next(size_t segment_index);
void clear();
};
}
#endif // __SD_CORE_LAYER_STREAM_PREFETCH_H__
src/core/runner_cache.cpp
+211
View File
@@ -0,0 +1,211 @@
#include "core/runner_cache.h"
#include <algorithm>
#include <iterator>
#include <unordered_set>
#include "core/ggml_graph_cut.h"
#include "core/util.h"
namespace sd {
static std::unordered_set<const ggml_tensor*> cache_graph_tensors(ggml_cgraph* graph) {
std::unordered_set<const ggml_tensor*> tensors;
for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
tensors.insert(ggml_graph_node(graph, i));
}
for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) {
tensors.insert(ggml_graph_cut::leaf_tensor(graph, i));
}
return tensors;
}
CachedTensor::~CachedTensor() {
ggml_backend_buffer_free(buffer);
ggml_free(context);
}
std::unique_ptr<CachedTensor> CachedTensor::copy(ggml_backend_t backend,
const std::string& name,
ggml_tensor* source) {
if (ggml_graph_cut::tensor_buffer(source) == nullptr) {
return nullptr;
}
auto entry = std::make_unique<CachedTensor>();
entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true});
if (entry->context == nullptr) {
return nullptr;
}
entry->tensor = ggml_dup_tensor(entry->context, source);
// Cut views are rebound with their original strides and offsets.
std::copy(std::begin(source->nb), std::end(source->nb), std::begin(entry->tensor->nb));
ggml_set_name(entry->tensor, name.c_str());
entry->buffer = ggml_backend_alloc_ctx_tensors(entry->context, backend);
if (entry->buffer == nullptr) {
return nullptr;
}
if (source->view_src != nullptr || !ggml_is_contiguous(source) || source->buffer == nullptr) {
std::vector<uint8_t> data(ggml_nbytes(source));
ggml_backend_tensor_get(source, data.data(), 0, data.size());
ggml_backend_tensor_set(entry->tensor, data.data(), 0, data.size());
} else {
ggml_backend_tensor_copy(source, entry->tensor);
}
return entry;
}
static ggml_tensor* cached_tensor(const CachedTensors& tensors, const std::string& name) {
auto entry = tensors.find(name);
return entry == tensors.end() ? nullptr : entry->second->tensor;
}
static size_t resident_bytes(const CachedTensors& tensors, ggml_backend_dev_t device) {
size_t bytes = 0;
for (const auto& entry : tensors) {
auto buffer = entry.second->buffer;
if (!ggml_backend_buffer_is_host(buffer) &&
ggml_backend_buft_get_device(ggml_backend_buffer_get_type(buffer)) == device) {
const size_t size = ggml_backend_buffer_get_size(buffer);
bytes = size > SIZE_MAX - bytes ? SIZE_MAX : bytes + size;
}
}
return bytes;
}
ggml_tensor* RunnerCache::get(const std::string& name) const {
return cached_tensor(committed_, name);
}
void RunnerCache::stage(const std::string& name, ggml_tensor* tensor) {
if (tensor != nullptr) {
ggml_set_output(tensor);
outputs_[name] = tensor;
}
}
size_t RunnerCache::pending_bytes(ggml_cgraph* graph) const {
if (outputs_.empty()) {
return 0;
}
auto tensors = cache_graph_tensors(graph);
auto buft = ggml_backend_get_default_buffer_type(backend_);
size_t bytes = 0;
for (const auto& output : outputs_) {
if (pending_.count(output.first) || !tensors.count(output.second)) {
continue;
}
const size_t size = GGML_PAD(ggml_backend_buft_get_alloc_size(buft, output.second),
ggml_backend_buft_get_alignment(buft));
bytes = size > SIZE_MAX - bytes ? SIZE_MAX : bytes + size;
}
return bytes;
}
size_t RunnerCache::resident_bytes(ggml_backend_dev_t device) const {
const size_t committed = sd::resident_bytes(committed_, device);
const size_t pending = sd::resident_bytes(pending_, device);
return pending > SIZE_MAX - committed ? SIZE_MAX : committed + pending;
}
bool RunnerCache::capture(ggml_cgraph* graph) {
if (outputs_.empty()) {
return true;
}
const auto tensors = cache_graph_tensors(graph);
for (const auto& output : outputs_) {
if (pending_.count(output.first) || !tensors.count(output.second)) {
continue;
}
GGML_ASSERT(ggml_is_contiguous(output.second));
auto entry = CachedTensor::copy(backend_, output.first, output.second);
if (entry == nullptr) {
return false;
}
pending_[output.first] = std::move(entry);
}
ggml_backend_synchronize(backend_);
return true;
}
void RunnerCache::graph_end(bool success) {
// Graph inputs can still reference the previous generation until graph end.
if (success) {
for (auto& entry : pending_) {
committed_[entry.first] = std::move(entry.second);
}
}
pending_.clear();
outputs_.clear();
}
void RunnerCache::clear() {
graph_end(false);
committed_.clear();
}
ggml_tensor* GraphCutTensorCache::get(const std::string& name) const {
return cached_tensor(tensors_, name);
}
size_t GraphCutTensorCache::resident_bytes(ggml_backend_dev_t device) const {
return sd::resident_bytes(tensors_, device);
}
size_t GraphCutTensorCache::estimate_output_bytes(
ggml_cgraph* graph,
const ggml_graph_cut::Segment& segment) const {
ggml_backend_buffer_type_t buffer_type =
ggml_backend_get_default_buffer_type(backend_);
if (buffer_type == nullptr) {
return SIZE_MAX;
}
const size_t alignment = ggml_backend_buft_get_alignment(buffer_type);
size_t total_size = 0;
for (size_t output_idx = 0; output_idx < segment.output_node_indices.size(); ++output_idx) {
ggml_tensor* output = ggml_graph_cut::output_tensor(graph, segment, output_idx);
if (output == nullptr || !ggml_graph_cut::is_graph_cut_tensor(output) ||
!segment.future_cut_names.count(output->name)) {
continue;
}
ggml_tensor* source = ggml_graph_cut::cache_source_tensor(output);
const size_t tensor_size = GGML_PAD(
ggml_backend_buft_get_alloc_size(buffer_type, source), alignment);
total_size = tensor_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + tensor_size;
}
return total_size;
}
void GraphCutTensorCache::prune(const std::unordered_set<std::string>& keep_names) {
for (auto it = tensors_.begin(); it != tensors_.end();) {
it = keep_names.count(it->first) ? std::next(it) : tensors_.erase(it);
}
}
bool GraphCutTensorCache::capture(ggml_cgraph* graph,
const ggml_graph_cut::Segment& segment,
const char* log_desc) {
size_t copied_bytes = 0;
size_t copied_count = 0;
for (int index : segment.output_node_indices) {
auto output = ggml_graph_node(graph, index);
if (!ggml_graph_cut::is_graph_cut_tensor(output) ||
!segment.future_cut_names.count(output->name)) {
continue;
}
auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output));
if (entry == nullptr) {
LOG_ERROR("%s failed to capture graph cut tensor: %s", log_desc, output->name);
return false;
}
const size_t size = ggml_backend_buffer_get_size(entry->buffer);
copied_bytes = size > SIZE_MAX - copied_bytes ? SIZE_MAX : copied_bytes + size;
++copied_count;
tensors_[output->name] = std::move(entry);
}
ggml_backend_synchronize(backend_);
if (copied_count > 0) {
LOG_DEBUG("%s graph cut cache added %6.2f MB (%zu tensors)",
log_desc, copied_bytes / (1024.f * 1024.f), copied_count);
}
return true;
}
}
src/core/runner_cache.h
+66
View File
@@ -0,0 +1,66 @@
#ifndef __SD_CORE_RUNNER_CACHE_H__
#define __SD_CORE_RUNNER_CACHE_H__
#include <map>
#include <memory>
#include <string>
#include <unordered_set>
#include "ggml-backend.h"
namespace sd::ggml_graph_cut {
struct Segment;
}
namespace sd {
struct CachedTensor {
ggml_context* context = nullptr;
ggml_backend_buffer_t buffer = nullptr;
ggml_tensor* tensor = nullptr;
~CachedTensor();
static std::unique_ptr<CachedTensor> copy(ggml_backend_t backend,
const std::string& name,
ggml_tensor* source);
};
using CachedTensors = std::map<std::string, std::unique_ptr<CachedTensor>>;
class RunnerCache {
ggml_backend_t backend_;
CachedTensors committed_;
CachedTensors pending_;
std::map<std::string, ggml_tensor*> outputs_;
public:
explicit RunnerCache(ggml_backend_t backend)
: backend_(backend) {}
RunnerCache(const RunnerCache&) = delete;
RunnerCache& operator=(const RunnerCache&) = delete;
ggml_tensor* get(const std::string& name) const;
void stage(const std::string& name, ggml_tensor* tensor);
const std::map<std::string, ggml_tensor*>& outputs() const { return outputs_; }
size_t pending_bytes(ggml_cgraph* graph) const;
size_t resident_bytes(ggml_backend_dev_t device) const;
bool capture(ggml_cgraph* graph);
void graph_end(bool success);
void clear();
};
class GraphCutTensorCache {
ggml_backend_t backend_;
CachedTensors tensors_;
public:
explicit GraphCutTensorCache(ggml_backend_t backend)
: backend_(backend) {}
ggml_tensor* get(const std::string& name) const;
size_t resident_bytes(ggml_backend_dev_t device) const;
size_t estimate_output_bytes(ggml_cgraph* graph,
const ggml_graph_cut::Segment& segment) const;
bool capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc);
void prune(const std::unordered_set<std::string>& keep_names);
void clear() { tensors_.clear(); }
};
}
#endif // __SD_CORE_RUNNER_CACHE_H__
src/core/segment_graph_bindings.cpp
+137
View File
@@ -0,0 +1,137 @@
#include "core/segment_graph_bindings.h"
#include <algorithm>
#include <iterator>
#include "core/ggml_graph_cut.h"
#include "core/runner_cache.h"
#include "core/util.h"
#include "ggml.h"
namespace sd {
SegmentGraphBindings::SegmentGraphBindings(GraphCutTensorCache& tensor_cache,
const ggml_graph_cut::Plan& plan,
ggml_cgraph* graph)
: tensor_cache_(tensor_cache),
graph_(graph) {
GGML_ASSERT(graph_ != nullptr);
for (int i = 0; i < ggml_graph_n_nodes(graph_); ++i) {
ggml_tensor* tensor = ggml_graph_node(graph_, i);
Topology topology{tensor->op, {}, tensor->view_src, tensor->flags};
std::copy(std::begin(tensor->src), std::end(tensor->src), topology.sources.begin());
topology_[tensor] = topology;
}
for (const auto& segment : plan.segments) {
for (const auto& input : segment.input_refs) {
if (input.type != ggml_graph_cut::Segment::INPUT_EXTERNAL) {
continue;
}
ggml_tensor* tensor = ggml_graph_cut::input_tensor(graph_, input);
if (tensor == nullptr || tensor->buffer == nullptr) {
continue;
}
external_bindings_[tensor] = {tensor->buffer, tensor->data, tensor->extra};
}
}
}
void SegmentGraphBindings::reset(const ggml_graph_cut::Segment& segment) {
restore();
for (const auto& input : segment.input_refs) {
ggml_tensor* tensor = ggml_graph_cut::input_tensor(graph_, input);
if (tensor == nullptr) {
continue;
}
switch (input.type) {
case ggml_graph_cut::Segment::INPUT_PREVIOUS_CUT:
tensor->buffer = nullptr;
tensor->data = nullptr;
tensor->extra = nullptr;
break;
case ggml_graph_cut::Segment::INPUT_EXTERNAL: {
auto binding = external_bindings_.find(tensor);
if (binding != external_bindings_.end()) {
tensor->buffer = binding->second.buffer;
tensor->data = binding->second.data;
tensor->extra = binding->second.extra;
} else {
tensor->buffer = nullptr;
tensor->data = nullptr;
tensor->extra = nullptr;
}
break;
}
case ggml_graph_cut::Segment::INPUT_PARAM:
break;
}
}
for (int node_index : segment.internal_node_indices) {
ggml_tensor* node = ggml_graph_node(graph_, node_index);
if (node == nullptr) {
continue;
}
node->buffer = nullptr;
node->data = nullptr;
node->extra = nullptr;
}
}
void SegmentGraphBindings::restore() {
for (const auto& entry : topology_) {
entry.first->op = entry.second.op;
entry.first->view_src = entry.second.view_source;
entry.first->flags = entry.second.flags;
std::copy(entry.second.sources.begin(), entry.second.sources.end(), std::begin(entry.first->src));
}
}
bool SegmentGraphBindings::bind_cached_inputs(
const ggml_graph_cut::Segment& segment,
const char* log_desc) {
std::unordered_map<ggml_tensor*, ggml_tensor*> cached_view_sources;
for (const auto& input : segment.input_refs) {
if (input.type != ggml_graph_cut::Segment::INPUT_PREVIOUS_CUT) {
continue;
}
ggml_tensor* input_tensor = ggml_graph_cut::input_tensor(graph_, input);
if (input_tensor == nullptr) {
continue;
}
ggml_tensor* cached_tensor = tensor_cache_.get(input.display_name);
if (cached_tensor == nullptr) {
LOG_ERROR("%s missing graph cut cache tensor: %s",
log_desc,
input.display_name.c_str());
return false;
}
if (input_tensor->view_src != nullptr) {
cached_view_sources[topology_.at(input_tensor).view_source] = cached_tensor;
input_tensor->view_src = cached_tensor;
input_tensor->buffer = nullptr;
input_tensor->data = cached_tensor->data == nullptr
? nullptr
: static_cast<void*>(static_cast<char*>(cached_tensor->data) +
input_tensor->view_offs);
input_tensor->extra = cached_tensor->extra;
} else {
input_tensor->buffer = cached_tensor->buffer;
input_tensor->data = cached_tensor->data;
input_tensor->extra = cached_tensor->extra;
}
for (int source_index = 0; source_index < GGML_MAX_SRC; ++source_index) {
input_tensor->src[source_index] = nullptr;
}
input_tensor->op = GGML_OP_NONE;
}
// ggml flattens view chains, so descendants also need the cached root.
for (int node_index : segment.internal_node_indices) {
ggml_tensor* node = ggml_graph_node(graph_, node_index);
auto cached_source = cached_view_sources.find(topology_.at(node).view_source);
if (cached_source != cached_view_sources.end()) {
node->view_src = cached_source->second;
}
}
return true;
}
}
src/core/segment_graph_bindings.h
+52
View File
@@ -0,0 +1,52 @@
#ifndef __SD_CORE_SEGMENT_GRAPH_BINDINGS_H__
#define __SD_CORE_SEGMENT_GRAPH_BINDINGS_H__
#include <array>
#include <unordered_map>
#include "ggml.h"
struct ggml_backend_buffer;
struct ggml_cgraph;
struct ggml_tensor;
namespace sd {
class GraphCutTensorCache;
namespace ggml_graph_cut {
struct Plan;
struct Segment;
}
class SegmentGraphBindings {
public:
SegmentGraphBindings(GraphCutTensorCache& tensor_cache,
const ggml_graph_cut::Plan& plan,
ggml_cgraph* graph);
void reset(const ggml_graph_cut::Segment& segment);
void restore();
~SegmentGraphBindings() { restore(); }
bool bind_cached_inputs(const ggml_graph_cut::Segment& segment,
const char* log_desc);
private:
struct ExternalBinding {
ggml_backend_buffer* buffer = nullptr;
void* data = nullptr;
void* extra = nullptr;
};
GraphCutTensorCache& tensor_cache_;
ggml_cgraph* graph_ = nullptr;
std::unordered_map<ggml_tensor*, ExternalBinding> external_bindings_;
struct Topology {
ggml_op op;
std::array<ggml_tensor*, GGML_MAX_SRC> sources;
ggml_tensor* view_source;
int flags;
};
std::unordered_map<ggml_tensor*, Topology> topology_;
};
}
#endif // __SD_CORE_SEGMENT_GRAPH_BINDINGS_H__
src/core/segment_weight_pipeline.cpp
+205
View File
@@ -0,0 +1,205 @@
#include "core/segment_weight_pipeline.h"
#include <utility>
#include "core/ggml_graph_cut.h"
#include "device_residency_manager.h"
namespace sd {
static ggml_tensor* canonical_param(
ggml_tensor* tensor,
const std::unordered_set<const ggml_tensor*>& params) {
for (ggml_tensor* current = tensor; current != nullptr; current = current->view_src) {
if (params.find(current) != params.end()) {
return current;
}
}
return nullptr;
}
SegmentWeightPipeline::SegmentWeightPipeline(
const std::shared_ptr<DeviceResidencyManager>& residency_manager,
ggml_backend_t compute_backend,
uintptr_t owner_id,
ggml_cgraph* graph,
const ggml_graph_cut::Plan& plan,
const std::unordered_set<const ggml_tensor*>& params,
bool enabled)
: residency_manager_(residency_manager),
compute_backend_(compute_backend),
owner_id_(owner_id),
enabled_(enabled && residency_manager != nullptr) {
segment_params_.resize(plan.segments.size());
for (size_t segment_index = 0; segment_index < plan.segments.size(); ++segment_index) {
std::unordered_set<ggml_tensor*> seen;
for (ggml_tensor* tensor :
ggml_graph_cut::param_tensors(graph, plan.segments[segment_index])) {
ggml_tensor* param = canonical_param(tensor, params);
if (param != nullptr && seen.insert(param).second) {
segment_params_[segment_index].push_back(param);
}
}
}
}
SegmentWeightPipeline::~SegmentWeightPipeline() {
segment_end();
clear();
}
size_t SegmentWeightPipeline::next_parameter_segment(size_t segment_index) const {
for (size_t next = segment_index + 1; next < segment_params_.size(); ++next) {
if (!segment_params_[next].empty()) {
return next;
}
}
return SIZE_MAX;
}
std::vector<std::vector<ggml_tensor*>> SegmentWeightPipeline::preferred_eviction_order() const {
return {segment_params_.rbegin(), segment_params_.rend()};
}
void SegmentWeightPipeline::disable() {
clear();
enabled_ = false;
}
void SegmentWeightPipeline::activate(size_t segment_index) {
if (!enabled_ || queued_segment_ == SIZE_MAX || queued_segment_ != segment_index) {
return;
}
auto manager = residency_manager_.lock();
if (manager == nullptr ||
!manager->activate_prefetched_params(owner_id_, queued_params_)) {
disable();
return;
}
queued_params_.clear();
queued_segment_ = SIZE_MAX;
}
bool SegmentWeightPipeline::ensure_segment_capacity(
size_t segment_index,
const std::vector<DeviceMemoryRequest>& requests) {
if (segment_index >= segment_params_.size()) {
return false;
}
auto manager = residency_manager_.lock();
if (manager == nullptr) {
return segment_params_[segment_index].empty();
}
std::vector<ggml_tensor*> protected_params = segment_params_[segment_index];
protected_params.insert(protected_params.end(), queued_params_.begin(), queued_params_.end());
for (const auto& request : requests) {
if (!manager->ensure_compute_backend_capacity(request, segment_params_[segment_index],
preferred_eviction_order(), protected_params)) {
return false;
}
}
return true;
}
bool SegmentWeightPipeline::segment_start(size_t segment_index, const std::function<bool()>& ensure_capacity) {
GGML_ASSERT(pinned_params_.empty());
activate(segment_index);
if (!ensure_capacity()) {
return false;
}
auto manager = residency_manager_.lock();
if (manager == nullptr) {
return segment_params_[segment_index].empty();
}
if (!manager->prepare_params(segment_params_[segment_index])) {
return false;
}
pinned_params_ = segment_params_[segment_index];
return true;
}
void SegmentWeightPipeline::segment_end() {
if (auto manager = residency_manager_.lock()) {
manager->release_compute_backend_params(pinned_params_);
}
pinned_params_.clear();
}
void SegmentWeightPipeline::enqueue_next(
size_t segment_index,
const DeviceMemoryRequest& request) {
if (!enabled_ || queued_segment_ != SIZE_MAX) {
return;
}
const size_t next_segment = next_parameter_segment(segment_index);
if (next_segment == SIZE_MAX) {
return;
}
std::unordered_set<ggml_tensor*> active_params(
segment_params_[segment_index].begin(),
segment_params_[segment_index].end());
std::vector<ggml_tensor*> params;
params.reserve(segment_params_[next_segment].size());
for (ggml_tensor* param : segment_params_[next_segment]) {
if (active_params.find(param) == active_params.end()) {
params.push_back(param);
}
}
if (params.empty()) {
return;
}
auto manager = residency_manager_.lock();
if (manager == nullptr) {
disable();
return;
}
const WeightResidencyInfo residency =
manager->inspect_compute_backend_params(params);
if (residency.missing_bytes == 0) {
return;
}
if (!residency.async_prefetch_supported) {
disable();
return;
}
DeviceMemoryRequest backend_request = request;
backend_request.compute_backend = compute_backend_;
backend_request.owner_id = owner_id_;
std::vector<ggml_tensor*> protected_params = segment_params_[segment_index];
protected_params.insert(protected_params.end(), params.begin(), params.end());
if (!manager->ensure_compute_backend_capacity(backend_request,
params,
preferred_eviction_order(),
protected_params)) {
disable();
return;
}
switch (manager->prefetch_params(owner_id_, params)) {
case WeightPrefetchResult::Scheduled:
queued_params_ = std::move(params);
queued_segment_ = next_segment;
return;
case WeightPrefetchResult::AlreadyResident:
return;
case WeightPrefetchResult::Unsupported:
disable();
return;
case WeightPrefetchResult::Failed:
disable();
return;
}
disable();
}
void SegmentWeightPipeline::clear() {
if (auto manager = residency_manager_.lock()) {
manager->clear_prefetched_params(owner_id_);
}
queued_params_.clear();
queued_segment_ = SIZE_MAX;
}
}
src/core/segment_weight_pipeline.h
+62
View File
@@ -0,0 +1,62 @@
#ifndef __SD_CORE_SEGMENT_WEIGHT_PIPELINE_H__
#define __SD_CORE_SEGMENT_WEIGHT_PIPELINE_H__
#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <unordered_set>
#include <vector>
#include "ggml-backend.h"
struct DeviceMemoryRequest;
struct DeviceResidencyManager;
struct ggml_cgraph;
struct ggml_tensor;
namespace sd::ggml_graph_cut {
struct Plan;
}
namespace sd {
class SegmentWeightPipeline {
private:
std::weak_ptr<DeviceResidencyManager> residency_manager_;
ggml_backend_t compute_backend_ = nullptr;
uintptr_t owner_id_ = 0;
std::vector<std::vector<ggml_tensor*>> segment_params_;
std::vector<ggml_tensor*> queued_params_;
std::vector<ggml_tensor*> pinned_params_;
size_t queued_segment_ = SIZE_MAX;
bool enabled_ = true;
size_t next_parameter_segment(size_t segment_index) const;
std::vector<std::vector<ggml_tensor*>> preferred_eviction_order() const;
void disable();
void activate(size_t segment_index);
void clear();
public:
SegmentWeightPipeline(
const std::shared_ptr<DeviceResidencyManager>& residency_manager,
ggml_backend_t compute_backend,
uintptr_t owner_id,
ggml_cgraph* graph,
const ggml_graph_cut::Plan& plan,
const std::unordered_set<const ggml_tensor*>& params,
bool enabled = true);
~SegmentWeightPipeline();
const std::vector<ggml_tensor*>& params(size_t index) const { return segment_params_[index]; }
bool ensure_segment_capacity(size_t segment_index,
const std::vector<DeviceMemoryRequest>& requests);
bool segment_start(size_t segment_index, const std::function<bool()>& ensure_capacity);
void segment_end();
// Prefetch is best effort; segment_start falls back to synchronous loading.
void enqueue_next(size_t segment_index,
const DeviceMemoryRequest& request);
};
}
#endif // __SD_CORE_SEGMENT_WEIGHT_PIPELINE_H__
src/core/tensor_ggml.hpp
+22 -3
View File
@@ -9,6 +9,7 @@
#include <type_traits>
#include "core/tensor.hpp"
#include "ggml-backend.h"
#include "ggml.h"
namespace sd {
@@ -54,10 +55,28 @@ namespace sd {
GGML_ABORT("ggml tensor type does not match sd::Tensor type");
}
Tensor<T> result(shape_from_ggml(tensor));
if (tensor->buffer != nullptr) {
ggml_backend_tensor_get(tensor, result.data(), 0, ggml_nbytes(tensor));
std::vector<uint8_t> strided_data;
void* destination = result.data();
if (!ggml_is_contiguous(tensor)) {
strided_data.resize(ggml_nbytes(tensor));
destination = strided_data.data();
}
auto buffer = tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer;
if (buffer != nullptr) {
ggml_backend_tensor_get(tensor, destination, 0, ggml_nbytes(tensor));
} else {
std::memcpy(result.data(), tensor->data, ggml_nbytes(tensor));
std::memcpy(destination, tensor->data, ggml_nbytes(tensor));
}
if (!strided_data.empty()) {
for (int64_t i = 0; i < result.numel(); ++i) {
int64_t index = i;
size_t offset = 0;
for (int d = 0; d < GGML_MAX_DIMS; ++d) {
offset += static_cast<size_t>(index % tensor->ne[d]) * tensor->nb[d];
index /= tensor->ne[d];
}
std::memcpy(result.data() + i, strided_data.data() + offset, sizeof(T));
}
}
return result;
}
src/detailer.cpp
+1 -1
View File
@@ -715,7 +715,7 @@ std::vector<ADetailerDetection> ADetailerGGML::predict(sd_image_t image,
LetterboxInput input = make_letterbox_input(image, params.input_size);
int64_t start = ggml_time_ms();
sd::Tensor<float> raw = detector->compute(n_threads, input.tensor);
detector->free_compute_buffer();
detector->runner_end();
if (raw.empty()) {
LOG_ERROR("YOLOv8 detector inference failed");
return {};
src/device_residency_manager.h
+76
View File
@@ -0,0 +1,76 @@
#ifndef __DEVICE_RESIDENCY_MANAGER_H__
#define __DEVICE_RESIDENCY_MANAGER_H__
#include <cstdint>
#include <functional>
#include <vector>
#include "ggml-backend.h"
struct ggml_tensor;
enum class WeightPrefetchResult {
Scheduled,
AlreadyResident,
Unsupported,
Failed,
};
struct WeightResidencyInfo {
bool async_prefetch_supported = false;
size_t missing_bytes = 0;
};
struct DeviceMemoryRequest {
ggml_backend_t compute_backend = nullptr;
uintptr_t owner_id = 0;
size_t pending_allocation_bytes = 0;
size_t runtime_resident_bytes = 0;
size_t max_backend_bytes = 0;
// Runtime buffers only; the manager accounts for weights separately.
size_t runtime_peak_bytes() const {
return pending_allocation_bytes > SIZE_MAX - runtime_resident_bytes
? SIZE_MAX
: runtime_resident_bytes + pending_allocation_bytes;
}
};
struct DeviceResidencyManager {
virtual ~DeviceResidencyManager() = default;
virtual bool segmented_compute_enabled() const = 0;
virtual bool prefetch_enabled() const = 0;
virtual void set_workspace_reclaimer(uintptr_t owner_id, std::function<bool()> reclaim) = 0;
virtual void remove_runtime_owner(uintptr_t owner_id) = 0;
// Capacity requests select their backend's weights; protection spans all backends.
virtual bool fits_compute_backend_capacity(const DeviceMemoryRequest& request,
const std::vector<ggml_tensor*>& required_params) const = 0;
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 evict_compute_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;
virtual WeightResidencyInfo inspect_compute_backend_params(
const std::vector<ggml_tensor*>& tensors) const = 0;
virtual void update_runtime_residency(uintptr_t owner_id,
ggml_backend_t compute_backend,
size_t resident_bytes) = 0;
virtual bool ensure_compute_backend_capacity(
const DeviceMemoryRequest& request,
const std::vector<ggml_tensor*>& required_params,
const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
const std::vector<ggml_tensor*>& protected_params) = 0;
virtual WeightPrefetchResult prefetch_params(
uintptr_t owner_id,
const std::vector<ggml_tensor*>& tensors) = 0;
virtual bool activate_prefetched_params(uintptr_t owner_id,
const std::vector<ggml_tensor*>& tensors) = 0;
virtual void clear_prefetched_params(uintptr_t owner_id) = 0;
};
// Transitional alias for model constructors that have not yet adopted the
// residency-oriented name. It does not introduce a second implementation.
using RunnerWeightManager = DeviceResidencyManager;
#endif // __DEVICE_RESIDENCY_MANAGER_H__
src/extensions/generation_extension.h
+1 -1
View File
@@ -49,7 +49,7 @@ struct GenerationExtension {
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>&) {}
virtual void collect_loras(std::vector<ModelManager::LoraSpec>&) {}
virtual void add_ignore_tensors(std::set<std::string>&) const {}
virtual void runner_done() {}
virtual void runner_end() {}
virtual void reset_runtime_condition() {}
virtual bool prepare_condition(GenerationExtensionConditionContext&) {
return false;
src/extensions/photomaker_extension.cpp
+2 -2
View File
@@ -175,9 +175,9 @@ struct PhotoMakerExtension : public GenerationExtension {
ignore_tensors.insert("pmid.unet.");
}
void runner_done() override {
void runner_end() override {
if (pmid_model != nullptr) {
pmid_model->runner_done();
pmid_model->runner_end();
}
}
src/model/adapter/ip_adapter.hpp
+1 -1
View File
@@ -200,7 +200,7 @@ namespace IPAdapter {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(image_embeds);
};
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true));
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
}
};
src/model/adapter/lora.hpp
+14 -12
View File
@@ -125,13 +125,12 @@ struct LoraModel : public GGMLRunner {
}
void release_loaded_tensors() {
runner_done();
free_compute_buffer();
runner_end();
model_manager.reset();
free_params_ctx();
alloc_params_ctx();
model_manager = std::make_shared<ModelManager>();
weight_manager = model_manager;
model_manager = std::make_shared<ModelManager>();
residency_manager = model_manager;
lora_tensors.clear();
original_tensor_to_final_tensor.clear();
applied_lora_tensors.clear();
@@ -952,16 +951,19 @@ struct LoraModel : public GGMLRunner {
auto get_graph = [&]() -> ggml_cgraph* {
return build_lora_graph(model_tensors, model_tensor_names, version);
};
GGMLRunner::compute<float>(get_graph, n_threads, false, false, false, true);
stat(!warn_unused);
for (auto item : original_tensor_to_final_tensor) {
ggml_tensor* original_tensor = item.first;
ggml_tensor* final_tensor = item.second;
ggml_backend_tensor_copy(final_tensor, original_tensor);
auto read_outputs = [&]() {
for (const auto& item : original_tensor_to_final_tensor) {
ggml_backend_tensor_copy(item.second, item.first);
}
return true;
};
auto result = GGMLRunner::compute<float>(get_graph, n_threads, false, true, read_outputs);
if (!result.has_value()) {
LOG_ERROR("LoRA graph execution failed");
}
stat(!warn_unused);
original_tensor_to_final_tensor.clear();
GGMLRunner::free_compute_buffer();
runner_end();
}
void apply(std::map<std::string, ggml_tensor*> model_tensors, SDVersion version, int n_threads, bool warn_unused = true) {
src/model/adapter/pmid.hpp
+1 -1
View File
@@ -558,7 +558,7 @@ public:
return build_graph(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds);
};
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true));
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
}
};
src/model/detector/yolov8.h
+1 -1
View File
@@ -355,7 +355,7 @@ struct YOLOv8Runner : public GGMLRunner {
sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& input) {
auto get_graph = [&]() { return build_graph(input); };
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false));
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, false));
}
};
src/model/diffusion/anima.hpp
+1 -1
View File
@@ -717,7 +717,7 @@ namespace Anima {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, t5_ids, t5_weights, ref_latents);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/boogu.hpp
+1 -1
View File
@@ -815,7 +815,7 @@ namespace Boogu {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, ref_latents);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/control.hpp
+18 -11
View File
@@ -423,19 +423,26 @@ struct ControlNet : public GGMLRunner {
return build_graph(x, hint, timesteps, context, y);
};
auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false, false, false, true);
auto read_outputs = [&]() {
controls.clear();
controls.reserve(control_outputs_ggml.size());
for (ggml_tensor* control : control_outputs_ggml) {
auto control_host = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml<float>(control), 4);
if (control_host.empty()) {
return false;
}
controls.push_back(std::move(control_host));
}
return true;
};
auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false, true, read_outputs);
control_outputs_ggml.clear();
guided_hint_output_ggml = nullptr;
if (!compute_result.has_value()) {
controls.clear();
return std::nullopt;
}
guided_hint_cached = get_cache_tensor_by_name(guided_hint_cache_name()) != nullptr;
controls.clear();
controls.reserve(control_outputs_ggml.size());
for (ggml_tensor* control : control_outputs_ggml) {
auto control_host = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml<float>(control), 4);
GGML_ASSERT(!control_host.empty());
controls.push_back(std::move(control_host));
}
return controls;
}
@@ -444,10 +451,10 @@ struct ControlNet : public GGMLRunner {
std::map<std::string, ggml_tensor*> tensors;
control_net.get_param_tensors(tensors);
auto manager = std::dynamic_pointer_cast<ModelManager>(weight_manager.lock());
auto manager = std::dynamic_pointer_cast<ModelManager>(residency_manager.lock());
if (manager == nullptr) {
owned_model_manager = std::make_shared<ModelManager>();
weight_manager = owned_model_manager;
residency_manager = owned_model_manager;
manager = owned_model_manager;
}
src/model/diffusion/ernie_image.hpp
+1 -1
View File
@@ -440,7 +440,7 @@ namespace ErnieImage {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/flux.hpp
+1 -1
View File
@@ -1626,7 +1626,7 @@ namespace Flux {
return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, ref_index_mode, skip_layers, pulid_id, pulid_id_weight);
};
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
return result;
}
src/model/diffusion/hidream_o1.hpp
+6 -8
View File
@@ -325,13 +325,11 @@ namespace HiDreamO1 {
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& image,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = true) {
bool auto_runner_end = true) {
auto get_graph = [&]() {
return build_graph(image);
};
auto output = GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params);
auto output = GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end);
return output.has_value() ? std::move(output.value()) : sd::Tensor<float>();
}
};
@@ -459,7 +457,7 @@ namespace HiDreamO1 {
auto get_graph = [&]() {
return build_graph(x, timestep, input_ids, input_pos, token_types, vinput_mask, image_embeds, ref_images);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
@@ -510,8 +508,8 @@ namespace HiDreamO1 {
vision_runner->set_weight_adapter(adapter);
}
void runner_done() override {
vision_runner->runner_done();
void runner_end() override {
vision_runner->runner_end();
}
SDCondition get_learned_condition(int n_threads,
@@ -659,7 +657,7 @@ namespace HiDreamO1 {
result.c_vinput_mask = sd::Tensor<int32_t>(vinput_mask_shape, std::move(vinput_mask));
result.c_image_embeds.reserve(vlm_images.size());
for (const auto& vlm_image : vlm_images) {
auto image_embed = vision_runner->compute(n_threads, vlm_image.second, false, true, true);
auto image_embed = vision_runner->compute(n_threads, vlm_image.second, false);
if (image_embed.empty()) {
LOG_ERROR("hidream_o1 conditioner: encode VLM image failed");
return SDCondition();
src/model/diffusion/hunyuan.hpp
+1 -1
View File
@@ -654,7 +654,7 @@ namespace Hunyuan {
return build_graph(x, timesteps, context, c_concat, y, guidance, byt5, vision, timestep_r);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/ideogram4.hpp
+1 -1
View File
@@ -537,7 +537,7 @@ namespace Ideogram4 {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, use_uncond_model);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/krea2.hpp
+1 -1
View File
@@ -775,7 +775,7 @@ namespace Krea2 {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, ref_latents, ref_image_params);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/lens.hpp
+1 -1
View File
@@ -408,7 +408,7 @@ namespace Lens {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/lingbot_video.hpp
+1 -1
View File
@@ -674,7 +674,7 @@ namespace LingBotVideo {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/ltxv.hpp
+1 -1
View File
@@ -1998,7 +1998,7 @@ namespace LTXV {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, audio_x, audio_timesteps, audio_length, frame_rate, video_positions);
};
auto out = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
auto out = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
return out;
}
src/model/diffusion/mage_flow.hpp
+1 -1
View File
@@ -142,7 +142,7 @@ namespace MageFlow {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, ref_latents);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/minimax_h3.hpp
-2
View File
@@ -1168,8 +1168,6 @@ namespace MiniMaxH3 {
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph,
n_threads,
false,
false,
false),
params.x->dim());
}
src/model/diffusion/minit2i.hpp
+1 -1
View File
@@ -589,7 +589,7 @@ namespace MiniT2I {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, mask);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/mmdit.hpp
+1 -1
View File
@@ -987,7 +987,7 @@ struct MMDiTRunner : public DiffusionModelRunner {
return build_graph(x, timesteps, context, y, skip_layers);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/pid.hpp
+1 -1
View File
@@ -938,7 +938,7 @@ namespace Pid {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, lq_latent, degrade_sigma);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/qwen_image.hpp
+1 -1
View File
@@ -707,7 +707,7 @@ namespace Qwen {
return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/unet.hpp
+1 -1
View File
@@ -835,7 +835,7 @@ struct UNetModelRunner : public DiffusionModelRunner {
return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength, ip_context, ip_scale);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/wan.hpp
+1 -1
View File
@@ -950,7 +950,7 @@ namespace WAN {
return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/diffusion/z_image.hpp
+1 -1
View File
@@ -636,7 +636,7 @@ namespace ZImage {
return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
src/model/te/clip.hpp
+2 -4
View File
@@ -568,13 +568,11 @@ struct CLIPTextModelRunner : public GGMLRunner {
size_t max_token_idx,
bool return_pooled,
int clip_skip,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = true) {
bool auto_runner_end = true) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input_ids, num_custom_embeddings, custom_embeddings_data, max_token_idx, return_pooled, clip_skip);
};
auto result = GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params);
auto result = GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end);
if (return_pooled) {
return take_or_empty(std::move(result));
}
src/model/te/llm.hpp
+8 -20
View File
@@ -2079,9 +2079,7 @@ namespace LLM {
const ImageEmbeds& image_embeds,
std::set<int> out_layers,
bool return_all_hidden_states = false,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = true,
bool auto_runner_end = true,
const DeepStackImageEmbeds& deepstack_image_embeds = {},
const std::vector<ImageGrid>& image_grids = {}) {
auto get_graph = [&]() -> ggml_cgraph* {
@@ -2093,7 +2091,7 @@ namespace LLM {
out_layers,
return_all_hidden_states);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params),
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end),
input_ids.dim() + 1);
}
@@ -2173,13 +2171,11 @@ namespace LLM {
sd::Tensor<float> encode_image(const int n_threads,
const sd::Tensor<float>& image,
bool auto_free = false,
bool free_compute_buffer = false,
bool free_compute_params = false) {
bool auto_runner_end = false) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_encode_image_graph(image);
};
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
}
ggml_cgraph* build_encode_image_outputs_graph(const sd::Tensor<float>& image_tensor) {
@@ -2287,13 +2283,11 @@ namespace LLM {
std::vector<sd::Tensor<float>> encode_image_outputs(const int n_threads,
const sd::Tensor<float>& image,
bool auto_free = false,
bool free_compute_buffer = false,
bool free_compute_params = false) {
bool auto_runner_end = false) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_encode_image_outputs_graph(image);
};
auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
if (combined.empty()) {
return {};
}
@@ -2312,20 +2306,14 @@ namespace LLM {
std::vector<sd::Tensor<float>> encode_video_block_outputs(const int n_threads,
const sd::Tensor<float>& frames,
bool auto_free = false,
bool free_compute_buffer = false,
bool free_compute_params = false) {
bool auto_runner_end = false) {
int grid_h = static_cast<int>(frames.shape()[1] / config.vision.patch_size);
int grid_w = static_cast<int>(frames.shape()[0] / config.vision.patch_size);
auto pixel_values = process_video_block_tensor(frames, config.vision);
auto get_graph = [&]() -> ggml_cgraph* {
return build_encode_video_block_outputs_graph(pixel_values, grid_h, grid_w);
};
auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph,
n_threads,
auto_free,
free_compute_buffer,
free_compute_params));
auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
if (combined.empty()) {
return {};
}
src/model/te/t5.hpp
+2 -4
View File
@@ -451,13 +451,11 @@ struct T5Runner : public GGMLRunner {
sd::Tensor<float> compute(const int n_threads,
const sd::Tensor<int32_t>& input_ids,
const sd::Tensor<float>& attention_mask,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = true) {
bool auto_runner_end = true) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input_ids, attention_mask);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params), 3);
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end), 3);
}
static std::vector<int> _relative_position_bucket(const std::vector<int>& relative_position,
src/model/upscaler/esrgan.hpp
+1 -1
View File
@@ -265,7 +265,7 @@ struct ESRGAN : public GGMLRunner {
sd::Tensor<float> compute(const int n_threads,
const sd::Tensor<float>& x) {
auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); };
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
return result;
}
};
src/model/upscaler/ltx_latent_upscaler.hpp
+1 -1
View File
@@ -499,7 +499,7 @@ namespace LTXVUpsampler {
}
size_t expected_dim = static_cast<size_t>(x.dim());
auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); };
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), expected_dim);
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), expected_dim);
}
};
src/model/vae/auto_encoder_kl.hpp
+1 -1
View File
@@ -744,7 +744,7 @@ struct AutoEncoderKL : public VAE {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(z, decode_graph);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), z.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z.dim());
}
sd::Tensor<float> gaussian_latent_sample(const sd::Tensor<float>& moments, std::shared_ptr<RNG> rng) {
src/model/vae/hunyuan_vae.hpp
+1 -3
View File
@@ -827,9 +827,7 @@ namespace Hunyuan {
};
auto output = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph,
n_threads,
true,
true,
true),
false),
graph_input.dim());
if (!output.empty() && input.dim() == 4) {
output.squeeze_(2);
src/model/vae/ltx_audio_vae.hpp
+1 -1
View File
@@ -1042,7 +1042,7 @@ namespace LTXV {
ggml_build_forward_expand(gf, waveform);
return gf;
};
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), 4);
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), 4);
int64_t t1 = ggml_time_ms();
LOG_INFO("ltx audio vae decode completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
return result;
src/model/vae/ltx_vae.hpp
+3 -5
View File
@@ -1334,7 +1334,6 @@ struct LTXVideoVAE : public VAE {
(int)plan.tiles.size());
free_cache_ctx_and_buffer();
cache_tensor_map.clear();
auto output = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& z_chunk, const VAETemporalTile& tile) {
LOG_DEBUG("LTX VAE temporal tile %lld/%d: latent frames [%lld, %lld), overlap=%d",
@@ -1349,12 +1348,11 @@ struct LTXVideoVAE : public VAE {
static_cast<int>(tile.start),
tile.overlap);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
expected_dim);
});
free_cache_ctx_and_buffer();
cache_tensor_map.clear();
return output;
}
@@ -1407,7 +1405,7 @@ struct LTXVideoVAE : public VAE {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input, decode_graph);
};
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), expected_dim);
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), expected_dim);
if (result.empty()) {
return {};
}
@@ -1420,7 +1418,7 @@ struct LTXVideoVAE : public VAE {
auto get_graph = [&]() -> ggml_cgraph* {
return build_latent_statistics_graph(z, normalize);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
static_cast<size_t>(z.dim()));
}
src/model/vae/mage_vae.hpp
+1 -1
View File
@@ -490,7 +490,7 @@ namespace MageVAE {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input, decode_graph);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), input.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), input.dim());
}
int get_encoder_output_channels(int input_channels) override {
src/model/vae/minimax_h3_audio_vae.hpp
+2 -2
View File
@@ -480,7 +480,7 @@ namespace MiniMaxH3 {
return graph;
};
auto result = restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
GGMLRunner::compute<float>(get_graph, n_threads, false),
4);
int64_t t1 = ggml_time_ms();
LOG_INFO("MiniMax-H3 audio VAE encode completed, taking %.2fs",
@@ -500,7 +500,7 @@ namespace MiniMaxH3 {
return graph;
};
auto result = restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
GGMLRunner::compute<float>(get_graph, n_threads, false),
4);
int64_t t1 = ggml_time_ms();
LOG_INFO("MiniMax-H3 audio VAE decode completed, taking %.2fs",
src/model/vae/minimax_h3_vae.hpp
-2
View File
@@ -793,8 +793,6 @@ namespace MiniMaxH3VAE {
return restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph,
n_threads,
false,
false,
false),
5);
}
src/model/vae/tae.hpp
+2 -2
View File
@@ -787,7 +787,7 @@ struct TinyImageAutoEncoder : public VAE {
return build_graph(z_tensor, decode_graph);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), z_tensor.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
}
};
@@ -872,7 +872,7 @@ struct TinyVideoAutoEncoder : public VAE {
return build_graph(z_tensor, decode_graph);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), z_tensor.dim());
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
}
};
src/model/vae/vae.hpp
+2 -2
View File
@@ -251,7 +251,7 @@ public:
tiling_params);
}
runner_done();
runner_end();
if (output.empty()) {
LOG_ERROR("vae encode compute failed");
@@ -305,7 +305,7 @@ public:
tiling_params);
}
runner_done();
runner_end();
if (output.empty()) {
LOG_ERROR("vae decode compute failed");
src/model/vae/wan_vae.hpp
+2 -4
View File
@@ -1415,7 +1415,6 @@ namespace WAN {
(int)plan.tiles.size());
free_cache_ctx_and_buffer();
cache_tensor_map.clear();
ae.clear_cache();
auto output = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& input_tile, const VAETemporalTile& tile) {
@@ -1428,12 +1427,11 @@ namespace WAN {
return build_temporal_tile_graph(input_tile, static_cast<int>(tile.start));
};
return restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
GGMLRunner::compute<float>(get_graph, n_threads, false),
static_cast<size_t>(input.dim()));
});
free_cache_ctx_and_buffer();
cache_tensor_map.clear();
ae.clear_cache();
return output;
}
@@ -1448,7 +1446,7 @@ namespace WAN {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input.empty() ? z : input, decode_graph);
};
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
input.empty() ? z.dim() : input.dim());
if (!result.empty() && z.dim() == 4) {
result.squeeze_(2);
src/model_manager.cpp
+588 -101
View File
@@ -143,7 +143,7 @@ size_t estimate_tensors_size(const std::map<std::string, ggml_tensor*>& tensors)
return size;
}
void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft) {
void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft, const std::vector<std::pair<ggml_backend_t, size_t>>& device_limits) {
if (compute_backend == nullptr) {
return;
}
@@ -152,6 +152,7 @@ void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_ba
return;
}
split_buffer_types_[compute_backend] = split_buft;
split_buffer_devices_[split_buft] = device_limits;
}
bool ModelManager::tensor_shape_supports_split_buffer(const ggml_tensor* tensor) {
@@ -164,11 +165,10 @@ bool ModelManager::tensor_shape_supports_split_buffer(const ggml_tensor* tensor)
}
ggml_backend_buffer_type_t ModelManager::split_buffer_type_for(const TensorState& state) const {
if (!state.allow_split_buffer || !tensor_shape_supports_split_buffer(state.tensor)) {
if (!tensor_shape_supports_split_buffer(state.tensor)) {
return nullptr;
}
auto it = split_buffer_types_.find(state.compute_backend);
return it != split_buffer_types_.end() ? it->second : nullptr;
return state.split_buffer_type;
}
bool ModelManager::register_param_tensors(const std::string& desc,
@@ -203,14 +203,17 @@ 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;
auto split_buffer = split_buffer_types_.find(compute_backend);
if (allow_split_buffer && split_buffer != split_buffer_types_.end()) {
state->split_buffer_type = split_buffer->second;
}
state->params_follow_compute_backend = params_follow_compute_backend;
if (tensor_ops != nullptr) {
auto op_it = tensor_ops->find(tensor);
@@ -240,7 +243,7 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg
if (state == nullptr || state->desc != desc) {
continue;
}
if (state->active_prepare_count > 0) {
if (state->pin_count > 0) {
LOG_ERROR("model manager cannot unregister active %s tensor '%s'",
desc.c_str(),
state->name.c_str());
@@ -287,7 +290,7 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg
if (state == nullptr) {
continue;
}
if (state->active_prepare_count > 0 || state->staged_to_compute_backend) {
if (state->pin_count > 0 || state->staged_to_compute_backend) {
LOG_ERROR("model manager cannot unregister %s while tensor '%s' is active",
desc.c_str(),
state->name.c_str());
@@ -403,14 +406,27 @@ bool ModelManager::load_tensors_to_params_backend(const std::vector<TensorState*
}
return false;
}
struct PrepareStats {
size_t bytes = 0;
size_t tensors = 0;
size_t blocks = 0;
};
std::map<ggml_backend_buffer_type_t, PrepareStats> prepared;
for (ParamsStorageBlock* block : created_storage_blocks) {
if (block != nullptr && block->buffer != nullptr) {
LOG_DEBUG("model manager prepared params backend buffer (%6.2f MB, %zu tensors, %s)",
ggml_backend_buffer_get_size(block->buffer) / (1024.f * 1024.f),
block->states.size(),
ggml_backend_buffer_is_host(block->buffer) ? "RAM" : "VRAM");
auto& stats = prepared[ggml_backend_buffer_get_type(block->buffer)];
stats.bytes += ggml_backend_buffer_get_size(block->buffer);
stats.tensors += block->states.size();
++stats.blocks;
}
}
for (const auto& entry : prepared) {
LOG_DEBUG("model manager prepared params backend buffers (%6.2f MB, %zu tensors, %zu blocks, %s) on %s",
entry.second.bytes / (1024.f * 1024.f),
entry.second.tensors, entry.second.blocks,
ggml_backend_buft_is_host(entry.first) ? "RAM" : "VRAM",
ggml_backend_buft_name(entry.first));
}
return true;
}
@@ -444,68 +460,99 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorStat
}
for (const auto& pair : states_by_staging_target) {
ggml_backend_t compute_backend = pair.first.first;
ggml_backend_buffer_type_t staging_buft = pair.first.second;
const std::vector<TensorState*>& states = pair.second;
if (states.empty()) {
ggml_backend_t compute_backend = pair.first.first;
ggml_backend_buffer_type_t staging_buft = pair.first.second;
const std::vector<TensorState*>& target_states = pair.second;
if (target_states.empty()) {
continue;
}
int64_t t0 = ggml_time_ms();
ggml_init_params init_params;
init_params.mem_size = std::max<size_t>(1, states.size()) * ggml_tensor_overhead();
init_params.mem_buffer = nullptr;
init_params.no_alloc = true;
ggml_context* staging_ctx = ggml_init(init_params);
GGML_ASSERT(staging_ctx != nullptr);
std::vector<std::pair<TensorState*, ggml_tensor*>> staged_tensors;
staged_tensors.reserve(states.size());
for (TensorState* state : states) {
ggml_tensor* staging_tensor = ggml_dup_tensor(staging_ctx, state->tensor);
ggml_set_name(staging_tensor, state->tensor->name);
staged_tensors.push_back({state, staging_tensor});
const size_t alignment = ggml_backend_buft_get_alignment(staging_buft);
size_t backend_limit = ggml_backend_buft_get_max_size(staging_buft);
if (!ggml_backend_buft_is_host(staging_buft) &&
(backend_limit == 0 || backend_limit > MAX_RESIDENCY_BLOCK_BYTES)) {
backend_limit = MAX_RESIDENCY_BLOCK_BYTES;
}
ggml_backend_buffer_t compute_buffer = ggml_backend_alloc_ctx_tensors_from_buft(staging_ctx, staging_buft);
if (compute_buffer == nullptr) {
LOG_ERROR("model manager alloc compute params backend buffer failed, num_tensors = %zu",
staged_tensors.size());
ggml_free(staging_ctx);
const int64_t t0 = ggml_time_ms();
size_t staged_bytes = 0;
size_t staged_blocks = 0;
auto stage_chunk = [&](const std::vector<TensorState*>& chunk) -> bool {
if (chunk.empty()) {
return true;
}
ggml_init_params init_params;
init_params.mem_size = std::max<size_t>(1, chunk.size()) * ggml_tensor_overhead();
init_params.mem_buffer = nullptr;
init_params.no_alloc = true;
ggml_context* staging_ctx = ggml_init(init_params);
GGML_ASSERT(staging_ctx != nullptr);
std::vector<std::pair<TensorState*, ggml_tensor*>> staged_tensors;
staged_tensors.reserve(chunk.size());
for (TensorState* state : chunk) {
ggml_tensor* staging_tensor = ggml_dup_tensor(staging_ctx, state->tensor);
ggml_set_name(staging_tensor, state->tensor->name);
staged_tensors.push_back({state, staging_tensor});
}
ggml_backend_buffer_t compute_buffer =
ggml_backend_alloc_ctx_tensors_from_buft(staging_ctx, staging_buft);
if (compute_buffer == nullptr) {
LOG_ERROR("model manager alloc compute params backend buffer failed, num_tensors = %zu",
staged_tensors.size());
ggml_free(staging_ctx);
return false;
}
ggml_backend_buffer_set_usage(compute_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
for (auto& staged_tensor : staged_tensors) {
TensorState* state = staged_tensor.first;
ggml_tensor* managed_tensor = state->tensor;
ggml_tensor* staging_tensor = staged_tensor.second;
ggml_backend_tensor_copy(managed_tensor, staging_tensor);
std::swap(managed_tensor->buffer, staging_tensor->buffer);
std::swap(managed_tensor->data, staging_tensor->data);
std::swap(managed_tensor->extra, staging_tensor->extra);
state->staged_to_compute_backend = true;
}
ggml_backend_synchronize(compute_backend);
auto block = std::make_unique<ComputeStagingBlock>();
block->compute_backend = compute_backend;
block->buffer = compute_buffer;
block->staging_ctx = staging_ctx;
block->staged_tensors = std::move(staged_tensors);
staged_bytes += ggml_backend_buffer_get_size(compute_buffer);
++staged_blocks;
compute_staging_blocks_.push_back(std::move(block));
return true;
};
std::vector<TensorState*> chunk;
size_t chunk_size = 0;
for (TensorState* state : target_states) {
const size_t tensor_size = GGML_PAD(
ggml_backend_buft_get_alloc_size(staging_buft, state->tensor), alignment);
if (!chunk.empty() && backend_limit > 0 &&
tensor_size > backend_limit - std::min(chunk_size, backend_limit)) {
if (!stage_chunk(chunk)) {
return false;
}
chunk.clear();
chunk_size = 0;
}
chunk.push_back(state);
chunk_size = tensor_size > SIZE_MAX - chunk_size ? SIZE_MAX : chunk_size + tensor_size;
}
if (!stage_chunk(chunk)) {
return false;
}
ggml_backend_buffer_set_usage(compute_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
for (auto& staged_tensor : staged_tensors) {
TensorState* state = staged_tensor.first;
ggml_tensor* managed_tensor = state->tensor;
ggml_tensor* staging_tensor = staged_tensor.second;
ggml_backend_tensor_copy(managed_tensor, staging_tensor);
std::swap(managed_tensor->buffer, staging_tensor->buffer);
std::swap(managed_tensor->data, staging_tensor->data);
std::swap(managed_tensor->extra, staging_tensor->extra);
}
ggml_backend_synchronize(compute_backend);
auto block = std::make_unique<ComputeStagingBlock>();
block->compute_backend = compute_backend;
block->buffer = compute_buffer;
block->staging_ctx = staging_ctx;
block->staged_tensors = std::move(staged_tensors);
for (auto& staged_tensor : block->staged_tensors) {
TensorState* state = staged_tensor.first;
state->staged_to_compute_backend = true;
}
compute_staging_blocks_.push_back(std::move(block));
int64_t t1 = ggml_time_ms();
LOG_DEBUG("model manager staged compute params (%6.2f MB, %zu tensors) to %s, taking %.2fs",
ggml_backend_buffer_get_size(compute_buffer) / (1024.f * 1024.f),
states.size(),
LOG_DEBUG("model manager staged compute params (%6.2f MB, %zu tensors, %zu blocks) to %s, taking %.2fs",
staged_bytes / (1024.f * 1024.f),
target_states.size(),
staged_blocks,
ggml_backend_name(compute_backend),
(t1 - t0) * 1.0f / 1000);
(ggml_time_ms() - t0) / 1000.f);
}
return true;
@@ -729,6 +776,10 @@ bool ModelManager::alloc_params_buffers(const std::vector<TensorState*>& states,
const std::vector<TensorState*>& states = pair.second;
size_t alignment = ggml_backend_buft_get_alignment(params_buft);
size_t max_size = ggml_backend_buft_get_max_size(params_buft);
if (!ggml_backend_buft_is_host(params_buft) &&
(max_size == 0 || max_size > MAX_RESIDENCY_BLOCK_BYTES)) {
max_size = MAX_RESIDENCY_BLOCK_BYTES;
}
auto alloc_chunk = [&](const std::vector<TensorState*>& chunk, size_t chunk_size) -> bool {
if (chunk.empty() || chunk_size == 0) {
@@ -935,10 +986,6 @@ void ModelManager::free_compute_staging_block(ComputeStagingBlock& block) {
}
if (block.buffer != nullptr) {
LOG_DEBUG("model manager releasing compute params (%6.2f MB, %zu tensors) from %s",
ggml_backend_buffer_get_size(block.buffer) / (1024.f * 1024.f),
block.staged_tensors.size(),
block.compute_backend != nullptr ? ggml_backend_name(block.compute_backend) : "unknown");
ggml_backend_buffer_free(block.buffer);
block.buffer = nullptr;
}
@@ -951,6 +998,12 @@ void ModelManager::free_compute_staging_block(ComputeStagingBlock& block) {
void ModelManager::release_compute_staging_blocks(bool force,
const std::unordered_set<TensorState*>* target_states) {
struct ReleaseStats {
size_t bytes = 0;
size_t tensors = 0;
size_t blocks = 0;
};
std::map<ggml_backend_t, ReleaseStats> released;
for (auto it = compute_staging_blocks_.begin(); it != compute_staging_blocks_.end();) {
ComputeStagingBlock* block = it->get();
bool can_release = force;
@@ -966,25 +1019,33 @@ void ModelManager::release_compute_staging_blocks(bool force,
target_states->find(state) == target_states->end()) {
return false;
}
return state->active_prepare_count == 0;
return state->pin_count == 0;
});
}
if (can_release) {
if (block->buffer != nullptr) {
auto& stats = released[block->compute_backend];
stats.bytes += ggml_backend_buffer_get_size(block->buffer);
stats.tensors += block->staged_tensors.size();
++stats.blocks;
}
free_compute_staging_block(*block);
it = compute_staging_blocks_.erase(it);
} else {
++it;
}
}
for (const auto& entry : released) {
LOG_DEBUG("model manager released compute params (%6.2f MB, %zu tensors, %zu blocks) from %s",
entry.second.bytes / (1024.f * 1024.f),
entry.second.tensors, entry.second.blocks,
entry.first != nullptr ? ggml_backend_name(entry.first) : "unknown");
}
}
void ModelManager::free_params_storage_block(ParamsStorageBlock& block) {
if (block.buffer != nullptr) {
LOG_DEBUG("model manager releasing params backend buffer (%6.2f MB, %zu tensors, %s)",
ggml_backend_buffer_get_size(block.buffer) / (1024.f * 1024.f),
block.states.size(),
ggml_backend_buffer_is_host(block.buffer) ? "RAM" : "VRAM");
ggml_backend_buffer_free(block.buffer);
block.buffer = nullptr;
}
@@ -1006,6 +1067,12 @@ void ModelManager::free_params_storage_block(ParamsStorageBlock& block) {
void ModelManager::release_params_storage_blocks(bool force,
const std::unordered_set<TensorState*>* target_states) {
struct ReleaseStats {
size_t bytes = 0;
size_t tensors = 0;
size_t blocks = 0;
};
std::map<ggml_backend_buffer_type_t, ReleaseStats> released;
for (auto it = params_storage_blocks_.begin(); it != params_storage_blocks_.end();) {
ParamsStorageBlock* block = it->get();
bool can_release = force;
@@ -1020,19 +1087,32 @@ void ModelManager::release_params_storage_blocks(bool force,
target_states->find(state) == target_states->end()) {
return false;
}
return state->active_prepare_count == 0 &&
return state->pin_count == 0 &&
!state->staged_to_compute_backend &&
state->residency_mode == ResidencyMode::Disk;
});
}
if (can_release) {
if (block->buffer != nullptr) {
auto& stats = released[ggml_backend_buffer_get_type(block->buffer)];
stats.bytes += ggml_backend_buffer_get_size(block->buffer);
stats.tensors += block->states.size();
++stats.blocks;
}
free_params_storage_block(*block);
it = params_storage_blocks_.erase(it);
} else {
++it;
}
}
for (const auto& entry : released) {
LOG_DEBUG("model manager released params backend buffers (%6.2f MB, %zu tensors, %zu blocks, %s) from %s",
entry.second.bytes / (1024.f * 1024.f),
entry.second.tensors, entry.second.blocks,
ggml_backend_buft_is_host(entry.first) ? "RAM" : "VRAM",
ggml_backend_buft_name(entry.first));
}
}
void ModelManager::erase_params_storage_block(ParamsStorageBlock* block) {
@@ -1048,16 +1128,19 @@ void ModelManager::erase_params_storage_block(ParamsStorageBlock* block) {
void ModelManager::release_all() {
clear_all_prefetched_params();
runtime_residencies_.clear();
workspace_reclaimers_.clear();
for (auto& state : tensor_states_) {
state->active_prepare_count = 0;
state->applied_lora_epoch = UINT64_MAX;
state->pin_count = 0;
state->applied_lora_epoch = UINT64_MAX;
}
release_compute_staging_blocks(true);
release_params_storage_blocks(true);
}
bool ModelManager::resolve_required_tensor_states(const std::vector<ggml_tensor*>& tensors,
std::vector<TensorState*>& required_states) const {
std::vector<TensorState*>& required_states,
ggml_backend_t compute_backend) const {
required_states.clear();
std::unordered_set<TensorState*> seen;
for (ggml_tensor* tensor : tensors) {
@@ -1079,7 +1162,9 @@ bool ModelManager::resolve_required_tensor_states(const std::vector<ggml_tensor*
LOG_ERROR("model manager tensor '%s' has no tensor state", raw_name);
return false;
}
if (seen.insert(state).second) {
if ((compute_backend == nullptr || state->compute_backend == nullptr ||
state->compute_backend == compute_backend) &&
seen.insert(state).second) {
required_states.push_back(state);
}
}
@@ -1115,7 +1200,7 @@ bool ModelManager::assign_compute_backend(const std::vector<ggml_tensor*>& tenso
continue;
}
if (state->active_prepare_count > 0 || state->staged_to_compute_backend) {
if (state->pin_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;
@@ -1135,6 +1220,131 @@ bool ModelManager::assign_compute_backend(const std::vector<ggml_tensor*>& tenso
return true;
}
size_t ModelManager::compute_backend_alloc_size(const std::vector<TensorState*>& states,
bool missing_only) const {
size_t total_size = 0;
std::unordered_set<TensorState*> seen;
for (TensorState* state : states) {
if (state == nullptr || state->tensor == nullptr || !seen.insert(state).second ||
should_ignore(*state) || is_optional_missing_tensor(state->name)) {
continue;
}
const bool compute_resident =
state->compute_backend == state->params_backend
? state->loaded_to_params_backend
: state->staged_to_compute_backend;
if (missing_only && compute_resident) {
continue;
}
ggml_backend_buffer_type_t buffer_type = nullptr;
if (state->compute_backend == state->params_backend) {
buffer_type = params_buffer_type_for(*state);
} else {
buffer_type = split_buffer_type_for(*state);
if (buffer_type == nullptr && state->compute_backend != nullptr) {
buffer_type = ggml_backend_get_default_buffer_type(state->compute_backend);
}
}
if (buffer_type == nullptr) {
continue;
}
const size_t alignment = ggml_backend_buft_get_alignment(buffer_type);
const size_t tensor_size = ggml_backend_buft_get_alloc_size(buffer_type, state->tensor);
const size_t alloc_size = GGML_PAD(tensor_size, alignment);
if (alloc_size > SIZE_MAX - total_size) {
return SIZE_MAX;
}
total_size += alloc_size;
}
return total_size;
}
size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backend) const {
if (compute_backend == nullptr) {
return 0;
}
ggml_backend_dev_t compute_device = ggml_backend_get_device(compute_backend);
if (compute_device == nullptr) {
return 0;
}
size_t total_size = 0;
auto add_buffer = [&](ggml_backend_buffer_t buffer) {
if (buffer == nullptr || ggml_backend_buffer_is_host(buffer)) {
return;
}
ggml_backend_buffer_type_t buffer_type = ggml_backend_buffer_get_type(buffer);
auto split_devices = split_buffer_devices_.find(buffer_type);
const bool on_device = split_devices == split_buffer_devices_.end()
? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device
: std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) {
return ggml_backend_get_device(entry.first) == compute_device;
});
if (!on_device) {
return;
}
const size_t buffer_size = ggml_backend_buffer_get_size(buffer);
total_size = buffer_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + buffer_size;
};
for (const auto& block : params_storage_blocks_) {
if (block != nullptr) {
add_buffer(block->buffer);
}
}
for (const auto& block : compute_staging_blocks_) {
if (block != nullptr) {
add_buffer(block->buffer);
}
}
for (const auto& entry : prefetch_blocks_) {
if (entry.second != nullptr) {
for (const auto& block : entry.second->staging_blocks) {
if (block != nullptr) {
add_buffer(block->buffer);
}
}
}
}
return total_size;
}
void ModelManager::update_runtime_residency(uintptr_t owner_id,
ggml_backend_t compute_backend,
size_t resident_bytes) {
if (owner_id == 0) {
return;
}
if (compute_backend == nullptr || resident_bytes == 0) {
runtime_residencies_.erase({owner_id, compute_backend});
return;
}
runtime_residencies_[{owner_id, compute_backend}] = {compute_backend, resident_bytes};
}
size_t ModelManager::other_runtime_resident_bytes(uintptr_t owner_id,
ggml_backend_t compute_backend) const {
if (compute_backend == nullptr) {
return 0;
}
ggml_backend_dev_t compute_device = ggml_backend_get_device(compute_backend);
if (compute_device == nullptr) {
return 0;
}
size_t total_size = 0;
for (const auto& entry : runtime_residencies_) {
if (entry.first.first == owner_id || entry.second.compute_backend == nullptr ||
ggml_backend_get_device(entry.second.compute_backend) != compute_device) {
continue;
}
total_size = entry.second.resident_bytes > SIZE_MAX - total_size
? SIZE_MAX
: total_size + entry.second.resident_bytes;
}
return total_size;
}
bool ModelManager::prepare_params(const std::vector<ggml_tensor*>& tensors) {
if (tensors.empty()) {
return true;
@@ -1155,18 +1365,20 @@ bool ModelManager::prepare_params(const std::vector<ggml_tensor*>& tensors) {
return false;
}
// LoRA execution may reclaim other residency blocks while these weights are in use.
const uint64_t use_epoch = ++residency_epoch_;
for (TensorState* state : required_states) {
if (state != nullptr) {
state->pin_count++;
state->last_use_epoch = use_epoch;
}
}
if (!apply_loras_to_params(required_states)) {
finish_compute_backend_usage(required_states);
release_compute_staging_blocks(false);
release_params_storage_blocks(false);
return false;
}
for (TensorState* state : required_states) {
if (state == nullptr) {
continue;
}
state->active_prepare_count++;
}
return true;
}
@@ -1180,11 +1392,10 @@ void ModelManager::finish_compute_backend_usage(const std::vector<TensorState*>&
if (state == nullptr || !target_states.insert(state).second) {
continue;
}
if (state->active_prepare_count > 0) {
state->active_prepare_count--;
if (state->pin_count > 0) {
state->pin_count--;
}
}
release_compute_staging_blocks(false, &target_states);
}
void ModelManager::release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) {
@@ -1198,7 +1409,7 @@ void ModelManager::release_compute_backend_params(const std::vector<ggml_tensor*
finish_compute_backend_usage(required_states);
}
void ModelManager::release_params_backend_params(const std::vector<ggml_tensor*>& tensors) {
void ModelManager::evict_compute_backend_params(const std::vector<ggml_tensor*>& tensors) {
if (tensors.empty()) {
return;
}
@@ -1206,9 +1417,285 @@ void ModelManager::release_params_backend_params(const std::vector<ggml_tensor*>
if (!resolve_required_tensor_states(tensors, required_states)) {
return;
}
if (required_states.empty()) {
return;
}
std::unordered_set<TensorState*> target_states(required_states.begin(), required_states.end());
for (const auto& block : compute_staging_blocks_) {
const bool intersects = std::any_of(
block->staged_tensors.begin(),
block->staged_tensors.end(),
[&](const std::pair<TensorState*, ggml_tensor*>& pair) {
return pair.first != nullptr && target_states.count(pair.first) > 0;
});
const bool fully_evictable = std::all_of(
block->staged_tensors.begin(),
block->staged_tensors.end(),
[](const std::pair<TensorState*, ggml_tensor*>& pair) {
return pair.first == nullptr || pair.first->pin_count == 0;
});
if (intersects && fully_evictable) {
for (const auto& pair : block->staged_tensors) {
if (pair.first != nullptr) {
target_states.insert(pair.first);
}
}
}
}
release_compute_staging_blocks(false, &target_states);
for (const auto& block : params_storage_blocks_) {
const bool intersects = std::any_of(
block->states.begin(),
block->states.end(),
[&](TensorState* state) {
return state != nullptr && target_states.count(state) > 0;
});
const bool fully_evictable = std::all_of(
block->states.begin(),
block->states.end(),
[](TensorState* state) {
return state == nullptr ||
(state->pin_count == 0 && !state->staged_to_compute_backend &&
state->residency_mode == ResidencyMode::Disk);
});
if (intersects && fully_evictable) {
target_states.insert(block->states.begin(), block->states.end());
}
}
release_params_storage_blocks(false, &target_states);
}
WeightResidencyInfo ModelManager::inspect_compute_backend_params(
const std::vector<ggml_tensor*>& tensors) const {
WeightResidencyInfo info;
std::vector<TensorState*> states;
if (!resolve_required_tensor_states(tensors, states)) {
return info;
}
ggml_backend_t prefetch_compute_backend = nullptr;
bool has_missing_params = false;
bool prefetch_candidate = true;
for (TensorState* state : states) {
if (state == nullptr || should_ignore(*state) ||
is_optional_missing_tensor(state->name)) {
continue;
}
const bool compute_resident =
state->compute_backend == state->params_backend
? state->loaded_to_params_backend
: state->staged_to_compute_backend;
if (compute_resident) {
continue;
}
has_missing_params = true;
if (split_buffer_type_for(*state) != nullptr) {
prefetch_candidate = false;
}
if (state->compute_backend == state->params_backend ||
state->compute_backend == nullptr || sd_backend_is_cpu(state->compute_backend)) {
prefetch_candidate = false;
continue;
}
if (prefetch_compute_backend == nullptr) {
prefetch_compute_backend = state->compute_backend;
} else if (prefetch_compute_backend != state->compute_backend) {
prefetch_candidate = false;
}
}
info.missing_bytes = compute_backend_alloc_size(states, true);
if (has_missing_params && prefetch_candidate && prefetch_compute_backend != nullptr) {
ggml_backend_dev_t device = ggml_backend_get_device(prefetch_compute_backend);
if (device != nullptr) {
ggml_backend_dev_props props{};
ggml_backend_dev_get_props(device, &props);
info.async_prefetch_supported = props.caps.async;
}
}
return info;
}
void ModelManager::set_workspace_reclaimer(uintptr_t owner_id, std::function<bool()> reclaim) {
workspace_reclaimers_[owner_id] = std::move(reclaim);
}
void ModelManager::remove_runtime_owner(uintptr_t owner_id) {
workspace_reclaimers_.erase(owner_id);
for (auto it = runtime_residencies_.begin(); it != runtime_residencies_.end();) {
if (it->first.first == owner_id) {
it = runtime_residencies_.erase(it);
} else {
++it;
}
}
}
ModelManager::CapacityCheck ModelManager::check_capacity(
const DeviceMemoryRequest& request,
const std::vector<TensorState*>& states) const {
CapacityCheck result;
if (request.compute_backend == nullptr || sd_backend_is_cpu(request.compute_backend)) {
return result;
}
auto add = [](size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; };
const size_t missing = compute_backend_alloc_size(states, true);
result.required_device_bytes = add(request.pending_allocation_bytes, missing);
result.required_budget_bytes = add(request.runtime_peak_bytes(), missing);
auto device = ggml_backend_get_device(request.compute_backend);
if (device != nullptr) {
size_t free_bytes = 0, total_bytes = 0;
ggml_backend_dev_memory(device, &free_bytes, &total_bytes);
if (free_bytes != 0 || total_bytes != 0) {
result.available_device_bytes = free_bytes;
}
}
if (request.max_backend_bytes > 0) {
const size_t resident = add(compute_backend_resident_bytes(request.compute_backend),
other_runtime_resident_bytes(request.owner_id, request.compute_backend));
result.available_budget_bytes = resident < request.max_backend_bytes
? request.max_backend_bytes - resident
: 0;
}
std::map<ggml_backend_t, size_t> split_devices;
for (auto state : states) {
auto placement = split_buffer_devices_.find(split_buffer_type_for(*state));
if (placement != split_buffer_devices_.end()) {
for (const auto& entry : placement->second) {
auto inserted = split_devices.emplace(entry);
if (!inserted.second && entry.second > 0) {
auto& limit = inserted.first->second;
limit = limit == 0 ? entry.second : std::min(limit, entry.second);
}
}
}
}
// GGML exposes only a split buffer's total size, not per-device allocations.
// Charge that upper bound on every participant instead of undercounting a shard.
for (const auto& entry : split_devices) {
size_t free_bytes = 0, total_bytes = 0;
ggml_backend_dev_memory(ggml_backend_get_device(entry.first), &free_bytes, &total_bytes);
if (free_bytes != 0 || total_bytes != 0) {
result.available_device_bytes = std::min(result.available_device_bytes, free_bytes);
}
if (entry.second > 0) {
const size_t resident = add(compute_backend_resident_bytes(entry.first),
other_runtime_resident_bytes(request.owner_id, entry.first));
const size_t available = resident < entry.second ? entry.second - resident : 0;
result.available_budget_bytes = std::min(result.available_budget_bytes, available);
}
}
return result;
}
bool ModelManager::fits_compute_backend_capacity(
const DeviceMemoryRequest& request,
const std::vector<ggml_tensor*>& required_params) const {
std::vector<TensorState*> states;
return resolve_required_tensor_states(required_params, states, request.compute_backend) &&
check_capacity(request, states).fits();
}
bool ModelManager::ensure_compute_backend_capacity(
const DeviceMemoryRequest& request,
const std::vector<ggml_tensor*>& required_params,
const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
const std::vector<ggml_tensor*>& protected_params) {
std::vector<TensorState*> required_states;
if (!resolve_required_tensor_states(required_params, required_states, request.compute_backend)) {
return false;
}
ggml_backend_t compute_backend = request.compute_backend;
if (compute_backend == nullptr) {
LOG_ERROR("model manager cannot reclaim memory for a null compute backend");
return false;
}
if (sd_backend_is_cpu(compute_backend)) {
return true;
}
auto fits = [&]() { return check_capacity(request, required_states).fits(); };
if (fits()) {
return true;
}
for (const auto& entry : workspace_reclaimers_) {
if (entry.first != request.owner_id) {
entry.second();
if (fits()) {
return true;
}
}
}
std::unordered_set<TensorState*> protected_states;
std::vector<TensorState*> resolved_protected;
if (!resolve_required_tensor_states(protected_params, resolved_protected)) {
return false;
}
protected_states.insert(resolved_protected.begin(), resolved_protected.end());
for (const auto& entry : prefetch_blocks_) {
if (entry.second != nullptr) {
protected_states.insert(entry.second->states.begin(), entry.second->states.end());
}
}
std::unordered_set<TensorState*> eviction_states;
auto add_evictable_state = [&](TensorState* state) {
if (state == nullptr || state->compute_backend != compute_backend ||
state->pin_count > 0 || protected_states.find(state) != protected_states.end()) {
return;
}
const bool reloadable = state->residency_mode == ResidencyMode::Disk ||
state->compute_backend != state->params_backend;
const bool resident = state->compute_backend == state->params_backend
? state->loaded_to_params_backend
: state->staged_to_compute_backend;
if (reloadable && resident) {
eviction_states.insert(state);
}
};
auto release_eviction_states = [&]() {
release_compute_staging_blocks(false, &eviction_states);
release_params_storage_blocks(false, &eviction_states);
return fits();
};
for (const auto& candidate_params : preferred_eviction_order) {
std::vector<TensorState*> candidate_states;
if (!resolve_required_tensor_states(candidate_params, candidate_states)) {
return false;
}
for (TensorState* state : candidate_states) {
add_evictable_state(state);
}
if (release_eviction_states()) {
return true;
}
}
std::vector<TensorState*> global_candidates;
global_candidates.reserve(tensor_states_.size());
for (const auto& state : tensor_states_) {
if (state != nullptr && eviction_states.find(state.get()) == eviction_states.end()) {
global_candidates.push_back(state.get());
}
}
std::stable_sort(global_candidates.begin(),
global_candidates.end(),
[](const TensorState* lhs, const TensorState* rhs) {
return lhs->last_use_epoch < rhs->last_use_epoch;
});
for (TensorState* state : global_candidates) {
add_evictable_state(state);
if (release_eviction_states()) {
return true;
}
}
const auto capacity = check_capacity(request, required_states);
LOG_WARN("model manager cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %.2f MB device / %.2f MB budget",
ggml_backend_name(compute_backend),
capacity.required_device_bytes / (1024.0 * 1024.0),
capacity.required_budget_bytes / (1024.0 * 1024.0),
capacity.available_device_bytes / (1024.0 * 1024.0),
capacity.available_budget_bytes / (1024.0 * 1024.0));
return false;
}
src/model_manager.h
+74 -23
View File
@@ -9,10 +9,10 @@
#include <unordered_set>
#include <vector>
#include "device_residency_manager.h"
#include "model_loader.h"
#include "weight_manager.h"
class ModelManager : public RunnerWeightManager {
class ModelManager : public DeviceResidencyManager {
public:
enum class ResidencyMode {
Disk,
@@ -28,24 +28,27 @@ public:
};
private:
static constexpr size_t MAX_RESIDENCY_BLOCK_BYTES = 64ULL * 1024ULL * 1024ULL;
struct TensorState {
std::string name;
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 params_follow_compute_backend = false;
bool metadata_validated = false;
enum ggml_op usage_op = GGML_OP_NONE;
ResidencyMode residency_mode = ResidencyMode::ParamBackend;
ggml_backend_t compute_backend = nullptr;
ggml_backend_t params_backend = nullptr;
ggml_backend_buffer_type_t split_buffer_type = nullptr;
bool params_follow_compute_backend = false;
bool metadata_validated = false;
enum ggml_op usage_op = GGML_OP_NONE;
int active_prepare_count = 0;
int pin_count = 0;
bool loaded_to_params_backend = false;
bool staged_to_compute_backend = false;
uint64_t applied_lora_epoch = UINT64_MAX;
uint64_t last_use_epoch = 0;
};
struct ParamsStorageBlock {
@@ -66,9 +69,12 @@ private:
ggml_backend_t compute_backend = nullptr;
ggml_backend_t transfer_backend = nullptr;
ggml_backend_event_t event = nullptr;
ggml_context* staging_ctx = nullptr;
ggml_backend_buffer_t buffer = nullptr;
std::vector<std::pair<TensorState*, ggml_tensor*>> staged_tensors;
std::vector<std::unique_ptr<ComputeStagingBlock>> staging_blocks;
};
struct RuntimeResidency {
ggml_backend_t compute_backend = nullptr;
size_t resident_bytes = 0;
};
ModelLoader model_loader_;
@@ -77,16 +83,22 @@ private:
std::vector<std::unique_ptr<ParamsStorageBlock>> params_storage_blocks_;
std::vector<std::unique_ptr<ComputeStagingBlock>> compute_staging_blocks_;
std::map<ggml_backend_t, ggml_backend_buffer_type_t> split_buffer_types_;
std::map<ggml_backend_buffer_type_t, std::vector<std::pair<ggml_backend_t, size_t>>> split_buffer_devices_;
std::map<uintptr_t, std::unique_ptr<PrefetchBlock>> prefetch_blocks_;
std::map<ggml_backend_t, ggml_backend_t> prefetch_backends_;
std::map<std::pair<uintptr_t, ggml_backend_t>, RuntimeResidency> runtime_residencies_;
std::map<uintptr_t, std::function<bool()>> workspace_reclaimers_;
bool warned_split_lora_skip_ = false;
std::set<std::string> common_ignore_tensors_;
std::vector<LoraSpec> loras_;
SDVersion lora_version_ = VERSION_COUNT;
uint64_t current_lora_epoch_ = 0;
int n_threads_ = 0;
bool enable_mmap_ = false;
bool writable_mmap_ = false;
SDVersion lora_version_ = VERSION_COUNT;
uint64_t current_lora_epoch_ = 0;
uint64_t residency_epoch_ = 0;
int n_threads_ = 0;
bool enable_mmap_ = false;
bool writable_mmap_ = false;
bool segmented_compute_disabled_ = false;
bool prefetch_disabled_ = false;
void finish_compute_backend_usage(const std::vector<TensorState*>& states);
void release_all();
@@ -99,7 +111,8 @@ private:
void release_prefetch();
bool resolve_required_tensor_states(const std::vector<ggml_tensor*>& tensors,
std::vector<TensorState*>& required_states) const;
std::vector<TensorState*>& required_states,
ggml_backend_t compute_backend = nullptr) const;
bool should_ignore(const TensorState& state) const;
bool is_optional_missing_tensor(const std::string& name) const;
bool validate_tensor(const TensorState& state) const;
@@ -113,6 +126,21 @@ private:
std::vector<ParamsStorageBlock*>& created_storage_blocks);
bool load_tensors(const std::vector<TensorState*>& states);
bool stage_tensors_to_compute_backend(const std::vector<TensorState*>& states);
size_t compute_backend_alloc_size(const std::vector<TensorState*>& states,
bool missing_only) const;
size_t compute_backend_resident_bytes(ggml_backend_t compute_backend) const;
struct CapacityCheck {
size_t required_device_bytes = 0;
size_t required_budget_bytes = 0;
size_t available_device_bytes = SIZE_MAX;
size_t available_budget_bytes = SIZE_MAX;
bool fits() const {
return required_device_bytes <= available_device_bytes &&
required_budget_bytes <= available_budget_bytes;
}
};
CapacityCheck check_capacity(const DeviceMemoryRequest& request,
const std::vector<TensorState*>& states) const;
ggml_backend_buffer_type_t params_buffer_type_for(const TensorState& state) const;
ggml_backend_buffer_type_t split_buffer_type_for(const TensorState& state) const;
@@ -124,6 +152,8 @@ private:
void free_params_storage_block(ParamsStorageBlock& block);
void erase_params_storage_block(ParamsStorageBlock* block);
void reset_lora_applied_params();
size_t other_runtime_resident_bytes(uintptr_t owner_id,
ggml_backend_t compute_backend) const;
public:
~ModelManager() override;
@@ -135,11 +165,15 @@ public:
n_threads_ = n_threads;
model_loader_.set_n_threads(n_threads);
}
void set_segmented_compute_disabled(bool disabled) {
segmented_compute_disabled_ = disabled;
}
void set_prefetch_disabled(bool disabled) { prefetch_disabled_ = disabled; }
void set_enable_mmap(bool enable_mmap) { enable_mmap_ = enable_mmap; }
void set_writable_mmap(bool writable_mmap) { writable_mmap_ = writable_mmap; }
void set_common_ignore_tensors(std::set<std::string> ignore_tensors);
void set_loras(std::vector<LoraSpec> loras, SDVersion version);
void set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft);
void set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft, const std::vector<std::pair<ggml_backend_t, size_t>>& device_limits);
static bool tensor_shape_supports_split_buffer(const ggml_tensor* tensor);
@@ -199,10 +233,27 @@ public:
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 set_workspace_reclaimer(uintptr_t owner_id, std::function<bool()> reclaim) override;
void remove_runtime_owner(uintptr_t owner_id) override;
bool fits_compute_backend_capacity(const DeviceMemoryRequest& request,
const std::vector<ggml_tensor*>& required_params) const override;
bool segmented_compute_enabled() const override { return !segmented_compute_disabled_; }
bool prefetch_enabled() const override { return !prefetch_disabled_; }
void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) override;
void release_params_backend_params(const std::vector<ggml_tensor*>& tensors) override;
bool prefetch_params(uintptr_t owner_id,
const std::vector<ggml_tensor*>& tensors) override;
void evict_compute_backend_params(const std::vector<ggml_tensor*>& tensors) override;
WeightResidencyInfo inspect_compute_backend_params(
const std::vector<ggml_tensor*>& tensors) const override;
void update_runtime_residency(uintptr_t owner_id,
ggml_backend_t compute_backend,
size_t resident_bytes) override;
bool ensure_compute_backend_capacity(
const DeviceMemoryRequest& request,
const std::vector<ggml_tensor*>& required_params,
const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
const std::vector<ggml_tensor*>& protected_params) override;
WeightPrefetchResult prefetch_params(
uintptr_t owner_id,
const std::vector<ggml_tensor*>& tensors) override;
bool activate_prefetched_params(uintptr_t owner_id,
const std::vector<ggml_tensor*>& tensors) override;
void clear_prefetched_params(uintptr_t owner_id) override;
src/model_manager_prefetch.cpp
+170 -88
View File
@@ -41,15 +41,21 @@ void ModelManager::synchronize_prefetch_block(PrefetchBlock& block) {
void ModelManager::free_prefetch_block(PrefetchBlock& block) {
synchronize_prefetch_block(block);
block.staged_tensors.clear();
if (block.buffer != nullptr) {
ggml_backend_buffer_free(block.buffer);
block.buffer = nullptr;
}
if (block.staging_ctx != nullptr) {
ggml_free(block.staging_ctx);
block.staging_ctx = nullptr;
for (auto& staging_block : block.staging_blocks) {
if (staging_block == nullptr) {
continue;
}
staging_block->staged_tensors.clear();
if (staging_block->buffer != nullptr) {
ggml_backend_buffer_free(staging_block->buffer);
staging_block->buffer = nullptr;
}
if (staging_block->staging_ctx != nullptr) {
ggml_free(staging_block->staging_ctx);
staging_block->staging_ctx = nullptr;
}
}
block.staging_blocks.clear();
}
bool ModelManager::populate_prefetch_block(PrefetchBlock& block) {
@@ -62,26 +68,13 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) {
return false;
}
ggml_init_params init_params;
init_params.mem_size = block.states.size() * ggml_tensor_overhead();
init_params.mem_buffer = nullptr;
init_params.no_alloc = true;
block.staging_ctx = ggml_init(init_params);
if (block.staging_ctx == nullptr) {
return false;
}
block.staged_tensors.reserve(block.states.size());
for (TensorState* state : block.states) {
if (state == nullptr || state->tensor == nullptr ||
state->tensor->buffer == nullptr || state->tensor->data == nullptr ||
state->params_backend == nullptr || state->staged_to_compute_backend ||
state->active_prepare_count > 0) {
state->pin_count > 0) {
return false;
}
ggml_tensor* staging_tensor = ggml_dup_tensor(block.staging_ctx, state->tensor);
ggml_set_name(staging_tensor, state->tensor->name);
block.staged_tensors.push_back({state, staging_tensor});
}
ggml_backend_buffer_type_t buffer_type =
@@ -89,38 +82,91 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) {
if (buffer_type == nullptr) {
return false;
}
block.buffer = ggml_backend_alloc_ctx_tensors_from_buft(block.staging_ctx, buffer_type);
if (block.buffer == nullptr) {
return false;
const size_t alignment = ggml_backend_buft_get_alignment(buffer_type);
size_t backend_limit = ggml_backend_buft_get_max_size(buffer_type);
if (!ggml_backend_buft_is_host(buffer_type) &&
(backend_limit == 0 || backend_limit > MAX_RESIDENCY_BLOCK_BYTES)) {
backend_limit = MAX_RESIDENCY_BLOCK_BYTES;
}
ggml_backend_buffer_set_usage(block.buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
for (const auto& pair : block.staged_tensors) {
TensorState* state = pair.first;
ggml_tensor* staging_tensor = pair.second;
const bool host_source = ggml_backend_buffer_is_host(state->tensor->buffer);
if (host_source &&
(!ggml_is_contiguous(state->tensor) || !ggml_is_contiguous(staging_tensor) ||
ggml_nbytes(state->tensor) != ggml_nbytes(staging_tensor))) {
auto enqueue_chunk = [&](const std::vector<TensorState*>& chunk) -> bool {
if (chunk.empty()) {
return true;
}
ggml_init_params init_params;
init_params.mem_size = chunk.size() * ggml_tensor_overhead();
init_params.mem_buffer = nullptr;
init_params.no_alloc = true;
ggml_context* staging_ctx = ggml_init(init_params);
if (staging_ctx == nullptr) {
return false;
}
}
for (const auto& pair : block.staged_tensors) {
TensorState* state = pair.first;
ggml_tensor* staging_tensor = pair.second;
if (ggml_backend_buffer_is_host(state->tensor->buffer)) {
ggml_backend_tensor_set_async(block.transfer_backend,
staging_tensor,
state->tensor->data,
0,
ggml_nbytes(state->tensor));
} else {
ggml_backend_tensor_copy_async(state->params_backend,
block.transfer_backend,
state->tensor,
staging_tensor);
auto staging_block = std::make_unique<ComputeStagingBlock>();
staging_block->compute_backend = block.compute_backend;
staging_block->staging_ctx = staging_ctx;
staging_block->staged_tensors.reserve(chunk.size());
for (TensorState* state : chunk) {
ggml_tensor* staging_tensor = ggml_dup_tensor(staging_ctx, state->tensor);
ggml_set_name(staging_tensor, state->tensor->name);
if (ggml_backend_buffer_is_host(state->tensor->buffer) &&
(!ggml_is_contiguous(state->tensor) || !ggml_is_contiguous(staging_tensor) ||
ggml_nbytes(state->tensor) != ggml_nbytes(staging_tensor))) {
ggml_free(staging_ctx);
staging_block->staging_ctx = nullptr;
return false;
}
staging_block->staged_tensors.push_back({state, staging_tensor});
}
staging_block->buffer =
ggml_backend_alloc_ctx_tensors_from_buft(staging_ctx, buffer_type);
if (staging_block->buffer == nullptr) {
ggml_free(staging_ctx);
staging_block->staging_ctx = nullptr;
return false;
}
ggml_backend_buffer_set_usage(staging_block->buffer,
GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
for (const auto& pair : staging_block->staged_tensors) {
TensorState* state = pair.first;
ggml_tensor* staging_tensor = pair.second;
const bool host_source = ggml_backend_buffer_is_host(state->tensor->buffer);
if (host_source) {
ggml_backend_tensor_set_async(block.transfer_backend,
staging_tensor,
state->tensor->data,
0,
ggml_nbytes(state->tensor));
} else {
ggml_backend_tensor_copy_async(state->params_backend,
block.transfer_backend,
state->tensor,
staging_tensor);
}
}
block.staging_blocks.push_back(std::move(staging_block));
return true;
};
std::vector<TensorState*> chunk;
size_t chunk_size = 0;
for (TensorState* state : block.states) {
const size_t tensor_size = GGML_PAD(
ggml_backend_buft_get_alloc_size(buffer_type, state->tensor), alignment);
if (!chunk.empty() && backend_limit > 0 &&
tensor_size > backend_limit - std::min(chunk_size, backend_limit)) {
if (!enqueue_chunk(chunk)) {
return false;
}
chunk.clear();
chunk_size = 0;
}
chunk.push_back(state);
chunk_size = tensor_size > SIZE_MAX - chunk_size ? SIZE_MAX : chunk_size + tensor_size;
}
if (!enqueue_chunk(chunk)) {
return false;
}
ggml_backend_dev_t device = ggml_backend_get_device(block.transfer_backend);
@@ -129,48 +175,79 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) {
ggml_backend_event_record(block.event, block.transfer_backend);
}
LOG_DEBUG("model manager queued layer prefetch (%6.2f MB, %zu tensors) to %s",
ggml_backend_buffer_get_size(block.buffer) / (1024.f * 1024.f),
size_t total_size = 0;
for (const auto& staging_block : block.staging_blocks) {
if (staging_block != nullptr && staging_block->buffer != nullptr) {
const size_t buffer_size = ggml_backend_buffer_get_size(staging_block->buffer);
total_size = buffer_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + buffer_size;
}
}
LOG_DEBUG("model manager queued segment prefetch (%6.2f MB, %zu tensors) to %s",
total_size / (1024.f * 1024.f),
block.states.size(),
ggml_backend_name(block.compute_backend));
return true;
}
bool ModelManager::prefetch_params(uintptr_t owner_id,
const std::vector<ggml_tensor*>& tensors) {
clear_prefetched_params(owner_id);
WeightPrefetchResult ModelManager::prefetch_params(
uintptr_t owner_id,
const std::vector<ggml_tensor*>& tensors) {
if (tensors.empty()) {
return true;
return WeightPrefetchResult::AlreadyResident;
}
std::vector<TensorState*> required_states;
if (!resolve_required_tensor_states(tensors, required_states) ||
!load_tensors_to_params_backend(required_states)) {
return false;
if (!resolve_required_tensor_states(tensors, required_states)) {
return WeightPrefetchResult::Failed;
}
std::vector<TensorState*> states;
states.reserve(required_states.size());
ggml_backend_t compute_backend = nullptr;
bool needs_synchronous_load = false;
for (TensorState* state : required_states) {
if (state == nullptr || should_ignore(*state) ||
is_optional_missing_tensor(state->name) ||
state->compute_backend == state->params_backend ||
state->staged_to_compute_backend || state->active_prepare_count > 0) {
is_optional_missing_tensor(state->name)) {
continue;
}
if (state->compute_backend == state->params_backend) {
needs_synchronous_load = needs_synchronous_load ||
!state->loaded_to_params_backend;
continue;
}
if (state->staged_to_compute_backend || state->pin_count > 0) {
continue;
}
// Split buffers cannot use the primary device's asynchronous upload path.
if (split_buffer_type_for(*state) != nullptr) {
return WeightPrefetchResult::Unsupported;
}
if (compute_backend == nullptr) {
compute_backend = state->compute_backend;
} else if (compute_backend != state->compute_backend) {
return false;
return WeightPrefetchResult::Failed;
}
states.push_back(state);
}
if (states.empty()) {
return true;
return needs_synchronous_load ? WeightPrefetchResult::Unsupported
: WeightPrefetchResult::AlreadyResident;
}
if (compute_backend == nullptr || sd_backend_is_cpu(compute_backend)) {
return false;
return WeightPrefetchResult::Unsupported;
}
ggml_backend_dev_t compute_device = ggml_backend_get_device(compute_backend);
ggml_backend_dev_props compute_props{};
if (compute_device == nullptr) {
return WeightPrefetchResult::Unsupported;
}
ggml_backend_dev_get_props(compute_device, &compute_props);
if (!compute_props.caps.async) {
return WeightPrefetchResult::Unsupported;
}
clear_prefetched_params(owner_id);
if (!load_tensors_to_params_backend(states)) {
return WeightPrefetchResult::Failed;
}
auto block = std::make_unique<PrefetchBlock>();
@@ -178,10 +255,10 @@ bool ModelManager::prefetch_params(uintptr_t owner_id,
block->compute_backend = compute_backend;
if (!populate_prefetch_block(*block)) {
free_prefetch_block(*block);
return false;
return WeightPrefetchResult::Failed;
}
prefetch_blocks_[owner_id] = std::move(block);
return true;
return WeightPrefetchResult::Scheduled;
}
bool ModelManager::activate_prefetched_params(
@@ -214,32 +291,37 @@ bool ModelManager::activate_prefetched_params(
prefetch_blocks_.erase(existing);
synchronize_prefetch_block(*block);
for (const auto& pair : block->staged_tensors) {
TensorState* state = pair.first;
ggml_tensor* staging_tensor = pair.second;
if (state == nullptr || state->tensor == nullptr || staging_tensor == nullptr ||
state->staged_to_compute_backend || state->active_prepare_count > 0) {
free_prefetch_block(*block);
return false;
for (const auto& staging_block : block->staging_blocks) {
if (staging_block == nullptr) {
continue;
}
for (const auto& pair : staging_block->staged_tensors) {
TensorState* state = pair.first;
ggml_tensor* staging_tensor = pair.second;
if (state == nullptr || state->tensor == nullptr || staging_tensor == nullptr ||
state->staged_to_compute_backend || state->pin_count > 0) {
free_prefetch_block(*block);
return false;
}
}
}
for (auto& pair : block->staged_tensors) {
TensorState* state = pair.first;
ggml_tensor* staging_tensor = pair.second;
std::swap(state->tensor->buffer, staging_tensor->buffer);
std::swap(state->tensor->data, staging_tensor->data);
std::swap(state->tensor->extra, staging_tensor->extra);
state->staged_to_compute_backend = true;
const uint64_t use_epoch = ++residency_epoch_;
for (auto& staging_block : block->staging_blocks) {
if (staging_block == nullptr) {
continue;
}
for (auto& pair : staging_block->staged_tensors) {
TensorState* state = pair.first;
ggml_tensor* staging_tensor = pair.second;
std::swap(state->tensor->buffer, staging_tensor->buffer);
std::swap(state->tensor->data, staging_tensor->data);
std::swap(state->tensor->extra, staging_tensor->extra);
state->staged_to_compute_backend = true;
state->last_use_epoch = use_epoch;
}
compute_staging_blocks_.push_back(std::move(staging_block));
}
auto staging_block = std::make_unique<ComputeStagingBlock>();
staging_block->compute_backend = block->compute_backend;
staging_block->buffer = block->buffer;
staging_block->staging_ctx = block->staging_ctx;
staging_block->staged_tensors = std::move(block->staged_tensors);
block->buffer = nullptr;
block->staging_ctx = nullptr;
compute_staging_blocks_.push_back(std::move(staging_block));
block->staging_blocks.clear();
return true;
}
src/stable-diffusion.cpp
+62 -68
View File
@@ -248,9 +248,9 @@ public:
sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr};
bool enable_mmap = false;
sd::ggml_graph_cut::MaxVramAssignment max_vram_assignment;
bool stream_layers = false;
bool disable_prefetch = false;
bool eager_load = false;
bool disable_prefetch = false;
bool disable_segmented_compute = false;
bool eager_load = false;
std::string backend_spec;
std::string params_backend_spec;
std::string split_mode_spec;
@@ -435,7 +435,11 @@ public:
if (split_buft == nullptr) {
return fall_back_to_layer_split("backend has no split buffer type");
}
model_manager->set_split_buffer_type(main_backend, split_buft);
std::vector<std::pair<ggml_backend_t, size_t>> split_device_limits;
for (auto backend : module_backends) {
split_device_limits.emplace_back(backend, max_vram_assignment.bytes_for_backend(backend));
}
model_manager->set_split_buffer_type(main_backend, split_buft, split_device_limits);
std::map<std::string, ggml_tensor*> split_tensors;
if constexpr (std::is_base_of_v<Conditioner, T>) {
@@ -599,6 +603,8 @@ public:
version,
"",
model_manager);
control_net->set_max_graph_vram_bytes(
max_graph_vram_bytes_for_module(SDBackendModule::CONTROL_NET));
if (diffusion_conv_direct) {
LOG_INFO("Using Conv2d direct in the control net");
control_net->set_conv2d_direct_enabled(true);
@@ -860,15 +866,15 @@ public:
}
bool init(const sd_ctx_params_t* sd_ctx_params) {
n_threads = sd_ctx_params->n_threads;
enable_mmap = sd_ctx_params->enable_mmap;
stream_layers = sd_ctx_params->stream_layers;
disable_prefetch = sd_ctx_params->disable_prefetch;
eager_load = sd_ctx_params->eager_load;
backend_spec = SAFE_STR(sd_ctx_params->backend);
params_backend_spec = SAFE_STR(sd_ctx_params->params_backend);
split_mode_spec = SAFE_STR(sd_ctx_params->split_mode);
auto_fit_enabled = sd_ctx_params->auto_fit;
n_threads = sd_ctx_params->n_threads;
enable_mmap = sd_ctx_params->enable_mmap;
disable_prefetch = sd_ctx_params->disable_prefetch;
disable_segmented_compute = sd_ctx_params->disable_segmented_compute;
eager_load = sd_ctx_params->eager_load;
backend_spec = SAFE_STR(sd_ctx_params->backend);
params_backend_spec = SAFE_STR(sd_ctx_params->params_backend);
split_mode_spec = SAFE_STR(sd_ctx_params->split_mode);
auto_fit_enabled = sd_ctx_params->auto_fit;
max_vram_assignment.reset(0.f);
{
std::string error;
@@ -897,6 +903,8 @@ public:
model_manager = std::make_shared<ModelManager>();
model_manager->set_n_threads(n_threads);
model_manager->set_enable_mmap(enable_mmap);
model_manager->set_segmented_compute_disabled(disable_segmented_compute);
model_manager->set_prefetch_disabled(disable_prefetch);
ModelLoader& model_loader = model_manager->loader();
if (!init_model_loader(model_loader, sd_ctx_params, use_tae, use_audio_vae, use_control_net)) {
@@ -931,10 +939,6 @@ public:
return false;
}
}
if (stream_layers && !backend_manager.params_backend_is_cpu(SDBackendModule::DIFFUSION)) {
LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring");
stream_layers = false;
}
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;
@@ -984,7 +988,8 @@ public:
}
// Avoid full-model LoRA merge buffers on constrained setups.
const bool params_offloaded = params_backend_for(SDBackendModule::DIFFUSION) != backend_for(SDBackendModule::DIFFUSION);
const bool streaming_constrained = stream_layers || params_offloaded;
const bool streaming_constrained = params_offloaded ||
backend_manager.params_backend_is_disk(SDBackendModule::DIFFUSION);
if (have_quantized_weight || streaming_constrained || row_split_active()) {
apply_lora_immediately = false;
} else {
@@ -1357,8 +1362,6 @@ public:
}
diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION));
diffusion_model->set_stream_layers_enabled(stream_layers);
diffusion_model->set_layer_prefetch_enabled(!disable_prefetch);
if (!register_runner_params("Diffusion model",
diffusion_model,
SDBackendModule::DIFFUSION,
@@ -1368,8 +1371,6 @@ public:
if (high_noise_diffusion_model) {
high_noise_diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION));
high_noise_diffusion_model->set_stream_layers_enabled(stream_layers);
high_noise_diffusion_model->set_layer_prefetch_enabled(!disable_prefetch);
if (!register_runner_params("High noise diffusion model",
high_noise_diffusion_model,
SDBackendModule::DIFFUSION,
@@ -1574,6 +1575,8 @@ public:
version,
"",
model_manager);
control_net->set_max_graph_vram_bytes(
max_graph_vram_bytes_for_module(SDBackendModule::CONTROL_NET));
if (sd_ctx_params->diffusion_conv_direct) {
LOG_INFO("Using Conv2d direct in the control net");
control_net->set_conv2d_direct_enabled(true);
@@ -1902,15 +1905,15 @@ public:
}
bool is_using_v_parameterization_for_sd2(bool is_inpaint = false) {
struct RunnerDoneOnExit {
struct RunnerEndOnExit {
GGMLRunner* runner = nullptr;
~RunnerDoneOnExit() {
~RunnerEndOnExit() {
if (runner != nullptr) {
runner->runner_done();
runner->runner_end();
}
}
};
RunnerDoneOnExit diffusion_runner_done{diffusion_model.get()};
RunnerEndOnExit diffusion_runner_end{diffusion_model.get()};
sd::Tensor<float> x_t = sd::full<float>({8, 8, 4, 1}, 0.5f);
sd::Tensor<float> c = sd::full<float>({1024, 2, 1, 1}, 0.5f);
@@ -2541,17 +2544,17 @@ public:
const sd_cache_params_t* cache_params,
bool preview_final_step,
const sd::Tensor<float>& video_positions = {}) {
struct RunnerDoneOnExit {
struct RunnerEndOnExit {
GGMLRunner* runner = nullptr;
~RunnerDoneOnExit() {
~RunnerEndOnExit() {
if (runner != nullptr) {
runner->runner_done();
runner->runner_end();
}
}
};
RunnerDoneOnExit sample_diffusion_runner_done{work_diffusion_model.get()};
RunnerEndOnExit sample_diffusion_runner_end{work_diffusion_model.get()};
RunnerDoneOnExit sample_control_runner_done{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr};
RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr};
std::vector<int> skip_layers(guidance.slg.layers, guidance.slg.layers + guidance.slg.layer_count);
float cfg_scale = guidance.txt_cfg;
@@ -2923,10 +2926,6 @@ public:
LOG_ERROR("Diffusion model sampling failed");
if (control_net) {
control_net->free_control_ctx();
control_net->free_compute_buffer();
}
if (work_diffusion_model) {
work_diffusion_model->free_compute_buffer();
}
return {};
}
@@ -2940,10 +2939,6 @@ public:
if (control_net) {
control_net->free_control_ctx();
control_net->free_compute_buffer();
}
if (work_diffusion_model) {
work_diffusion_model->free_compute_buffer();
}
return x0;
}
@@ -3092,7 +3087,6 @@ public:
while (decoded.empty() &&
auto_fit_enabled &&
sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) {
first_stage_model->free_compute_buffer();
decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
}
return decoded;
@@ -3567,27 +3561,27 @@ void sd_hires_params_init(sd_hires_params_t* hires_params) {
}
void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
*sd_ctx_params = {};
sd_ctx_params->n_threads = sd_get_num_physical_cores();
sd_ctx_params->wtype = SD_TYPE_COUNT;
sd_ctx_params->rng_type = CUDA_RNG;
sd_ctx_params->sampler_rng_type = RNG_TYPE_COUNT;
sd_ctx_params->prediction = PREDICTION_COUNT;
sd_ctx_params->lora_apply_mode = LORA_APPLY_AUTO;
sd_ctx_params->max_vram = nullptr;
sd_ctx_params->stream_layers = false;
sd_ctx_params->disable_prefetch = false;
sd_ctx_params->eager_load = false;
sd_ctx_params->enable_mmap = false;
sd_ctx_params->diffusion_flash_attn = false;
sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO;
sd_ctx_params->backend = nullptr;
sd_ctx_params->params_backend = nullptr;
sd_ctx_params->split_mode = nullptr;
sd_ctx_params->auto_fit = false;
sd_ctx_params->rpc_servers = nullptr;
sd_ctx_params->model_args = nullptr;
sd_ctx_params->pulid_weights_path = nullptr;
*sd_ctx_params = {};
sd_ctx_params->n_threads = sd_get_num_physical_cores();
sd_ctx_params->wtype = SD_TYPE_COUNT;
sd_ctx_params->rng_type = CUDA_RNG;
sd_ctx_params->sampler_rng_type = RNG_TYPE_COUNT;
sd_ctx_params->prediction = PREDICTION_COUNT;
sd_ctx_params->lora_apply_mode = LORA_APPLY_AUTO;
sd_ctx_params->max_vram = nullptr;
sd_ctx_params->disable_prefetch = false;
sd_ctx_params->disable_segmented_compute = false;
sd_ctx_params->eager_load = false;
sd_ctx_params->enable_mmap = false;
sd_ctx_params->diffusion_flash_attn = false;
sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO;
sd_ctx_params->backend = nullptr;
sd_ctx_params->params_backend = nullptr;
sd_ctx_params->split_mode = nullptr;
sd_ctx_params->auto_fit = false;
sd_ctx_params->rpc_servers = nullptr;
sd_ctx_params->model_args = nullptr;
sd_ctx_params->pulid_weights_path = nullptr;
}
char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
@@ -3621,8 +3615,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
"sampler_rng_type: %s\n"
"prediction: %s\n"
"max_vram: %s\n"
"stream_layers: %s\n"
"disable_prefetch: %s\n"
"disable_segmented_compute: %s\n"
"eager_load: %s\n"
"backend: %s\n"
"params_backend: %s\n"
@@ -3656,8 +3650,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
sd_rng_type_name(sd_ctx_params->sampler_rng_type),
sd_prediction_name(sd_ctx_params->prediction),
SAFE_STR(sd_ctx_params->max_vram),
BOOL_STR(sd_ctx_params->stream_layers),
BOOL_STR(sd_ctx_params->disable_prefetch),
BOOL_STR(sd_ctx_params->disable_segmented_compute),
BOOL_STR(sd_ctx_params->eager_load),
SAFE_STR(sd_ctx_params->backend),
SAFE_STR(sd_ctx_params->params_backend),
@@ -4809,11 +4803,11 @@ struct ImageGenerationEmbeds {
SDCondition img_uncond;
};
struct ConditionerRunnerDoneOnExit {
struct ConditionerRunnerEndOnExit {
Conditioner* conditioner = nullptr;
~ConditionerRunnerDoneOnExit() {
~ConditionerRunnerEndOnExit() {
if (conditioner != nullptr) {
conditioner->runner_done();
conditioner->runner_end();
}
}
};
@@ -5239,7 +5233,7 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
SamplePlan* plan,
ImageGenerationLatents* latents,
const RefImageParams& ref_image_params) {
ConditionerRunnerDoneOnExit conditioner_runner_done{sd_ctx->sd->cond_stage_model.get()};
ConditionerRunnerEndOnExit conditioner_runner_end{sd_ctx->sd->cond_stage_model.get()};
ConditionerParams condition_params;
condition_params.text = request->prompt;
@@ -6536,7 +6530,7 @@ static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx,
const sd_vid_gen_params_t* sd_vid_gen_params,
const GenerationRequest& request,
const ImageGenerationLatents& latents) {
ConditionerRunnerDoneOnExit conditioner_runner_done{sd_ctx->sd->cond_stage_model.get()};
ConditionerRunnerEndOnExit conditioner_runner_end{sd_ctx->sd->cond_stage_model.get()};
ImageGenerationEmbeds embeds;
ConditionerParams condition_params;
src/upscaler.cpp
+1 -9
View File
@@ -32,13 +32,6 @@ void UpscalerGGML::set_max_graph_vram_bytes(size_t max_vram_bytes) {
}
}
void UpscalerGGML::set_stream_layers_enabled(bool enabled) {
stream_layers_enabled = enabled;
if (esrgan_upscaler) {
esrgan_upscaler->set_stream_layers_enabled(enabled);
}
}
bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
int n_threads) {
ggml_log_set(ggml_log_callback_default, nullptr);
@@ -94,7 +87,6 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
return false;
}
esrgan_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes);
esrgan_upscaler->set_stream_layers_enabled(stream_layers_enabled);
if (direct) {
esrgan_upscaler->set_conv2d_direct_enabled(true);
}
@@ -139,7 +131,7 @@ sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_te
false,
on_processing);
}
esrgan_upscaler->free_compute_buffer();
esrgan_upscaler->runner_end();
if (upscaled.empty()) {
LOG_ERROR("esrgan compute failed");
return {};
src/upscaler.h
-2
View File
@@ -20,7 +20,6 @@ struct UpscalerGGML {
bool direct = false;
int tile_size = 128;
size_t max_graph_vram_bytes = 0;
bool stream_layers_enabled = false;
std::string backend_spec;
std::string params_backend_spec;
@@ -34,7 +33,6 @@ struct UpscalerGGML {
bool load_from_file(const std::string& esrgan_path,
int n_threads);
void set_max_graph_vram_bytes(size_t max_vram_bytes);
void set_stream_layers_enabled(bool enabled);
sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor);
sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor);
};
src/weight_manager.h
-25
View File
@@ -1,25 +0,0 @@
#ifndef __WEIGHT_MANAGER_H__
#define __WEIGHT_MANAGER_H__
#include <cstdint>
#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;
virtual bool prefetch_params(uintptr_t owner_id,
const std::vector<ggml_tensor*>& tensors) = 0;
virtual bool activate_prefetched_params(uintptr_t owner_id,
const std::vector<ggml_tensor*>& tensors) = 0;
virtual void clear_prefetched_params(uintptr_t owner_id) = 0;
};
#endif // __WEIGHT_MANAGER_H__