perf: accelerate INT8 diffusion inference with CUDA fusions (#2235)

src/core/compute_workspace.cpp

@@ -1,16 +1,13 @@
 #include "core/compute_workspace.h"
 
 #include <algorithm>
-#include <cstring>
 #include <exception>
 #include <map>
 #include <unordered_map>
 #include <unordered_set>
 
 #include "core/ggml_extend_backend.h"
-#include "core/ggml_graph_cut.h"
 #include "core/util.h"
-#include "ggml-cpu.h"
 #include "ggml-impl.h"
 
 namespace sd {
@@ -28,18 +25,6 @@ namespace sd {
         }
     }
 
-    bool ComputeWorkspace::needs_scheduler(ggml_cgraph* graph) const {
-        if (!extra_backends_.empty()) {
-            return true;
-        }
-        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
-            if (!ggml_backend_supports_op(backend_, ggml_graph_node(graph, i))) {
-                return true;
-            }
-        }
-        return false;
-    }
-
     ggml_backend_sched_t ComputeWorkspace::make_scheduler(size_t graph_size) {
         std::vector<ggml_backend_t> backends{backend_};
         backends.insert(backends.end(), extra_backends_.begin(), extra_backends_.end());
@@ -63,10 +48,6 @@ namespace sd {
                                       graph_size, false, false);
     }
 
-    bool ComputeWorkspace::measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const {
-        return measurement.scheduler == needs_scheduler(graph);
-    }
-
     bool ComputeWorkspace::prepare(const Measurement& measurement) {
         GGML_ASSERT(!active_);
         if (measurement.buffers.empty()) {
@@ -74,7 +55,7 @@ namespace sd {
         }
         const bool grows = std::any_of(measurement.buffers.begin(), measurement.buffers.end(),
                                        [&](const BackendBufferSize& size) { return size.bytes > bytes(size.backend); });
-        if (measurement.scheduler != (scheduler_ != nullptr) || grows) {
+        if (grows) {
             release();
         }
         return true;
@@ -88,47 +69,21 @@ namespace sd {
 
     bool ComputeWorkspace::allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes) {
         GGML_ASSERT(!active_);
-        const bool use_scheduler = needs_scheduler(graph);
-        if (use_scheduler) {
-            if (allocator_ != nullptr) {
-                release();
-            }
-            const size_t capacity = static_cast<size_t>(graph->n_nodes + graph->n_leafs) + 8;
-            if (scheduler_ == nullptr || capacity > scheduler_capacity_) {
-                release();
-                scheduler_          = make_scheduler(capacity);
-                scheduler_capacity_ = capacity;
-            }
-            if (scheduler_ == nullptr) {
-                return false;
-            }
-            ggml_backend_sched_reset(scheduler_);
-            assign_nodes(scheduler_, graph);
-            // Scheduler allocation rewrites sources. Split the execution graph only once.
-            if (!ggml_backend_sched_alloc_graph(scheduler_, graph)) {
-                release();
-                return false;
-            }
-        } else {
-            if (scheduler_ != nullptr) {
-                release();
-            }
-            if (allocator_ == nullptr) {
-                allocator_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));
-            }
-            auto signature = ggml_graph_cut::graph_layout(graph, true);
-            if (signature != reservation_) {
-                if (!ggml_gallocr_reserve(allocator_, graph)) {
-                    release();
-                    return false;
-                }
-                reservation_ = std::move(signature);
-                ++reservations_;
-            }
-            if (!ggml_gallocr_alloc_graph(allocator_, graph)) {
-                release();
-                return false;
-            }
+        const size_t capacity = static_cast<size_t>(graph->n_nodes + graph->n_leafs) + 8;
+        if (scheduler_ == nullptr || capacity > scheduler_capacity_) {
+            release();
+            scheduler_          = make_scheduler(capacity);
+            scheduler_capacity_ = capacity;
+        }
+        if (scheduler_ == nullptr) {
+            return false;
+        }
+        ggml_backend_sched_reset(scheduler_);
+        assign_nodes(scheduler_, graph);
+        // Scheduler allocation rewrites sources. Split the execution graph only once.
+        if (!ggml_backend_sched_alloc_graph(scheduler_, graph)) {
+            release();
+            return false;
         }
         active_ = true;
         return true;
@@ -136,12 +91,8 @@ namespace sd {
 
     ComputeWorkspace::Measurement ComputeWorkspace::measure(
         ggml_cgraph* graph,
-        size_t direct_bytes,
         const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
         const AssignNodes& assign_nodes) {
-        if (!needs_scheduler(graph)) {
-            return {{{backend_, direct_bytes}}, false};
-        }
         std::vector<const ggml_tensor*> tensors;
         std::unordered_set<const ggml_tensor*> seen;
         auto visit = [&](const ggml_tensor* tensor) {
@@ -202,9 +153,17 @@ namespace sd {
         for (int i = 0; i < graph->n_leafs; ++i) {
             copy_graph->leafs[i] = copies.at(graph->leafs[i]);
         }
+        copy_graph->order = graph->order;
+        // Preserve fusion metadata, including uses outside a graph view, without rebuilding its node order.
+        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);
+                copy_graph->use_counts[copy_index] = graph->use_counts[source_index];
+            }
+        }
         Measurement result;
-        result.scheduler = true;
-        auto scheduler   = make_scheduler(graph_size);
+        auto scheduler = make_scheduler(graph_size);
         if (scheduler != nullptr) {
             assign_nodes(scheduler, copy_graph);
             std::vector<size_t> sizes(extra_backends_.size() + 2);
@@ -253,12 +212,9 @@ namespace sd {
         if (active_) {
             return false;
         }
-        ggml_gallocr_free(allocator_);
-        allocator_ = nullptr;
         ggml_backend_sched_free(scheduler_);
         scheduler_          = nullptr;
         scheduler_capacity_ = 0;
-        reservation_.clear();
         return true;
     }
 
@@ -266,6 +222,6 @@ namespace sd {
         if (scheduler_ != nullptr) {
             return ggml_backend_sched_get_buffer_size(scheduler_, backend);
         }
-        return allocator_ != nullptr && backend == backend_ ? ggml_gallocr_get_buffer_size(allocator_, 0) : 0;
+        return 0;
     }
 }

src/core/compute_workspace.h

@@ -4,7 +4,6 @@
 #include <functional>
 #include <vector>
 
-#include "ggml-alloc.h"
 #include "ggml-backend.h"
 
 namespace sd {
@@ -17,20 +16,15 @@ namespace sd {
         ggml_backend_t backend_;
         std::vector<ggml_backend_t> extra_backends_;
         ggml_backend_t cpu_backend_     = nullptr;
-        ggml_gallocr_t allocator_       = nullptr;
         ggml_backend_sched_t scheduler_ = nullptr;
         size_t scheduler_capacity_      = 0;
-        std::vector<uint64_t> reservation_;
-        bool active_         = false;
-        size_t reservations_ = 0;
+        bool active_                    = false;
 
         ggml_backend_sched_t make_scheduler(size_t graph_size);
-        bool needs_scheduler(ggml_cgraph* graph) const;
 
     public:
         struct Measurement {
             std::vector<BackendBufferSize> buffers;
-            bool scheduler = false;
         };
         using AssignNodes = std::function<void(ggml_backend_sched_t, ggml_cgraph*)>;
 
@@ -41,13 +35,11 @@ namespace sd {
         ComputeWorkspace& operator=(const ComputeWorkspace&) = delete;
 
         void set_extra_backends(const std::vector<ggml_backend_t>& backends);
-        bool measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const;
         bool prepare(const Measurement& measurement);
         bool release_excess(const Measurement& measurement);
         bool allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes);
         Measurement measure(
             ggml_cgraph* graph,
-            size_t direct_bytes,
             const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
             const AssignNodes& assign_nodes);
         void synchronize() const;
@@ -57,7 +49,6 @@ namespace sd {
         ggml_backend_sched_t scheduler() const { return scheduler_; }
         ggml_backend_t cpu_backend() const { return cpu_backend_; }
         size_t bytes(ggml_backend_t backend) const;
-        size_t reservation_count() const { return reservations_; }
     };
 }
 

src/core/ggml_runner.cpp

@@ -52,8 +52,13 @@ void GGMLRunner::free_params_ctx() {
 
 void GGMLRunner::alloc_compute_ctx() {
     ggml_init_params params;
-    params.mem_size   = static_cast<size_t>(ggml_tensor_overhead() * MAX_GRAPH_SIZE + ggml_graph_overhead());
-    params.mem_buffer = nullptr;
+    params.mem_size = static_cast<size_t>(ggml_tensor_overhead() * MAX_GRAPH_SIZE + ggml_graph_overhead());
+    if (compute_ctx_buffer_ == nullptr) {
+        const size_t count = (params.mem_size + sizeof(std::max_align_t) - 1) / sizeof(std::max_align_t);
+        compute_ctx_buffer_.reset(new std::max_align_t[count]);
+    }
+    // Keep tensor addresses stable between steps so CUDA graphs can finish warming up.
+    params.mem_buffer = compute_ctx_buffer_.get();
     params.no_alloc   = true;
 
     compute_ctx = ggml_init(params);
@@ -482,6 +487,8 @@ void GGMLRunner::runner_end() {
         return;
     }
     workspace_.release();
+    free_compute_ctx();
+    compute_ctx_buffer_.reset();
     cache_.clear();
     logged_compute_bytes_.clear();
     logged_segment_count_ = 0;
@@ -705,7 +712,7 @@ static size_t add_bytes(size_t a, size_t b) {
     return b > SIZE_MAX - a ? SIZE_MAX : a + b;
 }
 
-ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, size_t direct_bytes) {
+ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph) {
     auto external_backend = [&](const ggml_tensor* tensor) -> ggml_backend_t {
         if (!params_tensor_set_.count(tensor)) {
             return nullptr;
@@ -716,7 +723,7 @@ ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, size_t dir
     auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) {
         pin_multi_device_nodes(scheduler, copy, graph);
     };
-    return workspace_.measure(graph, direct_bytes, external_backend, assign_nodes);
+    return workspace_.measure(graph, external_backend, assign_nodes);
 }
 
 std::vector<DeviceMemoryRequest> GGMLRunner::memory_requests(
@@ -767,18 +774,9 @@ bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) {
         sd_backend_cpu_set_n_threads(workspace_.cpu_backend(), n_threads);
     }
     auto scheduler = workspace_.scheduler();
-    ggml_status status;
-    if (scheduler != nullptr) {
-        if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) {
-            LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring", get_desc().c_str());
-            multi_device_eval_callback_warned = true;
-        }
-        status = ggml_backend_sched_graph_compute(scheduler, graph);
-    } else {
-        status = sd_backend_graph_compute_with_eval_callback(runtime_backend, graph,
-                                                             sd_get_backend_eval_callback(),
-                                                             sd_get_backend_eval_callback_data());
-    }
+    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);
     workspace_.synchronize();
     if (status != GGML_STATUS_SUCCESS) {
         last_compute_status_ = status;
@@ -832,7 +830,7 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
     }
     const auto params           = collect_used_param_tensors(graph);
     const auto& cached_plan     = resolve_graph_cut_plan(graph);
-    const auto full_measurement = measure(graph, cached_plan.compute_buffer_size);
+    const auto full_measurement = measure(graph);
     if (full_measurement.buffers.empty()) {
         last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
         return std::nullopt;
@@ -924,7 +922,7 @@ 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, segment.compute_buffer_size) : full_measurement;
+        auto measurement = segmented ? measure(segment_graph) : full_measurement;
         if (!workspace_.prepare(measurement)) {
             last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
             return fail_segment("workspace preparation");
@@ -951,9 +949,6 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
             return fail_segment("weight preparation");
         }
         // Preparing weights can execute LoRA graphs and reclaim an idle workspace.
-        if (!workspace_.measurement_matches(segment_graph, measurement)) {
-            measurement = measure(segment_graph, segment.compute_buffer_size);
-        }
         if (!workspace_.prepare(measurement)) {
             last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
             return fail_segment("workspace preparation");

src/core/ggml_runner.h

@@ -141,7 +141,7 @@ private:
     size_t logged_segment_count_     = 0;
     ggml_status last_compute_status_ = GGML_STATUS_SUCCESS;
 
-    sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, size_t direct_bytes);
+    sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph);
     std::vector<DeviceMemoryRequest> memory_requests(const std::vector<sd::BackendBufferSize>& sizes,
                                                      size_t pending_cache_bytes) const;
     bool fits(const std::vector<DeviceMemoryRequest>& requests,
@@ -161,6 +161,7 @@ protected:
     sd::GraphCutTensorCache cut_cache_;
     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;
 
@@ -169,7 +170,6 @@ protected:
     std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
 
     std::vector<ggml_backend_t> extra_runtime_backends;  // borrowed (SDBackendManager-owned)
-    bool multi_device_eval_callback_warned = false;
 
     std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
     std::weak_ptr<DeviceResidencyManager> residency_manager;

src/model/diffusion/qwen_image_2_1.hpp

@@ -266,16 +266,18 @@ namespace Qwen {
             blocks["img_mlp.out"] = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, false);
         }
 
-        static ggml_tensor* modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* params, int64_t prefix_length, bool gate = false) {
+        static ggml_tensor* modulate(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* params, int64_t prefix_length, bool gate = false) {
+            auto ctx   = runner_ctx->ggml_ctx;
             auto rows  = ggml_ext_chunk(ctx, params, 2, 1);
             auto apply = [&](ggml_tensor* part, ggml_tensor* row) {
                 row = gate ? ggml_tanh(ctx, row) : ggml_scale_bias(ctx, row, 1.f, 1.f);
+                runner_ctx->expand_graph(row);
                 return ggml_mul(ctx, part, row);
             };
-            auto target = apply(ggml_ext_slice(ctx, x, 1, prefix_length, x->ne[1]), rows[0]);
             if (prefix_length == 0) {
-                return target;
+                return apply(x, rows[0]);
             }
+            auto target = apply(ggml_ext_slice(ctx, x, 1, prefix_length, x->ne[1]), rows[0]);
             auto prefix = apply(ggml_ext_slice(ctx, x, 1, 0, prefix_length), rows[1]);
             return ggml_concat(ctx, prefix, target, 1);
         }
@@ -283,25 +285,22 @@ namespace Qwen {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector<ggml_tensor*>& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
             const int64_t prefix_length = cache.mode == QwenImage21PrefixCache::Mode::REUSE ? 0 : layout.prefix_length;
             auto h                      = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm1"])->forward(ctx, x);
-            h                           = modulate(ctx->ggml_ctx, h, modulation[0], prefix_length);
+            h                           = modulate(ctx, h, modulation[0], prefix_length);
             h                           = std::dynamic_pointer_cast<QwenImage21Attention>(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks, cache);
-            x                           = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], prefix_length, true));
+            x                           = ggml_add(ctx->ggml_ctx, x, modulate(ctx, h, modulation[1], prefix_length, true));
             h                           = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm2"])->forward(ctx, x);
-            h                           = modulate(ctx->ggml_ctx, h, modulation[2], prefix_length);
-            ggml_tensor* gate;
-            auto fused = blocks.find("img_mlp.gate_up");
+            h                           = modulate(ctx, h, modulation[2], prefix_length);
+            auto fused                  = blocks.find("img_mlp.gate_up");
             if (fused != blocks.end()) {
                 auto gate_up = std::dynamic_pointer_cast<Linear>(fused->second)->forward(ctx, h);
-                auto parts   = ggml_ext_chunk(ctx->ggml_ctx, gate_up, 2, 0);
-                gate         = parts[0];
-                h            = parts[1];
+                h            = ggml_swiglu(ctx->ggml_ctx, gate_up);
             } else {
-                gate = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.gate_layer"])->forward(ctx, h);
-                h    = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.proj"])->forward(ctx, h);
+                auto gate = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.gate_layer"])->forward(ctx, h);
+                h         = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.proj"])->forward(ctx, h);
+                h         = ggml_swiglu_split(ctx->ggml_ctx, gate, h);
             }
-            h = ggml_mul(ctx->ggml_ctx, h, ggml_silu(ctx->ggml_ctx, gate));
             h = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.out"])->forward(ctx, h);
-            return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], prefix_length, true));
+            return ggml_add(ctx->ggml_ctx, x, modulate(ctx, h, modulation[3], prefix_length, true));
         }
     };