refactor: split ggml extensions and move implementations to cpp files (#1945)

src/conditioning/conditioner.hpp

@@ -1,11 +1,13 @@
 #ifndef __SD_CONDITIONING_CONDITIONER_HPP__
 #define __SD_CONDITIONING_CONDITIONER_HPP__
 
+#include <cinttypes>
 #include <cmath>
 #include <iomanip>
 #include <limits>
 #include <optional>
 #include <sstream>
+#include "core/ggml_tensor_utils.h"
 
 #include "core/tensor_ggml.hpp"
 #include "core/util.h"
@@ -594,7 +596,7 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(pixel_values, return_pooled, clip_skip);
         };
-        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
+        return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true));
     }
 };
 
@@ -2876,7 +2878,7 @@ struct LTXAVTextProjectionRunner : public GGMLRunner {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(x);
         };
-        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
+        return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end));
     }
 };
 

src/core/ggml_extend.cpp

@@ -0,0 +1,727 @@
+#include "core/ggml_extend.h"
+
+#include <cmath>
+#include <utility>
+
+#include "core/ggml_extend_backend.h"
+
+ggml_tensor* ggml_ext_mul_n_mode(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b, int mode) {
+    // reshape A
+    // swap 0th and nth axis
+    a           = ggml_cont(ctx, ggml_permute(ctx, a, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0));
+    int64_t ne1 = a->ne[1];
+    int64_t ne2 = a->ne[2];
+    int64_t ne3 = a->ne[3];
+    // make 2D
+    a = ggml_cont(ctx, ggml_reshape_2d(ctx, a, a->ne[0], (ne3 * ne2 * ne1)));
+
+    ggml_tensor* result = ggml_cont(ctx, ggml_transpose(ctx, ggml_mul_mat(ctx, a, b)));
+
+    // reshape output (same shape as a after permutation except first dim)
+    result = ggml_reshape_4d(ctx, result, result->ne[0], ne1, ne2, ne3);
+    // swap back 0th and nth axis
+    result = ggml_permute(ctx, result, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0);
+    return result;
+}
+
+ggml_tensor* ggml_ext_kronecker(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b) {
+    return ggml_mul(ctx,
+                    ggml_interpolate(ctx,
+                                     a,
+                                     a->ne[0] * b->ne[0],
+                                     a->ne[1] * b->ne[1],
+                                     a->ne[2] * b->ne[2],
+                                     a->ne[3] * b->ne[3],
+                                     GGML_SCALE_MODE_NEAREST),
+                    b);
+}
+
+ggml_tensor* ggml_ext_cont(ggml_context* ctx,
+                           ggml_tensor* x) {
+    if (ggml_is_contiguous(x)) {
+        return x;
+    }
+    return ggml_cont(ctx, x);
+}
+
+ggml_tensor* ggml_ext_torch_permute(ggml_context* ctx,
+                                    ggml_tensor* x,
+                                    int axis0,
+                                    int axis1,
+                                    int axis2,
+                                    int axis3) {
+    int torch_axes[4] = {axis0, axis1, axis2, axis3};
+
+    int ggml_axes[4] = {0};
+    for (int i = 0; i < 4; ++i) {
+        int found = 0;
+        for (int j = 0; j < 4; ++j) {
+            if (torch_axes[j] == i) {
+                ggml_axes[i] = j;
+                found        = 1;
+                break;
+            }
+        }
+        GGML_ASSERT(found && "Invalid permute input: must be a permutation of 0-3");
+    }
+
+    return ggml_permute(ctx, x, ggml_axes[0], ggml_axes[1], ggml_axes[2], ggml_axes[3]);
+}
+
+ggml_tensor* ggml_ext_slice(ggml_context* ctx,
+                            ggml_tensor* x,
+                            int dim,
+                            int64_t start,
+                            int64_t end,
+                            bool cont) {
+    GGML_ASSERT(dim >= 0 && dim < 4);
+    if (x->ne[dim] == 1) {
+        return x;
+    }
+    while (start < 0) {
+        start = x->ne[dim] + start;
+    }
+    while (end < 0) {
+        end = x->ne[dim] + end;
+    }
+    GGML_ASSERT(end > start);
+    GGML_ASSERT(start >= 0 && start < x->ne[dim]);
+    GGML_ASSERT(end > start && end <= x->ne[dim]);
+
+    int64_t slice_size  = end - start;
+    int64_t slice_ne[4] = {x->ne[0], x->ne[1], x->ne[2], x->ne[3]};
+    slice_ne[dim]       = slice_size;
+
+    x = ggml_view_4d(ctx, x,
+                     slice_ne[0], slice_ne[1], slice_ne[2], slice_ne[3],
+                     x->nb[1], x->nb[2], x->nb[3], start * x->nb[dim]);
+
+    if (cont) {
+        x = ggml_cont(ctx, x);
+    }
+
+    return x;
+}
+
+std::vector<ggml_tensor*> ggml_ext_chunk(ggml_context* ctx,
+                                         ggml_tensor* x,
+                                         int num,
+                                         int64_t dim,
+                                         bool cont) {
+    GGML_ASSERT(dim >= 0 && dim < 4);
+    GGML_ASSERT(x->ne[dim] % num == 0);
+
+    std::vector<ggml_tensor*> chunks;
+    int64_t chunk_size  = x->ne[dim] / num;
+    int64_t stride      = chunk_size * x->nb[dim];
+    int64_t chunk_ne[4] = {x->ne[0], x->ne[1], x->ne[2], x->ne[3]};
+    chunk_ne[dim]       = chunk_size;
+    for (int i = 0; i < num; i++) {
+        auto chunk = ggml_view_4d(
+            ctx, x,
+            chunk_ne[0], chunk_ne[1], chunk_ne[2], chunk_ne[3],
+            x->nb[1], x->nb[2], x->nb[3], stride * i);
+        if (cont) {
+            chunk = ggml_cont(ctx, chunk);
+        }
+        chunks.push_back(chunk);
+    }
+
+    return chunks;
+}
+
+ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_first) {
+    // x: [ne3, ne2, ne1, ne0]
+    // return: [ne3, ne2, ne1, ne0/2]
+
+    auto x_vec = ggml_ext_chunk(ctx, x, 2, 0, false);
+    ggml_tensor* gate;
+    if (gate_first) {
+        gate = x_vec[0];
+        x    = x_vec[1];
+    } else {
+        x    = x_vec[0];
+        gate = x_vec[1];
+    }
+    gate = ggml_cont(ctx, gate);
+    gate = ggml_silu_inplace(ctx, gate);
+
+    x = ggml_mul(ctx, x, gate);  // [ne3, ne2, ne1, ne0/2]
+
+    return x;
+}
+
+ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx,
+                                    ggml_tensor* a) {
+    const float eps = 1e-6f;  // default eps parameter
+    return ggml_group_norm(ctx, a, 32, eps);
+}
+
+static bool ggml_ext_is_padded_1d(const ggml_tensor* x) {
+    return x->nb[0] == ggml_type_size(x->type) &&
+           x->nb[2] == x->nb[1] * x->ne[1] &&
+           x->nb[3] == x->nb[2] * x->ne[2];
+}
+
+ggml_tensor* ggml_ext_scale(ggml_context* ctx,
+                            ggml_tensor* x,
+                            float factor,
+                            bool inplace) {
+    if (!ggml_ext_is_padded_1d(x)) {
+        x = ggml_cont(ctx, x);
+    }
+    if (inplace) {
+        x = ggml_scale_inplace(ctx, x, factor);
+    } else {
+        x = ggml_scale(ctx, x, factor);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_gelu(ggml_context* ctx,
+                           ggml_tensor* x,
+                           bool inplace) {
+    if (!ggml_is_contiguous(x)) {
+        x = ggml_cont(ctx, x);
+    }
+    if (inplace) {
+        x = ggml_gelu_inplace(ctx, x);
+    } else {
+        x = ggml_gelu(ctx, x);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 bool inplace) {
+    if (!ggml_is_contiguous(x)) {
+        x = ggml_cont(ctx, x);
+    }
+    if (inplace) {
+        x = ggml_gelu_quick_inplace(ctx, x);
+    } else {
+        x = ggml_gelu_quick(ctx, x);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_linear(ggml_context* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* w,
+                             ggml_tensor* b,
+                             bool force_prec_f32,
+                             float scale) {
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, scale);
+    }
+    if (x->ne[2] * x->ne[3] > 1024) {
+        // workaround: avoid ggml cuda error
+        int64_t ne2 = x->ne[2];
+        int64_t ne3 = x->ne[3];
+        x           = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1] * x->ne[2] * x->ne[3]);
+        x           = ggml_mul_mat(ctx, w, x);
+        if (force_prec_f32) {
+            ggml_mul_mat_set_prec(x, GGML_PREC_F32);
+        }
+        x = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / ne2 / ne3, ne2, ne3);
+    } else {
+        x = ggml_mul_mat(ctx, w, x);
+        if (force_prec_f32) {
+            ggml_mul_mat_set_prec(x, GGML_PREC_F32);
+        }
+    }
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, 1.f / scale);
+    }
+    if (b != nullptr) {
+        x = ggml_add_inplace(ctx, x, b);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx,
+                                           ggml_tensor* x,
+                                           ggml_tensor* w,
+                                           ggml_tensor* weight_scale,
+                                           ggml_tensor* b,
+                                           int convrot_group_size,
+                                           float scale) {
+    GGML_ASSERT(x->type == GGML_TYPE_F32 || (x->type == GGML_TYPE_I8 && scale == 1.f));
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, scale);
+    }
+
+    ggml_tensor* fused_bias = scale == 1.f ? b : nullptr;
+    if (x->ne[2] * x->ne[3] > 1024) {
+        int64_t ne2 = x->ne[2];
+        int64_t ne3 = x->ne[3];
+        x           = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1] * x->ne[2] * x->ne[3]);
+        x           = ggml_mul_mat_i8_tensorwise(ctx, w, x, weight_scale, fused_bias, convrot_group_size);
+        x           = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / ne2 / ne3, ne2, ne3);
+    } else {
+        x = ggml_mul_mat_i8_tensorwise(ctx, w, x, weight_scale, fused_bias, convrot_group_size);
+    }
+
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, 1.f / scale);
+        if (b != nullptr) {
+            x = ggml_add_inplace(ctx, x, b);
+        }
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
+                              ggml_backend_t backend,
+                              ggml_tensor* x,
+                              int lp0,
+                              int rp0,
+                              int lp1,
+                              int rp1,
+                              int lp2,
+                              int rp2,
+                              int lp3,
+                              int rp3,
+                              bool circular_x,
+                              bool circular_y) {
+    if (circular_x && circular_y) {
+        return ggml_pad_ext_circular(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
+    }
+
+    if (circular_x && (lp0 != 0 || rp0 != 0)) {
+        x   = ggml_pad_ext_circular(ctx, x, lp0, rp0, 0, 0, 0, 0, 0, 0);
+        lp0 = rp0 = 0;
+    }
+    if (circular_y && (lp1 != 0 || rp1 != 0)) {
+        x   = ggml_pad_ext_circular(ctx, x, 0, 0, lp1, rp1, 0, 0, 0, 0);
+        lp1 = rp1 = 0;
+    }
+
+    if (lp0 != 0 || rp0 != 0 || lp1 != 0 || rp1 != 0 || lp2 != 0 || rp2 != 0 || lp3 != 0 || rp3 != 0) {
+        ggml_tensor* padded = ggml_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
+        if (backend == nullptr || ggml_backend_supports_op(backend, padded)) {
+            x = padded;
+        } else {
+            // Some backends (e.g. Metal) only implement right-padding for
+            // GGML_OP_PAD (see #850): pad right by lp+rp instead, then roll
+            // the padding around to the left. shift < ne always holds because
+            // ne grew by lp+rp.
+            x = ggml_pad_ext(ctx, x, 0, lp0 + rp0, 0, lp1 + rp1, 0, lp2 + rp2, 0, lp3 + rp3);
+            x = ggml_roll(ctx, x, lp0, lp1, lp2, lp3);
+        }
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_pad(ggml_context* ctx,
+                          ggml_tensor* x,
+                          int p0,
+                          int p1,
+                          int p2,
+                          int p3,
+                          bool circular_x,
+                          bool circular_y) {
+    return ggml_ext_pad_ext(ctx, nullptr, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y);
+}
+
+ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
+                              ggml_tensor* x,
+                              ggml_tensor* w,
+                              ggml_tensor* b,
+                              int s0,
+                              int s1,
+                              int p0,
+                              int p1,
+                              int d0,
+                              int d1,
+                              bool direct,
+                              bool circular_x,
+                              bool circular_y,
+                              float scale) {
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, scale);
+    }
+    if (w->ne[2] != x->ne[2] && ggml_n_dims(w) == 2) {
+        w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], w->ne[1]);
+    }
+
+    if ((p0 != 0 || p1 != 0) && (circular_x || circular_y)) {
+        x  = ggml_ext_pad_ext(ctx, nullptr, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y);
+        p0 = 0;
+        p1 = 0;
+    }
+
+    if (direct) {
+        x = ggml_conv_2d_direct(ctx, w, x, s0, s1, p0, p1, d0, d1);
+    } else {
+        x = ggml_conv_2d(ctx, w, x, s0, s1, p0, p1, d0, d1);
+    }
+    if (scale != 1.f) {
+        x = ggml_ext_scale(ctx, x, 1.f / scale);
+    }
+    if (b != nullptr) {
+        b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
+        x = ggml_add_inplace(ctx, x, b);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
+                              ggml_backend_t backend,
+                              ggml_tensor* x,
+                              ggml_tensor* w,
+                              ggml_tensor* b,
+                              int64_t IC,
+                              int s0,
+                              int s1,
+                              int s2,
+                              int p0,
+                              int p1,
+                              int p2,
+                              int d0,
+                              int d1,
+                              int d2,
+                              bool force_prec_f32) {
+    if (force_prec_f32) {
+        ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type);
+
+        int64_t OC = w->ne[3] / IC;
+        int64_t N  = x->ne[3] / IC;
+        x          = ggml_mul_mat(ctx,
+                                  ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]),
+                                  ggml_reshape_2d(ctx, w, w->ne[0] * w->ne[1] * w->ne[2] * IC, OC));
+        ggml_mul_mat_set_prec(x, GGML_PREC_F32);
+
+        int64_t OD = im2col->ne[3] / N;
+        x          = ggml_reshape_4d(ctx, x, im2col->ne[1] * im2col->ne[2], OD, N, OC);
+        x          = ggml_cont(ctx, ggml_permute(ctx, x, 0, 1, 3, 2));
+        x          = ggml_reshape_4d(ctx, x, im2col->ne[1], im2col->ne[2], OD, OC * N);
+    } else {
+        // ggml_conv_3d decomposes into GGML_OP_IM2COL_3D, which some backends
+        // (e.g. Metal, see #850) do not implement. Fall back to
+        // GGML_OP_CONV_3D on those backends.
+        bool im2col_3d_supported = true;
+        if (backend != nullptr) {
+            ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type);
+            im2col_3d_supported = ggml_backend_supports_op(backend, im2col);
+        }
+        if (im2col_3d_supported) {
+            x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2);
+        } else {
+            int64_t OC = w->ne[3] / IC;
+            int64_t N  = x->ne[3] / IC;
+            x          = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC);
+        }
+    }
+
+    if (b != nullptr) {
+        b = ggml_reshape_4d(ctx, b, 1, 1, 1, b->ne[0]);  // [OC, 1, 1, 1]
+        x = ggml_add_inplace(ctx, x, b);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_conv_3d_nx1x1(ggml_context* ctx,
+                                    ggml_tensor* x,
+                                    ggml_tensor* w,
+                                    ggml_tensor* b,
+                                    int s2,
+                                    int p2,
+                                    int d2) {
+    x = ggml_conv_2d(ctx, w, x, 1, s2, 0, p2, 1, d2);  // [N, OC, T, OH * OW]
+    if (b != nullptr) {
+        b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
+        x = ggml_add(ctx, x, b);
+    }
+    return x;  // [N, OC, T, OH * OW]
+}
+
+std::vector<ggml_tensor*> split_qkv(ggml_context* ctx,
+                                    ggml_tensor* qkv) {
+    qkv = ggml_reshape_4d(ctx, qkv, qkv->ne[0] / 3, 3, qkv->ne[1], qkv->ne[2]);  // [N, L, 3, C]
+    qkv = ggml_cont(ctx, ggml_permute(ctx, qkv, 0, 3, 1, 2));                    // [3, N, L, C]
+
+    int64_t offset = qkv->nb[2] * qkv->ne[2];
+    auto q         = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 0);  // [N, L, C]
+    auto k         = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 1);  // [N, L, C]
+    auto v         = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 2);  // [N, L, C]
+    return {q, k, v};
+}
+
+std::vector<ggml_tensor*> split_image_qkv(ggml_context* ctx,
+                                          ggml_tensor* qkv) {
+    int64_t W   = qkv->ne[0];
+    int64_t H   = qkv->ne[1];
+    int64_t C   = qkv->ne[2] / 3;
+    int64_t N   = qkv->ne[3];
+    int64_t nb1 = qkv->nb[1];
+    int64_t nb2 = qkv->nb[2];
+    qkv         = ggml_reshape_4d(ctx, qkv, W * H, C, 3, N);                     // [N, 3, C, H*W]
+    qkv         = ggml_cont(ctx, ggml_ext_torch_permute(ctx, qkv, 0, 1, 3, 2));  // [3, N, C, H*W]
+
+    int64_t offset = qkv->nb[2] * qkv->ne[2];
+    auto q         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 0);  // [N, C, H, W]
+    auto k         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 1);  // [N, C, H, W]
+    auto v         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 2);  // [N, C, H, W]
+    return {q, k, v};
+}
+
+ggml_tensor* ggml_ext_full(ggml_context* ctx,
+                           float value,
+                           int64_t ne0,
+                           int64_t ne1,
+                           int64_t ne2,
+                           int64_t ne3) {
+    auto one = ggml_get_tensor(ctx, "ggml_runner_build_in_tensor:one");
+    auto t   = ggml_ext_scale(ctx, one, value);             // [1,]
+    t        = ggml_repeat_4d(ctx, t, ne0, ne1, ne2, ne3);  // [ne0, ne1, ne2, ne3]
+    return t;
+}
+
+ggml_tensor* ggml_ext_zeros(ggml_context* ctx,
+                            int64_t ne0,
+                            int64_t ne1,
+                            int64_t ne2,
+                            int64_t ne3) {
+    return ggml_ext_full(ctx, 0.f, ne0, ne1, ne2, ne3);
+}
+
+ggml_tensor* ggml_ext_zeros_like(ggml_context* ctx,
+                                 ggml_tensor* x) {
+    return ggml_ext_zeros(ctx, x->ne[0], x->ne[1], x->ne[2], x->ne[3]);
+}
+
+ggml_tensor* ggml_ext_ones(ggml_context* ctx,
+                           int64_t ne0,
+                           int64_t ne1,
+                           int64_t ne2,
+                           int64_t ne3) {
+    return ggml_ext_full(ctx, 1.f, ne0, ne1, ne2, ne3);
+}
+
+ggml_tensor* ggml_ext_ones_like(ggml_context* ctx,
+                                ggml_tensor* x) {
+    return ggml_ext_ones(ctx, x->ne[0], x->ne[1], x->ne[2], x->ne[3]);
+}
+
+ggml_tensor* ggml_ext_cast_f32(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* a) {
+    if (sd_backend_is(backend, "Vulkan")) {
+        auto zero_index = ggml_get_tensor(ctx, "ggml_runner_build_in_tensor:zero_int");
+        auto out        = ggml_reshape_1d(ctx, a, ggml_nelements(a));
+        out             = ggml_get_rows(ctx, out, zero_index);
+        out             = ggml_reshape(ctx, out, a);
+        // auto out = ggml_cast(ctx, a, GGML_TYPE_F32);
+        return out;
+    } else {
+        auto out         = ggml_reshape_2d(ctx, a, 1, ggml_nelements(a));
+        ggml_tensor* one = ggml_ext_ones(ctx, 1, 1, 1, 1);  // [1,]
+        if (ggml_is_transposed(out)) {
+            out = ggml_mul_mat(ctx, one, out);
+        } else {
+            out = ggml_mul_mat(ctx, out, one);
+        }
+        out = ggml_reshape(ctx, out, a);
+        return out;
+    }
+}
+
+ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
+                                    ggml_backend_t backend,
+                                    ggml_tensor* q,
+                                    ggml_tensor* k,
+                                    ggml_tensor* v,
+                                    int64_t n_head,
+                                    ggml_tensor* mask,
+                                    bool skip_reshape,
+                                    bool flash_attn,
+                                    float kv_scale) {  // avoid overflow
+    int64_t L_q;
+    int64_t L_k;
+    int64_t C;
+    int64_t N;
+    int64_t d_head;
+    int64_t n_kv_head;
+    if (!skip_reshape) {
+        L_q       = q->ne[1];
+        L_k       = k->ne[1];
+        C         = q->ne[0];
+        N         = q->ne[2];
+        d_head    = C / n_head;
+        n_kv_head = k->ne[0] / d_head;
+
+        q = ggml_reshape_4d(ctx, q, d_head, n_head, L_q, N);       // [N, L_q, n_head, d_head]
+        q = ggml_ext_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3));  // [N, n_head, L_q, d_head]
+        q = ggml_reshape_3d(ctx, q, d_head, L_q, n_head * N);      // [N * n_head, L_q, d_head]
+
+        k = ggml_reshape_4d(ctx, k, d_head, n_kv_head, L_k, N);    // [N, L_k, n_kv_head, d_head]
+        k = ggml_ext_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3));  // [N, n_kv_head, L_k, d_head]
+        k = ggml_reshape_3d(ctx, k, d_head, L_k, n_kv_head * N);   // [N * n_kv_head, L_k, d_head]
+
+        v = ggml_reshape_4d(ctx, v, d_head, n_kv_head, L_k, N);  // [N, L_k, n_kv_head, d_head]
+    } else {
+        L_q       = q->ne[1];
+        L_k       = k->ne[1];
+        d_head    = v->ne[0];
+        N         = v->ne[3];
+        n_kv_head = k->ne[2] / N;
+        C         = d_head * n_head;
+    }
+
+    float scale = (1.0f / sqrt((float)d_head));
+
+    ggml_tensor* kqv = nullptr;
+
+    auto build_kqv = [&](ggml_tensor* q_in, ggml_tensor* k_in, ggml_tensor* v_in, ggml_tensor* mask_in) -> ggml_tensor* {
+        if (kv_scale != 1.0f) {
+            k_in = ggml_ext_scale(ctx, k_in, kv_scale);
+        }
+        k_in = ggml_cast(ctx, k_in, GGML_TYPE_F16);
+
+        v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v_in, 0, 2, 1, 3));
+        v_in = ggml_reshape_3d(ctx, v_in, d_head, L_k, n_kv_head * N);
+        if (kv_scale != 1.0f) {
+            v_in = ggml_ext_scale(ctx, v_in, kv_scale);
+        }
+        v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16);
+
+        if (mask_in != nullptr) {
+            // ggml_flash_attn_ext expects the mask as a contiguous F16 tensor shaped
+            // [n_kv, n_q, (heads), (batch)] (ne0 = key length, ne1 = query length) and,
+            // unlike the manual-attention path, does not broadcast the query dimension.
+            // Some callers (e.g. Chroma/T5) pass a per-key padding mask broadcast over
+            // queries ([n_kv, 1, ...]); materialize the query dimension to L_q so the
+            // kernel indexes it correctly. (A bare ggml_transpose here produced a
+            // [1, n_kv, ...] mask that the kernel silently misreads, yielding NaN/blank
+            // output for masked flash attention.)
+            if (mask_in->ne[1] != L_q) {
+                mask_in = ggml_repeat(ctx, mask_in,
+                                      ggml_new_tensor_4d(ctx, mask_in->type, mask_in->ne[0], L_q, mask_in->ne[2], mask_in->ne[3]));
+            }
+            mask_in = ggml_cast(ctx, mask_in, GGML_TYPE_F16);
+        }
+
+        auto out = ggml_flash_attn_ext(ctx, q_in, k_in, v_in, mask_in, scale / kv_scale, 0, 0);
+        if (!ggml_backend_supports_op(backend, out)) {
+            return nullptr;
+        }
+        ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32);
+        if (kv_scale != 1.0f) {
+            out = ggml_ext_scale(ctx, out, 1.0f / kv_scale);
+        }
+        return out;
+    };
+
+    if (flash_attn) {
+        // LOG_VERBOSE("attention_ext L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N);
+        bool can_use_flash_attn = true;
+        if (mask != nullptr) {
+            // TODO: figure out if we can bend t5 to work too
+            can_use_flash_attn = can_use_flash_attn && mask->ne[3] == 1;
+        }
+
+        if (can_use_flash_attn) {
+            kqv = build_kqv(q, k, v, mask);
+            if (kqv != nullptr) {
+                kqv = ggml_view_4d(ctx,
+                                   kqv,
+                                   d_head,
+                                   n_head,
+                                   L_q,
+                                   N,
+                                   kqv->nb[1],
+                                   kqv->nb[2],
+                                   kqv->nb[1] * n_head,
+                                   0);
+            }
+        }
+    }
+
+    if (kqv == nullptr) {
+        // if (flash_attn) {
+        //     LOG_VERBOSE("fallback to default attention, L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N);
+        // }
+        v = ggml_ext_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3));  // [N, n_kv_head, d_head, L_k]
+        v = ggml_reshape_3d(ctx, v, L_k, d_head, n_kv_head * N);   // [N * n_kv_head, d_head, L_k]
+
+        auto kq = ggml_mul_mat(ctx, k, q);  // [N * n_head, L_q, L_k]
+        ggml_mul_mat_set_prec(kq, GGML_PREC_F32);
+        kq = ggml_scale_inplace(ctx, kq, scale);
+        if (mask) {
+            kq = ggml_add_inplace(ctx, kq, mask);
+        }
+        kq = ggml_soft_max_inplace(ctx, kq);
+
+        kqv = ggml_mul_mat(ctx, v, kq);  // [N * n_head, L_q, d_head]
+
+        kqv = ggml_reshape_4d(ctx, kqv, d_head, L_q, n_head, N);  // [N, n_head, L_q, d_head]
+        kqv = ggml_permute(ctx, kqv, 0, 2, 1, 3);                 // [N, L_q, n_head, d_head]
+    }
+
+    kqv = ggml_ext_cont(ctx, kqv);
+    kqv = ggml_reshape_3d(ctx, kqv, d_head * n_head, L_q, N);  // [N, L_q, C]
+
+    return kqv;
+}
+
+ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 ggml_tensor* w,
+                                 ggml_tensor* b,
+                                 float eps) {
+    x = ggml_norm(ctx, x, eps);
+    if (w != nullptr) {
+        x = ggml_mul_inplace(ctx, x, w);
+        if (b != nullptr) {
+            x = ggml_add_inplace(ctx, x, b);
+        }
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_group_norm(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 ggml_tensor* w,
+                                 ggml_tensor* b,
+                                 int num_groups) {
+    if (ggml_n_dims(x) >= 3 && w != nullptr && b != nullptr) {
+        w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], 1);
+        b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1);
+    }
+
+    const float eps = 1e-6f;  // default eps parameter
+    x               = ggml_group_norm(ctx, x, num_groups, eps);
+    if (w != nullptr && b != nullptr) {
+        x = ggml_mul_inplace(ctx, x, w);
+        // b = ggml_repeat(ctx, b, x);
+        x = ggml_add_inplace(ctx, x, b);
+    }
+    return x;
+}
+
+ggml_tensor* ggml_ext_timestep_embedding(
+    ggml_context* ctx,
+    ggml_tensor* timesteps,
+    int dim,
+    int max_period,
+    float time_factor) {
+    timesteps = ggml_ext_scale(ctx, timesteps, time_factor);
+    return ggml_timestep_embedding(ctx, timesteps, dim, max_period);
+}
+
+ggml_tensor* ggml_ext_vec_concat(ggml_context* ctx,
+                                 std::vector<ggml_tensor*>& tensors,
+                                 int dim) {
+    while (tensors.size() > 1) {
+        std::vector<ggml_tensor*> next_level;
+        for (size_t i = 0; i < tensors.size(); i += 2) {
+            if (i + 1 < tensors.size()) {
+                next_level.push_back(ggml_concat(ctx, tensors[i], tensors[i + 1], dim));
+            } else {
+                next_level.push_back(tensors[i]);
+            }
+        }
+        tensors = std::move(next_level);
+    }
+    return tensors[0];
+}

src/core/ggml_extend.h

@@ -0,0 +1,235 @@
+#ifndef __SD_CORE_GGML_EXTEND_H__
+#define __SD_CORE_GGML_EXTEND_H__
+
+#include <cstdint>
+#include <vector>
+
+#include "ggml-backend.h"
+#include "ggml.h"
+
+#define EPS 1e-05f
+
+static_assert(GGML_MAX_NAME >= 128, "GGML_MAX_NAME must be at least 128");
+
+// n-mode tensor-matrix product
+// example: 2-mode product
+// A: [ne03, k, ne01, ne00]
+// B: k rows, m columns => [k, m]
+// result is [ne03, m, ne01, ne00]
+ggml_tensor* ggml_ext_mul_n_mode(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b, int mode = 0);
+
+// Kronecker product
+// [ne03,ne02,ne01,ne00] x [ne13,ne12,ne11,ne10] => [ne03*ne13,ne02*ne12,ne01*ne11,ne00*ne10]
+ggml_tensor* ggml_ext_kronecker(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b);
+
+ggml_tensor* ggml_ext_cont(ggml_context* ctx,
+                           ggml_tensor* x);
+
+// torch like permute
+ggml_tensor* ggml_ext_torch_permute(ggml_context* ctx,
+                                    ggml_tensor* x,
+                                    int axis0,
+                                    int axis1,
+                                    int axis2,
+                                    int axis3);
+
+ggml_tensor* ggml_ext_slice(ggml_context* ctx,
+                            ggml_tensor* x,
+                            int dim,
+                            int64_t start,
+                            int64_t end,
+                            bool cont = true);
+
+// example: [N, 3*C, H, W] => ([N, C, H, W], [N, C, H, W], [N, C, H, W])
+std::vector<ggml_tensor*> ggml_ext_chunk(ggml_context* ctx,
+                                         ggml_tensor* x,
+                                         int num,
+                                         int64_t dim,
+                                         bool cont = true);
+
+ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_first = true);
+
+ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx,
+                                    ggml_tensor* a);
+
+ggml_tensor* ggml_ext_scale(ggml_context* ctx,
+                            ggml_tensor* x,
+                            float factor,
+                            bool inplace = false);
+
+ggml_tensor* ggml_ext_gelu(ggml_context* ctx,
+                           ggml_tensor* x,
+                           bool inplace = false);
+
+ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 bool inplace = false);
+
+ggml_tensor* ggml_ext_linear(ggml_context* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* w,
+                             ggml_tensor* b,
+                             bool force_prec_f32 = false,
+                             float scale         = 1.f);
+
+ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx,
+                                           ggml_tensor* x,
+                                           ggml_tensor* w,
+                                           ggml_tensor* weight_scale,
+                                           ggml_tensor* b,
+                                           int convrot_group_size,
+                                           float scale = 1.f);
+
+ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
+                              ggml_backend_t backend,
+                              ggml_tensor* x,
+                              int lp0,
+                              int rp0,
+                              int lp1,
+                              int rp1,
+                              int lp2,
+                              int rp2,
+                              int lp3,
+                              int rp3,
+                              bool circular_x = false,
+                              bool circular_y = false);
+
+ggml_tensor* ggml_ext_pad(ggml_context* ctx,
+                          ggml_tensor* x,
+                          int p0,
+                          int p1,
+                          int p2          = 0,
+                          int p3          = 0,
+                          bool circular_x = false,
+                          bool circular_y = false);
+
+// w: [OC，IC, KH, KW]
+// x: [N, IC, IH, IW]
+// b: [OC,]
+// result: [N, OC, OH, OW]
+ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
+                              ggml_tensor* x,
+                              ggml_tensor* w,
+                              ggml_tensor* b,
+                              int s0          = 1,
+                              int s1          = 1,
+                              int p0          = 0,
+                              int p1          = 0,
+                              int d0          = 1,
+                              int d1          = 1,
+                              bool direct     = false,
+                              bool circular_x = false,
+                              bool circular_y = false,
+                              float scale     = 1.f);
+
+// w: [OC，IC, KD, 1 * 1]
+// x: [N, IC, IH, IW]
+// b: [OC,]
+// result: [N*OC, OD, OH, OW]
+ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
+                              ggml_backend_t backend,
+                              ggml_tensor* x,
+                              ggml_tensor* w,
+                              ggml_tensor* b,
+                              int64_t IC,
+                              int s0              = 1,
+                              int s1              = 1,
+                              int s2              = 1,
+                              int p0              = 0,
+                              int p1              = 0,
+                              int p2              = 0,
+                              int d0              = 1,
+                              int d1              = 1,
+                              int d2              = 1,
+                              bool force_prec_f32 = false);
+
+// w: [OC，IC, KD, 1 * 1]
+// x: [N, IC, ID, IH*IW]
+// b: [OC,]
+// result: [N, OC, OD, OH*OW]
+ggml_tensor* ggml_ext_conv_3d_nx1x1(ggml_context* ctx,
+                                    ggml_tensor* x,
+                                    ggml_tensor* w,
+                                    ggml_tensor* b,
+                                    int s2 = 1,
+                                    int p2 = 1,
+                                    int d2 = 1);
+
+// qkv: [N, L, 3*C]
+// return: ([N, L, C], [N, L, C], [N, L, C])
+std::vector<ggml_tensor*> split_qkv(ggml_context* ctx,
+                                    ggml_tensor* qkv);
+
+// qkv: [N, 3*C, H, W]
+// return: ([N, C, H, W], [N, C, H, W], [N, C, H, W])
+std::vector<ggml_tensor*> split_image_qkv(ggml_context* ctx,
+                                          ggml_tensor* qkv);
+
+// Constant and cast helpers require the built-in tensors initialized by GGMLRunner.
+ggml_tensor* ggml_ext_full(ggml_context* ctx,
+                           float value,
+                           int64_t ne0,
+                           int64_t ne1,
+                           int64_t ne2,
+                           int64_t ne3);
+
+ggml_tensor* ggml_ext_zeros(ggml_context* ctx,
+                            int64_t ne0,
+                            int64_t ne1,
+                            int64_t ne2,
+                            int64_t ne3);
+
+ggml_tensor* ggml_ext_zeros_like(ggml_context* ctx,
+                                 ggml_tensor* x);
+
+ggml_tensor* ggml_ext_ones(ggml_context* ctx,
+                           int64_t ne0,
+                           int64_t ne1,
+                           int64_t ne2,
+                           int64_t ne3);
+
+ggml_tensor* ggml_ext_ones_like(ggml_context* ctx,
+                                ggml_tensor* x);
+
+ggml_tensor* ggml_ext_cast_f32(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* a);
+
+// q: [N, L_q, C(n_head*d_head)] or [N*n_head, L_q, d_head]
+// k: [N, L_k, n_kv_head*d_head] or [N*n_kv_head, L_k, d_head]
+// v: [N, L_k, n_kv_head*d_head] or [N, L_k, n_kv_head, d_head]
+// mask: [N, L_q, L_k]
+// return: [N, L_q, C]
+ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
+                                    ggml_backend_t backend,
+                                    ggml_tensor* q,
+                                    ggml_tensor* k,
+                                    ggml_tensor* v,
+                                    int64_t n_head,
+                                    ggml_tensor* mask = nullptr,
+                                    bool skip_reshape = false,
+                                    bool flash_attn   = false,
+                                    float kv_scale    = 1.0f);
+
+ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 ggml_tensor* w,
+                                 ggml_tensor* b,
+                                 float eps = EPS);
+
+ggml_tensor* ggml_ext_group_norm(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 ggml_tensor* w,
+                                 ggml_tensor* b,
+                                 int num_groups = 32);
+
+ggml_tensor* ggml_ext_timestep_embedding(
+    ggml_context* ctx,
+    ggml_tensor* timesteps,
+    int dim,
+    int max_period    = 10000,
+    float time_factor = 1.0f);
+
+ggml_tensor* ggml_ext_vec_concat(ggml_context* ctx,
+                                 std::vector<ggml_tensor*>& tensors,
+                                 int dim);
+
+#endif  // __SD_CORE_GGML_EXTEND_H__

src/core/ggml_extend_backend.cpp

@@ -965,3 +965,35 @@ const char* sd_backend_module_name(SDBackendModule module) {
     }
     return "unknown";
 }
+
+void ggml_ext_backend_tensor_get_and_sync(ggml_backend_t backend, const ggml_tensor* tensor, void* data, size_t offset, size_t size) {
+    if ((sd_backend_is(backend, "ROCm") || sd_backend_is(backend, "CUDA") || sd_backend_is(backend, "SYCL")) &&
+        !sd_backend_is_cpu(backend)) {
+        ggml_backend_tensor_get_async(backend, tensor, data, offset, size);
+        ggml_backend_synchronize(backend);
+        return;
+    }
+
+    ggml_backend_tensor_get(tensor, data, offset, size);
+}
+
+float ggml_ext_backend_tensor_get_f32(ggml_tensor* tensor) {
+    GGML_ASSERT(tensor->type == GGML_TYPE_F32 || tensor->type == GGML_TYPE_F16 || tensor->type == GGML_TYPE_I32 || tensor->type == GGML_TYPE_BF16);
+    float value;
+    if (tensor->type == GGML_TYPE_F32) {
+        ggml_backend_tensor_get(tensor, &value, 0, sizeof(value));
+    } else if (tensor->type == GGML_TYPE_BF16) {
+        ggml_bf16_t bf16_value;
+        ggml_backend_tensor_get(tensor, &bf16_value, 0, sizeof(bf16_value));
+        value = ggml_bf16_to_fp32(bf16_value);
+    } else if (tensor->type == GGML_TYPE_F16) {
+        ggml_fp16_t f16_value;
+        ggml_backend_tensor_get(tensor, &f16_value, 0, sizeof(f16_value));
+        value = ggml_fp16_to_fp32(f16_value);
+    } else {  // GGML_TYPE_I32
+        int int32_value;
+        ggml_backend_tensor_get(tensor, &int32_value, 0, sizeof(int32_value));
+        value = (float)int32_value;
+    }
+    return value;
+}

src/core/ggml_extend_backend.h

@@ -96,4 +96,6 @@ std::string sd_backend_resolve_name(const std::string& name);
 const char* sd_backend_module_name(SDBackendModule module);
 void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value);
 bool add_rpc_devices(const std::string& servers);
+void ggml_ext_backend_tensor_get_and_sync(ggml_backend_t backend, const ggml_tensor* tensor, void* data, size_t offset, size_t size);
+float ggml_ext_backend_tensor_get_f32(ggml_tensor* tensor);
 #endif  // __SD_CORE_GGML_EXTEND_BACKEND_H__

src/core/ggml_runner.cpp

@@ -2,12 +2,666 @@
 #include <map>
 #include <utility>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/layer_split_partition.h"
 #include "core/segment_graph_bindings.h"
 #include "core/segment_weight_pipeline.h"
 
 using namespace sd;
 
+void GGMLRunner::alloc_params_ctx() {
+    ggml_init_params params;
+    params.mem_size   = static_cast<size_t>(MAX_PARAMS_TENSOR_NUM * ggml_tensor_overhead());
+    params.mem_buffer = nullptr;
+    params.no_alloc   = true;
+
+    params_ctx = ggml_init(params);
+    GGML_ASSERT(params_ctx != nullptr);
+    params_tensor_set_.clear();
+    params_tensor_set_dirty_ = true;
+}
+
+void GGMLRunner::free_params_ctx() {
+    if (params_ctx != nullptr) {
+        ggml_free(params_ctx);
+        params_ctx = nullptr;
+    }
+    params_tensor_set_.clear();
+    params_tensor_set_dirty_ = true;
+}
+
+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.no_alloc   = true;
+
+    compute_ctx = ggml_init(params);
+    GGML_ASSERT(compute_ctx != nullptr);
+}
+
+void GGMLRunner::free_compute_ctx() {
+    debug_tensors.clear();
+    if (compute_ctx != nullptr) {
+        ggml_free(compute_ctx);
+        compute_ctx = nullptr;
+    }
+    backend_tensor_data_map.clear();
+}
+
+void GGMLRunner::rebuild_params_tensor_set() {
+    if (!params_tensor_set_dirty_) {
+        return;
+    }
+    params_tensor_set_.clear();
+    if (params_ctx == nullptr) {
+        return;
+    }
+    for (ggml_tensor* t = ggml_get_first_tensor(params_ctx); t != nullptr; t = ggml_get_next_tensor(params_ctx, t)) {
+        params_tensor_set_.insert(t);
+    }
+    params_tensor_set_dirty_ = false;
+}
+
+ggml_tensor* GGMLRunner::canonical_param_tensor(ggml_tensor* tensor) {
+    if (tensor == nullptr) {
+        return nullptr;
+    }
+    if (params_tensor_set_.find(tensor) != params_tensor_set_.end()) {
+        return tensor;
+    }
+    if (tensor->view_src != nullptr &&
+        params_tensor_set_.find(tensor->view_src) != params_tensor_set_.end()) {
+        return tensor->view_src;
+    }
+    return nullptr;
+}
+
+std::vector<ggml_tensor*> GGMLRunner::collect_used_param_tensors(ggml_cgraph* gf) {
+    std::vector<ggml_tensor*> used_params;
+    rebuild_params_tensor_set();
+    if (gf == nullptr || params_tensor_set_.empty()) {
+        return used_params;
+    }
+
+    std::unordered_set<const ggml_tensor*> seen_params;
+    const int n_leafs = sd::ggml_graph_cut::leaf_count(gf);
+    seen_params.reserve(static_cast<size_t>(n_leafs));
+    for (int i = 0; i < n_leafs; ++i) {
+        ggml_tensor* leaf       = sd::ggml_graph_cut::leaf_tensor(gf, i);
+        ggml_tensor* param_leaf = canonical_param_tensor(leaf);
+        if (param_leaf != nullptr &&
+            seen_params.insert(param_leaf).second) {
+            used_params.push_back(param_leaf);
+        }
+    }
+    return used_params;
+}
+
+void GGMLRunner::evict_compute_backend_param_tensors(const std::vector<ggml_tensor*>& tensors) {
+    if (tensors.empty()) {
+        return;
+    }
+    auto manager = residency_manager.lock();
+    if (manager != nullptr) {
+        manager->evict_compute_backend_params(tensors);
+    }
+}
+
+void GGMLRunner::prepare_build_in_tensor_before() {
+    one_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_F32, 1);
+    ggml_set_name(one_tensor, "ggml_runner_build_in_tensor:one");
+    set_backend_tensor_data(one_tensor, one_vec.data());
+
+    zero_int_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, 1);
+    ggml_set_name(zero_int_tensor, "ggml_runner_build_in_tensor:zero_int");
+    set_backend_tensor_data(zero_int_tensor, zero_int_vec.data());
+}
+
+void GGMLRunner::prepare_build_in_tensor_after(ggml_cgraph* gf) {
+    ggml_build_forward_expand(gf, one_tensor);
+    ggml_build_forward_expand(gf, zero_int_tensor);
+}
+
+ggml_cgraph* GGMLRunner::new_graph_custom(size_t graph_size) {
+    if (weight_adapter) {
+        graph_size += weight_adapter->get_extra_graph_size();
+    }
+    return ggml_new_graph_custom(compute_ctx, graph_size, false);
+}
+
+ggml_cgraph* GGMLRunner::get_compute_graph(get_graph_cb_t get_graph) {
+    prepare_build_in_tensor_before();
+    ggml_cgraph* gf = get_graph();
+    if (gf == nullptr) {
+        return nullptr;
+    }
+    if (ggml_graph_n_nodes(gf) > 0) {
+        auto result = ggml_graph_node(gf, -1);
+        ggml_set_name(result, final_result_name.c_str());
+    }
+    for (const auto& entry : debug_tensors) {
+        if (entry.first != nullptr) {
+            ggml_build_forward_expand(gf, entry.first);
+        }
+    }
+    for (const auto& entry : cache_.outputs()) {
+        if (entry.second != nullptr) {
+            ggml_build_forward_expand(gf, entry.second);
+        }
+    }
+    prepare_build_in_tensor_after(gf);
+    return gf;
+}
+
+bool GGMLRunner::prepare_compute_graph(get_graph_cb_t get_graph,
+                                       ggml_cgraph** gf_out) {
+    GGML_ASSERT(gf_out != nullptr);
+
+    reset_compute_ctx();
+    ggml_cgraph* gf = get_compute_graph(get_graph);
+    if (gf == nullptr) {
+        free_compute_ctx();
+        return false;
+    }
+
+    *gf_out = gf;
+    return true;
+}
+
+ggml_backend_t GGMLRunner::backend_for_weight(const ggml_tensor* tensor) const {
+    if (tensor == nullptr || tensor->buffer == nullptr) {
+        return nullptr;
+    }
+    if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS ||
+        ggml_backend_buffer_is_host(tensor->buffer)) {
+        return nullptr;
+    }
+    ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer));
+    if (dev == nullptr) {
+        return nullptr;
+    }
+    if (ggml_backend_get_device(runtime_backend) == dev) {
+        return runtime_backend;
+    }
+    for (ggml_backend_t backend : extra_runtime_backends) {
+        if (ggml_backend_get_device(backend) == dev) {
+            return backend;
+        }
+    }
+    return nullptr;
+}
+
+void GGMLRunner::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph) {
+    if (sched == nullptr || gf == nullptr) {
+        return;
+    }
+    ggml_backend_t current = runtime_backend;
+    const int n_nodes      = ggml_graph_n_nodes(gf);
+    for (int i = 0; i < n_nodes; i++) {
+        ggml_tensor* node    = ggml_graph_node(gf, i);
+        auto node_assignment = graph_cut_layer_split_node_assignments_.find(original_graph == nullptr ? node : ggml_graph_node(original_graph, i));
+        if (node_assignment != graph_cut_layer_split_node_assignments_.end()) {
+            current = node_assignment->second;
+        }
+        for (int s = 0; s < GGML_MAX_SRC; s++) {
+            ggml_backend_t weight_backend = backend_for_weight(node->src[s]);
+            if (weight_backend != nullptr) {
+                if (node_assignment == graph_cut_layer_split_node_assignments_.end()) {
+                    current = weight_backend;
+                }
+            }
+        }
+        if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
+            node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) {
+            continue;
+        }
+        if (ggml_backend_supports_op(current, node)) {
+            ggml_backend_sched_set_tensor_backend(sched, node, current);
+        }
+    }
+}
+
+size_t GGMLRunner::retained_runtime_buffer_bytes(ggml_backend_t backend) const {
+    backend      = backend == nullptr ? runtime_backend : backend;
+    size_t bytes = workspace_.bytes(backend);
+    if (backend == runtime_backend) {
+        const size_t cache_bytes = cache_.resident_bytes(ggml_backend_get_device(backend));
+        bytes                    = cache_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cache_bytes;
+        const size_t cut_bytes   = cut_cache_.resident_bytes(ggml_backend_get_device(backend));
+        bytes                    = cut_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cut_bytes;
+    }
+    return bytes;
+}
+
+void GGMLRunner::sync_runtime_residency() {
+    if (auto manager = residency_manager.lock()) {
+        manager->update_runtime_residency(reinterpret_cast<uintptr_t>(this),
+                                          runtime_backend, retained_runtime_buffer_bytes());
+        for (auto backend : extra_runtime_backends) {
+            manager->update_runtime_residency(reinterpret_cast<uintptr_t>(this),
+                                              backend, retained_runtime_buffer_bytes(backend));
+        }
+    }
+}
+
+std::optional<sd::Tensor<float>> GGMLRunner::read_graph_tensor(ggml_tensor* tensor, const char* label) {
+    if (tensor == nullptr) {
+        LOG_ERROR("%s %s tensor is null", get_desc().c_str(), label);
+        return std::nullopt;
+    }
+    if (tensor->type != GGML_TYPE_F32) {
+        LOG_ERROR("%s %s tensor type mismatch: got %s",
+                  get_desc().c_str(),
+                  label,
+                  ggml_type_name(tensor->type));
+        return std::nullopt;
+    }
+    ggml_backend_buffer_t buf = sd::ggml_graph_cut::tensor_buffer(tensor);
+    if (buf == nullptr) {
+        LOG_ERROR("%s %s tensor buffer missing: name=%s op=%s buffer=%p view_src=%p view_src_buffer=%p data=%p",
+                  get_desc().c_str(),
+                  label,
+                  tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
+                  ggml_op_name(tensor->op),
+                  tensor->buffer,
+                  tensor->view_src,
+                  tensor->view_src ? tensor->view_src->buffer : nullptr,
+                  tensor->data);
+        return std::nullopt;
+    }
+
+    return sd::make_sd_tensor_from_ggml<float>(tensor);
+}
+
+void GGMLRunner::copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy) {
+    GGML_ASSERT(gf != nullptr);
+    std::unordered_set<const ggml_tensor*> graph_tensor_set;
+    const int n_leafs = sd::ggml_graph_cut::leaf_count(gf);
+    const int n_nodes = ggml_graph_n_nodes(gf);
+    graph_tensor_set.reserve(static_cast<size_t>(n_leafs + n_nodes));
+    for (int i = 0; i < n_leafs; ++i) {
+        graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i));
+    }
+    for (int i = 0; i < n_nodes; ++i) {
+        graph_tensor_set.insert(ggml_graph_node(gf, i));
+    }
+
+    for (auto& kv : backend_tensor_data_map) {
+        auto tensor = kv.first;
+        auto data   = kv.second;
+        if (tensor == nullptr || data == nullptr) {
+            continue;
+        }
+        const char* name = ggml_get_name(tensor);
+        if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) {
+            continue;
+        }
+        if (tensor->buffer == nullptr) {
+            LOG_WARN("%s skip backend tensor copy: tensor buffer not set, name='%s', ne=[%lld,%lld,%lld,%lld], type=%s",
+                     get_desc().c_str(),
+                     name != nullptr ? name : "",
+                     (long long)tensor->ne[0],
+                     (long long)tensor->ne[1],
+                     (long long)tensor->ne[2],
+                     (long long)tensor->ne[3],
+                     ggml_type_name(tensor->type));
+            continue;
+        }
+
+        ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
+        if (buf == nullptr) {
+            LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p",
+                     get_desc().c_str(),
+                     tensor && tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
+                     tensor ? ggml_op_name(tensor->op) : "<null>",
+                     data,
+                     tensor ? tensor->view_src : nullptr,
+                     (tensor && tensor->view_src) ? tensor->view_src->buffer : nullptr);
+            continue;
+        }
+
+        ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor));
+    }
+
+    if (clear_after_copy) {
+        backend_tensor_data_map.clear();
+    }
+}
+
+bool GGMLRunner::resolve_graph_cut_plan(ggml_cgraph* gf,
+                                        GraphCutPlan* plan_out) {
+    GGML_ASSERT(plan_out != nullptr);
+    GGML_ASSERT(gf != nullptr);
+    *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend,
+                                                 gf,
+                                                 &graph_cut_plan_cache_,
+                                                 params_tensor_set_,
+                                                 get_desc().c_str());
+    return true;
+}
+
+bool GGMLRunner::resolve_graph_cut_layer_split_plan(ggml_cgraph* gf,
+                                                    GraphCutPlan* plan_out) {
+    return resolve_graph_cut_plan(gf, plan_out);
+}
+
+bool GGMLRunner::assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
+    graph_cut_layer_split_node_assignments_.clear();
+    if (!graph_cut_layer_split_enabled) {
+        return true;
+    }
+    if (!is_multi_device()) {
+        LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", get_desc().c_str());
+        return false;
+    }
+
+    GraphCutPlan plan;
+    if (!resolve_graph_cut_layer_split_plan(gf, &plan)) {
+        return false;
+    }
+    if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) {
+        auto manager = residency_manager.lock();
+        if (manager == nullptr) {
+            LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
+            return false;
+        }
+        std::vector<ggml_tensor*> graph_params = collect_used_param_tensors(gf);
+        if (!graph_params.empty() &&
+            !manager->assign_compute_backend(graph_params, runtime_backend)) {
+            LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s",
+                      get_desc().c_str(),
+                      sd::layer_split_backend_device_display_name(runtime_backend).c_str());
+            return false;
+        }
+        for (ggml_tensor* param : graph_params) {
+            if (param != nullptr) {
+                graph_cut_layer_split_assignments_[param] = runtime_backend;
+            }
+        }
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        for (int i = 0; i < n_nodes; i++) {
+            ggml_tensor* node = ggml_graph_node(gf, i);
+            if (node != nullptr) {
+                graph_cut_layer_split_node_assignments_[node] = runtime_backend;
+            }
+        }
+        if (!graph_cut_layer_split_primary_notice_logged_) {
+            LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
+                     get_desc().c_str(),
+                     sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
+                     graph_params.size());
+            graph_cut_layer_split_primary_notice_logged_ = true;
+        } else {
+            LOG_VERBOSE("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
+                        get_desc().c_str(),
+                        sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
+                        graph_params.size());
+        }
+        return true;
+    }
+
+    std::vector<ggml_backend_t> split_backends;
+    split_backends.reserve(extra_runtime_backends.size() + 1);
+    split_backends.push_back(runtime_backend);
+    for (ggml_backend_t backend : extra_runtime_backends) {
+        if (backend != nullptr) {
+            split_backends.push_back(backend);
+        }
+    }
+
+    auto manager = residency_manager.lock();
+    if (manager == nullptr) {
+        LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
+        return false;
+    }
+
+    sd::GraphCutLayerSplitAssignment assignment;
+    auto canonicalize_param = [this](ggml_tensor* tensor) {
+        return canonical_param_tensor(tensor);
+    };
+    if (!sd::partition_graph_cut_layer_split(get_desc().c_str(),
+                                             gf,
+                                             plan,
+                                             split_backends,
+                                             graph_cut_layer_split_backend_vram_limits_,
+                                             max_graph_vram_bytes,
+                                             graph_cut_layer_split_assignments_,
+                                             canonicalize_param,
+                                             &assignment)) {
+        return false;
+    }
+
+    for (size_t i = 0; i < split_backends.size(); i++) {
+        if (assignment.tensors_by_backend[i].empty()) {
+            continue;
+        }
+        if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) {
+            LOG_ERROR("%s graph-cut layer split failed to assign params to %s",
+                      get_desc().c_str(),
+                      sd::layer_split_backend_device_display_name(split_backends[i]).c_str());
+            return false;
+        }
+    }
+
+    graph_cut_layer_split_node_assignments_ = std::move(assignment.node_assignments);
+    sd::log_graph_cut_layer_split_assignment(get_desc().c_str(), split_backends, assignment);
+
+    return true;
+}
+
+bool GGMLRunner::runner_start() {
+    if (runner_started_) {
+        return true;
+    }
+    cache_.clear();
+    workspace_.set_extra_backends(extra_runtime_backends);
+    if (auto manager = residency_manager.lock()) {
+        manager->set_workspace_reclaimer(reinterpret_cast<uintptr_t>(this), [this]() {
+            if (!workspace_.release()) {
+                return false;
+            }
+            sync_runtime_residency();
+            return true;
+        });
+    }
+    runner_started_ = true;
+    return true;
+}
+
+void GGMLRunner::runner_end() {
+    GGML_ASSERT(!graph_active_);
+    if (!runner_started_) {
+        return;
+    }
+    workspace_.release();
+    cache_.clear();
+    logged_compute_bytes_.clear();
+    logged_segment_count_ = 0;
+    if (auto manager = residency_manager.lock()) {
+        manager->clear_prefetched_params(reinterpret_cast<uintptr_t>(this));
+        std::vector<ggml_tensor*> tensors;
+        for (auto tensor = ggml_get_first_tensor(params_ctx); tensor != nullptr;
+             tensor      = ggml_get_next_tensor(params_ctx, tensor)) {
+            tensors.push_back(tensor);
+        }
+        manager->evict_compute_backend_params(tensors);
+        manager->remove_runtime_owner(reinterpret_cast<uintptr_t>(this));
+    }
+    runner_started_ = false;
+}
+
+GGMLRunner::GGMLRunner(ggml_backend_t backend,
+                       std::shared_ptr<DeviceResidencyManager> manager)
+    : runtime_backend(backend),
+      cache_(backend),
+      cut_cache_(backend),
+      workspace_(backend),
+      residency_manager(manager) {
+    GGML_ASSERT(runtime_backend != nullptr);
+    alloc_params_ctx();
+}
+
+GGMLRunner::~GGMLRunner() {
+    runner_end();
+    free_compute_ctx();
+    free_params_ctx();
+}
+
+GGMLRunnerContext GGMLRunner::get_context() {
+    GGMLRunnerContext runner_ctx;
+    runner_ctx.ggml_ctx              = compute_ctx;
+    runner_ctx.backend               = runtime_backend;
+    runner_ctx.flash_attn_enabled    = flash_attn_enabled;
+    runner_ctx.conv2d_direct_enabled = conv2d_direct_enabled;
+    runner_ctx.circular_x_enabled    = circular_x_enabled;
+    runner_ctx.circular_y_enabled    = circular_y_enabled;
+    runner_ctx.weight_adapter        = weight_adapter;
+    runner_ctx.debug_tensors         = &debug_tensors;
+    runner_ctx.get_cache_tensor      = [this](const std::string& name) {
+        return this->get_cache_tensor_by_name(name);
+    };
+    runner_ctx.cache_tensor = [this](const std::string& name, ggml_tensor* tensor) {
+        this->cache(name, tensor);
+    };
+    runner_ctx.set_backend_tensor_data = [this](ggml_tensor* tensor, const void* data) {
+        this->set_backend_tensor_data(tensor, data);
+    };
+    return runner_ctx;
+}
+
+void GGMLRunner::reset_compute_ctx() {
+    free_compute_ctx();
+    alloc_compute_ctx();
+}
+
+void GGMLRunner::free_cache_ctx_and_buffer() {
+    cache_.clear();
+    sync_runtime_residency();
+}
+
+void GGMLRunner::set_backend_tensor_data(ggml_tensor* tensor, const void* data) {
+    // The scheduler only allocates standalone data tensors when they are
+    // marked as graph inputs. The flag is harmless for single-backend graphs.
+    ggml_set_input(tensor);
+    backend_tensor_data_map[tensor] = data;
+}
+
+ggml_tensor* GGMLRunner::to_backend(ggml_tensor* tensor) {
+    GGML_ASSERT(compute_ctx != nullptr);
+    if (tensor == nullptr) {
+        return nullptr;
+    }
+    // it's performing a compute, check if backend isn't cpu
+    if (!sd_backend_is_cpu(runtime_backend) && (tensor->buffer == nullptr || ggml_backend_buffer_is_host(tensor->buffer))) {
+        // pass input tensors to gpu memory
+        auto backend_tensor = ggml_dup_tensor(compute_ctx, tensor);
+
+        set_backend_tensor_data(backend_tensor, tensor->data);
+        return backend_tensor;
+    } else {
+        return tensor;
+    }
+}
+
+void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) {
+    if (tensor != nullptr && tensor->view_src != nullptr) {
+        tensor = ggml_cont(compute_ctx, tensor);
+    }
+    if (tensor != nullptr) {
+        ggml_set_output(tensor);
+    }
+    cache_.stage(name, tensor);
+}
+
+std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
+                                                     int n_threads,
+                                                     bool auto_runner_end,
+                                                     bool no_return,
+                                                     const std::function<bool()>& read_outputs) {
+    if (graph_active_) {
+        LOG_ERROR("%s does not support reentrant graph execution", get_desc().c_str());
+        return std::nullopt;
+    }
+    if (!runner_start()) {
+        runner_end();
+        return std::nullopt;
+    }
+    struct RunnerEndGuard {
+        GGMLRunner& runner;
+        bool enabled;
+        ~RunnerEndGuard() {
+            if (enabled) {
+                runner.runner_end();
+            }
+        }
+    } runner_guard{*this, auto_runner_end};
+    graph_active_ = true;
+    bool success  = false;
+    struct GraphEndGuard {
+        GGMLRunner& runner;
+        const bool& success;
+        ~GraphEndGuard() {
+            runner.workspace_.segment_end();
+            runner.cache_.graph_end(false);
+            runner.cut_cache_.clear();
+            runner.free_compute_ctx();
+            runner.graph_active_ = false;
+            if (!success) {
+                runner.workspace_.release();
+            }
+            runner.sync_runtime_residency();
+        }
+    } graph_guard{*this, success};
+
+    ggml_cgraph* graph = nullptr;
+    if (!prepare_compute_graph(get_graph, &graph)) {
+        return std::nullopt;
+    }
+    rebuild_params_tensor_set();
+    auto output = execute_graph(graph, n_threads, no_return, read_outputs);
+    success     = output.has_value();
+    if (success) {
+        cache_.graph_end(true);
+    }
+    return output;
+}
+
+void GGMLRunner::set_graph_cut_layer_split_enabled(bool enabled) {
+    graph_cut_layer_split_enabled = enabled;
+    if (!enabled) {
+        graph_cut_layer_split_assignments_.clear();
+        graph_cut_layer_split_node_assignments_.clear();
+        graph_cut_layer_split_primary_notice_logged_ = false;
+    }
+}
+
+void GGMLRunner::set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {
+    graph_cut_layer_split_backend_vram_limits_ = limits;
+    graph_cut_layer_split_assignments_.clear();
+    graph_cut_layer_split_node_assignments_.clear();
+    graph_cut_layer_split_primary_notice_logged_ = false;
+}
+
+void GGMLRunner::set_runtime_backends(const std::vector<ggml_backend_t>& backends) {
+    extra_runtime_backends.clear();
+    for (ggml_backend_t backend : backends) {
+        if (backend == nullptr || backend == runtime_backend) {
+            continue;
+        }
+        if (std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), backend) ==
+            extra_runtime_backends.end()) {
+            extra_runtime_backends.push_back(backend);
+        }
+    }
+    workspace_.set_extra_backends(extra_runtime_backends);
+    graph_cut_layer_split_assignments_.clear();
+    graph_cut_layer_split_node_assignments_.clear();
+    graph_cut_layer_split_primary_notice_logged_ = false;
+}
+
 static size_t add_bytes(size_t a, size_t b) {
     return b > SIZE_MAX - a ? SIZE_MAX : a + b;
 }
@@ -275,7 +929,7 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
             }
             if (!no_return) {
                 auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str());
-                output      = read_graph_tensor<float>(result, "output");
+                output      = read_graph_tensor(result, "output");
                 if (!output.has_value()) {
                     return fail_segment("output readback");
                 }

src/core/ggml_runner.h

@@ -0,0 +1,350 @@
+#ifndef __SD_CORE_GGML_RUNNER_H__
+#define __SD_CORE_GGML_RUNNER_H__
+
+#include <cstddef>
+#include <functional>
+#include <map>
+#include <memory>
+#include <optional>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "core/compute_workspace.h"
+#include "core/ggml_graph_cut.h"
+#include "core/runner_cache.h"
+#include "core/tensor_ggml.hpp"
+#include "core/util.h"
+#include "device_residency_manager.h"
+
+/* SDXL with LoRA requires more space */
+#define MAX_PARAMS_TENSOR_NUM 32768
+#define MAX_GRAPH_SIZE 327680
+
+struct WeightAdapter {
+    struct ForwardParams {
+        enum class op_type_t {
+            OP_LINEAR,
+            OP_CONV2D,
+        } op_type;
+        struct {
+            bool force_prec_f32 = false;
+            float scale         = 1.f;
+        } linear;
+        struct conv2d_params_t {
+            int s0          = 1;
+            int s1          = 1;
+            int p0          = 0;
+            int p1          = 0;
+            int d0          = 1;
+            int d1          = 1;
+            bool direct     = false;
+            bool circular_x = false;
+            bool circular_y = false;
+            float scale     = 1.f;
+        } conv2d;
+    };
+    virtual ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name) = 0;
+    virtual ggml_tensor* forward_with_lora(ggml_context* ctx,
+                                           ggml_backend_t backend,
+                                           ggml_tensor* x,
+                                           ggml_tensor* w,
+                                           ggml_tensor* b,
+                                           const std::string& prefix,
+                                           ForwardParams forward_params)                                                              = 0;
+    virtual ggml_tensor* add_lora_to_output(ggml_context* ctx,
+                                            ggml_backend_t backend,
+                                            ggml_tensor* x,
+                                            ggml_tensor* w,
+                                            ggml_tensor* output,
+                                            const std::string& prefix,
+                                            ForwardParams forward_params)                                                             = 0;
+    virtual size_t get_extra_graph_size()                                                                                             = 0;
+};
+
+struct GGMLRunnerContext {
+    ggml_backend_t backend                                           = nullptr;
+    ggml_context* ggml_ctx                                           = nullptr;
+    bool flash_attn_enabled                                          = false;
+    bool conv2d_direct_enabled                                       = false;
+    bool circular_x_enabled                                          = false;
+    bool circular_y_enabled                                          = false;
+    ggml_tensor* ip_context                                          = nullptr;
+    float ip_scale                                                   = 1.0f;
+    std::shared_ptr<WeightAdapter> weight_adapter                    = nullptr;
+    std::vector<std::pair<ggml_tensor*, std::string>>* debug_tensors = nullptr;
+    std::function<ggml_tensor*(const std::string&)> get_cache_tensor;
+    std::function<void(const std::string&, ggml_tensor*)> cache_tensor;
+    std::function<void(ggml_tensor*, const void*)> set_backend_tensor_data;
+    std::map<std::pair<ggml_tensor*, int>, ggml_tensor*> int8_convrot_cache;
+
+    void capture_tensor(const std::string& name, ggml_tensor* tensor) {
+        if (debug_tensors == nullptr || tensor == nullptr) {
+            return;
+        }
+        ggml_tensor* snapshot = ggml_cont(ggml_ctx, tensor);
+        ggml_tensor* dst      = ggml_dup_tensor(ggml_ctx, snapshot);
+        snapshot              = ggml_cpy(ggml_ctx, snapshot, dst);
+        ggml_set_output(snapshot);
+        debug_tensors->push_back({snapshot, name});
+    }
+
+    ggml_tensor* load_cache_tensor(const std::string& name) const {
+        if (!get_cache_tensor) {
+            return nullptr;
+        }
+        return get_cache_tensor(name);
+    }
+
+    void persist_cache_tensor(const std::string& name, ggml_tensor* tensor) const {
+        if (!cache_tensor || tensor == nullptr) {
+            return;
+        }
+        cache_tensor(name, tensor);
+    }
+
+    void bind_backend_tensor_data(ggml_tensor* tensor, const void* data) const {
+        if (!set_backend_tensor_data || tensor == nullptr || data == nullptr) {
+            return;
+        }
+        set_backend_tensor_data(tensor, data);
+    }
+};
+
+struct GGMLRunner {
+private:
+    std::map<ggml_backend_t, size_t> logged_compute_bytes_;
+    size_t logged_segment_count_ = 0;
+
+    sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, size_t direct_bytes);
+    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);
+    std::optional<sd::Tensor<float>> execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs);
+
+protected:
+    typedef std::function<ggml_cgraph*()> get_graph_cb_t;
+    using GraphCutPlan = sd::ggml_graph_cut::Plan;
+
+    ggml_backend_t runtime_backend = nullptr;
+
+    ggml_context* params_ctx = nullptr;
+
+    sd::RunnerCache cache_;
+    sd::GraphCutTensorCache cut_cache_;
+    sd::ComputeWorkspace workspace_;
+    ggml_context* compute_ctx = nullptr;
+    bool runner_started_      = false;
+    bool graph_active_        = false;
+
+    size_t max_graph_vram_bytes        = 0;
+    bool graph_cut_layer_split_enabled = false;
+    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;
+    bool params_tensor_set_dirty_ = true;
+
+    std::vector<float> one_vec = {1.f};
+    ggml_tensor* one_tensor    = nullptr;
+
+    std::vector<int> zero_int_vec = {0};
+    ggml_tensor* zero_int_tensor  = nullptr;
+
+    std::map<ggml_tensor*, const void*> backend_tensor_data_map;
+    std::vector<std::pair<ggml_tensor*, std::string>> debug_tensors;
+    const std::string final_result_name = "ggml_runner_final_result_tensor";
+
+    bool flash_attn_enabled    = false;
+    bool conv2d_direct_enabled = false;
+    bool circular_x_enabled    = false;
+    bool circular_y_enabled    = false;
+
+    sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_;
+    std::unordered_set<const ggml_tensor*> params_tensor_set_;
+    std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_assignments_;
+    std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_node_assignments_;
+    bool graph_cut_layer_split_primary_notice_logged_ = false;
+
+    template <typename T>
+    static sd::Tensor<T> take_or_empty(std::optional<sd::Tensor<T>> tensor) {
+        if (!tensor.has_value()) {
+            return {};
+        }
+        return std::move(*tensor);
+    }
+
+    template <typename T>
+    static sd::Tensor<T> restore_trailing_singleton_dims(std::optional<sd::Tensor<T>> tensor,
+                                                         size_t expected_dim) {
+        return restore_trailing_singleton_dims(take_or_empty(std::move(tensor)), expected_dim);
+    }
+
+    template <typename T>
+    static sd::Tensor<T> restore_trailing_singleton_dims(sd::Tensor<T> tensor,
+                                                         size_t expected_dim) {
+        if (tensor.empty()) {
+            return tensor;
+        }
+        while (static_cast<size_t>(tensor.dim()) < expected_dim) {
+            tensor.unsqueeze_(tensor.dim());
+        }
+        return tensor;
+    }
+
+    void alloc_params_ctx();
+
+    void free_params_ctx();
+
+    void alloc_compute_ctx();
+
+    void free_compute_ctx();
+
+    void rebuild_params_tensor_set();
+
+    ggml_tensor* canonical_param_tensor(ggml_tensor* tensor);
+
+    std::vector<ggml_tensor*> collect_used_param_tensors(ggml_cgraph* gf);
+
+    void evict_compute_backend_param_tensors(const std::vector<ggml_tensor*>& tensors);
+
+    void prepare_build_in_tensor_before();
+
+    void prepare_build_in_tensor_after(ggml_cgraph* gf);
+
+    ggml_cgraph* new_graph_custom(size_t graph_size);
+
+    ggml_cgraph* get_compute_graph(get_graph_cb_t get_graph);
+
+    bool prepare_compute_graph(get_graph_cb_t get_graph,
+                               ggml_cgraph** gf_out);
+
+    ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const;
+
+    // Weightless ops have no scheduler anchor, so pin them to the most recent
+    // weight device. Views must stay unpinned or cross-device copies can be
+    // skipped for their consumers.
+    void pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph = nullptr);
+
+    bool is_multi_device() const {
+        return !extra_runtime_backends.empty();
+    }
+
+    size_t reusable_compute_buffer_bytes() const {
+        return workspace_.bytes(runtime_backend);
+    }
+
+    size_t retained_runtime_buffer_bytes(ggml_backend_t backend = nullptr) const;
+
+    void sync_runtime_residency();
+
+    std::optional<sd::Tensor<float>> read_graph_tensor(ggml_tensor* tensor, const char* label);
+
+    void copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy = true);
+
+    bool resolve_graph_cut_plan(ggml_cgraph* gf,
+                                GraphCutPlan* plan_out);
+
+    bool resolve_graph_cut_layer_split_plan(ggml_cgraph* gf,
+                                            GraphCutPlan* plan_out);
+
+    bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf);
+
+public:
+    bool runner_start();
+
+    bool runner_started() const { return runner_started_; }
+
+    void runner_end();
+
+public:
+    virtual std::string get_desc() = 0;
+
+    GGMLRunner(ggml_backend_t backend,
+               std::shared_ptr<DeviceResidencyManager> manager = nullptr);
+
+    virtual ~GGMLRunner();
+
+    virtual GGMLRunnerContext get_context();
+
+    void reset_compute_ctx();
+
+public:
+    void free_cache_ctx_and_buffer();
+
+    // do copy after alloc graph
+    void set_backend_tensor_data(ggml_tensor* tensor, const void* data);
+
+    template <typename T>
+    ggml_tensor* make_input(const sd::Tensor<T>& tensor) {
+        ggml_tensor* input = sd::make_ggml_tensor(compute_ctx, tensor, false);
+        set_backend_tensor_data(input, tensor.data());
+        return input;
+    }
+
+    template <typename T>
+    ggml_tensor* make_optional_input(const sd::Tensor<T>& tensor) {
+        if (tensor.empty()) {
+            return nullptr;
+        }
+        return make_input(tensor);
+    }
+
+    template <typename T>
+    ggml_tensor* make_optional_input(const sd::Tensor<T>* tensor) {
+        if (tensor == nullptr) {
+            return nullptr;
+        }
+        return make_input(*tensor);
+    }
+
+    ggml_tensor* to_backend(ggml_tensor* tensor);
+
+    void cache(const std::string name, ggml_tensor* tensor);
+
+    ggml_tensor* get_cache_tensor_by_name(const std::string& name) {
+        return cache_.get(name);
+    }
+
+    std::optional<sd::Tensor<float>> compute(get_graph_cb_t get_graph,
+                                             int n_threads,
+                                             bool auto_runner_end                      = true,
+                                             bool no_return                            = false,
+                                             const std::function<bool()>& read_outputs = {});
+
+    void set_flash_attention_enabled(bool enabled) {
+        flash_attn_enabled = enabled;
+    }
+
+    void set_conv2d_direct_enabled(bool enabled) {
+        conv2d_direct_enabled = enabled;
+    }
+
+    void set_circular_axes(bool circular_x, bool circular_y) {
+        circular_x_enabled = circular_x;
+        circular_y_enabled = circular_y;
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {
+        weight_adapter = adapter;
+    }
+
+    void set_max_graph_vram_bytes(size_t max_vram_bytes) {
+        max_graph_vram_bytes = max_vram_bytes;
+    }
+
+    void set_graph_cut_layer_split_enabled(bool enabled);
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits);
+
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends);
+};
+
+#endif  // __SD_CORE_GGML_RUNNER_H__

src/core/ggml_tensor_utils.cpp

@@ -0,0 +1,428 @@
+#include "core/ggml_tensor_utils.h"
+
+#include <algorithm>
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+
+#include "core/ggml_extend_backend.h"
+#include "core/rng.hpp"
+
+void ggml_ext_im_set_randn_f32(ggml_tensor* tensor, std::shared_ptr<RNG> rng) {
+    uint32_t n                        = (uint32_t)ggml_nelements(tensor);
+    std::vector<float> random_numbers = rng->randn(n);
+    for (uint32_t i = 0; i < n; i++) {
+        ggml_ext_im_set_f32_1d(tensor, i, random_numbers[i]);
+    }
+}
+
+void print_ggml_tensor(ggml_tensor* tensor, bool shape_only, const char* mark) {
+    printf("%s (%s): shape(%zu, %zu, %zu, %zu)\n", mark, ggml_type_name(tensor->type), tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]);
+    fflush(stdout);
+    if (shape_only) {
+        return;
+    }
+    int range = 3;
+    for (int i3 = 0; i3 < tensor->ne[3]; i3++) {
+        if (i3 >= range && i3 + range < tensor->ne[3]) {
+            continue;
+        }
+        for (int i2 = 0; i2 < tensor->ne[2]; i2++) {
+            if (i2 >= range && i2 + range < tensor->ne[2]) {
+                continue;
+            }
+            for (int i1 = 0; i1 < tensor->ne[1]; i1++) {
+                if (i1 >= range && i1 + range < tensor->ne[1]) {
+                    continue;
+                }
+                for (int i0 = 0; i0 < tensor->ne[0]; i0++) {
+                    if (i0 >= range && i0 + range < tensor->ne[0]) {
+                        continue;
+                    }
+                    if (tensor->type == GGML_TYPE_F32) {
+                        printf("  [%d, %d, %d, %d] = %f\n", i3, i2, i1, i0, ggml_ext_tensor_get_f32(tensor, i0, i1, i2, i3));
+                    } else if (tensor->type == GGML_TYPE_F16) {
+                        printf("  [%d, %d, %d, %d] = %f\n", i3, i2, i1, i0, ggml_fp16_to_fp32(ggml_ext_tensor_get_f16(tensor, i0, i1, i2, i3)));
+                    } else if (tensor->type == GGML_TYPE_I32) {
+                        printf("  [%d, %d, %d, %d] = %i3\n", i3, i2, i1, i0, ggml_ext_tensor_get_i32(tensor, i0, i1, i2, i3));
+                    }
+                    fflush(stdout);
+                }
+            }
+        }
+    }
+}
+
+void ggml_ext_tensor_iter(
+    ggml_tensor* tensor,
+    const std::function<void(ggml_tensor*, int64_t, int64_t, int64_t, int64_t)>& fn) {
+    int64_t n0 = tensor->ne[0];
+    int64_t n1 = tensor->ne[1];
+    int64_t n2 = tensor->ne[2];
+    int64_t n3 = tensor->ne[3];
+
+    for (int64_t i3 = 0; i3 < n3; i3++) {
+        for (int64_t i2 = 0; i2 < n2; i2++) {
+            for (int64_t i1 = 0; i1 < n1; i1++) {
+                for (int64_t i0 = 0; i0 < n0; i0++) {
+                    fn(tensor, i0, i1, i2, i3);
+                }
+            }
+        }
+    }
+}
+
+void ggml_ext_tensor_iter(
+    ggml_tensor* tensor,
+    const std::function<void(ggml_tensor*, int64_t)>& fn) {
+    int64_t n0 = tensor->ne[0];
+    int64_t n1 = tensor->ne[1];
+    int64_t n2 = tensor->ne[2];
+    int64_t n3 = tensor->ne[3];
+
+    for (int64_t i = 0; i < ggml_nelements(tensor); i++) {
+        fn(tensor, i);
+    }
+}
+
+void ggml_ext_tensor_diff(
+    ggml_tensor* a,
+    ggml_tensor* b,
+    float gap) {
+    GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b));
+    ggml_ext_tensor_iter(a, [&](ggml_tensor* a, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+        float a_value = ggml_ext_tensor_get_f32(a, i0, i1, i2, i3);
+        float b_value = ggml_ext_tensor_get_f32(b, i0, i1, i2, i3);
+        if (abs(a_value - b_value) > gap) {
+            LOG_WARN("[%ld, %ld, %ld, %ld] %f %f", i3, i2, i1, i0, a_value, b_value);
+        }
+    });
+}
+
+ggml_tensor* load_tensor_from_file(ggml_context* ctx, const std::string& file_path) {
+    std::ifstream file(file_path, std::ios::binary);
+    if (!file.is_open()) {
+        LOG_ERROR("failed to open '%s'", file_path.c_str());
+        return nullptr;
+    }
+    int32_t n_dims;
+    int32_t length;
+    int32_t ttype;
+
+    file.read(reinterpret_cast<char*>(&n_dims), sizeof(n_dims));
+    file.read(reinterpret_cast<char*>(&length), sizeof(length));
+    file.read(reinterpret_cast<char*>(&ttype), sizeof(ttype));
+
+    LOG_VERBOSE("load_tensor_from_file %d %d %d", n_dims, length, ttype);
+
+    if (file.eof()) {
+        LOG_ERROR("incomplete file '%s'", file_path.c_str());
+        return nullptr;
+    }
+
+    int32_t nelements = 1;
+    int32_t ne[4]     = {1, 1, 1, 1};
+    for (int i = 0; i < n_dims; ++i) {
+        file.read(reinterpret_cast<char*>(&ne[i]), sizeof(ne[i]));
+        nelements *= ne[i];
+    }
+    std::string name(length, 0);
+    file.read(&name[0], length);
+    ggml_tensor* tensor = ggml_new_tensor_4d(ctx, (ggml_type)ttype, ne[0], ne[1], ne[2], ne[3]);
+    const size_t bpe    = ggml_type_size(ggml_type(ttype));
+    file.read(reinterpret_cast<char*>(tensor->data), ggml_nbytes(tensor));
+    return tensor;
+}
+
+// __STATIC_INLINE__ void save_tensor_to_file(const std::string& file_name, ggml_tensor* tensor, const std::string & name) {
+//     std::string file_name_ = file_name + ".tensor";
+//     std::string name_ = name;
+//     std::ofstream file("./" + file_name_, std::ios::binary);
+//     file.write(reinterpret_cast<char*>(&tensor->n_dims), sizeof(tensor->n_dims));
+//     int len = (int)name_.size();
+//     file.write(reinterpret_cast<char*>(&len), sizeof(len));
+//     int ttype = (int)tensor->type;
+//     file.write(reinterpret_cast<char*>(&ttype), sizeof(ttype));
+//     for (int i = 0; i < tensor->n_dims; ++i) {
+//         int ne_ = (int) tensor->ne[i];
+//         file.write(reinterpret_cast<char*>(&ne_), sizeof(ne_));
+//     }
+//     file.write(&name_[0], len);
+//     char* data = nullptr;
+//     file.write((char*)tensor->data, ggml_nbytes(tensor));
+//     file.close();
+// }
+
+uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, uint8_t* image_data) {
+    int64_t width    = input->ne[0];
+    int64_t height   = input->ne[1];
+    int64_t channels = input->ne[2];
+    GGML_ASSERT(input->type == GGML_TYPE_F32);
+    if (image_data == nullptr) {
+        image_data = (uint8_t*)malloc(width * height * channels);
+    }
+    for (int iy = 0; iy < height; iy++) {
+        for (int ix = 0; ix < width; ix++) {
+            for (int k = 0; k < channels; k++) {
+                float value                                               = ggml_ext_tensor_get_f32(input, ix, iy, k);
+                *(image_data + iy * width * channels + ix * channels + k) = (uint8_t)(value * 255.0f);
+            }
+        }
+    }
+    return image_data;
+}
+
+uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, int idx, bool video) {
+    int64_t width  = input->ne[0];
+    int64_t height = input->ne[1];
+    int64_t channels;
+    if (video) {
+        channels = input->ne[3];
+    } else {
+        channels = input->ne[2];
+    }
+    GGML_ASSERT(channels == 3 && input->type == GGML_TYPE_F32);
+    uint8_t* image_data = (uint8_t*)malloc(width * height * channels);
+    for (int ih = 0; ih < height; ih++) {
+        for (int iw = 0; iw < width; iw++) {
+            for (int ic = 0; ic < channels; ic++) {
+                float value;
+                if (video) {
+                    value = ggml_ext_tensor_get_f32(input, iw, ih, idx, ic);
+                } else {
+                    value = ggml_ext_tensor_get_f32(input, iw, ih, ic, idx);
+                }
+                *(image_data + ih * width * channels + iw * channels + ic) = (uint8_t)(value * 255.0f);
+            }
+        }
+    }
+    return image_data;
+}
+
+void sd_image_to_ggml_tensor(sd_image_t image,
+                             ggml_tensor* tensor,
+                             bool scale) {
+    GGML_ASSERT(image.width == tensor->ne[0]);
+    GGML_ASSERT(image.height == tensor->ne[1]);
+    GGML_ASSERT(image.channel == tensor->ne[2]);
+    GGML_ASSERT(1 == tensor->ne[3]);
+    GGML_ASSERT(tensor->type == GGML_TYPE_F32);
+    ggml_ext_tensor_iter(tensor, [&](ggml_tensor* tensor, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+        float value = sd_image_get_f32(image, i0, i1, i2, scale);
+        ggml_ext_tensor_set_f32(tensor, value, i0, i1, i2, i3);
+    });
+}
+
+void ggml_ext_tensor_apply_mask(ggml_tensor* image_data,
+                                ggml_tensor* mask,
+                                ggml_tensor* output,
+                                float masked_value) {
+    int64_t width    = output->ne[0];
+    int64_t height   = output->ne[1];
+    int64_t channels = output->ne[2];
+    float rescale_mx = 1.f * mask->ne[0] / output->ne[0];
+    float rescale_my = 1.f * mask->ne[1] / output->ne[1];
+    GGML_ASSERT(output->type == GGML_TYPE_F32);
+    for (int ix = 0; ix < width; ix++) {
+        for (int iy = 0; iy < height; iy++) {
+            int mx  = (int)(ix * rescale_mx);
+            int my  = (int)(iy * rescale_my);
+            float m = ggml_ext_tensor_get_f32(mask, mx, my);
+            m       = round(m);  // inpaint models need binary masks
+            ggml_ext_tensor_set_f32(mask, m, mx, my);
+            for (int k = 0; k < channels; k++) {
+                float value = ggml_ext_tensor_get_f32(image_data, ix, iy, k);
+                value       = (1 - m) * (value - masked_value) + masked_value;
+                ggml_ext_tensor_set_f32(output, value, ix, iy, k);
+            }
+        }
+    }
+}
+
+float ggml_ext_tensor_mean(ggml_tensor* src) {
+    float mean        = 0.0f;
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        mean += data[i] / nelements * 1.0f;
+    }
+    return mean;
+}
+
+void ggml_ext_tensor_add_inplace(ggml_tensor* a, ggml_tensor* b) {
+    GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b));
+    int64_t nelements = ggml_nelements(a);
+    float* vec_a      = (float*)a->data;
+    float* vec_b      = (float*)b->data;
+    for (int i = 0; i < nelements; i++) {
+        vec_a[i] = vec_a[i] + vec_b[i];
+    }
+}
+
+void ggml_ext_tensor_scale_inplace(ggml_tensor* src, float scale) {
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        data[i] = data[i] * scale;
+    }
+}
+
+void ggml_ext_tensor_clamp_inplace(ggml_tensor* src, float min, float max) {
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        float val = data[i];
+        data[i]   = val < min ? min : (val > max ? max : val);
+    }
+}
+
+ggml_tensor* ggml_ext_tensor_concat(ggml_context* ctx,
+                                    ggml_tensor* a,
+                                    ggml_tensor* b,
+                                    int dim) {
+    int64_t ne[GGML_MAX_DIMS];
+    for (int d = 0; d < GGML_MAX_DIMS; ++d) {
+        if (d == dim) {
+            ne[d] = a->ne[d] + b->ne[d];
+            continue;
+        }
+        GGML_ASSERT(a->ne[d] == b->ne[d]);
+        ne[d] = a->ne[d];
+    }
+    ggml_tensor* result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, ne);
+    int64_t o[4]        = {0, 0, 0, 0};
+    o[dim]              = a->ne[dim];
+
+    float v;
+    for (int i3 = 0; i3 < result->ne[3]; i3++) {
+        for (int i2 = 0; i2 < result->ne[2]; i2++) {
+            for (int i1 = 0; i1 < result->ne[1]; i1++) {
+                for (int i0 = 0; i0 < result->ne[0]; i0++) {
+                    if (i0 < a->ne[0] && i1 < a->ne[1] && i2 < a->ne[2] && i3 < a->ne[3]) {
+                        v = ggml_ext_tensor_get_f32(a, i0, i1, i2, i3);
+                    } else {
+                        v = ggml_ext_tensor_get_f32(b, i0 - o[0], i1 - o[1], i2 - o[2], i3 - o[3]);
+                    }
+
+                    ggml_ext_tensor_set_f32(result, v, i0, i1, i2, i3);
+                }
+            }
+        }
+    }
+    return result;
+}
+
+void scale_to_minus1_1(ggml_tensor* src) {
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        float val = data[i];
+        data[i]   = val * 2.0f - 1.0f;
+    }
+}
+
+void scale_to_0_1(ggml_tensor* src) {
+    int64_t nelements = ggml_nelements(src);
+    float* data       = (float*)src->data;
+    for (int i = 0; i < nelements; i++) {
+        float val = data[i];
+        data[i]   = (val + 1.0f) * 0.5f;
+    }
+}
+
+ggml_tensor* vector_to_ggml_tensor(ggml_context* ctx,
+                                   const std::vector<float>& vec) {
+    ggml_tensor* t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, vec.size());
+    memcpy(t->data, (const void*)vec.data(), ggml_nbytes(t));
+    return t;
+}
+
+ggml_tensor* vector_to_ggml_tensor_i32(ggml_context* ctx,
+                                       const std::vector<int>& vec) {
+    ggml_tensor* t = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, vec.size());
+    memcpy(t->data, (const void*)vec.data(), ggml_nbytes(t));
+    return t;
+}
+
+std::vector<float> arange(float start, float end, float step) {
+    std::vector<float> result;
+
+    for (float value = start; value < end; value += step) {
+        result.push_back(value);
+    }
+
+    return result;
+}
+
+std::vector<float> timestep_embedding(std::vector<float> timesteps,
+                                      int dim,
+                                      int max_period,
+                                      bool flip_sin_to_cos,
+                                      float scale) {
+    // timesteps: [N,]
+    // embedding: [N, dim]
+    size_t N = timesteps.size();
+    std::vector<float> embedding(N * dim, 0.f);
+    int half = dim / 2;
+    std::vector<float> freqs(half);
+    for (int i = 0; i < half; ++i) {
+        freqs[i] = (float)std::exp(-std::log(max_period) * i / half);
+    }
+    for (int i = 0; i < N; ++i) {
+        for (int j = 0; j < half; ++j) {
+            float arg = timesteps[i] * freqs[j] * scale;
+            if (flip_sin_to_cos) {
+                embedding[i * dim + j]        = std::cos(arg);
+                embedding[i * dim + j + half] = std::sin(arg);
+            } else {
+                embedding[i * dim + j]        = std::sin(arg);
+                embedding[i * dim + j + half] = std::cos(arg);
+            }
+        }
+    }
+    return embedding;
+}
+
+void set_timestep_embedding(std::vector<float> timesteps,
+                            ggml_tensor* embedding,
+                            int dim,
+                            int max_period) {
+    std::vector<float> embedding_vec = timestep_embedding(timesteps, dim, max_period);
+    memcpy(((char*)embedding->data), ((char*)embedding_vec.data()), ggml_nbytes(embedding));
+}
+
+void set_timestep_embedding(std::vector<float> timesteps,
+                            sd::Tensor<float>* embedding,
+                            int dim,
+                            int max_period) {
+    GGML_ASSERT(embedding != nullptr);
+    std::vector<float> embedding_vec = timestep_embedding(timesteps, dim, max_period);
+    if (embedding->numel() != static_cast<int64_t>(embedding_vec.size())) {
+        embedding->resize({dim, static_cast<int64_t>(timesteps.size())});
+    }
+    std::copy(embedding_vec.begin(), embedding_vec.end(), embedding->values().begin());
+}
+
+ggml_tensor* new_timestep_embedding(ggml_context* ctx,
+                                    std::vector<float> timesteps,
+                                    int dim,
+                                    int max_period) {
+    // timesteps: [N,]
+    // embedding: [N, dim]
+    std::vector<float> embedding_vec = timestep_embedding(timesteps, dim, max_period);
+    ggml_tensor* embedding           = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, timesteps.size());
+    if (embedding->data != nullptr) {
+        memcpy(((char*)embedding->data), ((char*)embedding_vec.data()), ggml_nbytes(embedding));
+    } else {
+        ggml_backend_tensor_set(embedding, embedding_vec.data(), 0, ggml_nbytes(embedding));
+    }
+    return embedding;
+}
+
+size_t ggml_tensor_num(ggml_context* ctx) {
+    size_t num = 0;
+    for (ggml_tensor* t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+        num++;
+    }
+    return num;
+}

src/core/ggml_tensor_utils.h

@@ -0,0 +1,210 @@
+#ifndef __SD_CORE_GGML_TENSOR_UTILS_H__
+#define __SD_CORE_GGML_TENSOR_UTILS_H__
+
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <functional>
+#include <memory>
+#include <string>
+#include <type_traits>
+#include <vector>
+
+#include "core/tensor.hpp"
+#include "core/util.h"
+#include "ggml-backend.h"
+#include "ggml.h"
+#include "stable-diffusion.h"
+
+class RNG;
+
+__STATIC_INLINE__ int align_up_offset(int n, int multiple) {
+    return (multiple - n % multiple) % multiple;
+}
+
+__STATIC_INLINE__ int align_up(int n, int multiple) {
+    return n + align_up_offset(n, multiple);
+}
+
+void ggml_ext_im_set_randn_f32(ggml_tensor* tensor, std::shared_ptr<RNG> rng);
+
+__STATIC_INLINE__ void ggml_ext_tensor_set_f32(ggml_tensor* tensor, float value, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
+    GGML_ASSERT(tensor->nb[0] == sizeof(float));
+    *(float*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]) = value;
+}
+
+__STATIC_INLINE__ float ggml_ext_tensor_get_f32(const ggml_tensor* tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
+    if (tensor->buffer != nullptr) {
+        float value;
+        ggml_backend_tensor_get(tensor, &value, i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0], sizeof(float));
+        return value;
+    }
+    GGML_ASSERT(tensor->nb[0] == sizeof(float));
+    return *(float*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]);
+}
+
+__STATIC_INLINE__ int ggml_ext_tensor_get_i32(const ggml_tensor* tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
+    if (tensor->buffer != nullptr) {
+        int value;
+        ggml_backend_tensor_get(tensor, &value, i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0], sizeof(int));
+        return value;
+    }
+    GGML_ASSERT(tensor->nb[0] == sizeof(int));
+    return *(int*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]);
+}
+
+__STATIC_INLINE__ ggml_fp16_t ggml_ext_tensor_get_f16(const ggml_tensor* tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) {
+    GGML_ASSERT(tensor->nb[0] == sizeof(ggml_fp16_t));
+    return *(ggml_fp16_t*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]);
+}
+
+__STATIC_INLINE__ float sd_image_get_f32(sd_image_t image, int64_t iw, int64_t ih, int64_t ic, bool scale = true) {
+    float value = *(image.data + ih * image.width * image.channel + iw * image.channel + ic);
+    if (scale) {
+        value /= 255.f;
+    }
+    return value;
+}
+
+void print_ggml_tensor(ggml_tensor* tensor, bool shape_only = false, const char* mark = "");
+
+template <typename T>
+__STATIC_INLINE__ void print_sd_tensor(const sd::Tensor<T>& tensor, bool shape_only = false, const char* mark = "") {
+    printf("%s: shape(", mark);
+    for (size_t i = 0; i < static_cast<size_t>(tensor.dim()); ++i) {
+        printf("%s%lld", i == 0 ? "" : ", ", static_cast<long long>(tensor.shape()[i]));
+    }
+    printf(")\n");
+    fflush(stdout);
+    if (shape_only) {
+        return;
+    }
+    if (tensor.empty()) {
+        return;
+    }
+    int range                  = 3;
+    std::vector<int64_t> shape = tensor.shape();
+    while (shape.size() < 4) {
+        shape.push_back(1);
+    }
+    for (int64_t i3 = 0; i3 < shape[3]; i3++) {
+        if (i3 >= range && i3 + range < shape[3]) {
+            continue;
+        }
+        for (int64_t i2 = 0; i2 < shape[2]; i2++) {
+            if (i2 >= range && i2 + range < shape[2]) {
+                continue;
+            }
+            for (int64_t i1 = 0; i1 < shape[1]; i1++) {
+                if (i1 >= range && i1 + range < shape[1]) {
+                    continue;
+                }
+                for (int64_t i0 = 0; i0 < shape[0]; i0++) {
+                    if (i0 >= range && i0 + range < shape[0]) {
+                        continue;
+                    }
+                    size_t offset = static_cast<size_t>(i0 + shape[0] * (i1 + shape[1] * (i2 + shape[2] * i3)));
+                    printf("  [%lld, %lld, %lld, %lld] = ", static_cast<long long>(i3), static_cast<long long>(i2), static_cast<long long>(i1), static_cast<long long>(i0));
+                    if constexpr (std::is_same_v<T, float>) {
+                        printf("%f\n", tensor[static_cast<int64_t>(offset)]);
+                    } else if constexpr (std::is_same_v<T, ggml_fp16_t>) {
+                        printf("%f\n", ggml_fp16_to_fp32(tensor[static_cast<int64_t>(offset)]));
+                    } else if constexpr (std::is_same_v<T, int32_t>) {
+                        printf("%d\n", tensor[static_cast<int64_t>(offset)]);
+                    } else if constexpr (std::is_same_v<T, int64_t>) {
+                        printf("%lld\n", static_cast<long long>(tensor[static_cast<int64_t>(offset)]));
+                    }
+                    fflush(stdout);
+                }
+            }
+        }
+    }
+}
+
+void ggml_ext_tensor_iter(
+    ggml_tensor* tensor,
+    const std::function<void(ggml_tensor*, int64_t, int64_t, int64_t, int64_t)>& fn);
+
+void ggml_ext_tensor_iter(
+    ggml_tensor* tensor,
+    const std::function<void(ggml_tensor*, int64_t)>& fn);
+
+void ggml_ext_tensor_diff(
+    ggml_tensor* a,
+    ggml_tensor* b,
+    float gap = 0.1f);
+
+ggml_tensor* load_tensor_from_file(ggml_context* ctx, const std::string& file_path);
+
+__STATIC_INLINE__ float sigmoid(float x) {
+    return 1 / (1.0f + expf(-x));
+}
+
+// SPECIAL OPERATIONS WITH TENSORS
+
+uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, uint8_t* image_data = nullptr);
+
+uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, int idx, bool video = false);
+
+void sd_image_to_ggml_tensor(sd_image_t image,
+                             ggml_tensor* tensor,
+                             bool scale = true);
+
+void ggml_ext_tensor_apply_mask(ggml_tensor* image_data,
+                                ggml_tensor* mask,
+                                ggml_tensor* output,
+                                float masked_value = 0.5f);
+
+float ggml_ext_tensor_mean(ggml_tensor* src);
+
+// a = a+b
+void ggml_ext_tensor_add_inplace(ggml_tensor* a, ggml_tensor* b);
+
+void ggml_ext_tensor_scale_inplace(ggml_tensor* src, float scale);
+
+void ggml_ext_tensor_clamp_inplace(ggml_tensor* src, float min, float max);
+
+ggml_tensor* ggml_ext_tensor_concat(ggml_context* ctx,
+                                    ggml_tensor* a,
+                                    ggml_tensor* b,
+                                    int dim);
+
+// convert values from [0, 1] to [-1, 1]
+void scale_to_minus1_1(ggml_tensor* src);
+
+// convert values from [-1, 1] to [0, 1]
+void scale_to_0_1(ggml_tensor* src);
+
+ggml_tensor* vector_to_ggml_tensor(ggml_context* ctx,
+                                   const std::vector<float>& vec);
+
+ggml_tensor* vector_to_ggml_tensor_i32(ggml_context* ctx,
+                                       const std::vector<int>& vec);
+
+std::vector<float> arange(float start, float end, float step = 1.f);
+
+// Ref: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/util.py#L151
+std::vector<float> timestep_embedding(std::vector<float> timesteps,
+                                      int dim,
+                                      int max_period       = 10000,
+                                      bool flip_sin_to_cos = true,
+                                      float scale          = 1.f);
+
+void set_timestep_embedding(std::vector<float> timesteps,
+                            ggml_tensor* embedding,
+                            int dim,
+                            int max_period = 10000);
+
+void set_timestep_embedding(std::vector<float> timesteps,
+                            sd::Tensor<float>* embedding,
+                            int dim,
+                            int max_period = 10000);
+
+ggml_tensor* new_timestep_embedding(ggml_context* ctx,
+                                    std::vector<float> timesteps,
+                                    int dim,
+                                    int max_period = 10000);
+
+size_t ggml_tensor_num(ggml_context* ctx);
+
+#endif  // __SD_CORE_GGML_TENSOR_UTILS_H__

src/core/util.cpp

@@ -15,6 +15,7 @@
 #include <thread>
 #include <unordered_set>
 #include <vector>
+#include "core/ggml_tensor_utils.h"
 #include "runtime/preprocessing.hpp"
 
 #if defined(__APPLE__) && defined(__MACH__)
@@ -618,6 +619,25 @@ void log_printf(sd_log_level_t level, const char* file, int line, const char* fo
     va_end(args);
 }
 
+void sd_ggml_log_callback(ggml_log_level level, const char* text, void*) {
+    switch (level) {
+        case GGML_LOG_LEVEL_DEBUG:
+            LOG_VERBOSE(text);
+            break;
+        case GGML_LOG_LEVEL_INFO:
+            LOG_INFO(text);
+            break;
+        case GGML_LOG_LEVEL_WARN:
+            LOG_WARN(text);
+            break;
+        case GGML_LOG_LEVEL_ERROR:
+            LOG_ERROR(text);
+            break;
+        default:
+            LOG_VERBOSE(text);
+    }
+}
+
 void sd_set_log_callback(sd_log_cb_t cb, void* data) {
     sd_log_cb      = cb;
     sd_log_cb_data = data;

src/core/util.h

@@ -11,6 +11,14 @@
 #include "ggml-backend.h"
 #include "stable-diffusion.h"
 
+#ifndef __STATIC_INLINE__
+#define __STATIC_INLINE__ static inline
+#endif
+
+#ifndef SD_UNUSED
+#define SD_UNUSED(x) (void)(x)
+#endif
+
 #define SAFE_STR(s) ((s) ? (s) : "")
 #define BOOL_STR(b) ((b) ? "true" : "false")
 
@@ -79,6 +87,7 @@ void pretty_progress(int step, int steps, float time);
 void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float elapsed_seconds);
 
 void log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...);
+void sd_ggml_log_callback(ggml_log_level level, const char* text, void*);
 
 ggml_type sd_type_to_ggml_type(sd_type_t sdtype);
 

src/extensions/photomaker_extension.cpp

@@ -1,3 +1,4 @@
+#include <cinttypes>
 #include "extensions/generation_extension.h"
 
 #include <algorithm>

src/model/adapter/ip_adapter.hpp

@@ -1,8 +1,10 @@
 #ifndef __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__
 #define __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 #include "model/common/block.hpp"
+#include "model/common/ggml_block.hpp"
 #include "model_loader.h"
 
 namespace IPAdapter {
@@ -200,7 +202,7 @@ namespace IPAdapter {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(image_embeds);
             };
-            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
+            return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true));
         }
     };
 

src/model/adapter/lora.hpp

@@ -2,7 +2,13 @@
 #define __SD_MODEL_ADAPTER_LORA_HPP__
 
 #include <mutex>
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
+#include "model.h"
+#include "model/adapter/lora_ops.h"
 #include "model_loader.h"
 #include "model_manager.h"
 
@@ -957,7 +963,7 @@ struct LoraModel : public GGMLRunner {
             }
             return true;
         };
-        auto result = GGMLRunner::compute<float>(get_graph, n_threads, false, true, read_outputs);
+        auto result = GGMLRunner::compute(get_graph, n_threads, false, true, read_outputs);
         if (!result.has_value()) {
             LOG_ERROR("LoRA graph execution failed");
         }

src/model/adapter/lora_ops.cpp

@@ -0,0 +1,207 @@
+#include "model/adapter/lora_ops.h"
+
+#include <cmath>
+
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+
+ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx,
+                                 ggml_tensor* lora_down,
+                                 ggml_tensor* lora_up,
+                                 ggml_tensor* lora_mid) {
+    ggml_tensor* updown;
+    // flat lora tensors to multiply it
+    int64_t lora_up_rows  = lora_up->ne[ggml_n_dims(lora_up) - 1];
+    lora_up               = ggml_reshape_2d(ctx, lora_up, ggml_nelements(lora_up) / lora_up_rows, lora_up_rows);
+    auto lora_down_n_dims = ggml_n_dims(lora_down);
+    // assume n_dims should always be a multiple of 2 (otherwise rank 1 doesn't work)
+    lora_down_n_dims       = (lora_down_n_dims + lora_down_n_dims % 2);
+    int64_t lora_down_rows = lora_down->ne[lora_down_n_dims - 1];
+    lora_down              = ggml_reshape_2d(ctx, lora_down, ggml_nelements(lora_down) / lora_down_rows, lora_down_rows);
+
+    // ggml_mul_mat requires tensor b transposed
+    lora_down = ggml_cont(ctx, ggml_transpose(ctx, lora_down));
+    if (lora_mid == nullptr) {
+        updown = ggml_mul_mat(ctx, lora_up, lora_down);
+        updown = ggml_cont(ctx, ggml_transpose(ctx, updown));
+    } else {
+        // undoing tucker decomposition for conv layers.
+        // lora_mid  has shape (3,    3,   Rank, Rank)
+        // lora_down has shape (Rank, In,  1,    1)
+        // lora_up   has shape (Rank, Out, 1,    1)
+        // conv layer shape is (3,    3,   Out,  In)
+        updown = ggml_ext_mul_n_mode(ctx, ggml_ext_mul_n_mode(ctx, lora_mid, lora_down, 3), lora_up, 2);
+        updown = ggml_cont(ctx, updown);
+    }
+    return updown;
+}
+
+ggml_tensor* ggml_ext_lokr_forward(
+    ggml_context* ctx,
+    ggml_backend_t backend,
+    ggml_tensor* h,    // Input: [q, batch] or [W, H, q, batch]
+    ggml_tensor* w1,   // Outer C (Full rank)
+    ggml_tensor* w1a,  // Outer A (Low rank part 1)
+    ggml_tensor* w1b,  // Outer B (Low rank part 2)
+    ggml_tensor* w2,   // Inner BA (Full rank)
+    ggml_tensor* w2a,  // Inner A (Low rank part 1)
+    ggml_tensor* w2b,  // Inner B (Low rank part 2)
+    bool is_conv,
+    WeightAdapter::ForwardParams::conv2d_params_t conv_params,
+    float scale) {
+    GGML_ASSERT((w1 != nullptr || (w1a != nullptr && w1b != nullptr)));
+    GGML_ASSERT((w2 != nullptr || (w2a != nullptr && w2b != nullptr)));
+
+    int uq = (w1 != nullptr) ? (int)w1->ne[0] : (int)w1a->ne[0];
+    int up = (w1 != nullptr) ? (int)w1->ne[1] : (int)w1b->ne[1];
+
+    int q_actual = is_conv ? (int)h->ne[2] : (int)h->ne[0];
+    int vq       = q_actual / uq;
+
+    int vp = (w2 != nullptr) ? (is_conv ? (int)w2->ne[3] : (int)w2->ne[1])
+                             : (int)w2a->ne[1];
+    GGML_ASSERT(q_actual == (uq * vq) && "Input dimension mismatch for LoKR split");
+
+    ggml_tensor* hb;
+
+    if (!is_conv) {
+        int batch          = (int)h->ne[1];
+        int merge_batch_uq = batch;
+        int merge_batch_vp = batch;
+
+        if (sd_backend_is(backend, "Vulkan")) {
+            if (batch > 1) {
+                // no access to backend here, worst case is slightly worse perfs for other backends when built alongside Vulkan backend
+                int max_batch    = 65535;
+                int max_batch_uq = max_batch / uq;
+                merge_batch_uq   = 1;
+                for (int i = max_batch_uq; i > 0; i--) {
+                    if (batch % i == 0) {
+                        merge_batch_uq = i;
+                        break;
+                    }
+                }
+
+                int max_batch_vp = max_batch / vp;
+                merge_batch_vp   = 1;
+                for (int i = max_batch_vp; i > 0; i--) {
+                    if (batch % i == 0) {
+                        merge_batch_vp = i;
+                        break;
+                    }
+                }
+            }
+        }
+
+        ggml_tensor* h_split = ggml_reshape_3d(ctx, h, vq, uq * merge_batch_uq, batch / merge_batch_uq);
+        if (w2 != nullptr) {
+            hb = ggml_mul_mat(ctx, w2, h_split);
+        } else {
+            hb = ggml_mul_mat(ctx, w2b, ggml_mul_mat(ctx, w2a, h_split));
+        }
+
+        if (batch > 1) {
+            hb = ggml_reshape_3d(ctx, hb, vp, uq, batch);
+        }
+        ggml_tensor* hb_t = ggml_cont(ctx, ggml_transpose(ctx, hb));
+        hb_t              = ggml_reshape_3d(ctx, hb_t, uq, vp * merge_batch_vp, batch / merge_batch_vp);
+
+        ggml_tensor* hc_t;
+        if (w1 != nullptr) {
+            hc_t = ggml_mul_mat(ctx, w1, hb_t);
+        } else {
+            hc_t = ggml_mul_mat(ctx, w1b, ggml_mul_mat(ctx, w1a, hb_t));
+        }
+
+        if (batch > 1) {
+            hc_t = ggml_reshape_3d(ctx, hc_t, up, vp, batch);
+        }
+
+        ggml_tensor* hc  = ggml_transpose(ctx, hc_t);
+        ggml_tensor* out = ggml_reshape_2d(ctx, ggml_cont(ctx, hc), up * vp, batch);
+        return ggml_ext_scale(ctx, out, scale);
+    } else {
+        int batch = (int)h->ne[3];
+        // 1. Reshape input: [W, H, vq*uq, batch] -> [W, H, vq, uq * batch]
+        ggml_tensor* h_split = ggml_reshape_4d(ctx, h, h->ne[0], h->ne[1], vq, uq * batch);
+
+        if (w2 != nullptr) {
+            hb = ggml_ext_conv_2d(ctx, h_split, w2, nullptr,
+                                  conv_params.s0,
+                                  conv_params.s1,
+                                  conv_params.p0,
+                                  conv_params.p1,
+                                  conv_params.d0,
+                                  conv_params.d1,
+                                  conv_params.direct,
+                                  conv_params.circular_x,
+                                  conv_params.circular_y,
+                                  conv_params.scale);
+        } else {
+            // swap a and b order for conv lora
+            ggml_tensor* a = w2b;
+            ggml_tensor* b = w2a;
+
+            // unpack conv2d weights if needed
+            if (ggml_n_dims(a) < 4) {
+                int k = (int)sqrt(a->ne[0] / h_split->ne[2]);
+                GGML_ASSERT(k * k * h_split->ne[2] == a->ne[0]);
+                a = ggml_reshape_4d(ctx, a, k, k, a->ne[0] / (k * k), a->ne[1]);
+            } else if (a->ne[2] != h_split->ne[2]) {
+                int k = (int)sqrt(a->ne[2] / h_split->ne[2]);
+                GGML_ASSERT(k * k * h_split->ne[2] == a->ne[2]);
+                a = ggml_reshape_4d(ctx, a, a->ne[0] * k, a->ne[1] * k, a->ne[2] / (k * k), a->ne[3]);
+            }
+            ggml_tensor* ha = ggml_ext_conv_2d(ctx, h_split, a, nullptr,
+                                               conv_params.s0,
+                                               conv_params.s1,
+                                               conv_params.p0,
+                                               conv_params.p1,
+                                               conv_params.d0,
+                                               conv_params.d1,
+                                               conv_params.direct,
+                                               conv_params.circular_x,
+                                               conv_params.circular_y,
+                                               conv_params.scale);
+
+            // not supporting lora_mid here
+            hb = ggml_ext_conv_2d(ctx,
+                                  ha,
+                                  b,
+                                  nullptr,
+                                  1,
+                                  1,
+                                  0,
+                                  0,
+                                  1,
+                                  1,
+                                  conv_params.direct,
+                                  conv_params.circular_x,
+                                  conv_params.circular_y,
+                                  conv_params.scale);
+        }
+
+        // Current hb shape: [W_out, H_out, vp, uq * batch]
+        int w_out = (int)hb->ne[0];
+        int h_out = (int)hb->ne[1];
+
+        // ggml_tensor* hb_cat = ggml_reshape_4d(ctx, hb, w_out , h_out , vp * uq, batch);
+        // [W_out, H_out, vp * uq,  batch]
+        // Now left to compute (W1 kr Id) * hb_cat == (W1 kr W2) cv h
+
+        // merge the uq groups of size vp*w_out*h_out
+        ggml_tensor* hb_merged = ggml_reshape_2d(ctx, hb, w_out * h_out * vp, uq * batch);
+        ggml_tensor* hc_t;
+        ggml_tensor* hb_merged_t = ggml_cont(ctx, ggml_transpose(ctx, hb_merged));
+        if (w1 != nullptr) {
+            // Would be great to be able to transpose w1 instead to avoid transposing both hb and hc
+            hc_t = ggml_mul_mat(ctx, w1, hb_merged_t);
+        } else {
+            hc_t = ggml_mul_mat(ctx, w1b, ggml_mul_mat(ctx, w1a, hb_merged_t));
+        }
+        ggml_tensor* hc = ggml_transpose(ctx, hc_t);
+        // ungroup
+        ggml_tensor* out = ggml_reshape_4d(ctx, ggml_cont(ctx, hc), w_out, h_out, up * vp, batch);
+        return ggml_ext_scale(ctx, out, scale);
+    }
+}

src/model/adapter/lora_ops.h

@@ -0,0 +1,25 @@
+#ifndef __SD_MODEL_ADAPTER_LORA_OPS_H__
+#define __SD_MODEL_ADAPTER_LORA_OPS_H__
+
+#include "core/ggml_runner.h"
+
+ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx,
+                                 ggml_tensor* lora_down,
+                                 ggml_tensor* lora_up,
+                                 ggml_tensor* lora_mid = nullptr);
+
+ggml_tensor* ggml_ext_lokr_forward(
+    ggml_context* ctx,
+    ggml_backend_t backend,
+    ggml_tensor* h,    // Input: [q, batch] or [W, H, q, batch]
+    ggml_tensor* w1,   // Outer C (Full rank)
+    ggml_tensor* w1a,  // Outer A (Low rank part 1)
+    ggml_tensor* w1b,  // Outer B (Low rank part 2)
+    ggml_tensor* w2,   // Inner BA (Full rank)
+    ggml_tensor* w2a,  // Inner A (Low rank part 1)
+    ggml_tensor* w2b,  // Inner B (Low rank part 2)
+    bool is_conv,
+    WeightAdapter::ForwardParams::conv2d_params_t conv_params,
+    float scale);
+
+#endif  // __SD_MODEL_ADAPTER_LORA_OPS_H__

src/model/adapter/pmid.hpp

@@ -1,7 +1,10 @@
 #ifndef __SD_MODEL_ADAPTER_PMID_HPP__
 #define __SD_MODEL_ADAPTER_PMID_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 
 #include "model/adapter/lora.hpp"
 #include "model/common/block.hpp"
@@ -558,7 +561,7 @@ public:
             return build_graph(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds);
         };
 
-        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
+        return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true));
     }
 };
 

src/model/adapter/pulid.hpp

@@ -1,8 +1,10 @@
 #ifndef __PULID_HPP__
 #define __PULID_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 #include "model/common/block.hpp"
+#include "model/common/ggml_block.hpp"
 
 class PuLIDPerceiverAttentionCA : public GGMLBlock {
 public:

src/model/common/block.hpp

@@ -1,9 +1,11 @@
 #ifndef __SD_MODEL_COMMON_BLOCK_HPP__
 #define __SD_MODEL_COMMON_BLOCK_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 #include "core/util.h"
 #include "ggml-backend.h"
+#include "model/common/ggml_block.hpp"
 
 class DownSampleBlock : public GGMLBlock {
 protected:

src/model/common/ggml_block.hpp

@@ -0,0 +1,868 @@
+#ifndef __SD_MODEL_COMMON_GGML_BLOCK_HPP__
+#define __SD_MODEL_COMMON_GGML_BLOCK_HPP__
+
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <set>
+#include <string>
+#include <tuple>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "model.h"
+
+class GGMLBlock {
+protected:
+    typedef std::unordered_map<std::string, ggml_tensor*> ParameterMap;
+    typedef std::unordered_map<std::string, std::shared_ptr<GGMLBlock>> GGMLBlockMap;
+    GGMLBlockMap blocks;
+    ParameterMap params;
+
+    ggml_type get_type(const std::string& name, const String2TensorStorage& tensor_storage_map, ggml_type default_type) {
+        ggml_type wtype = default_type;
+        auto iter       = tensor_storage_map.find(name);
+        if (iter != tensor_storage_map.end()) {
+            const TensorStorage& tensor_storage = iter->second;
+            if (tensor_storage.expected_type != GGML_TYPE_COUNT) {
+                wtype = tensor_storage.expected_type;
+            } else {
+                wtype = tensor_storage.type;
+            }
+        }
+        return wtype;
+    }
+
+    void init_blocks(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {
+        for (auto& pair : blocks) {
+            auto& block = pair.second;
+            block->init(ctx, tensor_storage_map, prefix + pair.first);
+        }
+    }
+
+    virtual void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {}
+
+    virtual enum ggml_op param_usage_op(const std::string& name) const {
+        (void)name;
+        return GGML_OP_NONE;
+    }
+
+public:
+    void init(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") {
+        if (prefix.size() > 0) {
+            prefix = prefix + ".";
+        }
+        init_params(ctx, tensor_storage_map, prefix);
+        init_blocks(ctx, tensor_storage_map, prefix);
+    }
+
+    size_t get_params_num() {
+        size_t num_tensors = params.size();
+        for (auto& pair : blocks) {
+            auto& block = pair.second;
+
+            num_tensors += block->get_params_num();
+        }
+        return num_tensors;
+    };
+
+    size_t get_params_mem_size() {
+        size_t mem_size = 0;
+        for (auto& pair : blocks) {
+            auto& block = pair.second;
+
+            mem_size += block->get_params_mem_size();
+        }
+
+        for (auto& pair : params) {
+            mem_size += ggml_nbytes(pair.second);
+        }
+
+        return mem_size;
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, std::string prefix = "") {
+        if (prefix.size() > 0) {
+            prefix = prefix + ".";
+        }
+        for (auto& pair : blocks) {
+            auto& block = pair.second;
+            block->get_param_tensors(tensors, prefix + pair.first);
+        }
+
+        for (auto& pair : params) {
+            ggml_tensor* param           = pair.second;
+            tensors[prefix + pair.first] = pair.second;
+            ggml_set_name(param, (prefix + pair.first).c_str());
+        }
+    }
+
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
+        for (auto& pair : blocks) {
+            pair.second->get_param_tensor_ops(tensor_ops);
+        }
+        for (auto& pair : params) {
+            enum ggml_op op = param_usage_op(pair.first);
+            if (op != GGML_OP_NONE) {
+                tensor_ops[pair.second] = op;
+            }
+        }
+    }
+
+    virtual std::string get_desc() {
+        return "GGMLBlock";
+    }
+
+    void get_all_blocks(std::vector<GGMLBlock*>& result) {
+        result.push_back(this);
+        for (auto& block_iter : blocks) {
+            if (block_iter.second) {
+                block_iter.second->get_all_blocks(result);
+            }
+        }
+    }
+};
+
+class UnaryBlock : public GGMLBlock {
+public:
+    virtual ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) = 0;
+};
+
+class Identity : public UnaryBlock {
+public:
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        return x;
+    }
+};
+
+class Linear : public UnaryBlock {
+protected:
+    int64_t in_features;
+    int64_t out_features;
+    bool bias;
+    bool force_f32;
+    bool force_prec_f32;
+    bool has_weight_scale       = false;
+    bool int8_convrot           = false;
+    int int8_convrot_group_size = 0;
+    float scale;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        this->prefix            = prefix;
+        has_weight_scale        = false;
+        int8_convrot            = false;
+        int8_convrot_group_size = 0;
+        enum ggml_type wtype    = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
+        if (in_features % ggml_blck_size(wtype) != 0 || force_f32) {
+            wtype = GGML_TYPE_F32;
+        }
+        params["weight"] = ggml_new_tensor_2d(ctx, wtype, in_features, out_features);
+        if (bias) {
+            enum ggml_type wtype = GGML_TYPE_F32;
+            params["bias"]       = ggml_new_tensor_1d(ctx, wtype, out_features);
+        }
+        auto weight_storage           = tensor_storage_map.find(prefix + "weight");
+        const bool is_int8_tensorwise = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise;
+        auto weight_scale_storage     = tensor_storage_map.find(prefix + "weight_scale");
+        if (weight_scale_storage != tensor_storage_map.end()) {
+            const int64_t scale_nelements = weight_scale_storage->second.nelements();
+            GGML_ASSERT(scale_nelements == 1 || scale_nelements == out_features);
+            params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scale_nelements);
+            has_weight_scale       = true;
+        }
+        if (is_int8_tensorwise) {
+            GGML_ASSERT(wtype == GGML_TYPE_I8);
+            GGML_ASSERT(has_weight_scale);
+            int8_convrot            = weight_storage->second.int8_convrot;
+            int8_convrot_group_size = weight_storage->second.int8_convrot_group_size;
+        }
+    }
+
+public:
+    Linear(int64_t in_features,
+           int64_t out_features,
+           bool bias           = true,
+           bool force_f32      = false,
+           bool force_prec_f32 = false,
+           float scale         = 1.f)
+        : in_features(in_features),
+          out_features(out_features),
+          bias(bias),
+          force_f32(force_f32),
+          force_prec_f32(force_prec_f32),
+          scale(scale) {}
+
+    void set_scale(float scale_) {
+        scale = scale_;
+    }
+
+    void set_force_prec_f32(bool force_prec_f32_) {
+        force_prec_f32 = force_prec_f32_;
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w            = params["weight"];
+        ggml_tensor* weight_scale = has_weight_scale ? params["weight_scale"] : nullptr;
+        if (w->type == GGML_TYPE_F8_E4M3 || w->type == GGML_TYPE_F8_E5M2) {
+            bool supports_fp8_matmul = false;
+            if (ctx->backend != nullptr) {
+                ggml_tensor* fp8_matmul = ggml_mul_mat(ctx->ggml_ctx, w, x);
+                if (force_prec_f32) {
+                    ggml_mul_mat_set_prec(fp8_matmul, GGML_PREC_F32);
+                }
+                supports_fp8_matmul = ggml_backend_supports_op(ctx->backend, fp8_matmul);
+            }
+            if (!supports_fp8_matmul) {
+                w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_BF16);
+            }
+        }
+        ggml_tensor* b = nullptr;
+        if (bias) {
+            b = params["bias"];
+        }
+        ggml_tensor* linear_bias = has_weight_scale ? nullptr : b;
+        ggml_tensor* out         = nullptr;
+        if (w->type == GGML_TYPE_I8) {
+            if (x->type != GGML_TYPE_F32) {
+                x = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, x);
+            }
+            if (!ggml_is_contiguous(x)) {
+                x = ggml_cont(ctx->ggml_ctx, x);
+            }
+            ggml_tensor* lora_input = x;
+            if (ctx->weight_adapter && b != nullptr) {
+                b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+            }
+            if (int8_convrot && scale == 1.f) {
+                const auto cache_key = std::make_pair(x, int8_convrot_group_size);
+                auto cached          = ctx->int8_convrot_cache.find(cache_key);
+                if (cached == ctx->int8_convrot_cache.end()) {
+                    x = ggml_quantize_i8_convrot(ctx->ggml_ctx, x, int8_convrot_group_size);
+                    ctx->int8_convrot_cache.emplace(cache_key, x);
+                } else {
+                    x = cached->second;
+                }
+            }
+            out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx,
+                                                x,
+                                                w,
+                                                weight_scale,
+                                                b,
+                                                int8_convrot ? int8_convrot_group_size : 0,
+                                                scale);
+            if (ctx->weight_adapter) {
+                WeightAdapter::ForwardParams forward_params;
+                forward_params.op_type               = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
+                forward_params.linear.force_prec_f32 = force_prec_f32;
+                forward_params.linear.scale          = scale;
+                out                                  = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx,
+                                                                                               ctx->backend,
+                                                                                               lora_input,
+                                                                                               w,
+                                                                                               out,
+                                                                                               prefix,
+                                                                                               forward_params);
+            }
+            return out;
+        }
+        if (has_weight_scale) {
+            out = ggml_ext_linear(ctx->ggml_ctx, x, w, nullptr, force_prec_f32, scale);
+            out = ggml_mul(ctx->ggml_ctx, out, weight_scale);
+            if (ctx->weight_adapter) {
+                WeightAdapter::ForwardParams forward_params;
+                forward_params.op_type               = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
+                forward_params.linear.force_prec_f32 = force_prec_f32;
+                forward_params.linear.scale          = scale;
+                out                                  = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx,
+                                                                                               ctx->backend,
+                                                                                               x,
+                                                                                               w,
+                                                                                               out,
+                                                                                               prefix,
+                                                                                               forward_params);
+                if (b != nullptr) {
+                    b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+                }
+            }
+            if (b != nullptr) {
+                out = ggml_add_inplace(ctx->ggml_ctx, out, b);
+            }
+            return out;
+        }
+        if (ctx->weight_adapter) {
+            WeightAdapter::ForwardParams forward_params;
+            forward_params.op_type               = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
+            forward_params.linear.force_prec_f32 = force_prec_f32;
+            forward_params.linear.scale          = scale;
+            out                                  = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, linear_bias, prefix, forward_params);
+        } else {
+            out = ggml_ext_linear(ctx->ggml_ctx, x, w, linear_bias, force_prec_f32, scale);
+        }
+        return out;
+    }
+};
+
+__STATIC_INLINE__ bool support_get_rows(ggml_type wtype) {
+    std::set<ggml_type> allow_types = {GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0};
+    if (allow_types.find(wtype) != allow_types.end()) {
+        return true;
+    }
+    return false;
+}
+
+class Embedding : public UnaryBlock {
+protected:
+    int64_t embedding_dim;
+    int64_t num_embeddings;
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override {
+        enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
+        if (!support_get_rows(wtype)) {
+            wtype = GGML_TYPE_F32;
+        }
+        params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings);
+    }
+
+    enum ggml_op param_usage_op(const std::string& name) const override {
+        return name == "weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE;
+    }
+
+public:
+    Embedding(int64_t num_embeddings, int64_t embedding_dim)
+        : embedding_dim(embedding_dim),
+          num_embeddings(num_embeddings) {
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* input_ids) override {
+        // input_ids: [N, n_token]
+        auto weight = params["weight"];
+
+        // There are issues with ggml batch inference, so we are expanding it here first.
+        // TODO: fix ggml batch inference
+        int64_t n = input_ids->ne[1];
+        input_ids = ggml_reshape_1d(ctx->ggml_ctx, input_ids, input_ids->ne[0] * input_ids->ne[1]);
+
+        input_ids      = ggml_reshape_3d(ctx->ggml_ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]);
+        auto embedding = ggml_get_rows(ctx->ggml_ctx, weight, input_ids);
+        embedding      = ggml_reshape_3d(ctx->ggml_ctx, embedding, embedding->ne[0], embedding->ne[1] / n, n);
+
+        // [N, n_token, embedding_dim]
+        return embedding;
+    }
+};
+
+class Conv2d : public UnaryBlock {
+protected:
+    int64_t in_channels;
+    int64_t out_channels;
+    std::pair<int, int> kernel_size;
+    std::pair<int, int> stride;
+    std::pair<int, int> padding;
+    std::pair<int, int> dilation;
+    bool bias;
+    float scale = 1.f;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override {
+        this->prefix         = prefix;
+        enum ggml_type wtype = GGML_TYPE_F16;
+        params["weight"]     = ggml_new_tensor_4d(ctx, wtype, kernel_size.second, kernel_size.first, in_channels, out_channels);
+        if (bias) {
+            enum ggml_type wtype = GGML_TYPE_F32;
+            params["bias"]       = ggml_new_tensor_1d(ctx, wtype, out_channels);
+        }
+    }
+
+public:
+    Conv2d(int64_t in_channels,
+           int64_t out_channels,
+           std::pair<int, int> kernel_size,
+           std::pair<int, int> stride   = {1, 1},
+           std::pair<int, int> padding  = {0, 0},
+           std::pair<int, int> dilation = {1, 1},
+           bool bias                    = true)
+        : in_channels(in_channels),
+          out_channels(out_channels),
+          kernel_size(kernel_size),
+          stride(stride),
+          padding(padding),
+          dilation(dilation),
+          bias(bias) {}
+
+    void set_scale(float scale_value) {
+        scale = scale_value;
+    }
+
+    std::string get_desc() override {
+        return "Conv2d";
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w = params["weight"];
+        ggml_tensor* b = nullptr;
+        if (bias) {
+            b = params["bias"];
+        }
+        if (ctx->weight_adapter) {
+            WeightAdapter::ForwardParams forward_params;
+            forward_params.op_type           = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
+            forward_params.conv2d.s0         = stride.second;
+            forward_params.conv2d.s1         = stride.first;
+            forward_params.conv2d.p0         = padding.second;
+            forward_params.conv2d.p1         = padding.first;
+            forward_params.conv2d.d0         = dilation.second;
+            forward_params.conv2d.d1         = dilation.first;
+            forward_params.conv2d.direct     = ctx->conv2d_direct_enabled;
+            forward_params.conv2d.circular_x = ctx->circular_x_enabled;
+            forward_params.conv2d.circular_y = ctx->circular_y_enabled;
+            forward_params.conv2d.scale      = scale;
+            return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params);
+        }
+        return ggml_ext_conv_2d(ctx->ggml_ctx,
+                                x,
+                                w,
+                                b,
+                                stride.second,
+                                stride.first,
+                                padding.second,
+                                padding.first,
+                                dilation.second,
+                                dilation.first,
+                                ctx->conv2d_direct_enabled,
+                                ctx->circular_x_enabled,
+                                ctx->circular_y_enabled,
+                                scale);
+    }
+};
+
+class Conv2d_grouped : public UnaryBlock {
+protected:
+    int64_t in_channels;
+    int64_t out_channels;
+    int groups;
+    std::pair<int, int> kernel_size;
+    std::pair<int, int> stride;
+    std::pair<int, int> padding;
+    std::pair<int, int> dilation;
+    bool bias;
+    float scale = 1.f;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override {
+        this->prefix         = prefix;
+        enum ggml_type wtype = GGML_TYPE_F16;
+        params["weight"]     = ggml_new_tensor_4d(ctx, wtype, kernel_size.second, kernel_size.first, in_channels / groups, out_channels);
+        if (bias) {
+            enum ggml_type wtype = GGML_TYPE_F32;
+            params["bias"]       = ggml_new_tensor_1d(ctx, wtype, out_channels);
+        }
+    }
+
+public:
+    Conv2d_grouped(int64_t in_channels,
+                   int64_t out_channels,
+                   int groups,
+                   std::pair<int, int> kernel_size,
+                   std::pair<int, int> stride   = {1, 1},
+                   std::pair<int, int> padding  = {0, 0},
+                   std::pair<int, int> dilation = {1, 1},
+                   bool bias                    = true)
+        : in_channels(in_channels),
+          out_channels(out_channels),
+          groups(groups),
+          kernel_size(kernel_size),
+          stride(stride),
+          padding(padding),
+          dilation(dilation),
+          bias(bias) {}
+
+    void set_scale(float scale_value) {
+        scale = scale_value;
+    }
+
+    std::string get_desc() override {
+        return "Conv2d_grouped";
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w = params["weight"];
+        ggml_tensor* b = nullptr;
+        if (bias) {
+            b = params["bias"];
+        }
+
+        if (groups == 1) {
+            if (ctx->weight_adapter) {
+                WeightAdapter::ForwardParams forward_params;
+                forward_params.op_type           = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
+                forward_params.conv2d.s0         = stride.second;
+                forward_params.conv2d.s1         = stride.first;
+                forward_params.conv2d.p0         = padding.second;
+                forward_params.conv2d.p1         = padding.first;
+                forward_params.conv2d.d0         = dilation.second;
+                forward_params.conv2d.d1         = dilation.first;
+                forward_params.conv2d.direct     = ctx->conv2d_direct_enabled;
+                forward_params.conv2d.circular_x = ctx->circular_x_enabled;
+                forward_params.conv2d.circular_y = ctx->circular_y_enabled;
+                forward_params.conv2d.scale      = scale;
+                return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params);
+            }
+            return ggml_ext_conv_2d(ctx->ggml_ctx, x, w, b,
+                                    stride.second, stride.first,
+                                    padding.second, padding.first,
+                                    dilation.second, dilation.first,
+                                    ctx->conv2d_direct_enabled,
+                                    ctx->circular_x_enabled,
+                                    ctx->circular_y_enabled,
+                                    scale);
+        }
+
+        if (groups == in_channels && groups == out_channels) {
+            ggml_tensor* res;
+            if (ctx->conv2d_direct_enabled) {
+                res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, w, x,
+                                             stride.second, stride.first,
+                                             padding.second, padding.first,
+                                             dilation.second, dilation.first);
+            } else {
+                res = ggml_conv_2d_dw(ctx->ggml_ctx, w, x,
+                                      stride.second, stride.first,
+                                      padding.second, padding.first,
+                                      dilation.second, dilation.first);
+            }
+            if (b) {
+                b   = ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, b->ne[0], 1);
+                res = ggml_add_inplace(ctx->ggml_ctx, res, b);
+            }
+            return res;
+        }
+
+        int64_t ic_g = in_channels / groups;
+        int64_t oc_g = out_channels / groups;
+
+        std::vector<ggml_tensor*> out_slices(groups);
+
+        for (int i = 0; i < groups; ++i) {
+            size_t x_offset  = i * ic_g * x->nb[2];
+            ggml_tensor* x_i = ggml_view_4d(ctx->ggml_ctx, x,
+                                            x->ne[0], x->ne[1], ic_g, x->ne[3],
+                                            x->nb[1], x->nb[2], x->nb[3],
+                                            x_offset);
+
+            size_t w_offset  = i * oc_g * w->nb[3];
+            ggml_tensor* w_i = ggml_view_4d(ctx->ggml_ctx, w,
+                                            w->ne[0], w->ne[1], w->ne[2], oc_g,
+                                            w->nb[1], w->nb[2], w->nb[3],
+                                            w_offset);
+
+            ggml_tensor* b_i = nullptr;
+            if (b) {
+                size_t b_offset = i * oc_g * b->nb[0];
+                b_i             = ggml_view_1d(ctx->ggml_ctx, b, oc_g, b_offset);
+            }
+
+            if (ctx->weight_adapter) {
+                WeightAdapter::ForwardParams forward_params;
+                forward_params.op_type           = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
+                forward_params.conv2d.s0         = stride.second;
+                forward_params.conv2d.s1         = stride.first;
+                forward_params.conv2d.p0         = padding.second;
+                forward_params.conv2d.p1         = padding.first;
+                forward_params.conv2d.d0         = dilation.second;
+                forward_params.conv2d.d1         = dilation.first;
+                forward_params.conv2d.direct     = ctx->conv2d_direct_enabled;
+                forward_params.conv2d.circular_x = ctx->circular_x_enabled;
+                forward_params.conv2d.circular_y = ctx->circular_y_enabled;
+                forward_params.conv2d.scale      = scale;
+                out_slices[i]                    = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x_i, w_i, b_i, prefix, forward_params);
+            } else {
+                out_slices[i] = ggml_ext_conv_2d(ctx->ggml_ctx, x_i, w_i, b_i,
+                                                 stride.second, stride.first,
+                                                 padding.second, padding.first,
+                                                 dilation.second, dilation.first,
+                                                 ctx->conv2d_direct_enabled,
+                                                 ctx->circular_x_enabled,
+                                                 ctx->circular_y_enabled,
+                                                 scale);
+            }
+        }
+
+        ggml_tensor* out = ggml_ext_vec_concat(ctx->ggml_ctx, out_slices, 2);
+
+        return out;
+    }
+};
+
+class Conv3d : public UnaryBlock {
+protected:
+    int64_t in_channels;
+    int64_t out_channels;
+    std::tuple<int, int, int> kernel_size;
+    std::tuple<int, int, int> stride;
+    std::tuple<int, int, int> padding;
+    std::tuple<int, int, int> dilation;
+    bool bias;
+    bool force_prec_f32;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override {
+        this->prefix         = prefix;
+        enum ggml_type wtype = GGML_TYPE_F16;
+        params["weight"]     = ggml_new_tensor_4d(ctx,
+                                                  wtype,
+                                                  std::get<2>(kernel_size),
+                                                  std::get<1>(kernel_size),
+                                                  std::get<0>(kernel_size),
+                                                  in_channels * out_channels);
+        if (bias) {
+            params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
+        }
+    }
+
+public:
+    Conv3d(int64_t in_channels,
+           int64_t out_channels,
+           std::tuple<int, int, int> kernel_size,
+           std::tuple<int, int, int> stride   = {1, 1, 1},
+           std::tuple<int, int, int> padding  = {0, 0, 0},
+           std::tuple<int, int, int> dilation = {1, 1, 1},
+           bool bias                          = true,
+           bool force_prec_f32                = false)
+        : in_channels(in_channels),
+          out_channels(out_channels),
+          kernel_size(kernel_size),
+          stride(stride),
+          padding(padding),
+          dilation(dilation),
+          bias(bias),
+          force_prec_f32(force_prec_f32) {}
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w = params["weight"];
+        ggml_tensor* b = nullptr;
+        if (ctx->weight_adapter) {
+            w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+            if (w->type != GGML_TYPE_F16) {
+                w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_F16);
+            }
+        }
+        if (bias) {
+            b = params["bias"];
+            if (ctx->weight_adapter) {
+                b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+            }
+        }
+        return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
+                                std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
+                                std::get<2>(padding), std::get<1>(padding), std::get<0>(padding),
+                                std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation),
+                                force_prec_f32);
+    }
+};
+
+class LayerNorm : public UnaryBlock {
+protected:
+    int64_t normalized_shape;
+    float eps;
+    bool elementwise_affine;
+    bool bias;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        this->prefix = prefix;
+        if (elementwise_affine) {
+            enum ggml_type wtype = GGML_TYPE_F32;
+            params["weight"]     = ggml_new_tensor_1d(ctx, wtype, normalized_shape);
+            if (bias) {
+                enum ggml_type wtype = GGML_TYPE_F32;
+                params["bias"]       = ggml_new_tensor_1d(ctx, wtype, normalized_shape);
+            }
+        }
+    }
+
+public:
+    LayerNorm(int64_t normalized_shape,
+              float eps               = 1e-05f,
+              bool elementwise_affine = true,
+              bool bias               = true)
+        : normalized_shape(normalized_shape),
+          eps(eps),
+          elementwise_affine(elementwise_affine),
+          bias(bias) {}
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w = nullptr;
+        ggml_tensor* b = nullptr;
+
+        if (elementwise_affine) {
+            w = params["weight"];
+            if (ctx->weight_adapter) {
+                w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+            }
+            if (bias) {
+                b = params["bias"];
+                if (ctx->weight_adapter) {
+                    b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+                }
+            }
+        }
+        return ggml_ext_layer_norm(ctx->ggml_ctx, x, w, b, eps);
+    }
+};
+
+class GroupNorm : public GGMLBlock {
+protected:
+    int num_groups;
+    int64_t num_channels;
+    float eps;
+    bool affine;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        this->prefix = prefix;
+        if (affine) {
+            enum ggml_type wtype      = GGML_TYPE_F32;
+            enum ggml_type bias_wtype = GGML_TYPE_F32;
+            params["weight"]          = ggml_new_tensor_1d(ctx, wtype, num_channels);
+            params["bias"]            = ggml_new_tensor_1d(ctx, bias_wtype, num_channels);
+        }
+    }
+
+public:
+    GroupNorm(int num_groups,
+              int64_t num_channels,
+              float eps   = 1e-05f,
+              bool affine = true)
+        : num_groups(num_groups),
+          num_channels(num_channels),
+          eps(eps),
+          affine(affine) {}
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        ggml_tensor* w = nullptr;
+        ggml_tensor* b = nullptr;
+        if (affine) {
+            w = params["weight"];
+            b = params["bias"];
+            if (ctx->weight_adapter) {
+                w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+                b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+            }
+        }
+        return ggml_ext_group_norm(ctx->ggml_ctx, x, w, b, num_groups);
+    }
+};
+
+class GroupNorm32 : public GroupNorm {
+public:
+    GroupNorm32(int64_t num_channels)
+        : GroupNorm(32, num_channels, 1e-06f) {}
+};
+
+class RMSNorm : public UnaryBlock {
+protected:
+    int64_t hidden_size;
+    float eps;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
+        this->prefix         = prefix;
+        enum ggml_type wtype = GGML_TYPE_F32;
+        params["weight"]     = ggml_new_tensor_1d(ctx, wtype, hidden_size);
+    }
+
+public:
+    RMSNorm(int64_t hidden_size,
+            float eps = 1e-06f)
+        : hidden_size(hidden_size),
+          eps(eps) {}
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        ggml_tensor* w = params["weight"];
+        if (ctx->weight_adapter) {
+            w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+        }
+        x = ggml_rms_norm(ctx->ggml_ctx, x, eps);
+        x = ggml_mul_inplace(ctx->ggml_ctx, x, w);
+        return x;
+    }
+};
+
+class MultiheadAttention : public GGMLBlock {
+protected:
+    int64_t embed_dim;
+    int64_t n_head;
+    bool proj_in;
+    std::string q_proj_name;
+    std::string k_proj_name;
+    std::string v_proj_name;
+    std::string in_proj_name;
+    std::string out_proj_name;
+
+public:
+    MultiheadAttention(int64_t embed_dim,
+                       int64_t n_head,
+                       bool qkv_proj_bias        = true,
+                       bool out_proj_bias        = true,
+                       bool proj_in              = false,
+                       std::string q_proj_name   = "q_proj",
+                       std::string k_proj_name   = "k_proj",
+                       std::string v_proj_name   = "v_proj",
+                       std::string in_proj_name  = "in_proj",
+                       std::string out_proj_name = "out_proj")
+        : embed_dim(embed_dim),
+          n_head(n_head),
+          proj_in(proj_in),
+          q_proj_name(q_proj_name),
+          k_proj_name(k_proj_name),
+          v_proj_name(v_proj_name),
+          in_proj_name(in_proj_name),
+          out_proj_name(out_proj_name) {
+        if (proj_in) {
+            blocks[in_proj_name] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim * 3, qkv_proj_bias));
+        } else {
+            blocks[q_proj_name] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, qkv_proj_bias));
+            blocks[k_proj_name] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, qkv_proj_bias));
+            blocks[v_proj_name] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, qkv_proj_bias));
+        }
+        blocks[out_proj_name] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, embed_dim, out_proj_bias));
+    }
+
+    // x: [N, n_token, embed_dim]
+    ggml_tensor* forward(GGMLRunnerContext* ctx,
+                         ggml_tensor* x,
+                         ggml_tensor* mask = nullptr) {
+        auto out_proj = std::dynamic_pointer_cast<Linear>(blocks[out_proj_name]);
+
+        ggml_tensor* q;
+        ggml_tensor* k;
+        ggml_tensor* v;
+        if (proj_in) {
+            auto in_proj = std::dynamic_pointer_cast<Linear>(blocks[in_proj_name]);
+            auto qkv     = in_proj->forward(ctx, x);
+            auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
+            q            = qkv_vec[0];
+            k            = qkv_vec[1];
+            v            = qkv_vec[2];
+        } else {
+            auto q_proj = std::dynamic_pointer_cast<Linear>(blocks[q_proj_name]);
+            auto k_proj = std::dynamic_pointer_cast<Linear>(blocks[k_proj_name]);
+            auto v_proj = std::dynamic_pointer_cast<Linear>(blocks[v_proj_name]);
+
+            q = q_proj->forward(ctx, x);
+            k = k_proj->forward(ctx, x);
+            v = v_proj->forward(ctx, x);
+        }
+
+        x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, false);  // [N, n_token, embed_dim]
+
+        x = out_proj->forward(ctx, x);  // [N, n_token, embed_dim]
+        return x;
+    }
+};
+
+#endif  // __SD_MODEL_COMMON_GGML_BLOCK_HPP__

src/model/common/rope.hpp

@@ -2,9 +2,13 @@
 #define __SD_MODEL_COMMON_ROPE_HPP__
 
 #include <algorithm>
+#include <cassert>
 #include <cmath>
+#include <set>
 #include <vector>
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
 
 namespace Rope {
     enum class EmbedNDLayout {

src/model/detector/yolov8.h

@@ -8,8 +8,9 @@
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_runner.h"
 #include "core/util.h"
+#include "model/common/ggml_block.hpp"
 
 struct YOLOv8Config {
     std::array<int, 23> out_channels{};
@@ -355,7 +356,7 @@ struct YOLOv8Runner : public GGMLRunner {
 
     sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& input) {
         auto get_graph = [&]() { return build_graph(input); };
-        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, false));
+        return take_or_empty(GGMLRunner::compute(get_graph, n_threads, false));
     }
 };
 

src/model/diffusion/anima.hpp

@@ -2,6 +2,7 @@
 #define __SD_MODEL_DIFFUSION_ANIMA_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <memory>
 #include <utility>
@@ -717,7 +718,7 @@ namespace Anima {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, t5_ids, t5_weights, ref_latents);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/animatediff.hpp

@@ -1,8 +1,10 @@
 #ifndef __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__
 #define __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 #include "model/common/block.hpp"
+#include "model/common/ggml_block.hpp"
 
 // AnimateDiff (https://arxiv.org/abs/2307.04725) SD 1.5 motion modules.
 namespace AnimateDiff {

src/model/diffusion/boogu.hpp

@@ -2,11 +2,15 @@
 #define __SD_MODEL_DIFFUSION_BOOGU_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <tuple>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model/diffusion/model.hpp"
@@ -815,7 +819,7 @@ namespace Boogu {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, ref_latents);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/control.hpp

@@ -435,7 +435,7 @@ struct ControlNet : public GGMLRunner {
             }
             return true;
         };
-        auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false, true, read_outputs);
+        auto compute_result = GGMLRunner::compute(get_graph, n_threads, false, true, read_outputs);
         control_outputs_ggml.clear();
         guided_hint_output_ggml = nullptr;
         if (!compute_result.has_value()) {

src/model/diffusion/dit.hpp

@@ -1,7 +1,8 @@
 #ifndef __SD_MODEL_DIFFUSION_DIT_HPP__
 #define __SD_MODEL_DIFFUSION_DIT_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 
 namespace DiT {
     inline ggml_tensor* patchify(ggml_context* ctx,

src/model/diffusion/ernie_image.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__
 #define __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__
 
+#include <cinttypes>
 #include <memory>
 #include <vector>
 
@@ -440,7 +441,7 @@ namespace ErnieImage {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/flux.hpp

@@ -1,8 +1,11 @@
 #ifndef __SD_MODEL_DIFFUSION_FLUX_HPP__
 #define __SD_MODEL_DIFFUSION_FLUX_HPP__
 
+#include <cinttypes>
 #include <memory>
 #include <vector>
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
 
 #include "core/util.h"
 #include "model/adapter/pulid.hpp"
@@ -1626,7 +1629,7 @@ namespace Flux {
                 return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, ref_index_mode, skip_layers, pulid_id, pulid_id_weight);
             };
 
-            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
             return result;
         }
 

src/model/diffusion/hidream_o1.hpp

@@ -329,7 +329,7 @@ namespace HiDreamO1 {
             auto get_graph = [&]() {
                 return build_graph(image);
             };
-            auto output = GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end);
+            auto output = GGMLRunner::compute(get_graph, n_threads, auto_runner_end);
             return output.has_value() ? std::move(output.value()) : sd::Tensor<float>();
         }
     };
@@ -457,7 +457,7 @@ namespace HiDreamO1 {
             auto get_graph = [&]() {
                 return build_graph(x, timestep, input_ids, input_pos, token_types, vinput_mask, image_embeds, ref_images);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/hunyuan.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_HUNYUAN_HPP__
 #define __SD_MODEL_DIFFUSION_HUNYUAN_HPP__
 
+#include <cinttypes>
 #include <memory>
 
 #include "model/common/block.hpp"
@@ -654,7 +655,7 @@ namespace Hunyuan {
                 return build_graph(x, timesteps, context, c_concat, y, guidance, byt5, vision, timestep_r);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/ideogram4.hpp

@@ -2,14 +2,18 @@
 #define __SD_MODEL_DIFFUSION_IDEOGRAM4_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <cstdlib>
 #include <memory>
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
 #include "core/ggml_graph_cut.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/model.hpp"
 
@@ -537,7 +541,7 @@ namespace Ideogram4 {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, use_uncond_model);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/krea2.hpp

@@ -12,8 +12,11 @@
 #include <utility>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
 #include "core/ggml_graph_cut.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model/diffusion/flux.hpp"
@@ -775,7 +778,7 @@ namespace Krea2 {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, ref_latents, ref_image_params);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/lens.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_LENS_HPP__
 #define __SD_MODEL_DIFFUSION_LENS_HPP__
 
+#include <cinttypes>
 #include <memory>
 #include <vector>
 
@@ -408,7 +409,7 @@ namespace Lens {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/lingbot_video.hpp

@@ -674,7 +674,7 @@ namespace LingBotVideo {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/ltxv.hpp

@@ -2,12 +2,15 @@
 #define __SD_MODEL_DIFFUSION_LTXV_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <memory>
 #include <string>
 #include <tuple>
 #include <utility>
 #include <vector>
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
 
 #include "model/common/block.hpp"
 #include "model/common/rope.hpp"
@@ -1998,7 +2001,7 @@ namespace LTXV {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, audio_x, audio_timesteps, audio_length, frame_rate, video_positions);
             };
-            auto out = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            auto out = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
             return out;
         }
 

src/model/diffusion/mage_flow.hpp

@@ -142,7 +142,7 @@ namespace MageFlow {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, ref_latents);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/minimax_h3.hpp

@@ -2,12 +2,14 @@
 #define __SD_MODEL_DIFFUSION_MINIMAX_H3_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <set>
 #include <string>
 #include <tuple>
 #include <utility>
 #include <vector>
+#include "core/ggml_tensor_utils.h"
 
 #include "core/ggml_graph_cut.h"
 #include "model/diffusion/dit.hpp"
@@ -1166,9 +1168,9 @@ namespace MiniMaxH3 {
                                    extra->video_sigma_shift,
                                    extra->audio_sigma_shift);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph,
-                                                                              n_threads,
-                                                                              false),
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph,
+                                                                       n_threads,
+                                                                       false),
                                                    params.x->dim());
         }
     };

src/model/diffusion/minit2i.hpp

@@ -2,6 +2,7 @@
 #define __SD_MODEL_DIFFUSION_MINIT2I_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <cstdint>
 #include <cstdlib>
@@ -9,7 +10,10 @@
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model/diffusion/model.hpp"
@@ -589,7 +593,7 @@ namespace MiniT2I {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, mask);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/mmdit.hpp

@@ -2,12 +2,18 @@
 #define __SD_MODEL_DIFFUSION_MMDIT_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 #include <memory>
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
 #include "model/common/block.hpp"
+#include "model/common/ggml_block.hpp"
 #include "model/diffusion/model.hpp"
 #include "model_loader.h"
 
@@ -987,7 +993,7 @@ struct MMDiTRunner : public DiffusionModelRunner {
             return build_graph(x, timesteps, context, y, skip_layers);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
     }
 
     sd::Tensor<float> compute(int n_threads,

src/model/diffusion/model.hpp

@@ -5,7 +5,7 @@
 #include <utility>
 #include <variant>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_runner.h"
 #include "core/tensor_ggml.hpp"
 #include "model/common/rope.hpp"
 #include "model_manager.h"

src/model/diffusion/pid.hpp

@@ -1,13 +1,18 @@
 #ifndef __SD_MODEL_DIFFUSION_PID_HPP__
 #define __SD_MODEL_DIFFUSION_PID_HPP__
 
+#include <cinttypes>
 #include <cmath>
 #include <cstdlib>
 #include <memory>
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model/diffusion/mmdit.hpp"
@@ -938,7 +943,7 @@ namespace Pid {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, lq_latent, degrade_sigma);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/qwen_image.hpp

@@ -2,6 +2,8 @@
 #define __SD_MODEL_DIFFUSION_QWEN_IMAGE_HPP__
 
 #include <memory>
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
 
 #include "core/util.h"
 #include "model/common/block.hpp"
@@ -707,7 +709,7 @@ namespace Qwen {
                 return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/sefi_image.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__
 #define __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__
 
+#include <cinttypes>
 #include <memory>
 
 #include "model/common/block.hpp"

src/model/diffusion/unet.hpp

@@ -3,6 +3,7 @@
 
 #include <algorithm>
 #include <vector>
+#include "core/ggml_tensor_utils.h"
 
 #include "model.h"
 #include "model/common/block.hpp"
@@ -835,7 +836,7 @@ struct UNetModelRunner : public DiffusionModelRunner {
             return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength, ip_context, ip_scale);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
     }
 
     sd::Tensor<float> compute(int n_threads,

src/model/diffusion/wan.hpp

@@ -1,9 +1,12 @@
 #ifndef __SD_MODEL_DIFFUSION_WAN_HPP__
 #define __SD_MODEL_DIFFUSION_WAN_HPP__
 
+#include <cinttypes>
 #include <map>
 #include <memory>
 #include <utility>
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
 
 #include "model/common/block.hpp"
 #include "model/common/rope.hpp"
@@ -950,7 +953,7 @@ namespace WAN {
                 return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/z_image.hpp

@@ -2,8 +2,14 @@
 #define __SD_MODEL_DIFFUSION_Z_IMAGE_HPP__
 
 #include <algorithm>
+#include <cinttypes>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/diffusion/flux.hpp"
 #include "model/diffusion/mmdit.hpp"
 #include "model/diffusion/model.hpp"
@@ -636,7 +642,7 @@ namespace ZImage {
                 return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/te/clip.hpp

@@ -1,8 +1,11 @@
 #ifndef __SD_MODEL_TE_CLIP_HPP__
 #define __SD_MODEL_TE_CLIP_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/util.h"
 #include "model.h"
+#include "model/common/ggml_block.hpp"
 #include "tokenizers/clip_tokenizer.h"
 
 /*================================================ FrozenCLIPEmbedder ================================================*/
@@ -572,7 +575,7 @@ struct CLIPTextModelRunner : public GGMLRunner {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input_ids, num_custom_embeddings, custom_embeddings_data, max_token_idx, return_pooled, clip_skip);
         };
-        auto result = GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end);
+        auto result = GGMLRunner::compute(get_graph, n_threads, auto_runner_end);
         if (return_pooled) {
             return take_or_empty(std::move(result));
         }

src/model/te/llm.hpp

@@ -3,6 +3,7 @@
 
 #include <algorithm>
 #include <array>
+#include <cinttypes>
 #include <cmath>
 #include <fstream>
 #include <functional>
@@ -18,8 +19,13 @@
 #include <utility>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
 #include "json.hpp"
+#include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
 #include "model_loader.h"
 #include "model_manager.h"
@@ -2091,7 +2097,7 @@ namespace LLM {
                                    out_layers,
                                    return_all_hidden_states);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end),
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_runner_end),
                                                    input_ids.dim() + 1);
         }
 
@@ -2175,7 +2181,7 @@ namespace LLM {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_encode_image_graph(image);
             };
-            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
+            return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end));
         }
 
         ggml_cgraph* build_encode_image_outputs_graph(const sd::Tensor<float>& image_tensor) {
@@ -2287,7 +2293,7 @@ namespace LLM {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_encode_image_outputs_graph(image);
             };
-            auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
+            auto combined = take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end));
             if (combined.empty()) {
                 return {};
             }
@@ -2313,7 +2319,7 @@ namespace LLM {
             auto get_graph    = [&]() -> ggml_cgraph* {
                 return build_encode_video_block_outputs_graph(pixel_values, grid_h, grid_w);
             };
-            auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
+            auto combined = take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end));
             if (combined.empty()) {
                 return {};
             }

src/model/te/t5.hpp

@@ -10,7 +10,12 @@
 #include <string>
 #include <unordered_map>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model_loader.h"
 #include "model_manager.h"
 #include "tokenizers/t5_unigram_tokenizer.h"
@@ -455,7 +460,7 @@ struct T5Runner : public GGMLRunner {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input_ids, attention_mask);
         };
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end), 3);
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_runner_end), 3);
     }
 
     static std::vector<int> _relative_position_bucket(const std::vector<int>& relative_position,

src/model/upscaler/esrgan.hpp

@@ -7,8 +7,10 @@
 #include <utility>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
 #include "core/util.h"
+#include "model/common/ggml_block.hpp"
 
 /*
     ===================================    ESRGAN  ===================================
@@ -265,7 +267,7 @@ struct ESRGAN : public GGMLRunner {
     sd::Tensor<float> compute(const int n_threads,
                               const sd::Tensor<float>& x) {
         auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); };
-        auto result    = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        auto result    = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
         return result;
     }
 };

src/model/upscaler/ltx_latent_upscaler.hpp

@@ -11,9 +11,11 @@
 #include <utility>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
 #include "core/ggml_graph_cut.h"
+#include "core/ggml_runner.h"
 #include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model_loader.h"
 
@@ -499,7 +501,7 @@ namespace LTXVUpsampler {
             }
             size_t expected_dim = static_cast<size_t>(x.dim());
             auto get_graph      = [&]() -> ggml_cgraph* { return build_graph(x); };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), expected_dim);
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), expected_dim);
         }
     };
 

src/model/vae/audio_vae.hpp

@@ -1,7 +1,8 @@
 #ifndef __SD_MODEL_VAE_AUDIO_VAE_HPP__
 #define __SD_MODEL_VAE_AUDIO_VAE_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_runner.h"
+#include "core/util.h"
 
 struct AudioVAERunner : public GGMLRunner {
     AudioVAERunner(ggml_backend_t backend,

src/model/vae/auto_encoder_kl.hpp

@@ -1,6 +1,8 @@
 #ifndef __SD_MODEL_VAE_AUTO_ENCODER_KL_HPP__
 #define __SD_MODEL_VAE_AUTO_ENCODER_KL_HPP__
 
+#include <cinttypes>
+#include "core/ggml_tensor_utils.h"
 #include "model/vae/vae.hpp"
 
 /*================================================== AutoEncoderKL ===================================================*/
@@ -744,7 +746,7 @@ struct AutoEncoderKL : public VAE {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(z, decode_graph);
         };
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), z.dim());
     }
 
     sd::Tensor<float> gaussian_latent_sample(const sd::Tensor<float>& moments, std::shared_ptr<RNG> rng) {

src/model/vae/hunyuan_vae.hpp

@@ -825,9 +825,9 @@ namespace Hunyuan {
             auto get_graph          = [&]() -> ggml_cgraph* {
                 return build_graph(graph_input, decode_graph);
             };
-            auto output = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph,
-                                                                                     n_threads,
-                                                                                     false),
+            auto output = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph,
+                                                                              n_threads,
+                                                                              false),
                                                           graph_input.dim());
             if (!output.empty() && input.dim() == 4) {
                 output.squeeze_(2);

src/model/vae/ltx_audio_vae.hpp

@@ -7,7 +7,12 @@
 #include <string>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
+#include "core/util.h"
+#include "model/common/ggml_block.hpp"
 #include "model/vae/audio_vae.hpp"
 #include "model_loader.h"
 #include "model_manager.h"
@@ -1042,7 +1047,7 @@ namespace LTXV {
                 ggml_build_forward_expand(gf, waveform);
                 return gf;
             };
-            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), 4);
+            auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), 4);
             int64_t t1  = ggml_time_ms();
             LOG_INFO("ltx audio vae decode completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
             return result;

src/model/vae/ltx_vae.hpp

@@ -8,6 +8,8 @@
 #include <tuple>
 #include <utility>
 #include <vector>
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
 
 #include "model/diffusion/ltxv.hpp"
 #include "model/vae/vae.hpp"
@@ -1348,7 +1350,7 @@ struct LTXVideoVAE : public VAE {
                                                  static_cast<int>(tile.start),
                                                  tile.overlap);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false),
                                                    expected_dim);
         });
 
@@ -1405,7 +1407,7 @@ struct LTXVideoVAE : public VAE {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input, decode_graph);
         };
-        auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), expected_dim);
+        auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), expected_dim);
         if (result.empty()) {
             return {};
         }
@@ -1418,7 +1420,7 @@ struct LTXVideoVAE : public VAE {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_latent_statistics_graph(z, normalize);
         };
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false),
                                                static_cast<size_t>(z.dim()));
     }
 

src/model/vae/mage_vae.hpp

@@ -490,7 +490,7 @@ namespace MageVAE {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(input, decode_graph);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), input.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), input.dim());
         }
 
         int get_encoder_output_channels(int input_channels) override {

src/model/vae/minimax_h3_audio_vae.hpp

@@ -480,7 +480,7 @@ namespace MiniMaxH3 {
                 return graph;
             };
             auto result = restore_trailing_singleton_dims(
-                GGMLRunner::compute<float>(get_graph, n_threads, false),
+                GGMLRunner::compute(get_graph, n_threads, false),
                 4);
             int64_t t1 = ggml_time_ms();
             LOG_INFO("MiniMax-H3 audio VAE encode completed, taking %.2fs",
@@ -500,7 +500,7 @@ namespace MiniMaxH3 {
                 return graph;
             };
             auto result = restore_trailing_singleton_dims(
-                GGMLRunner::compute<float>(get_graph, n_threads, false),
+                GGMLRunner::compute(get_graph, n_threads, false),
                 4);
             int64_t t1 = ggml_time_ms();
             LOG_INFO("MiniMax-H3 audio VAE decode completed, taking %.2fs",

src/model/vae/minimax_h3_vae.hpp

@@ -791,9 +791,9 @@ namespace MiniMaxH3VAE {
                 return graph;
             };
             return restore_trailing_singleton_dims(
-                GGMLRunner::compute<float>(get_graph,
-                                           n_threads,
-                                           false),
+                GGMLRunner::compute(get_graph,
+                                    n_threads,
+                                    false),
                 5);
         }
     };

src/model/vae/tae.hpp

@@ -1,8 +1,12 @@
 #ifndef __SD_MODEL_VAE_TAE_HPP__
 #define __SD_MODEL_VAE_TAE_HPP__
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend.h"
+#include "core/ggml_runner.h"
+#include "core/rng.hpp"
+#include "core/util.h"
 #include "model.h"
+#include "model/common/ggml_block.hpp"
 
 /*
     ===================================    TinyAutoEncoder  ===================================
@@ -787,7 +791,7 @@ struct TinyImageAutoEncoder : public VAE {
             return build_graph(z_tensor, decode_graph);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), z_tensor.dim());
     }
 };
 
@@ -872,7 +876,7 @@ struct TinyVideoAutoEncoder : public VAE {
             return build_graph(z_tensor, decode_graph);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), z_tensor.dim());
     }
 };
 

src/model/vae/vae.hpp

@@ -5,6 +5,7 @@
 #include "model/common/block.hpp"
 #include "model/vae/vae_tiling.hpp"
 #include "model_manager.h"
+#include "runtime/tiling.h"
 
 struct VAE : public GGMLRunner {
 protected:

src/model/vae/wan_vae.hpp

@@ -4,6 +4,8 @@
 #include <map>
 #include <memory>
 #include <utility>
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_tensor_utils.h"
 
 #include "model/common/block.hpp"
 #include "model/vae/vae.hpp"
@@ -1427,7 +1429,7 @@ namespace WAN {
                     return build_temporal_tile_graph(input_tile, static_cast<int>(tile.start));
                 };
                 return restore_trailing_singleton_dims(
-                    GGMLRunner::compute<float>(get_graph, n_threads, false),
+                    GGMLRunner::compute(get_graph, n_threads, false),
                     static_cast<size_t>(input.dim()));
             });
 
@@ -1446,7 +1448,7 @@ namespace WAN {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(input.empty() ? z : input, decode_graph);
             };
-            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
+            auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false),
                                                           input.empty() ? z.dim() : input.dim());
             if (!result.empty() && z.dim() == 4) {
                 result.squeeze_(2);

src/runtime/cache_dit.hpp

@@ -8,8 +8,8 @@
 #include <unordered_map>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
 #include "core/tensor.hpp"
+#include "core/util.h"
 #include "runtime/condition_cache_utils.hpp"
 
 struct DBCacheConfig {

src/runtime/denoiser.hpp

@@ -11,8 +11,10 @@
 #include <string>
 #include <utility>
 
-#include "core/ggml_extend.hpp"
+#include "core/rng.hpp"
 #include "core/tensor.hpp"
+#include "core/util.h"
+#include "model.h"
 #include "runtime/gits_noise.h"
 #include "runtime/guidance.h"
 

src/runtime/easycache.hpp

@@ -6,8 +6,8 @@
 #include <unordered_map>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
 #include "core/tensor.hpp"
+#include "core/util.h"
 #include "runtime/condition_cache_utils.hpp"
 #include "runtime/denoiser.hpp"
 

src/runtime/preprocessing.hpp

@@ -4,7 +4,7 @@
 #include <cmath>
 #include <limits>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_tensor_utils.h"
 
 #define M_PI_ 3.14159265358979323846f
 

src/runtime/spectrum.hpp

@@ -5,8 +5,8 @@
 #include <cstring>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
 #include "core/tensor.hpp"
+#include "core/util.h"
 
 struct SpectrumConfig {
     float w            = 0.40f;

src/runtime/tiling.cpp

@@ -0,0 +1,268 @@
+#include "runtime/tiling.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "core/util.h"
+#include "ggml.h"
+
+static void sd_tiling_calc_tiles(int& num_tiles_dim,
+                                 float& tile_overlap_factor_dim,
+                                 int small_dim,
+                                 int tile_size,
+                                 const float tile_overlap_factor,
+                                 bool circular) {
+    int tile_overlap     = static_cast<int>(tile_size * tile_overlap_factor);
+    int non_tile_overlap = tile_size - tile_overlap;
+
+    if (circular) {
+        // circular means the last and first tile are overlapping (wraping around)
+        num_tiles_dim = small_dim / non_tile_overlap;
+
+        if (num_tiles_dim < 1) {
+            num_tiles_dim = 1;
+        }
+
+        tile_overlap_factor_dim = (tile_size - small_dim / num_tiles_dim) / (float)tile_size;
+
+        // if single tile and tile_overlap_factor is not 0, add one to ensure we have at least two overlapping tiles
+        if (num_tiles_dim == 1 && tile_overlap_factor_dim > 0) {
+            num_tiles_dim++;
+            tile_overlap_factor_dim = 0.5;
+        }
+
+        return;
+    }
+    // else, non-circular means the last and first tile are not overlapping
+
+    num_tiles_dim     = (small_dim - tile_overlap) / non_tile_overlap;
+    int overshoot_dim = ((num_tiles_dim + 1) * non_tile_overlap + tile_overlap) % small_dim;
+
+    if ((overshoot_dim != non_tile_overlap) && (overshoot_dim <= num_tiles_dim * (tile_size / 2 - tile_overlap))) {
+        // if tiles don't fit perfectly using the desired overlap
+        // and there is enough room to squeeze an extra tile without overlap becoming >0.5
+        num_tiles_dim++;
+    }
+
+    tile_overlap_factor_dim = (float)(tile_size * num_tiles_dim - small_dim) / (float)(tile_size * (num_tiles_dim - 1));
+    if (num_tiles_dim <= 2) {
+        if (small_dim <= tile_size) {
+            num_tiles_dim           = 1;
+            tile_overlap_factor_dim = 0;
+        } else {
+            num_tiles_dim           = 2;
+            tile_overlap_factor_dim = (2 * tile_size - small_dim) / (float)tile_size;
+        }
+    }
+}
+
+static int64_t sd_tensor_plane_size(const sd::Tensor<float>& tensor) {
+    GGML_ASSERT(tensor.dim() >= 2);
+    return tensor.shape()[0] * tensor.shape()[1];
+}
+
+static sd::Tensor<float> sd_tensor_split_2d(const sd::Tensor<float>& input, int width, int height, int x, int y) {
+    GGML_ASSERT(input.dim() >= 4);
+    std::vector<int64_t> output_shape = input.shape();
+    output_shape[0]                   = width;
+    output_shape[1]                   = height;
+    sd::Tensor<float> output(std::move(output_shape));
+    int64_t input_width  = input.shape()[0];
+    int64_t input_height = input.shape()[1];
+    int64_t input_plane  = sd_tensor_plane_size(input);
+    int64_t output_plane = sd_tensor_plane_size(output);
+    int64_t plane_count  = input.numel() / input_plane;
+    for (int iy = 0; iy < height; iy++) {
+        for (int ix = 0; ix < width; ix++) {
+            int64_t src_xy = (ix + x) % input_width + input_width * ((iy + y) % input_height);
+            int64_t dst_xy = ix + width * iy;
+            for (int64_t plane = 0; plane < plane_count; ++plane) {
+                output[plane * output_plane + dst_xy] = input[plane * input_plane + src_xy];
+            }
+        }
+    }
+    return output;
+}
+
+static void sd_tensor_merge_2d(const sd::Tensor<float>& input,
+                               sd::Tensor<float>* output,
+                               int x,
+                               int y,
+                               int overlap_x,
+                               int overlap_y,
+                               bool circular_x,
+                               bool circular_y,
+                               int x_skip,
+                               int y_skip) {
+    GGML_ASSERT(output != nullptr);
+    int64_t width        = input.shape()[0];
+    int64_t height       = input.shape()[1];
+    int64_t img_width    = output->shape()[0];
+    int64_t img_height   = output->shape()[1];
+    int64_t input_plane  = sd_tensor_plane_size(input);
+    int64_t output_plane = sd_tensor_plane_size(*output);
+    int64_t plane_count  = input.numel() / input_plane;
+    GGML_ASSERT(output->numel() / output_plane == plane_count);
+
+    // unclamped -> expects x in the range [0-1]
+    auto smootherstep_f32 = [](const float x) -> float {
+        GGML_ASSERT(x >= 0.f && x <= 1.f);
+        return x * x * x * (x * (6.0f * x - 15.0f) + 10.0f);
+    };
+
+    for (int iy = y_skip; iy < height; iy++) {
+        for (int ix = x_skip; ix < width; ix++) {
+            int64_t src_xy = ix + width * iy;
+            int64_t ox     = (x + ix) % img_width;
+            int64_t oy     = (y + iy) % img_height;
+            int64_t dst_xy = ox + img_width * oy;
+            for (int64_t plane = 0; plane < plane_count; ++plane) {
+                float new_value = input[plane * input_plane + src_xy];
+                if (overlap_x > 0 || overlap_y > 0) {
+                    float old_value   = (*output)[plane * output_plane + dst_xy];
+                    const float x_f_0 = (circular_x || (overlap_x > 0 && x > 0)) ? (ix - x_skip) / float(overlap_x) : 1.f;
+                    const float x_f_1 = (circular_x || (overlap_x > 0 && x < (img_width - width))) ? (width - ix) / float(overlap_x) : 1.f;
+                    const float y_f_0 = (circular_y || (overlap_y > 0 && y > 0)) ? (iy - y_skip) / float(overlap_y) : 1.f;
+                    const float y_f_1 = (circular_y || (overlap_y > 0 && y < (img_height - height))) ? (height - iy) / float(overlap_y) : 1.f;
+                    const float x_f   = std::min(std::min(x_f_0, x_f_1), 1.f);
+                    const float y_f   = std::min(std::min(y_f_0, y_f_1), 1.f);
+                    (*output)[plane * output_plane + dst_xy] =
+                        old_value + new_value * smootherstep_f32(y_f) * smootherstep_f32(x_f);
+                } else {
+                    (*output)[plane * output_plane + dst_xy] = new_value;
+                }
+            }
+        }
+    }
+}
+
+sd::Tensor<float> process_tiles_2d(const sd::Tensor<float>& input,
+                                   int output_width,
+                                   int output_height,
+                                   int scale,
+                                   int p_tile_size_x,
+                                   int p_tile_size_y,
+                                   float tile_overlap_factor,
+                                   bool circular_x,
+                                   bool circular_y,
+                                   const TileProcessCallback& on_processing,
+                                   bool silent) {
+    sd::Tensor<float> output;
+    int input_width  = static_cast<int>(input.shape()[0]);
+    int input_height = static_cast<int>(input.shape()[1]);
+
+    GGML_ASSERT(((input_width / output_width) == (input_height / output_height)) &&
+                ((output_width / input_width) == (output_height / input_height)));
+    GGML_ASSERT(((input_width / output_width) == scale) ||
+                ((output_width / input_width) == scale));
+
+    int small_width  = output_width;
+    int small_height = output_height;
+    bool decode      = output_width > input_width;
+    if (decode) {
+        small_width  = input_width;
+        small_height = input_height;
+    }
+
+    int num_tiles_x;
+    float tile_overlap_factor_x;
+    sd_tiling_calc_tiles(num_tiles_x, tile_overlap_factor_x, small_width, p_tile_size_x, tile_overlap_factor, circular_x);
+
+    int num_tiles_y;
+    float tile_overlap_factor_y;
+    sd_tiling_calc_tiles(num_tiles_y, tile_overlap_factor_y, small_height, p_tile_size_y, tile_overlap_factor, circular_y);
+
+    int tile_overlap_x     = static_cast<int32_t>(p_tile_size_x * tile_overlap_factor_x);
+    int non_tile_overlap_x = p_tile_size_x - tile_overlap_x;
+    int tile_overlap_y     = static_cast<int32_t>(p_tile_size_y * tile_overlap_factor_y);
+    int non_tile_overlap_y = p_tile_size_y - tile_overlap_y;
+    int tile_size_x        = p_tile_size_x < small_width ? p_tile_size_x : small_width;
+    int tile_size_y        = p_tile_size_y < small_height ? p_tile_size_y : small_height;
+    int input_tile_size_x  = tile_size_x;
+    int input_tile_size_y  = tile_size_y;
+    int output_tile_size_x = tile_size_x;
+    int output_tile_size_y = tile_size_y;
+    if (decode) {
+        output_tile_size_x *= scale;
+        output_tile_size_y *= scale;
+    } else {
+        input_tile_size_x *= scale;
+        input_tile_size_y *= scale;
+    }
+
+    int num_tiles   = num_tiles_x * num_tiles_y;
+    int tile_count  = 1;
+    bool last_y     = false;
+    bool last_x     = false;
+    float last_time = 0.0f;
+    if (!silent) {
+        LOG_VERBOSE("num tiles : %d, %d ", num_tiles_x, num_tiles_y);
+        LOG_VERBOSE("optimal overlap : %f, %f (targeting %f)", tile_overlap_factor_x, tile_overlap_factor_y, tile_overlap_factor);
+        LOG_VERBOSE("processing %i tiles", num_tiles);
+        pretty_progress(0, num_tiles, 0.0f);
+    }
+    for (int y = 0; y < small_height && !last_y; y += non_tile_overlap_y) {
+        int dy = 0;
+        if (!circular_y && y + tile_size_y >= small_height) {
+            int original_y = y;
+            y              = small_height - tile_size_y;
+            dy             = original_y - y;
+            if (decode) {
+                dy *= scale;
+            }
+            last_y = true;
+        }
+        for (int x = 0; x < small_width && !last_x; x += non_tile_overlap_x) {
+            int dx = 0;
+            if (!circular_x && x + tile_size_x >= small_width) {
+                int original_x = x;
+                x              = small_width - tile_size_x;
+                dx             = original_x - x;
+                if (decode) {
+                    dx *= scale;
+                }
+                last_x = true;
+            }
+
+            int x_in  = decode ? x : scale * x;
+            int y_in  = decode ? y : scale * y;
+            int x_out = decode ? x * scale : x;
+            int y_out = decode ? y * scale : y;
+
+            int overlap_x_out = decode ? tile_overlap_x * scale : tile_overlap_x;
+            int overlap_y_out = decode ? tile_overlap_y * scale : tile_overlap_y;
+
+            int64_t t1       = ggml_time_ms();
+            auto input_tile  = sd_tensor_split_2d(input, input_tile_size_x, input_tile_size_y, x_in, y_in);
+            auto output_tile = on_processing(input_tile);
+            if (output_tile.empty()) {
+                return {};
+            }
+            GGML_ASSERT(output_tile.shape()[0] == output_tile_size_x && output_tile.shape()[1] == output_tile_size_y);
+            if (output.empty()) {
+                std::vector<int64_t> output_shape = output_tile.shape();
+                output_shape[0]                   = output_width;
+                output_shape[1]                   = output_height;
+                output                            = sd::Tensor<float>::zeros(std::move(output_shape));
+            }
+            sd_tensor_merge_2d(output_tile, &output, x_out, y_out, overlap_x_out, overlap_y_out, circular_x, circular_y, dx, dy);
+
+            if (!silent) {
+                int64_t t2 = ggml_time_ms();
+                last_time  = (t2 - t1) / 1000.0f;
+                pretty_progress(tile_count, num_tiles, last_time);
+            }
+            tile_count++;
+        }
+        last_x = false;
+    }
+    if (!silent && tile_count < num_tiles) {
+        pretty_progress(num_tiles, num_tiles, last_time);
+    }
+    if (output.empty()) {
+        return {};
+    }
+    return output;
+}

src/runtime/tiling.h

@@ -0,0 +1,22 @@
+#ifndef __SD_RUNTIME_TILING_H__
+#define __SD_RUNTIME_TILING_H__
+
+#include <functional>
+
+#include "core/tensor.hpp"
+
+using TileProcessCallback = std::function<sd::Tensor<float>(const sd::Tensor<float>&)>;
+
+sd::Tensor<float> process_tiles_2d(const sd::Tensor<float>& input,
+                                   int output_width,
+                                   int output_height,
+                                   int scale,
+                                   int p_tile_size_x,
+                                   int p_tile_size_y,
+                                   float tile_overlap_factor,
+                                   bool circular_x,
+                                   bool circular_y,
+                                   const TileProcessCallback& on_processing,
+                                   bool silent = false);
+
+#endif  // __SD_RUNTIME_TILING_H__

src/runtime/ucache.hpp

@@ -6,8 +6,8 @@
 #include <unordered_map>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
 #include "core/tensor.hpp"
+#include "core/util.h"
 #include "runtime/condition_cache_utils.hpp"
 #include "runtime/denoiser.hpp"
 

src/stable-diffusion.cpp

@@ -1,4 +1,5 @@
 #include <algorithm>
+#include <cinttypes>
 #include <cmath>
 #include <cstdlib>
 #include <set>
@@ -7,9 +8,12 @@
 #include <utility>
 #include <vector>
 
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend_backend.h"
 #include "core/ggml_graph_cut.h"
+#include "core/ggml_runner.h"
+#include "core/ggml_tensor_utils.h"
 #include "core/layer_split_partition.h"
+#include "model.h"
 
 #include "core/rng.hpp"
 #include "core/rng_mt19937.hpp"
@@ -898,7 +902,7 @@ public:
             sampler_rng = rng;
         }
 
-        ggml_log_set(ggml_log_callback_default, nullptr);
+        ggml_log_set(sd_ggml_log_callback, nullptr);
 
         model_manager = std::make_shared<ModelManager>();
         model_manager->set_n_threads(n_threads);

src/upscaler.cpp

@@ -1,7 +1,8 @@
 #include "upscaler.h"
-#include "core/ggml_extend.hpp"
+#include "core/ggml_extend_backend.h"
 #include "core/util.h"
 #include "model_loader.h"
+#include "runtime/tiling.h"
 #include "stable-diffusion.h"
 
 #include <cstdlib>
@@ -34,7 +35,7 @@ void UpscalerGGML::set_max_graph_vram_bytes(size_t max_vram_bytes) {
 
 bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
                                   int n_threads) {
-    ggml_log_set(ggml_log_callback_default, nullptr);
+    ggml_log_set(sd_ggml_log_callback, nullptr);
 
     std::string error;
     if (!backend_manager.init(backend_spec.c_str(),