refactor: register RoPE kernels and fusion outside ggml

docs/ggml_extensions.md

@@ -0,0 +1,107 @@
+# ggml operator extensions
+
+The patched ggml exposes `ggml-ext.h` for application-owned operators, kernel
+overrides, and graph fusion rules. Registrations are in-process and scoped to a
+backend device. Register them before capability queries or graph allocation;
+callback code and user data must remain alive until all dependent graphs and
+backends have been destroyed. Registration names are unique, and duplicate
+registrations fail rather than replacing live implementations.
+
+## Operator definitions
+
+`ggml_ext_register_op` registers a versioned name, attribute size, and an output
+metadata inference callback. `ggml_ext_op_create` creates an ordinary ggml tensor
+with explicit input dependencies. The initial interface supports one contiguous,
+out-of-place output with an existing `ggml_type`.
+
+Attributes are copied into the unused portion of `op_params`. They are small,
+plain values; use `memcpy` when reading attributes that need stricter alignment.
+Auxiliary tensors must appear in `src`, not in an attribute pointer or user data.
+Registered operators use `GGML_OP_CUSTOM` internally, with an identifiable
+callback marker. Unsupported devices do not mistake them for CPU callbacks.
+These process-local descriptors are not a graph serialization format.
+
+## Kernel selection
+
+`ggml_ext_register_kernel(device, &desc)` associates a kernel with a built-in
+`ggml_op`, or with a registered custom operator. `supports` describes correctness
+constraints; `preferred` optionally decides whether to override an available
+native implementation. Higher priorities are considered first. When native
+execution is unsupported, a supported extension can run regardless of its
+preference callback.
+
+Capability queries include registered single-node implementations, so the
+scheduler can select a device even when its native implementation rejects the
+operator. Callbacks inspect metadata only and must not recursively call the
+public capability query for the same device and operator. A fusion-only
+registration does not claim single-node support.
+
+`prepare` creates optional host-side plan state and reports workspace bytes.
+It can run during memory measurement, before tensor addresses are assigned, and
+must not launch work or allocate device buffers. Workspace is allocated by ggml
+before execution and included in scheduler memory estimates and buffer-size
+reports. `execute` receives the backend, workspace buffer, and current tensor
+addresses. `destroy` releases prepared state. A preparation or execution error
+is propagated; execution never retries a native implementation after launch.
+
+Overriding a built-in operator preserves its numerical and aliasing contract,
+including allocator reuse of input storage. Explicit views/in-place nodes are
+not intercepted by this first interface. Rebuild the scheduler allocation after
+changing graph metadata, selection policy, or registrations.
+
+## Fusion
+
+A `match` callback returns the number of consecutive nodes it can replace.
+The scheduler freezes this selection before allocating tensors. ggml checks
+compute flags, external uses, output flags, and views; the initial implementation
+accepts a conservative set of pure operators with only the final node exposed
+as the fused output. The rule remains responsible for mathematical equivalence,
+dtype, shape, layout, and hardware constraints.
+
+External inputs stay allocated through the final node. Native graph optimization
+and execution operate on the remaining contiguous ranges, so native fusion cannot
+consume nodes selected by an extension. Intermediate allocations are retained
+in this initial implementation; only their computation is removed.
+
+Automatic extension fusion is enabled in scheduler execution with one pipeline
+copy and without evaluation callbacks. Multiple pipeline copies and evaluation
+callbacks retain individual operations. Direct backend graph computation supports
+single-node extensions; it does not fuse an already allocated graph. Direct calls
+use temporary plans and synchronize before releasing them; scheduler plans and
+workspaces are reused across executions.
+
+Use the backend graph APIs or the scheduler. The legacy CPU `ggml_graph_compute`
+and backend `graph_plan_*` interfaces do not dispatch these extensions.
+
+## Backend services
+
+Services are discovered through `ggml_backend_reg_get_proc_address`:
+
+| Service | Purpose |
+| --- | --- |
+| `ggml_ext_cpu_parallel_v1` | Run an application callback on the CPU backend's configured thread pool, with its abort callback. |
+| `ggml_ext_cuda_context_v1` | Borrow the current CUDA device and stream, ordered with native work. |
+| `ggml_ext_cuda_device_v1` | Get the physical CUDA ordinal for capability checks before backend creation, including virtual ggml devices. |
+
+Workspace is a backend buffer. Its base is a CUDA device pointer on CUDA; callers
+must use the representation appropriate to their backend. A Vulkan command
+recording service is not implemented yet.
+
+## sd.cpp integration
+
+`src/core/ggml_extensions.cpp` registers the CPU implementation of
+`sd.rope_apply.v1`, an override for existing `GGML_OP_ROPE_APPLY` nodes, and a
+`RMSNorm -> Mul -> RoPE` fusion. The fusion accepts a one-dimensional weight
+vector. Both interleaved and split-half RoPE, partial rotation, and strided input
+rows are supported. Model construction uses the custom operator on CPU and
+retains existing backend implementations elsewhere.
+
+`SD_USE_UPSTREAM_GGML` excludes these registrations. The stable-diffusion public C
+API is unchanged. INT8 ConvRot retains its existing packed activation layout and
+separate weight scale tensor.
+
+CuTe Flash Attention still resides in ggml, including its existing fusion paths.
+Migrating it requires moving its pool/conversion dependencies onto the extension
+workspace and CUDA service, then registering its selection and fusion rules in
+sd.cpp. This initial implementation establishes and validates the mechanism;
+it does not yet externalize all model-specific kernels or add dynamic dtypes.

src/core/ggml_extend.h

@@ -6,6 +6,7 @@
 
 #include "ggml-backend.h"
 #include "ggml.h"
+#include "core/ggml_extensions.h"
 
 #define EPS 1e-05f
 

src/core/ggml_extend_backend.cpp

@@ -1,4 +1,5 @@
 #include "core/ggml_extend_backend.h"
+#include "core/ggml_extensions.h"
 
 #include <algorithm>
 #include <cctype>
@@ -596,7 +597,9 @@ bool sd_backend_supports_conv_implicit_gemm(ggml_backend_t backend, bool conv3d)
 
 ggml_backend_t sd_backend_cpu_init() {
     ggml_backend_load_all_once();
-    return ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
+    auto backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
+    sd_ggml_register_extensions(backend);
+    return backend;
 }
 
 bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) {
@@ -1207,6 +1210,7 @@ ggml_backend_t SDBackendManager::init_cached_backend(const std::string& name) {
     }
 
     SDBackendHandle handle(backend);
+    sd_ggml_register_extensions(backend);
     backends_.emplace(actual_key, std::move(handle));
     return backend;
 }

src/core/ggml_extensions.cpp

@@ -0,0 +1,207 @@
+#include "core/ggml_extensions.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+#include "ggml-ext.h"
+
+#include <cmath>
+#include <cstring>
+#include <mutex>
+#include <unordered_set>
+
+static bool rope_infer(ggml_tensor* const* inputs, int count, const void*, ggml_ext_tensor_desc* out) {
+    if (count != 2) {
+        return false;
+    }
+    const auto* x  = inputs[0];
+    const auto* pe = inputs[1];
+    if (x->type != GGML_TYPE_F32 || pe->type != GGML_TYPE_F32 || x->ne[0] % 2 != 0 ||
+        pe->ne[0] != 2 || pe->ne[1] != 2 || pe->ne[2] <= 0 || pe->ne[2] > x->ne[0] / 2 ||
+        pe->ne[3] != x->ne[2]) {
+        return false;
+    }
+    out->type  = GGML_TYPE_F32;
+    out->ne[0] = x->ne[0];
+    out->ne[1] = x->ne[2];
+    out->ne[2] = x->ne[1] * x->ne[3];
+    out->ne[3] = 1;
+    return true;
+}
+
+static ggml_ext_op_t rope_op() {
+    static const ggml_ext_op_t op = [] {
+        const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.rope_apply.v1", sizeof(int32_t), rope_infer};
+        auto result                 = ggml_ext_register_op(&desc);
+        GGML_ASSERT(result);
+        return result;
+    }();
+    return op;
+}
+
+static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+    if (node->type != GGML_TYPE_F32 || !ggml_is_contiguous(node) || !node->src[0] || !node->src[1]) {
+        return false;
+    }
+    ggml_tensor* inputs[] = {node->src[0], node->src[1]};
+    ggml_ext_tensor_desc output;
+    return rope_infer(inputs, 2, nullptr, &output) &&
+           std::memcmp(output.ne, node->ne, sizeof(output.ne)) == 0;
+}
+
+struct RopeExecution {
+    const ggml_tensor* rope;
+    const ggml_tensor* norm;
+    const ggml_tensor* weight;
+};
+
+static void rope_compute(ggml_tensor* dst, int ith, int nth, void* user_data) {
+    const auto& execution = *static_cast<const RopeExecution*>(user_data);
+    const auto* rope      = execution.rope;
+    const auto* x         = execution.norm ? execution.norm->src[0] : rope->src[0];
+    const auto* pe        = rope->src[1];
+    int32_t interleaved;
+    std::memcpy(&interleaved, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(interleaved));
+    float eps = 0.0f;
+    if (execution.norm) {
+        std::memcpy(&eps, execution.norm->op_params, sizeof(eps));
+    }
+    const int64_t half_dim = pe->ne[2];
+    const int64_t rows     = x->ne[1] * x->ne[2] * x->ne[3];
+    for (int64_t row = ith; row < rows; row += nth) {
+        const int64_t token = row % x->ne[2];
+        const int64_t head  = (row / x->ne[2]) % x->ne[1];
+        const int64_t batch = row / (x->ne[2] * x->ne[1]);
+        const char* input   = static_cast<const char*>(x->data) + head * x->nb[1] + token * x->nb[2] + batch * x->nb[3];
+        const char* pos     = static_cast<const char*>(pe->data) + token * pe->nb[3];
+        float* output       = static_cast<float*>(dst->data) + row * x->ne[0];
+        float scale         = 1.0f;
+        if (execution.norm) {
+            double sum = 0.0;
+            for (int64_t col = 0; col < x->ne[0]; ++col) {
+                const float v = *reinterpret_cast<const float*>(input + col * x->nb[0]);
+                sum += static_cast<double>(v * v);
+            }
+            const float mean = static_cast<float>(sum / x->ne[0]);
+            scale            = 1.0f / std::sqrt(mean + eps);
+        }
+        auto value = [&](int64_t col) {
+            float v = *reinterpret_cast<const float*>(input + col * x->nb[0]);
+            if (execution.norm) {
+                const auto* weight = static_cast<const char*>(execution.weight->data);
+                v                  = (v * scale) * *reinterpret_cast<const float*>(weight + col * execution.weight->nb[0]);
+            }
+            return v;
+        };
+        for (int64_t j = 0; j < half_dim; ++j) {
+            const int64_t i0 = interleaved ? 2 * j : j;
+            const int64_t i1 = interleaved ? 2 * j + 1 : j + half_dim;
+            const float a    = value(i0);
+            const float b    = value(i1);
+            const char* m    = pos + j * pe->nb[2];
+            const float c    = *reinterpret_cast<const float*>(m);
+            const float s    = *reinterpret_cast<const float*>(m + pe->nb[1]);
+            if (interleaved) {
+                const float neg_s = *reinterpret_cast<const float*>(m + pe->nb[0]);
+                const float c1    = *reinterpret_cast<const float*>(m + pe->nb[0] + pe->nb[1]);
+                output[i0]        = a * c + b * neg_s;
+                output[i1]        = a * s + b * c1;
+            } else {
+                output[i0] = a * c - b * s;
+                output[i1] = a * s + b * c;
+            }
+        }
+        for (int64_t col = 2 * half_dim; col < x->ne[0]; ++col) {
+            output[col] = value(col);
+        }
+    }
+}
+
+static int rope_fusion_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+    auto* view = const_cast<ggml_cgraph*>(graph);
+    if (first + 2 >= ggml_graph_n_nodes(view)) {
+        return 0;
+    }
+    const auto* norm = ggml_graph_node(view, first);
+    const auto* mul  = ggml_graph_node(view, first + 1);
+    const auto* rope = ggml_graph_node(view, first + 2);
+    if (norm->op != GGML_OP_RMS_NORM || norm->type != GGML_TYPE_F32 || !norm->src[0] ||
+        norm->src[0]->type != GGML_TYPE_F32 || norm->src[0]->nb[0] != sizeof(float) ||
+        mul->op != GGML_OP_MUL || mul->type != GGML_TYPE_F32 || rope->src[0] != mul ||
+        (rope->op != GGML_OP_ROPE_APPLY && ggml_ext_op_get(rope) != rope_op()) ||
+        !rope_supports(device, rope, nullptr) || !ggml_are_same_shape(norm, mul)) {
+        return 0;
+    }
+    const auto* weight = mul->src[0] == norm ? mul->src[1] : (mul->src[1] == norm ? mul->src[0] : nullptr);
+    if (!weight || weight->type != GGML_TYPE_F32 || weight->ne[0] != norm->ne[0] ||
+        weight->ne[1] != 1 || weight->ne[2] != 1 || weight->ne[3] != 1) {
+        return 0;
+    }
+    return 3;
+}
+
+static ggml_status rope_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int count, void*, void*) {
+    auto device = ggml_backend_get_device(context->backend);
+    auto reg    = ggml_backend_dev_backend_reg(device);
+    auto run    = reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1"));
+    GGML_ASSERT(run);
+    RopeExecution execution = {nodes[count - 1], nullptr, nullptr};
+    if (count == 3) {
+        execution.norm   = nodes[0];
+        execution.weight = nodes[1]->src[0] == nodes[0] ? nodes[1]->src[1] : nodes[1]->src[0];
+    }
+    return run(context->backend, nodes[count - 1], rope_compute, &execution);
+}
+
+void sd_ggml_register_extensions(ggml_backend_t backend) {
+    if (!backend) {
+        return;
+    }
+    auto device = ggml_backend_get_device(backend);
+    auto reg    = ggml_backend_dev_backend_reg(device);
+    if (!ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1")) {
+        return;
+    }
+    static std::mutex mutex;
+    static std::unordered_set<ggml_backend_dev_t> registered;
+    std::lock_guard<std::mutex> lock(mutex);
+    if (!registered.insert(device).second) {
+        return;
+    }
+    ggml_ext_kernel_desc desc = {};
+    desc.abi_version          = GGML_EXT_ABI_VERSION;
+    desc.name                 = "sd.rope_apply.cpu.v1";
+    desc.op                   = GGML_OP_CUSTOM;
+    desc.custom_op            = rope_op();
+    desc.priority             = 100;
+    desc.supports             = rope_supports;
+    desc.execute              = rope_execute;
+    GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+
+    desc.name      = "sd.rope_apply.cpu.override.v1";
+    desc.op        = GGML_OP_ROPE_APPLY;
+    desc.custom_op = nullptr;
+    GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+
+    desc.name     = "sd.rms_norm_mul_rope.cpu.v1";
+    desc.op       = GGML_OP_RMS_NORM;
+    desc.priority = 200;
+    desc.supports = nullptr;
+    desc.match    = rope_fusion_match;
+    GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+}
+
+ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved) {
+    sd_ggml_register_extensions(backend);
+    if (backend) {
+        auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
+        if (ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1")) {
+            ggml_tensor* inputs[] = {x, pe};
+            int32_t attributes    = interleaved ? 1 : 0;
+            return ggml_ext_op_create(ctx, rope_op(), inputs, 2, &attributes);
+        }
+    }
+    return ggml_rope_apply(ctx, x, pe, interleaved);
+}
+
+#else
+void sd_ggml_register_extensions(ggml_backend_t) {}
+#endif

src/core/ggml_extensions.h

@@ -0,0 +1,12 @@
+#ifndef __SD_CORE_GGML_EXTENSIONS_H__
+#define __SD_CORE_GGML_EXTENSIONS_H__
+
+#include "ggml-backend.h"
+
+void sd_ggml_register_extensions(ggml_backend_t backend);
+
+#ifndef SD_USE_UPSTREAM_GGML
+ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved);
+#endif
+
+#endif

src/model/common/rope.hpp

@@ -1077,8 +1077,8 @@ namespace Rope {
         if (backend != nullptr && x->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32 &&
             x->ne[0] % 2 == 0 && pe->ne[0] == 2 && pe->ne[1] == 2 &&
             pe->ne[2] == x->ne[0] / 2 && pe->ne[3] == x->ne[2]) {
-            auto out = ggml_rope_apply(ctx, x, pe, rope_interleaved);
-            if (ggml_backend_supports_op(backend, out)) {
+            auto out = sd_ggml_rope_apply(ctx, backend, x, pe, rope_interleaved);
+            if (out && ggml_backend_supports_op(backend, out)) {
                 return out;
             }
         }

src/model/diffusion/minimax_h3.hpp

@@ -180,8 +180,8 @@ namespace MiniMaxH3 {
         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)) {
+            auto out = sd_ggml_rope_apply(ctx, runner_ctx->backend, x, pe, false);
+            if (out && ggml_backend_supports_op(runner_ctx->backend, out)) {
                 return out;
             }
         }