perf: optimize H3 inference and segmented weight prefetch (#2240)

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;

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(&copy_graph->visited_hash_set, entry.second);
+                const size_t copy_index            = ggml_hash_insert(&copy_graph->visited_hash_set, entry.second);
                 copy_graph->use_counts[copy_index] = graph->use_counts[source_index];
             }
         }

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,

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,

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;
     }

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);

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;

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) {

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(); }

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_);

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();
     };
 }
 

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;

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"]);
+        ggml_tensor* project_qkv(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            return std::dynamic_pointer_cast<Linear>(blocks["qkv_proj"])->forward(ctx, x);
+        }
 
-            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_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,

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;

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;

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,11 +128,45 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) {
         ggml_backend_buffer_set_usage(staging_block->buffer,
                                       GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
 
+        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 (!allocate_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 (!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;
-            const bool host_source      = ggml_backend_buffer_is_host(state->tensor->buffer);
-            if (host_source) {
+            if (ggml_backend_buffer_is_host(state->tensor->buffer)) {
                 ggml_backend_tensor_set_async(block.transfer_backend,
                                               staging_tensor,
                                               state->tensor->data,
@@ -145,28 +179,6 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) {
                                                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);
@@ -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);