Commit 37bc88b for stable-diffusion.cpp
commit 37bc88b5e97a213821949732ea66c58ae8ffc784
Author: leejet <leejet714@gmail.com>
Date: Fri Oct 30 00:38:14 2026 +0800
perf: optimize H3 inference and segmented weight prefetch (#2240)
diff --git a/examples/server/model_catalog.cpp b/examples/server/model_catalog.cpp
index 7daed53..f665567 100644
--- a/examples/server/model_catalog.cpp
+++ b/examples/server/model_catalog.cpp
@@ -671,7 +671,7 @@ void prepare_model_load(ServerRuntime& runtime, json& body, SDContextParams& con
check_keys(body, {"recipe", "models"});
const auto recipe_id = body.at("recipe").get<std::string>();
// `models` is optional; unselected slots may fall back to recipe defaults below.
- auto selection = body.value("models", json::object());
+ auto selection = body.value("models", json::object());
if (!selection.is_object())
throw std::runtime_error("models must be an object");
auto& models = *runtime.models;
diff --git a/ggml b/ggml
index 643e8d9..7b2b435 160000
--- a/ggml
+++ b/ggml
@@ -1 +1 @@
-Subproject commit 643e8d970cd2d926a0074f5cb12a68b38078157a
+Subproject commit 7b2b435da087614be772ee6f5ea80e68eca81de6
diff --git a/src/core/compute_workspace.cpp b/src/core/compute_workspace.cpp
index 03ca9bb..462a5e3 100644
--- a/src/core/compute_workspace.cpp
+++ b/src/core/compute_workspace.cpp
@@ -158,7 +158,7 @@ namespace sd {
for (const auto& entry : copies) {
const size_t source_index = ggml_hash_find(&graph->visited_hash_set, entry.first);
if (source_index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, source_index)) {
- const size_t copy_index = ggml_hash_insert(©_graph->visited_hash_set, entry.second);
+ const size_t copy_index = ggml_hash_insert(©_graph->visited_hash_set, entry.second);
copy_graph->use_counts[copy_index] = graph->use_counts[source_index];
}
}
diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp
index 9e57efa..8586d87 100644
--- a/src/core/ggml_extend_backend.cpp
+++ b/src/core/ggml_extend_backend.cpp
@@ -615,6 +615,32 @@ bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) {
return false;
}
+void sd_backend_set_buffer_cache_enabled(ggml_backend_t backend, bool enabled) {
+ auto dev = backend != nullptr ? ggml_backend_get_device(backend) : nullptr;
+ if (dev == nullptr) {
+ return;
+ }
+ auto reg = ggml_backend_dev_backend_reg(dev);
+ using SetBufferCacheEnabled = void (*)(ggml_backend_t, bool);
+ auto fn = reinterpret_cast<SetBufferCacheEnabled>(ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_buffer_cache_enabled"));
+ if (fn != nullptr) {
+ fn(backend, enabled);
+ }
+}
+
+void sd_backend_clear_buffer_cache(ggml_backend_t backend) {
+ auto dev = backend != nullptr ? ggml_backend_get_device(backend) : nullptr;
+ if (dev == nullptr) {
+ return;
+ }
+ auto reg = ggml_backend_dev_backend_reg(dev);
+ using ClearBufferCache = void (*)(ggml_backend_t);
+ auto fn = reinterpret_cast<ClearBufferCache>(ggml_backend_reg_get_proc_address(reg, "ggml_backend_clear_buffer_cache"));
+ if (fn != nullptr) {
+ fn(backend);
+ }
+}
+
static ggml_cgraph sd_ggml_graph_view(ggml_cgraph* cgraph0, int i0, int i1) {
ggml_cgraph cgraph = {
/*.size =*/0,
diff --git a/src/core/ggml_extend_backend.h b/src/core/ggml_extend_backend.h
index 295ab16..037ef0d 100644
--- a/src/core/ggml_extend_backend.h
+++ b/src/core/ggml_extend_backend.h
@@ -96,6 +96,9 @@ ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t dev
size_t max_tensor_size);
ggml_backend_t sd_backend_cpu_init();
bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
+void sd_backend_set_buffer_cache_enabled(ggml_backend_t backend, bool enabled);
+void sd_backend_clear_buffer_cache(ggml_backend_t backend);
+
ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
ggml_cgraph* gf,
sd_graph_eval_callback_t callback_eval,
diff --git a/src/core/ggml_graph_cut.cpp b/src/core/ggml_graph_cut.cpp
index b10e7a1..8454d43 100644
--- a/src/core/ggml_graph_cut.cpp
+++ b/src/core/ggml_graph_cut.cpp
@@ -753,6 +753,14 @@ namespace sd::ggml_graph_cut {
for (int node_idx : segment.internal_node_indices) {
ggml_graph_add_node(segment_graph, ggml_graph_node(gf, node_idx));
}
+ // Fusion checks need the original use counts, including consumers outside this segment.
+ for (ggml_tensor* tensor : represented) {
+ const size_t source_index = ggml_hash_find(&gf->visited_hash_set, tensor);
+ if (source_index != GGML_HASHSET_FULL && ggml_bitset_get(gf->visited_hash_set.used, source_index)) {
+ const size_t index = ggml_hash_insert(&segment_graph->visited_hash_set, tensor);
+ segment_graph->use_counts[index] = gf->use_counts[source_index];
+ }
+ }
*graph_ctx_out = graph_ctx;
return segment_graph;
}
diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp
index e47b013..7d8aeb0 100644
--- a/src/core/ggml_runner.cpp
+++ b/src/core/ggml_runner.cpp
@@ -466,6 +466,10 @@ bool GGMLRunner::runner_start() {
if (runner_started_) {
return true;
}
+ sd_backend_set_buffer_cache_enabled(runtime_backend, true);
+ for (auto backend : extra_runtime_backends) {
+ sd_backend_set_buffer_cache_enabled(backend, true);
+ }
cache_.clear();
workspace_.set_extra_backends(extra_runtime_backends);
if (auto manager = residency_manager.lock()) {
@@ -490,6 +494,7 @@ void GGMLRunner::runner_end() {
free_compute_ctx();
compute_ctx_buffer_.reset();
cache_.clear();
+ cut_cache_.clear();
logged_compute_bytes_.clear();
logged_segment_count_ = 0;
if (auto manager = residency_manager.lock()) {
@@ -503,6 +508,10 @@ void GGMLRunner::runner_end() {
manager->evict_compute_backend_params(tensors);
manager->remove_runtime_owner(reinterpret_cast<uintptr_t>(this));
}
+ sd_backend_clear_buffer_cache(runtime_backend);
+ for (auto backend : extra_runtime_backends) {
+ sd_backend_clear_buffer_cache(backend);
+ }
runner_started_ = false;
}
@@ -712,8 +721,11 @@ 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) {
+ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, const std::unordered_set<const ggml_tensor*>* external_outputs) {
auto external_backend = [&](const ggml_tensor* tensor) -> ggml_backend_t {
+ if (external_outputs != nullptr && external_outputs->count(tensor)) {
+ return runtime_backend;
+ }
if (!params_tensor_set_.count(tensor)) {
return nullptr;
}
@@ -766,7 +778,7 @@ bool GGMLRunner::fits(const std::vector<DeviceMemoryRequest>& requests,
return true;
}
-bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) {
+bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads, const std::function<void()>& on_before_compute) {
if (sd_backend_is_cpu(runtime_backend)) {
sd_backend_cpu_set_n_threads(runtime_backend, n_threads);
}
@@ -776,7 +788,27 @@ bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) {
auto scheduler = workspace_.scheduler();
GGML_ASSERT(scheduler != nullptr);
ggml_backend_sched_set_eval_callback(scheduler, sd_get_backend_eval_callback(), sd_get_backend_eval_callback_data());
- const auto status = ggml_backend_sched_graph_compute(scheduler, graph);
+#ifndef SD_USE_UPSTREAM_GGML
+ struct BeforeComputeCallback {
+ ggml_backend_t backend;
+ const std::function<void()>& callback;
+ bool invoked = false;
+ } before_compute{runtime_backend, on_before_compute};
+ ggml_backend_sched_set_compute_callback(
+ scheduler, [](ggml_backend_t backend, void* data) {
+ auto& hook = *static_cast<BeforeComputeCallback*>(data);
+ if (backend == hook.backend && !hook.invoked) {
+ hook.invoked = true;
+ hook.callback();
+ }
+ },
+ &before_compute);
+ const auto status = ggml_backend_sched_graph_compute_async(scheduler, graph);
+ ggml_backend_sched_set_compute_callback(scheduler, nullptr, nullptr);
+#else
+ on_before_compute();
+ const auto status = ggml_backend_sched_graph_compute_async(scheduler, graph);
+#endif
workspace_.synchronize();
if (status != GGML_STATUS_SUCCESS) {
last_compute_status_ = status;
@@ -828,27 +860,31 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
if (!assign_graph_cut_layer_split_backends(graph)) {
return std::nullopt;
}
- const auto params = collect_used_param_tensors(graph);
- const auto& cached_plan = resolve_graph_cut_plan(graph);
- const auto full_measurement = measure(graph);
- if (full_measurement.buffers.empty()) {
- last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
- return std::nullopt;
- }
- auto fits_monolithic = [&]() {
- // Planning headroom absorbs allocation estimate drift; execution keeps the normal limits.
- constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL;
- auto requests = memory_requests(full_measurement.buffers, cache_.pending_bytes(graph));
- for (auto& request : requests) {
- request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom);
+ const auto params = collect_used_param_tensors(graph);
+ const auto& cached_plan = resolve_graph_cut_plan(graph);
+ auto manager = residency_manager.lock();
+ const bool can_segment = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) &&
+ manager != nullptr && manager->segmented_compute_enabled() &&
+ cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1;
+ // An empty workspace request is a lower bound on the full graph's memory demand.
+ bool segmented = can_segment && !fits(memory_requests({{runtime_backend, 0}}, 0), params);
+ ComputeWorkspace::Measurement full_measurement;
+ if (!segmented) {
+ full_measurement = measure(graph);
+ if (full_measurement.buffers.empty()) {
+ last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+ return std::nullopt;
}
- return fits(requests, params);
- };
- auto manager = residency_manager.lock();
- const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) &&
- manager != nullptr && manager->segmented_compute_enabled() &&
- cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1 &&
- !fits_monolithic();
+ if (can_segment) {
+ // Planning headroom absorbs allocation estimate drift; execution keeps the normal limits.
+ constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL;
+ auto requests = memory_requests(full_measurement.buffers, cache_.pending_bytes(graph));
+ for (auto& request : requests) {
+ request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom);
+ }
+ segmented = !fits(requests, params);
+ }
+ }
ggml_graph_cut::Plan monolithic_plan;
if (!segmented) {
monolithic_plan.segments.emplace_back();
@@ -922,14 +958,24 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
}
} segment_cleanup{*this, weights, bindings, segment_context};
- auto measurement = segmented ? measure(segment_graph) : full_measurement;
+ std::unordered_set<const ggml_tensor*> cache_outputs;
+ if (segmented) {
+ for (int output_index : segment.output_node_indices) {
+ auto tensor = ggml_graph_node(graph, output_index);
+ if (ggml_graph_cut::is_graph_cut_tensor(tensor) && segment.future_cut_names.count(tensor->name) &&
+ GraphCutTensorCache::can_bind_output(runtime_backend, tensor)) {
+ cache_outputs.insert(tensor);
+ }
+ }
+ }
+ auto measurement = segmented ? measure(segment_graph, &cache_outputs) : full_measurement;
if (!workspace_.prepare(measurement)) {
last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
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 = [&]() {
+ const size_t cut_bytes = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment);
+ 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))) {
@@ -956,6 +1002,11 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
if (!ensure_capacity()) {
return fail_segment("workspace capacity check");
}
+ if (!cut_cache_.bind_outputs(cache_outputs)) {
+ last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+ return fail_segment("output cache allocation");
+ }
+ new_cache_bytes = add_bytes(last ? 0 : cut_cache_.estimate_output_bytes(graph, segment), cache_.pending_bytes(segment_graph));
if (!workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) {
pin_multi_device_nodes(scheduler, current);
})) {
@@ -978,7 +1029,7 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
}
LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(),
index + 1, plan.segments.size(), segment.group_name.c_str());
- if (!execute_segment(segment_graph, n_threads)) {
+ if (!execute_segment(segment_graph, n_threads, [&]() { weights.submit_next(); })) {
return fail_segment("execution");
}
auto cache_status = cache_.capture(segment_graph);
diff --git a/src/core/ggml_runner.h b/src/core/ggml_runner.h
index a146312..5c21b1d 100644
--- a/src/core/ggml_runner.h
+++ b/src/core/ggml_runner.h
@@ -141,12 +141,12 @@ private:
size_t logged_segment_count_ = 0;
ggml_status last_compute_status_ = GGML_STATUS_SUCCESS;
- sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph);
+ sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, const std::unordered_set<const ggml_tensor*>* external_outputs = nullptr);
std::vector<DeviceMemoryRequest> memory_requests(const std::vector<sd::BackendBufferSize>& sizes,
size_t pending_cache_bytes) const;
bool fits(const std::vector<DeviceMemoryRequest>& requests,
const std::vector<ggml_tensor*>& params) const;
- bool execute_segment(ggml_cgraph* graph, int n_threads);
+ bool execute_segment(ggml_cgraph* graph, int n_threads, const std::function<void()>& on_before_compute);
std::optional<sd::Tensor<float>> execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs);
protected:
@@ -162,8 +162,8 @@ protected:
sd::ComputeWorkspace workspace_;
ggml_context* compute_ctx = nullptr;
std::unique_ptr<std::max_align_t[]> compute_ctx_buffer_;
- bool runner_started_ = false;
- bool graph_active_ = false;
+ bool runner_started_ = false;
+ bool graph_active_ = false;
size_t max_graph_vram_bytes = 0;
bool graph_cut_layer_split_enabled = false;
diff --git a/src/core/runner_cache.cpp b/src/core/runner_cache.cpp
index b640bb4..9bbcae0 100644
--- a/src/core/runner_cache.cpp
+++ b/src/core/runner_cache.cpp
@@ -24,15 +24,7 @@ namespace sd {
ggml_free(context);
}
- std::unique_ptr<CachedTensor> CachedTensor::copy(ggml_backend_t backend,
- const std::string& name,
- ggml_tensor* source,
- ggml_status& status) {
- status = GGML_STATUS_FAILED;
- if (ggml_graph_cut::tensor_buffer(source) == nullptr) {
- return nullptr;
- }
- status = GGML_STATUS_ALLOC_FAILED;
+ std::unique_ptr<CachedTensor> CachedTensor::allocate(ggml_backend_t backend, const std::string& name, ggml_tensor* source) {
auto entry = std::make_unique<CachedTensor>();
entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true});
if (entry->context == nullptr) {
@@ -46,6 +38,22 @@ namespace sd {
if (entry->buffer == nullptr) {
return nullptr;
}
+ return entry;
+ }
+
+ std::unique_ptr<CachedTensor> CachedTensor::copy(ggml_backend_t backend,
+ const std::string& name,
+ ggml_tensor* source,
+ ggml_status& status) {
+ status = GGML_STATUS_FAILED;
+ if (ggml_graph_cut::tensor_buffer(source) == nullptr) {
+ return nullptr;
+ }
+ status = GGML_STATUS_ALLOC_FAILED;
+ auto entry = allocate(backend, name, source);
+ if (entry == 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());
@@ -155,6 +163,26 @@ namespace sd {
return sd::resident_bytes(tensors_, device);
}
+ bool GraphCutTensorCache::can_bind_output(ggml_backend_t backend, const ggml_tensor* tensor) {
+ return tensor != nullptr && tensor->view_src == nullptr && tensor->op != GGML_OP_NONE &&
+ ggml_is_contiguous(tensor) && ggml_backend_supports_op(backend, tensor);
+ }
+
+ bool GraphCutTensorCache::bind_outputs(const std::unordered_set<const ggml_tensor*>& outputs) {
+ for (const ggml_tensor* output : outputs) {
+ auto tensor = const_cast<ggml_tensor*>(output);
+ auto entry = CachedTensor::allocate(backend_, tensor->name, tensor);
+ if (entry == nullptr) {
+ return false;
+ }
+ tensor->buffer = entry->tensor->buffer;
+ tensor->data = entry->tensor->data;
+ tensor->extra = entry->tensor->extra;
+ tensors_[tensor->name] = std::move(entry);
+ }
+ return true;
+ }
+
size_t GraphCutTensorCache::estimate_output_bytes(
ggml_cgraph* graph,
const ggml_graph_cut::Segment& segment) const {
@@ -171,6 +199,10 @@ namespace sd {
!segment.future_cut_names.count(output->name)) {
continue;
}
+ const ggml_tensor* cached = get(output->name);
+ if (cached != nullptr && output->buffer == cached->buffer && output->data == cached->data) {
+ 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);
@@ -196,6 +228,10 @@ namespace sd {
!segment.future_cut_names.count(output->name)) {
continue;
}
+ const ggml_tensor* cached = get(output->name);
+ if (cached != nullptr && output->buffer == cached->buffer && output->data == cached->data) {
+ continue;
+ }
ggml_status status;
auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output), status);
if (entry == nullptr) {
diff --git a/src/core/runner_cache.h b/src/core/runner_cache.h
index f1b5cf9..b6fe1c7 100644
--- a/src/core/runner_cache.h
+++ b/src/core/runner_cache.h
@@ -18,6 +18,7 @@ namespace sd {
ggml_backend_buffer_t buffer = nullptr;
ggml_tensor* tensor = nullptr;
~CachedTensor();
+ static std::unique_ptr<CachedTensor> allocate(ggml_backend_t backend, const std::string& name, ggml_tensor* source);
static std::unique_ptr<CachedTensor> copy(ggml_backend_t backend,
const std::string& name,
ggml_tensor* source,
@@ -59,6 +60,8 @@ namespace sd {
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;
+ static bool can_bind_output(ggml_backend_t backend, const ggml_tensor* tensor);
+ bool bind_outputs(const std::unordered_set<const ggml_tensor*>& outputs);
ggml_status 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(); }
diff --git a/src/core/segment_weight_pipeline.cpp b/src/core/segment_weight_pipeline.cpp
index f769794..2df9a89 100644
--- a/src/core/segment_weight_pipeline.cpp
+++ b/src/core/segment_weight_pipeline.cpp
@@ -178,7 +178,7 @@ namespace sd {
disable();
return;
}
- switch (manager->prefetch_params(owner_id_, params)) {
+ switch (manager->prepare_prefetch_params(owner_id_, params)) {
case WeightPrefetchResult::Scheduled:
queued_params_ = std::move(params);
queued_segment_ = next_segment;
@@ -195,6 +195,14 @@ namespace sd {
disable();
}
+ void SegmentWeightPipeline::submit_next() {
+ if (queued_segment_ != SIZE_MAX) {
+ if (auto manager = residency_manager_.lock()) {
+ manager->submit_prefetched_params(owner_id_);
+ }
+ }
+ }
+
void SegmentWeightPipeline::clear() {
if (auto manager = residency_manager_.lock()) {
manager->clear_prefetched_params(owner_id_);
diff --git a/src/core/segment_weight_pipeline.h b/src/core/segment_weight_pipeline.h
index 598bf9f..67949ed 100644
--- a/src/core/segment_weight_pipeline.h
+++ b/src/core/segment_weight_pipeline.h
@@ -56,6 +56,7 @@ namespace sd {
// Prefetch is best effort; segment_start falls back to synchronous loading.
void enqueue_next(size_t segment_index,
const DeviceMemoryRequest& request);
+ void submit_next();
};
}
diff --git a/src/device_residency_manager.h b/src/device_residency_manager.h
index 0a2a889..faf7be2 100644
--- a/src/device_residency_manager.h
+++ b/src/device_residency_manager.h
@@ -64,7 +64,13 @@ struct DeviceResidencyManager {
const std::vector<ggml_tensor*>& protected_params) = 0;
virtual WeightPrefetchResult prefetch_params(
uintptr_t owner_id,
- const std::vector<ggml_tensor*>& tensors) = 0;
+ const std::vector<ggml_tensor*>& tensors) = 0;
+ virtual WeightPrefetchResult prepare_prefetch_params(
+ uintptr_t owner_id,
+ const std::vector<ggml_tensor*>& tensors) {
+ return prefetch_params(owner_id, tensors);
+ }
+ virtual void submit_prefetched_params(uintptr_t owner_id) {}
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;
diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp
index e93c70f..739ef77 100644
--- a/src/model/diffusion/minimax_h3.hpp
+++ b/src/model/diffusion/minimax_h3.hpp
@@ -146,13 +146,24 @@ namespace MiniMaxH3 {
blocks["fc2"] = std::make_shared<Linear>(ffn_hidden_size, hidden_size, false, false, true, 1.f / 128.f);
}
+ void init_params(ggml_context* ctx,
+ const String2TensorStorage& tensors = {},
+ const std::string prefix = "") override {
+ // INT8 accumulates in INT32 and returns FP32, so it needs no FP16 overflow scaling.
+ for (const auto* name : {"fc1", "fc2"}) {
+ auto linear = std::dynamic_pointer_cast<Linear>(blocks[name]);
+ linear->set_scale(get_type(prefix + name + ".weight", tensors, GGML_TYPE_F32) == GGML_TYPE_I8
+ ? 1.f
+ : 1.f / 128.f);
+ }
+ }
+
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
- auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
- auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
- auto uv = ggml_ext_chunk(ctx->ggml_ctx, fc1->forward(ctx, x), 2, 0);
- return fc2->forward(ctx, ggml_mul(ctx->ggml_ctx,
- ggml_silu(ctx->ggml_ctx, uv[0]),
- uv[1]));
+ auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
+ auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
+ auto projected = fc1->forward(ctx, x);
+ auto activated = ggml_swiglu(ctx->ggml_ctx, projected);
+ return fc2->forward(ctx, activated);
}
};
@@ -167,6 +178,14 @@ namespace MiniMaxH3 {
auto ctx = runner_ctx->ggml_ctx;
int64_t rot_dim = pe->ne[2] * 2;
GGML_ASSERT(rot_dim <= x->ne[0]);
+#ifndef SD_USE_UPSTREAM_GGML
+ if (runner_ctx->backend != nullptr && x->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32) {
+ auto out = ggml_rope_apply(ctx, x, pe, false);
+ if (ggml_backend_supports_op(runner_ctx->backend, out)) {
+ return out;
+ }
+ }
+#endif
auto rotated = Rope::apply_rope(runner_ctx,
ggml_ext_slice(ctx, x, 0, 0, rot_dim),
pe,
@@ -194,22 +213,31 @@ namespace MiniMaxH3 {
blocks["out_proj"] = std::make_shared<Linear>(inner, hidden_size, false);
}
- ggml_tensor* forward(GGMLRunnerContext* ctx,
- ggml_tensor* x,
- ggml_tensor* pe = nullptr) {
- auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv_proj"]);
- auto q_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
- auto k_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
- auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out_proj"]);
-
- int64_t sequence = x->ne[1];
- int64_t batch = x->ne[2] * x->ne[3];
- auto qkv = ggml_ext_chunk(ctx->ggml_ctx, qkv_proj->forward(ctx, x), 3, 0);
- auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv[0], head_dim, heads, sequence, batch);
- auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv[1], head_dim, heads, sequence, batch);
- auto v = ggml_reshape_4d(ctx->ggml_ctx, qkv[2], head_dim, heads, sequence, batch);
- q = q_norm->forward(ctx, q);
- k = k_norm->forward(ctx, k);
+ ggml_tensor* project_qkv(GGMLRunnerContext* ctx, ggml_tensor* x) {
+ return std::dynamic_pointer_cast<Linear>(blocks["qkv_proj"])->forward(ctx, x);
+ }
+
+ ggml_tensor* project_output(GGMLRunnerContext* ctx, ggml_tensor* x) {
+ return std::dynamic_pointer_cast<Linear>(blocks["out_proj"])->forward(ctx, x);
+ }
+
+ ggml_tensor* attend(GGMLRunnerContext* ctx,
+ ggml_tensor* projected,
+ ggml_tensor* pe = nullptr) {
+ auto q_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
+ auto k_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
+ int64_t sequence = projected->ne[1];
+ int64_t batch = projected->ne[2] * projected->ne[3];
+ auto view = [&](int index) {
+ return ggml_view_4d(ctx->ggml_ctx, projected, head_dim, heads, sequence, batch,
+ head_dim * projected->nb[0], projected->nb[1], projected->nb[2],
+ index * heads * head_dim * projected->nb[0]);
+ };
+ auto q = view(0);
+ auto k = view(1);
+ auto v = view(2);
+ q = q_norm->forward(ctx, q);
+ k = k_norm->forward(ctx, k);
if (pe != nullptr) {
q = apply_partial_rope(ctx, q, pe);
k = apply_partial_rope(ctx, k, pe);
@@ -226,7 +254,13 @@ namespace MiniMaxH3 {
true,
ctx->flash_attn_enabled,
1.f / 128.f);
- return out_proj->forward(ctx, out);
+ return out;
+ }
+
+ ggml_tensor* forward(GGMLRunnerContext* ctx,
+ ggml_tensor* x,
+ ggml_tensor* pe = nullptr) {
+ return project_output(ctx, attend(ctx, project_qkv(ctx, x), pe));
}
};
@@ -340,50 +374,6 @@ namespace MiniMaxH3 {
return ggml_ext_chunk(ctx, selected, expand, 0);
}
- static ggml_tensor* modulate_segments(ggml_context* ctx,
- ggml_tensor* x,
- ggml_tensor* projection,
- const std::vector<TokenModulationSpan>& segments,
- int64_t hidden_size,
- int expand,
- int modalities,
- int shift_index,
- int scale_index) {
- ggml_tensor* out = nullptr;
- for (const auto& segment : segments) {
- auto mods = modulation_row(ctx,
- projection,
- hidden_size,
- expand,
- modalities,
- segment.modulation_row);
- auto part = ggml_ext_slice(ctx, x, 1, segment.start, segment.end);
- part = ggml_add(ctx,
- ggml_add(ctx, part, ggml_mul(ctx, part, mods[scale_index])),
- mods[shift_index]);
- out = out == nullptr ? part : ggml_concat(ctx, out, part, 1);
- }
- return out;
- }
-
- static ggml_tensor* gated_residual_segments(ggml_context* ctx,
- ggml_tensor* x,
- ggml_tensor* update,
- ggml_tensor* projection,
- const std::vector<TokenModulationSpan>& segments,
- int64_t hidden_size,
- int gate_index) {
- ggml_tensor* out = nullptr;
- for (const auto& segment : segments) {
- auto mods = modulation_row(ctx, projection, hidden_size, 6, 3, segment.modulation_row);
- auto base = ggml_ext_slice(ctx, x, 1, segment.start, segment.end);
- auto add = ggml_ext_slice(ctx, update, 1, segment.start, segment.end);
- auto part = ggml_add(ctx, base, ggml_mul(ctx, add, mods[gate_index]));
- out = out == nullptr ? part : ggml_concat(ctx, out, part, 1);
- }
- return out;
- }
-
struct TransformerBlock : public GGMLBlock {
Config config;
@@ -417,38 +407,43 @@ namespace MiniMaxH3 {
auto adaln = std::dynamic_pointer_cast<AdaLayerNormModulation>(blocks["adaln_proj"]);
auto mods = adaln->forward(ctx, t_emb);
- auto h = modulate_segments(ctx->ggml_ctx,
- norm1->forward(ctx, x),
- mods,
- segments,
- config.hidden_size,
- 6,
- 3,
- 0,
- 1);
- x = gated_residual_segments(ctx->ggml_ctx,
- x,
- attn->forward(ctx, h, pe),
- mods,
- segments,
- config.hidden_size,
- 2);
- h = modulate_segments(ctx->ggml_ctx,
- norm2->forward(ctx, x),
- mods,
- segments,
- config.hidden_size,
- 6,
- 3,
- 3,
- 4);
- return gated_residual_segments(ctx->ggml_ctx,
- x,
- mlp->forward(ctx, h),
- mods,
- segments,
- config.hidden_size,
- 5);
+ std::vector<std::vector<ggml_tensor*>> rows;
+ std::vector<ggml_tensor*> inputs;
+ for (const auto& segment : segments) {
+ rows.push_back(modulation_row(ctx->ggml_ctx, mods, config.hidden_size,
+ 6, 3, segment.modulation_row));
+ inputs.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 1, segment.start, segment.end, false));
+ }
+ auto modulate = [&](ggml_tensor* value, const std::vector<ggml_tensor*>& row, int shift, int scale) {
+ return ggml_add(ctx->ggml_ctx,
+ ggml_add(ctx->ggml_ctx, value, ggml_mul(ctx->ggml_ctx, value, row[scale])),
+ row[shift]);
+ };
+ ggml_tensor* qkv = nullptr;
+ for (size_t i = 0; i < segments.size(); ++i) {
+ auto h = modulate(norm1->forward(ctx, inputs[i]), rows[i], 0, 1);
+ h = attn->project_qkv(ctx, h);
+ if (qkv == nullptr) {
+ qkv = h;
+ continue;
+ }
+ qkv = ggml_concat(ctx->ggml_ctx, qkv, h, 1);
+ }
+ auto attended = attn->attend(ctx, qkv, pe);
+ std::vector<ggml_tensor*> projections;
+ for (const auto& segment : segments) {
+ auto part = ggml_ext_slice(ctx->ggml_ctx, attended, 1, segment.start, segment.end, false);
+ projections.push_back(attn->project_output(ctx, part));
+ }
+ ggml_tensor* output = nullptr;
+ for (size_t i = 0; i < segments.size(); ++i) {
+ auto residual = ggml_add(ctx->ggml_ctx, inputs[i],
+ ggml_mul(ctx->ggml_ctx, projections[i], rows[i][2]));
+ auto h = modulate(norm2->forward(ctx, residual), rows[i], 3, 4);
+ h = ggml_add(ctx->ggml_ctx, residual, ggml_mul(ctx->ggml_ctx, mlp->forward(ctx, h), rows[i][5]));
+ output = output == nullptr ? h : ggml_concat(ctx->ggml_ctx, output, h, 1);
+ }
+ return output;
}
};
@@ -956,8 +951,6 @@ namespace MiniMaxH3 {
struct MiniMaxH3Runner : public DiffusionModelRunner {
Config config;
MiniMaxH3Transformer3DModel model;
- sd::Tensor<float> video_input_cache;
- sd::Tensor<float> audio_input_cache;
sd::Tensor<float> position_input_cache;
sd::Tensor<float> timestep_feature_input_cache;
sd::Tensor<int32_t> curve_index_input_cache;
@@ -983,29 +976,6 @@ namespace MiniMaxH3 {
model.get_param_tensors(tensors, prefix);
}
- std::pair<sd::Tensor<float>, sd::Tensor<float>> split_av_latents(const sd::Tensor<float>& packed,
- int audio_length) const {
- GGML_ASSERT(packed.dim() == 4 || packed.dim() == 5);
- int64_t spatial = packed.shape()[0] * packed.shape()[1] * packed.shape()[2];
- int64_t video_values = spatial * config.video_latent_channels;
- sd::Tensor<float> video({packed.shape()[0],
- packed.shape()[1],
- packed.shape()[2],
- config.video_latent_channels,
- 1});
- std::copy_n(packed.data(), static_cast<size_t>(video_values), video.data());
- if (audio_length <= 0) {
- return {video, {}};
- }
- int64_t audio_values = audio_length * 2 * config.audio_latent_channels;
- GGML_ASSERT(packed.numel() >= video_values + audio_values);
- sd::Tensor<float> audio({audio_length, 2, config.audio_latent_channels, 1});
- std::copy_n(packed.data() + video_values,
- static_cast<size_t>(audio_values),
- audio.data());
- return {video, audio};
- }
-
ggml_tensor* merge_av_latents(ggml_context* ctx,
ggml_tensor* video,
ggml_tensor* audio) const {
@@ -1036,14 +1006,26 @@ namespace MiniMaxH3 {
int audio_length,
float video_shift,
float audio_shift) {
- auto split = split_av_latents(packed, audio_length);
- video_input_cache = std::move(split.first);
- audio_input_cache = std::move(split.second);
- GGML_ASSERT(!audio_input_cache.empty());
+ GGML_ASSERT(packed.dim() == 4 || packed.dim() == 5);
+ GGML_ASSERT(audio_length > 0);
GGML_ASSERT(!context_tensor.empty());
-
- auto video = make_input(video_input_cache);
- auto audio_carrier = make_input(audio_input_cache);
+ const int64_t spatial = packed.shape()[0] * packed.shape()[1] * packed.shape()[2];
+ const int64_t video_values = spatial * config.video_latent_channels;
+ const int64_t audio_values = int64_t(audio_length) * 2 * config.audio_latent_channels;
+ GGML_ASSERT(packed.numel() >= video_values + audio_values);
+ auto packed_input = make_input(packed);
+ auto video = ggml_view_4d(compute_ctx, packed_input,
+ packed.shape()[0], packed.shape()[1], packed.shape()[2],
+ config.video_latent_channels,
+ packed.shape()[0] * sizeof(float),
+ packed.shape()[0] * packed.shape()[1] * sizeof(float),
+ spatial * sizeof(float), 0);
+ auto audio_carrier = ggml_view_4d(compute_ctx, packed_input,
+ audio_length, 2, config.audio_latent_channels, 1,
+ audio_length * sizeof(float),
+ int64_t(audio_length) * 2 * sizeof(float),
+ audio_values * sizeof(float),
+ video_values * sizeof(float));
auto context = make_input(context_tensor);
std::vector<ggml_tensor*> condition_inputs;
condition_inputs.reserve(condition_videos.size());
@@ -1065,9 +1047,9 @@ namespace MiniMaxH3 {
// latent follows one sigma schedule. Restore x_a for the H3 network.
auto audio = ggml_ext_scale(compute_ctx, audio_carrier, sigma_a / sigma_v);
auto layout = build_layout(context_tensor.shape()[1],
- video_input_cache.shape()[2],
- video_input_cache.shape()[1],
- video_input_cache.shape()[0],
+ packed.shape()[2],
+ packed.shape()[1],
+ packed.shape()[0],
audio_length,
condition_videos,
condition_audios,
diff --git a/src/model_manager.cpp b/src/model_manager.cpp
index 243811d..d07d63a 100644
--- a/src/model_manager.cpp
+++ b/src/model_manager.cpp
@@ -1767,10 +1767,16 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
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);
// Backend scratch buffers and pipelines are not included in graph measurements.
- constexpr size_t safety_margin = 512ULL * 1024ULL * 1024ULL;
- result.required_device_bytes = add(add(request.pending_allocation_bytes, missing), safety_margin);
- result.required_budget_bytes = add(request.runtime_peak_bytes(), missing);
- auto available_device_bytes = [&](ggml_backend_t backend) {
+ size_t safety_margin = 512ULL * 1024ULL * 1024ULL;
+#ifdef _WIN32
+ if (sd_backend_is(request.compute_backend, "CUDA")) {
+ // Avoid WDDM allocation churn when transient buffers compete with reusable staging buffers.
+ safety_margin = 1024ULL * 1024ULL * 1024ULL;
+ }
+#endif
+ result.required_device_bytes = add(add(request.pending_allocation_bytes, missing), safety_margin);
+ result.required_budget_bytes = add(request.runtime_peak_bytes(), missing);
+ auto available_device_bytes = [&](ggml_backend_t backend) {
auto device = ggml_backend_get_device(backend);
if (device == nullptr) {
return SIZE_MAX;
diff --git a/src/model_manager.h b/src/model_manager.h
index fdb1108..81d2546 100644
--- a/src/model_manager.h
+++ b/src/model_manager.h
@@ -76,6 +76,7 @@ private:
ggml_backend_t compute_backend = nullptr;
ggml_backend_t transfer_backend = nullptr;
ggml_backend_event_t event = nullptr;
+ bool submitted = false;
std::vector<std::unique_ptr<ComputeStagingBlock>> staging_blocks;
};
@@ -292,6 +293,10 @@ public:
WeightPrefetchResult prefetch_params(
uintptr_t owner_id,
const std::vector<ggml_tensor*>& tensors) override;
+ WeightPrefetchResult prepare_prefetch_params(
+ uintptr_t owner_id,
+ const std::vector<ggml_tensor*>& tensors) override;
+ void submit_prefetched_params(uintptr_t owner_id) 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;
diff --git a/src/model_manager_prefetch.cpp b/src/model_manager_prefetch.cpp
index de96d5a..4056b8f 100644
--- a/src/model_manager_prefetch.cpp
+++ b/src/model_manager_prefetch.cpp
@@ -33,7 +33,7 @@ void ModelManager::synchronize_prefetch_block(PrefetchBlock& block) {
ggml_backend_event_synchronize(block.event);
ggml_backend_event_free(block.event);
block.event = nullptr;
- } else if (block.transfer_backend != nullptr) {
+ } else if (block.submitted && block.transfer_backend != nullptr) {
ggml_backend_synchronize(block.transfer_backend);
}
block.transfer_backend = nullptr;
@@ -89,7 +89,7 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) {
backend_limit = MAX_RESIDENCY_BLOCK_BYTES;
}
- auto enqueue_chunk = [&](const std::vector<TensorState*>& chunk) -> bool {
+ auto allocate_chunk = [&](const std::vector<TensorState*>& chunk) -> bool {
if (chunk.empty()) {
return true;
}
@@ -128,23 +128,6 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) {
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;
};
@@ -156,7 +139,7 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) {
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)) {
+ if (!allocate_chunk(chunk)) {
return false;
}
chunk.clear();
@@ -165,10 +148,39 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) {
chunk.push_back(state);
chunk_size = tensor_size > SIZE_MAX - chunk_size ? SIZE_MAX : chunk_size + tensor_size;
}
- if (!enqueue_chunk(chunk)) {
+ if (!allocate_chunk(chunk)) {
return false;
}
+ return true;
+}
+
+void ModelManager::submit_prefetched_params(uintptr_t owner_id) {
+ auto existing = prefetch_blocks_.find(owner_id);
+ if (existing == prefetch_blocks_.end() || existing->second->submitted) {
+ return;
+ }
+ auto& block = *existing->second;
+ block.submitted = true;
+ for (const auto& staging_block : block.staging_blocks) {
+ for (const auto& pair : staging_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);
+ }
+ }
+ }
+
ggml_backend_dev_t device = ggml_backend_get_device(block.transfer_backend);
block.event = ggml_backend_event_new(device);
if (block.event != nullptr) {
@@ -186,10 +198,19 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) {
total_size / (1024.f * 1024.f),
block.states.size(),
ggml_backend_name(block.compute_backend));
- return true;
}
WeightPrefetchResult ModelManager::prefetch_params(
+ uintptr_t owner_id,
+ const std::vector<ggml_tensor*>& tensors) {
+ auto result = prepare_prefetch_params(owner_id, tensors);
+ if (result == WeightPrefetchResult::Scheduled) {
+ submit_prefetched_params(owner_id);
+ }
+ return result;
+}
+
+WeightPrefetchResult ModelManager::prepare_prefetch_params(
uintptr_t owner_id,
const std::vector<ggml_tensor*>& tensors) {
if (tensors.empty()) {
@@ -288,6 +309,7 @@ bool ModelManager::activate_prefetched_params(
if (existing == prefetch_blocks_.end()) {
return false;
}
+ submit_prefetched_params(owner_id);
std::unique_ptr<PrefetchBlock> block = std::move(existing->second);
prefetch_blocks_.erase(existing);
synchronize_prefetch_block(*block);