Skip to content

refactor: route all runner params through model manager #1649

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
leejet merged 1 commit into from
Jun 13, 2026
Merged

refactor: route all runner params through model manager #1649

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
146 changes: 48 additions & 98 deletions src/conditioning/conditioner.hpp
View file
Open in desktop

Large diffs are not rendered by default.

102 changes: 5 additions & 97 deletions src/core/ggml_extend.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1696,11 +1696,9 @@ struct GGMLRunner {
using GraphCutSegment = sd::ggml_graph_cut::Segment;
using GraphCutPlan = sd::ggml_graph_cut::Plan;

ggml_backend_t params_backend = nullptr;
ggml_backend_t runtime_backend = nullptr;

ggml_context* params_ctx = nullptr;
ggml_backend_buffer_t params_buffer = nullptr;
ggml_context* params_ctx = nullptr;

ggml_context* cache_ctx = nullptr;
ggml_backend_buffer_t cache_buffer = nullptr;
Expand Down Expand Up @@ -1880,9 +1878,6 @@ struct GGMLRunner {
auto manager = weight_manager.lock();
if (manager == nullptr) {
if (!params_to_prepare.empty()) {
if (params_buffer != nullptr) {
return true;
}
LOG_ERROR("%s weight manager is not set for graph params", get_desc().c_str());
return false;
}
Expand Down Expand Up @@ -2194,13 +2189,11 @@ struct GGMLRunner {
plan.valid &&
max_graph_vram_bytes > 0 &&
plan.segments.size() > 1 &&
params_backend != runtime_backend &&
!sd_backend_is_cpu(runtime_backend);
}

bool can_attempt_graph_cut_segmented_compute() const {
return max_graph_vram_bytes > 0 &&
params_backend != runtime_backend &&
!sd_backend_is_cpu(runtime_backend);
}

Expand Down Expand Up @@ -2631,16 +2624,15 @@ struct GGMLRunner {
public:
virtual std::string get_desc() = 0;

GGMLRunner(ggml_backend_t backend, ggml_backend_t params_backend)
: params_backend(params_backend),
runtime_backend(backend) {
GGMLRunner(ggml_backend_t backend,
std::shared_ptr<RunnerWeightManager> manager = nullptr)
: runtime_backend(backend),
weight_manager(manager) {
GGML_ASSERT(runtime_backend != nullptr);
GGML_ASSERT(params_backend != nullptr);
alloc_params_ctx();
}

virtual ~GGMLRunner() {
free_params_buffer();
free_compute_buffer();
free_params_ctx();
free_compute_ctx();
Expand Down Expand Up @@ -2674,73 +2666,6 @@ struct GGMLRunner {
alloc_compute_ctx();
}

bool alloc_params_buffer() {
size_t num_tensors = ggml_tensor_num(params_ctx);
if (num_tensors > 0) {
// ggml_backend_alloc_ctx_tensors fails when all tensors are already allocated
// (typical for memory-mapped weights). See ggml-alloc.c n_buffers==0 branch.
bool all_have_data = true;
for (ggml_tensor* t = ggml_get_first_tensor(params_ctx); t != nullptr; t = ggml_get_next_tensor(params_ctx, t)) {
if (t->data == nullptr) {
all_have_data = false;
break;
}
}
if (all_have_data) {
LOG_DEBUG("%s all params already mmap-allocated (no separate buffer needed)", get_desc().c_str());
params_buffer = nullptr;
rebuild_params_tensor_set();
return true;
}
} else {
LOG_DEBUG("%s skipping params allocation (no tensors)", get_desc().c_str());
return true;
}
// Pinned host buffer when CPU-offloaded for DMA-direct H2D.
ggml_backend_buffer_type_t params_buft = nullptr;
if (params_backend != runtime_backend) {
ggml_backend_dev_t runtime_dev = ggml_backend_get_device(runtime_backend);
if (runtime_dev != nullptr) {
params_buft = ggml_backend_dev_host_buffer_type(runtime_dev);
}
}
if (params_buft == nullptr) {
params_buft = ggml_backend_get_default_buffer_type(params_backend);
}
params_buffer = ggml_backend_alloc_ctx_tensors_from_buft(params_ctx, params_buft);
if (params_buffer == nullptr) {
LOG_ERROR("%s alloc params backend buffer failed, num_tensors = %i",
get_desc().c_str(),
num_tensors);
return false;
}
rebuild_params_tensor_set();
ggml_backend_buffer_set_usage(params_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
size_t params_buffer_size = ggml_backend_buffer_get_size(params_buffer);
LOG_DEBUG("%s params backend buffer size = % 6.2f MB(%s) (%i tensors)",
get_desc().c_str(),
params_buffer_size / (1024.f * 1024.f),
sd_backend_is_cpu(params_backend) ? "RAM" : "VRAM",
num_tensors);
return true;
}

protected:
void free_params_buffer() {
if (params_buffer != nullptr) {
ggml_backend_buffer_free(params_buffer);
params_buffer = nullptr;
}
observed_max_effective_budget_ = 0;
}

size_t get_params_buffer_size() {
if (params_buffer != nullptr) {
return ggml_backend_buffer_get_size(params_buffer);
}
return 0;
}

public:
void free_cache_ctx_and_buffer() {
free_cache_buffer();
Expand Down Expand Up @@ -2886,30 +2811,13 @@ struct GGMLRunner {
weight_adapter = adapter;
}

void set_weight_manager(const std::shared_ptr<RunnerWeightManager>& manager) {
weight_manager = manager;
}

void set_weight_manager(const std::shared_ptr<RunnerWeightManager>& manager,
const std::string&) {
set_weight_manager(manager);
}

void set_max_graph_vram_bytes(size_t max_vram_bytes) {
max_graph_vram_bytes = max_vram_bytes;
}

void set_stream_layers_enabled(bool enabled) {
stream_layers_enabled = enabled;
}

ggml_backend_t get_runtime_backend() {
return runtime_backend;
}

ggml_backend_t get_params_backend() {
return params_backend;
}
};

class GGMLBlock {
Expand Down
2 changes: 1 addition & 1 deletion src/extensions/generation_extension.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ struct GenerationExtensionInitContext {
SDVersion version;
const String2TensorStorage& tensor_storage_map;
ModelLoader& model_loader;
std::shared_ptr<ModelManager> model_manager;
int n_threads;
std::function<bool(SDBackendModule)> ensure_backend_pair;
std::function<ggml_backend_t(SDBackendModule)> backend_for;
Expand Down Expand Up @@ -46,7 +47,6 @@ 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 set_weight_manager(const std::shared_ptr<RunnerWeightManager>&) {}
virtual void runner_done() {}
virtual void reset_runtime_condition() {}
virtual bool prepare_condition(GenerationExtensionConditionContext&) {
Expand Down
11 changes: 3 additions & 8 deletions src/extensions/photomaker_extension.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,11 +134,12 @@ struct PhotoMakerExtension : public GenerationExtension {
}

pmid_model = std::make_shared<PhotoMakerIDEncoder>(ctx.backend_for(SDBackendModule::PHOTOMAKER),
ctx.params_backend_for(SDBackendModule::PHOTOMAKER),
ctx.tensor_storage_map,
"pmid",
ctx.version,
pm_version);
pm_version,
20.f,
ctx.model_manager);
if (pm_version == PM_VERSION_2) {
LOG_INFO("using PhotoMaker Version 2");
}
Expand Down Expand Up @@ -174,12 +175,6 @@ struct PhotoMakerExtension : public GenerationExtension {
ignore_tensors.insert("pmid.unet.");
}

void set_weight_manager(const std::shared_ptr<RunnerWeightManager>& manager) override {
if (pmid_model != nullptr) {
pmid_model->set_weight_manager(manager);
}
}

void runner_done() override {
if (pmid_model != nullptr) {
pmid_model->runner_done();
Expand Down
59 changes: 42 additions & 17 deletions src/model/adapter/lora.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
#include <mutex>
#include "core/ggml_extend.hpp"
#include "model_loader.h"
#include "model_manager.h"

#define LORA_GRAPH_BASE_SIZE 10240

Expand All @@ -14,22 +15,24 @@ struct LoraModel : public GGMLRunner {
std::map<ggml_tensor*, ggml_tensor*> original_tensor_to_final_tensor;
std::set<std::string> applied_lora_tensors;
std::string file_path;
ModelLoader model_loader;
bool load_failed = false;
bool applied = false;
bool tensor_preprocessed = false;
std::shared_ptr<ModelManager> model_manager;
ggml_backend_t params_backend = nullptr;
bool load_failed = false;
bool applied = false;
bool tensor_preprocessed = false;

typedef std::function<bool(const std::string&)> filter_t;

LoraModel(const std::string& lora_id,
ggml_backend_t backend,
ggml_backend_t params_backend,
const std::string& file_path = "",
std::string prefix = "",
SDVersion version = VERSION_COUNT)
: lora_id(lora_id), file_path(file_path), GGMLRunner(backend, params_backend) {
ggml_backend_t params_backend_,
const std::string& file_path = "",
std::string prefix = "",
SDVersion version = VERSION_COUNT,
std::shared_ptr<ModelManager> manager = std::make_shared<ModelManager>())
: GGMLRunner(backend, manager), lora_id(lora_id), file_path(file_path), model_manager(std::move(manager)), params_backend(params_backend_) {
prefix = "lora." + prefix;
if (!model_loader.init_from_file_and_convert_name(file_path, prefix, version)) {
if (model_manager == nullptr || !model_manager->loader().init_from_file_and_convert_name(file_path, prefix, version)) {
load_failed = true;
}
}
Expand Down Expand Up @@ -71,7 +74,10 @@ struct LoraModel : public GGMLRunner {
return true;
};

model_loader.set_n_threads(n_threads);
if (model_manager != nullptr) {
model_manager->set_n_threads(n_threads);
}
ModelLoader& model_loader = model_manager->loader();
model_loader.load_tensors(on_new_tensor_cb);

if (tensors_to_create.empty()) {
Expand All @@ -88,23 +94,42 @@ struct LoraModel : public GGMLRunner {
lora_tensors[name] = real;
}

if (!alloc_params_buffer()) {
LOG_ERROR("lora model buffer allocation failed");
std::map<std::string, ggml_tensor*> tensors;
for (const auto& pair : lora_tensors) {
tensors[pair.first] = pair.second;
}
if (model_manager == nullptr ||
!model_manager->register_param_tensors("LoRA",
std::move(tensors),
ModelManager::ResidencyMode::Resident,
runtime_backend,
params_backend) ||
!model_manager->validate_registered_tensors()) {
LOG_ERROR("lora model manager registration failed");
return false;
}
std::vector<ggml_tensor*> lora_params;
lora_params.reserve(lora_tensors.size());
for (const auto& pair : lora_tensors) {
lora_params.push_back(pair.second);
}
if (!model_manager->prepare_params(lora_params)) {
LOG_ERROR("lora model manager prepare params failed");
return false;
}

dry_run = false;
model_loader.load_tensors(on_new_tensor_cb);

LOG_DEBUG("finished loaded lora");
return true;
}

void release_loaded_tensors() {
runner_done();
free_compute_buffer();
free_params_buffer();
model_manager.reset();
free_params_ctx();
alloc_params_ctx();
model_manager = std::make_shared<ModelManager>();
weight_manager = model_manager;
lora_tensors.clear();
original_tensor_to_final_tensor.clear();
applied_lora_tensors.clear();
Expand Down
55 changes: 34 additions & 21 deletions src/model/adapter/pmid.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -413,13 +413,13 @@ struct PhotoMakerIDEncoder : public GGMLRunner {

public:
PhotoMakerIDEncoder(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map,
const std::string prefix,
SDVersion version = VERSION_SDXL,
PMVersion pm_v = PM_VERSION_1,
float sty = 20.f)
: GGMLRunner(backend, params_backend),
SDVersion version = VERSION_SDXL,
PMVersion pm_v = PM_VERSION_1,
float sty = 20.f,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager),
version(version),
pm_version(pm_v),
style_strength(sty) {
Expand Down Expand Up @@ -565,17 +565,18 @@ struct PhotoMakerIDEncoder : public GGMLRunner {
struct PhotoMakerIDEmbed : public GGMLRunner {
std::map<std::string, ggml_tensor*> tensors;
std::string file_path;
ModelLoader* model_loader;
bool load_failed = false;
bool applied = false;
std::shared_ptr<ModelManager> model_manager;
ggml_backend_t params_backend = nullptr;
bool load_failed = false;
bool applied = false;

PhotoMakerIDEmbed(ggml_backend_t backend,
ggml_backend_t params_backend,
ModelLoader* ml,
const std::string& file_path = "",
const std::string& prefix = "")
: file_path(file_path), GGMLRunner(backend, params_backend), model_loader(ml) {
if (!model_loader->init_from_file_and_convert_name(file_path, prefix)) {
ggml_backend_t params_backend_,
std::shared_ptr<ModelManager> manager = std::make_shared<ModelManager>(),
const std::string& file_path = "",
const std::string& prefix = "")
: GGMLRunner(backend, manager), file_path(file_path), model_manager(std::move(manager)), params_backend(params_backend_) {
if (model_manager == nullptr || !model_manager->loader().init_from_file_and_convert_name(file_path, prefix)) {
load_failed = true;
}
}
Expand Down Expand Up @@ -616,15 +617,27 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
return true;
};

model_loader->set_n_threads(n_threads);
model_loader->load_tensors(on_new_tensor_cb);
if (!alloc_params_buffer()) {
LOG_ERROR("PhotoMaker ID embeds buffer allocation failed");
model_manager->set_n_threads(n_threads);
ModelLoader& model_loader = model_manager->loader();
model_loader.load_tensors(on_new_tensor_cb);
if (!model_manager->register_param_tensors("PhotoMaker ID embeds",
tensors,
ModelManager::ResidencyMode::Resident,
runtime_backend,
params_backend) ||
!model_manager->validate_registered_tensors()) {
LOG_ERROR("PhotoMaker ID embeds model manager registration failed");
return false;
}
std::vector<ggml_tensor*> id_embed_params;
id_embed_params.reserve(tensors.size());
for (const auto& pair : tensors) {
id_embed_params.push_back(pair.second);
}
if (!model_manager->prepare_params(id_embed_params)) {
LOG_ERROR("PhotoMaker ID embeds model manager prepare params failed");
return false;
}

dry_run = false;
model_loader->load_tensors(on_new_tensor_cb);

LOG_DEBUG("finished loading PhotoMaker ID Embeds ");
return true;
Expand Down
Loading
Loading