refactor: centralize runner weight staging and cleanup (#1644)

This commit is contained in:
leejet authored and GitHub committed 2026-06-13 13:19:13 +08:00
1 parent 3a54597776
commit 563137a592
43 files changed
+2335 -1729

No files matched your search

+5 -2
View File
@@ -567,11 +567,14 @@ struct CLIPTextModelRunner : public GGMLRunner {
void* custom_embeddings_data,
size_t max_token_idx,
bool return_pooled,
int clip_skip) {
int clip_skip,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = 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, true);
auto result = GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params);
if (return_pooled) {
return take_or_empty(std::move(result));
}
+10 -4
View File
@@ -1733,7 +1733,10 @@ namespace LLM {
const sd::Tensor<float>& attention_mask,
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
std::set<int> out_layers,
bool return_all_hidden_states = false) {
bool return_all_hidden_states = false,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = true) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input_ids,
attention_mask,
@@ -1741,7 +1744,7 @@ namespace LLM {
out_layers,
return_all_hidden_states);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true),
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params),
input_ids.dim() + 1);
}
@@ -1802,11 +1805,14 @@ namespace LLM {
}
sd::Tensor<float> encode_image(const int n_threads,
const sd::Tensor<float>& image) {
const sd::Tensor<float>& image,
bool auto_free = false,
bool free_compute_buffer = false,
bool free_compute_params = false) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_encode_image_graph(image);
};
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, false));
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
}
};
+5 -2
View File
@@ -394,11 +394,14 @@ 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) {
const sd::Tensor<float>& attention_mask,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = 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, true), 3);
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params), 3);
}
static std::vector<int> _relative_position_bucket(const std::vector<int>& relative_position,