refactor: unify ggml operation helpers around runner context (#2314)

src/conditioning/conditioner.hpp

@@ -3448,8 +3448,8 @@ struct LTXAVTextProjection : public GGMLBlock {
 
         auto video_projection = std::dynamic_pointer_cast<Linear>(blocks["video_aggregate_embed"]);
         auto audio_projection = std::dynamic_pointer_cast<Linear>(blocks["audio_aggregate_embed"]);
-        auto video_in         = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(4096.f / static_cast<float>(kHiddenSize)));
-        auto audio_in         = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(2048.f / static_cast<float>(kHiddenSize)));
+        auto video_in         = ggml_ext_scale(ctx, x, std::sqrt(4096.f / static_cast<float>(kHiddenSize)));
+        auto audio_in         = ggml_ext_scale(ctx, x, std::sqrt(2048.f / static_cast<float>(kHiddenSize)));
         auto video            = video_projection->forward(ctx, video_in);
         auto audio            = audio_projection->forward(ctx, audio_in);
         return ggml_concat(ctx->ggml_ctx, video, audio, 0);

src/core/ggml_extend.h

@@ -10,33 +10,24 @@
 
 struct GGMLRunnerContext;
 
+constexpr uint32_t GGML_PRECISION_F32_VULKAN = 1u << 0;
+constexpr uint32_t GGML_PRECISION_F32_ROCM   = 1u << 1;
+
 #define EPS 1e-05f
 
 static_assert(GGML_MAX_NAME >= 160, "GGML_MAX_NAME must be at least 160");
 
-// 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);
+ggml_tensor* ggml_ext_cont(GGMLRunnerContext* runner_ctx, ggml_tensor* x);
 
 // torch like permute
-ggml_tensor* ggml_ext_torch_permute(ggml_context* ctx,
+ggml_tensor* ggml_ext_torch_permute(GGMLRunnerContext* runner_ctx,
                                     ggml_tensor* x,
                                     int axis0,
                                     int axis1,
                                     int axis2,
                                     int axis3);
 
-ggml_tensor* ggml_ext_slice(ggml_context* ctx,
+ggml_tensor* ggml_ext_slice(GGMLRunnerContext* runner_ctx,
                             ggml_tensor* x,
                             int dim,
                             int64_t start,
@@ -44,48 +35,38 @@ ggml_tensor* ggml_ext_slice(ggml_context* ctx,
                             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,
+std::vector<ggml_tensor*> ggml_ext_chunk(GGMLRunnerContext* runner_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_silu_act(GGMLRunnerContext* runner_ctx, ggml_tensor* x, bool gate_first = true);
 
 // Without gate, hidden contains [hidden, gate] packed along ne[0].
-ggml_tensor* ggml_ext_situ(GGMLRunnerContext* ctx, ggml_tensor* hidden, ggml_tensor* gate = nullptr);
+ggml_tensor* ggml_ext_situ(GGMLRunnerContext* runner_ctx, ggml_tensor* hidden, ggml_tensor* gate = nullptr);
 
-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_scale(GGMLRunnerContext* runner_ctx, ggml_tensor* x, float factor, bool inplace = false);
 
 // scale and optional shift must already be broadcastable to x in ggml layout.
 // Computes (x + x * scale) + shift; F32 uses sd.modulate.v1 when supported by ctx->backend.
-ggml_tensor* ggml_ext_modulate(GGMLRunnerContext* ctx,
+ggml_tensor* ggml_ext_modulate(GGMLRunnerContext* runner_ctx,
                                ggml_tensor* x,
                                ggml_tensor* shift,
                                ggml_tensor* scale);
 
-ggml_tensor* ggml_ext_gelu(ggml_context* ctx,
-                           ggml_tensor* x,
-                           bool inplace = false);
+ggml_tensor* ggml_ext_gelu(GGMLRunnerContext* runner_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_gelu_quick(GGMLRunnerContext* runner_ctx, ggml_tensor* x, bool inplace = false);
 
-ggml_tensor* ggml_ext_linear(ggml_context* ctx,
+ggml_tensor* ggml_ext_linear(GGMLRunnerContext* runner_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* ggml_ext_linear_i8_tensorwise(GGMLRunnerContext* runner_ctx,
                                            ggml_tensor* x,
                                            ggml_tensor* w,
                                            ggml_tensor* weight_scale,
@@ -93,8 +74,7 @@ ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx,
                                            int convrot_group_size,
                                            float scale = 1.f);
 
-ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
-                              ggml_backend_t backend,
+ggml_tensor* ggml_ext_pad_ext(GGMLRunnerContext* runner_ctx,
                               ggml_tensor* x,
                               int lp0,
                               int rp0,
@@ -107,7 +87,7 @@ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
                               bool circular_x = false,
                               bool circular_y = false);
 
-ggml_tensor* ggml_ext_pad(ggml_context* ctx,
+ggml_tensor* ggml_ext_pad(GGMLRunnerContext* runner_ctx,
                           ggml_tensor* x,
                           int p0,
                           int p1,
@@ -118,7 +98,7 @@ ggml_tensor* ggml_ext_pad(ggml_context* ctx,
 
 // ggml layout: x [L, IC, N], w [K, IC/groups, OC], b [OC], result [OL, OC, N].
 // force_prec_f32 keeps both input patches and weights in F32.
-ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx,
+ggml_tensor* ggml_ext_conv_1d(GGMLRunnerContext* runner_ctx,
                               ggml_tensor* x,
                               ggml_tensor* w,
                               ggml_tensor* b,
@@ -132,7 +112,7 @@ ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx,
 // x: [N, IC, IH, IW]
 // b: [OC,]
 // result: [N, OC, OH, OW]
-ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
+ggml_tensor* ggml_ext_conv_2d(GGMLRunnerContext* runner_ctx,
                               ggml_tensor* x,
                               ggml_tensor* w,
                               ggml_tensor* b,
@@ -151,8 +131,7 @@ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
 // 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* ggml_ext_conv_3d(GGMLRunnerContext* runner_ctx,
                               ggml_tensor* x,
                               ggml_tensor* w,
                               ggml_tensor* b,
@@ -168,65 +147,41 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
                               int d2              = 1,
                               bool force_prec_f32 = false,
                               bool direct         = false,
-                              float scale         = 1.f);
-
-// 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);
+                              float scale         = 1.f,
+                              bool allow_conv2d   = false);
 
 // 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);
+std::vector<ggml_tensor*> ggml_ext_split_qkv(GGMLRunnerContext* runner_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);
+std::vector<ggml_tensor*> ggml_ext_split_image_qkv(GGMLRunnerContext* runner_ctx, ggml_tensor* qkv);
 
 // Constant and cast helpers require the built-in tensors initialized by GGMLRunner.
-ggml_tensor* ggml_ext_full(ggml_context* ctx,
+ggml_tensor* ggml_ext_full(GGMLRunnerContext* runner_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(GGMLRunnerContext* runner_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_zeros_like(GGMLRunnerContext* runner_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(GGMLRunnerContext* runner_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_ones_like(GGMLRunnerContext* runner_ctx, ggml_tensor* x);
 
-ggml_tensor* ggml_ext_cast_f32(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* a);
+ggml_tensor* ggml_ext_cast_f32(GGMLRunnerContext* runner_ctx, 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* ggml_ext_attention_ext(GGMLRunnerContext* runner_ctx,
                                     ggml_tensor* q,
                                     ggml_tensor* k,
                                     ggml_tensor* v,
@@ -235,31 +190,138 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
                                     bool skip_reshape     = false,
                                     bool flash_attn       = false,
                                     float kv_scale        = 1.0f,
-                                    bool sage_attn        = false,
                                     bool* used_flash_attn = nullptr);
 
-ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
+ggml_tensor* ggml_ext_layer_norm(GGMLRunnerContext* runner_ctx,
                                  ggml_tensor* x,
                                  ggml_tensor* w,
                                  ggml_tensor* b,
                                  float eps = EPS);
 
-ggml_tensor* ggml_ext_group_norm(ggml_context* ctx,
+ggml_tensor* ggml_ext_group_norm(GGMLRunnerContext* runner_ctx,
                                  ggml_tensor* x,
                                  ggml_tensor* w,
                                  ggml_tensor* b,
                                  int num_groups = 32,
                                  float eps      = 1e-6f);
 
-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_timestep_embedding(GGMLRunnerContext* runner_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);
+ggml_tensor* ggml_ext_vec_concat(GGMLRunnerContext* runner_ctx, std::vector<ggml_tensor*>& tensors, int dim);
+
+#ifndef SD_USE_UPSTREAM_GGML
+ggml_tensor* ggml_ext_mul_mat_i8_tensorwise(GGMLRunnerContext* runner_ctx,
+                                            ggml_tensor* weight,
+                                            ggml_tensor* input,
+                                            ggml_tensor* weight_scale,
+                                            ggml_tensor* bias,
+                                            int group_size,
+                                            ggml_tensor* input_scale = nullptr);
+
+ggml_tensor* ggml_ext_quantize_i8_convrot(GGMLRunnerContext* runner_ctx, ggml_tensor* input, int group_size);
+
+// Packed I8 [K/2, N], low nibble first. W4A8 uses a 16-entry F32 codebook
+// and F8_E4M3 [K/16, N] relative scales. F32 inputs get an explicit quantize node.
+ggml_tensor* ggml_ext_mul_mat_w4_convrot(GGMLRunnerContext* runner_ctx,
+                                         ggml_tensor* weight,
+                                         ggml_tensor* input,
+                                         ggml_tensor* weight_scale,
+                                         ggml_tensor* s_channel,
+                                         ggml_tensor* s_rel,
+                                         ggml_tensor* bias,
+                                         SDW4Kind kind,
+                                         int group_size);
+
+// Q/K: contiguous F32 [D, tokens, heads, batch], V: contiguous F16 in the same layout.
+// Output: F32 [D, query heads, query tokens, batch]. CUDA-only, unmasked inference.
+ggml_tensor* ggml_ext_sage_attn(GGMLRunnerContext* runner_ctx,
+                                ggml_tensor* q,
+                                ggml_tensor* k,
+                                ggml_tensor* v,
+                                float scale,
+                                ggml_sage_attn_mode mode);
+
+// Noncausal Sol-Attn: contiguous F32 Q/K/V [128, tokens, heads, batch], output [128, heads, tokens, batch].
+ggml_tensor* ggml_ext_sol_attn(GGMLRunnerContext* runner_ctx,
+                               ggml_tensor* q,
+                               ggml_tensor* k,
+                               ggml_tensor* v,
+                               float scale,
+                               float tau);
+
+// Contiguous F32 Q [D*heads, W*H], K/V [D*kv_heads, W*H*planes].
+// Returns Q's layout, or nullptr when the backend has no compatible kernel.
+ggml_tensor* ggml_ext_window_attn(GGMLRunnerContext* runner_ctx,
+                                  ggml_tensor* q,
+                                  ggml_tensor* k,
+                                  ggml_tensor* v,
+                                  const SDWindowAttentionParams& params);
+#endif
+
+// x: F32 [head_dim, heads, tokens, batch], pe: F32 [2, 2, rot_dim/2, tokens].
+// Returns [head_dim, tokens, heads*batch]; channels beyond rot_dim are unchanged.
+ggml_tensor* ggml_ext_rope_apply(GGMLRunnerContext* runner_ctx,
+                                 ggml_tensor* x,
+                                 ggml_tensor* pe,
+                                 bool interleaved = true);
+
+ggml_tensor* ggml_ext_repeat_4d(GGMLRunnerContext* runner_ctx,
+                                ggml_tensor* x,
+                                int64_t ne0,
+                                int64_t ne1,
+                                int64_t ne2,
+                                int64_t ne3);
+
+ggml_tensor* ggml_ext_repeat(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* shape);
+
+ggml_tensor* ggml_ext_embedding(GGMLRunnerContext* runner_ctx,
+                                ggml_tensor* weight,
+                                ggml_tensor* input_ids,
+                                ggml_tensor* weight_scale = nullptr,
+                                int group_size            = 0);
+
+ggml_tensor* ggml_ext_prepare_linear_weight(GGMLRunnerContext* runner_ctx,
+                                            ggml_tensor* w,
+                                            ggml_tensor* x,
+                                            bool force_prec_f32);
+
+ggml_tensor* ggml_ext_prepare_i8_input(GGMLRunnerContext* runner_ctx, ggml_tensor* x);
+
+bool ggml_ext_force_prec_f32(GGMLRunnerContext* runner_ctx,
+                             bool requested,
+                             uint32_t flags);
+
+#ifndef SD_USE_UPSTREAM_GGML
+ggml_tensor* ggml_ext_linear_w4_convrot(GGMLRunnerContext* runner_ctx,
+                                        ggml_tensor* x,
+                                        ggml_tensor* w,
+                                        ggml_tensor* ws,
+                                        ggml_tensor* sc,
+                                        ggml_tensor* sr,
+                                        ggml_tensor* bias,
+                                        SDW4Kind kind,
+                                        int group,
+                                        float scale = 1.f);
+#endif
+
+ggml_tensor* ggml_ext_conv_2d_grouped(GGMLRunnerContext* runner_ctx,
+                                      ggml_tensor* x,
+                                      ggml_tensor* w,
+                                      ggml_tensor* b,
+                                      int groups,
+                                      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);
 
 #endif  // __SD_CORE_GGML_EXTEND_H__

src/core/ggml_extensions.cpp

@@ -131,20 +131,12 @@ bool sd_ggml_rope_describe(const ggml_tensor* node) {
            std::memcmp(output.ne, node->ne, sizeof(output.ne)) == 0;
 }
 
-struct I8Attributes {
-    int64_t shape[3];
-    int32_t group_size;
-    int32_t flags;
-};
-enum I8Flags { I8_PACKED       = 1,
-               I8_INPUT_SCALE  = 2,
-               I8_WEIGHT_SCALE = 4,
-               I8_BIAS         = 8 };
-
-static bool valid_group(int group) {
+bool sd_ggml_valid_convrot_group(int group) {
     return group == 1 || group == 4 || group == 16 || group == 64 || group == 256;
 }
+
 static constexpr int64_t i8_size_limit = int64_t(std::min(uint64_t(INT64_MAX), uint64_t(SIZE_MAX)));
+
 static bool packed_height(int64_t k, int64_t rows, int64_t* height) {
     if (k <= 0 || rows <= 0 || rows > i8_size_limit / 4)
         return false;
@@ -155,15 +147,17 @@ static bool packed_height(int64_t k, int64_t rows, int64_t* height) {
     *height = padded + scales;
     return true;
 }
+
 static bool scale_matches(const ggml_tensor* scale, int64_t count) {
     return !scale || (scale->type == GGML_TYPE_F32 && ggml_is_contiguous(scale) && ggml_nelements(scale) == count);
 }
+
 static bool i8_shape(const SDI8Operation& op, const int64_t* shape) {
     const auto* w = op.weight;
     const auto* x = op.input;
     if (!w || !x || w->type != GGML_TYPE_I8 || w->ne[2] != 1 || w->ne[3] != 1 ||
         !ggml_is_contiguous(w) || !ggml_is_contiguous(x) || x->ne[0] != w->ne[0] ||
-        (op.group_size != 0 && (!valid_group(op.group_size) || w->ne[0] % op.group_size)))
+        (op.group_size != 0 && (!sd_ggml_valid_convrot_group(op.group_size) || w->ne[0] % op.group_size)))
         return false;
     int64_t rows = 1;
     for (int i = 0; i < 3; ++i) {
@@ -177,7 +171,7 @@ static bool i8_shape(const SDI8Operation& op, const int64_t* shape) {
         rows > (i8_size_limit - 3) / (k + 4 + n * 4) || !packed_height(k, rows, &height))
         return false;
     if (op.packed) {
-        if (x->type != GGML_TYPE_I8 || !valid_group(op.group_size) || op.input_scale ||
+        if (x->type != GGML_TYPE_I8 || !sd_ggml_valid_convrot_group(op.group_size) || op.input_scale ||
             x->ne[1] != height || x->ne[2] != 1 || x->ne[3] != 1)
             return false;
     } else {
@@ -191,14 +185,15 @@ static bool i8_shape(const SDI8Operation& op, const int64_t* shape) {
     }
     return scale_matches(op.input_scale, rows) && scale_matches(op.weight_scale, w->ne[1]) && scale_matches(op.bias, w->ne[1]);
 }
+
 static bool i8_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* output) {
-    I8Attributes attr;
+    SDI8Attributes attr;
     std::memcpy(&attr, attributes, sizeof(attr));
     if (count < 2 || (attr.flags & ~15))
         return false;
-    SDI8Operation op = {inputs[0], inputs[1], nullptr, nullptr, nullptr, attr.group_size, (attr.flags & I8_PACKED) != 0};
+    SDI8Operation op = {inputs[0], inputs[1], nullptr, nullptr, nullptr, attr.group_size, (attr.flags & SD_I8_PACKED) != 0};
     int index        = 2;
-    for (auto entry : {std::make_pair(I8_INPUT_SCALE, &op.input_scale), std::make_pair(I8_WEIGHT_SCALE, &op.weight_scale), std::make_pair(I8_BIAS, &op.bias)}) {
+    for (auto entry : {std::make_pair(SD_I8_INPUT_SCALE, &op.input_scale), std::make_pair(SD_I8_WEIGHT_SCALE, &op.weight_scale), std::make_pair(SD_I8_BIAS, &op.bias)}) {
         if (attr.flags & entry.first) {
             if (index >= count)
                 return false;
@@ -212,10 +207,11 @@ static bool i8_infer(ggml_tensor* const* inputs, int count, const void* attribut
     std::memcpy(output->ne + 1, attr.shape, sizeof(attr.shape));
     return true;
 }
+
 static bool convrot_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* output) {
     int32_t group;
     std::memcpy(&group, attributes, sizeof(group));
-    if (count != 1 || inputs[0]->type != GGML_TYPE_F32 || !ggml_is_contiguous(inputs[0]) || !valid_group(group) || inputs[0]->ne[0] <= 0 || inputs[0]->ne[0] % group)
+    if (count != 1 || inputs[0]->type != GGML_TYPE_F32 || !ggml_is_contiguous(inputs[0]) || !sd_ggml_valid_convrot_group(group) || inputs[0]->ne[0] <= 0 || inputs[0]->ne[0] % group)
         return false;
     const int64_t k = inputs[0]->ne[0], rows = ggml_nrows(inputs[0]);
     int64_t height;
@@ -227,38 +223,42 @@ static bool convrot_infer(ggml_tensor* const* inputs, int count, const void* att
     output->ne[2] = output->ne[3] = 1;
     return true;
 }
+
 ggml_ext_op_t sd_ggml_i8_op() {
-    static const auto op = [] { ggml_ext_op_desc d = {GGML_EXT_ABI_VERSION, "sd.linear_i8.v1", sizeof(I8Attributes), i8_infer}; auto result = ggml_ext_register_op(&d); GGML_ASSERT(result); return result; }();
+    static const auto op = [] { ggml_ext_op_desc d = {GGML_EXT_ABI_VERSION, "sd.linear_i8.v1", sizeof(SDI8Attributes), i8_infer}; auto result = ggml_ext_register_op(&d); GGML_ASSERT(result); return result; }();
     return op;
 }
+
 ggml_ext_op_t sd_ggml_convrot_op() {
     static const auto op = [] { ggml_ext_op_desc d = {GGML_EXT_ABI_VERSION, "sd.quantize_i8_convrot.v1", sizeof(int32_t), convrot_infer}; auto result = ggml_ext_register_op(&d); GGML_ASSERT(result); return result; }();
     return op;
 }
+
 bool sd_ggml_i8_describe(const ggml_tensor* node, SDI8Operation* op) {
     if (!node || node->type != GGML_TYPE_F32 || !ggml_is_contiguous(node))
         return false;
     *op = {node->src[0], node->src[1], nullptr, nullptr, nullptr, 0, false};
     if (ggml_ext_op_get(node) == sd_ggml_i8_op()) {
-        I8Attributes attr;
+        SDI8Attributes attr;
         std::memcpy(&attr, ggml_ext_op_attributes(node), sizeof(attr));
         if (attr.flags & ~15 || std::memcmp(attr.shape, node->ne + 1, sizeof(attr.shape)))
             return false;
         op->group_size = attr.group_size;
-        op->packed     = (attr.flags & I8_PACKED) != 0;
+        op->packed     = (attr.flags & SD_I8_PACKED) != 0;
         int index      = 2;
-        if (attr.flags & I8_INPUT_SCALE)
+        if (attr.flags & SD_I8_INPUT_SCALE)
             op->input_scale = node->src[index++];
-        if (attr.flags & I8_WEIGHT_SCALE)
+        if (attr.flags & SD_I8_WEIGHT_SCALE)
             op->weight_scale = node->src[index++];
-        if (attr.flags & I8_BIAS)
+        if (attr.flags & SD_I8_BIAS)
             op->bias = node->src[index++];
-        if (((attr.flags & I8_INPUT_SCALE) && !op->input_scale) || ((attr.flags & I8_WEIGHT_SCALE) && !op->weight_scale) || ((attr.flags & I8_BIAS) && !op->bias))
+        if (((attr.flags & SD_I8_INPUT_SCALE) && !op->input_scale) || ((attr.flags & SD_I8_WEIGHT_SCALE) && !op->weight_scale) || ((attr.flags & SD_I8_BIAS) && !op->bias))
             return false;
     } else
         return false;
     return i8_shape(*op, node->ne + 1) && node->ne[0] == op->weight->ne[1];
 }
+
 bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group) {
     if (!node || !node->src[0] || !ggml_is_contiguous(node))
         return false;
@@ -269,48 +269,11 @@ bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group) {
     ggml_ext_tensor_desc output;
     return convrot_infer(node->src, 1, group, &output) && output.type == node->type && !std::memcmp(output.ne, node->ne, sizeof(output.ne));
 }
-ggml_tensor* ggml_mul_mat_i8_tensorwise(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* ws, ggml_tensor* bias, int group, ggml_tensor* xs) {
-    if (!w || !x)
-        return nullptr;
-    const auto* logical = x;
-    if (x->type == GGML_TYPE_I8 && !xs) {
-        int packed_group;
-        if (!sd_ggml_convrot_describe(x, &packed_group) || packed_group != group)
-            return nullptr;
-        logical = x->src[0];
-    }
-    I8Attributes attr      = {{logical->ne[1], logical->ne[2], logical->ne[3]}, group, logical != x ? I8_PACKED : 0};
-    ggml_tensor* inputs[5] = {w, x};
-    int count              = 2;
-    if (xs) {
-        attr.flags |= I8_INPUT_SCALE;
-        inputs[count++] = xs;
-    }
-    if (ws) {
-        attr.flags |= I8_WEIGHT_SCALE;
-        inputs[count++] = ws;
-    }
-    if (bias) {
-        attr.flags |= I8_BIAS;
-        inputs[count++] = bias;
-    }
-    return ggml_ext_op_create(ctx, sd_ggml_i8_op(), inputs, count, &attr);
-}
-ggml_tensor* ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* x, int group) {
-    int32_t attr = group;
-    return ggml_ext_op_create(ctx, sd_ggml_convrot_op(), &x, 1, &attr);
-}
-
-struct W4Attributes {
-    int64_t shape[3];
-    SDW4Kind kind;
-    int32_t has_bias;
-};
 
 static bool w4_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
     if (!attributes || count < 3)
         return false;
-    W4Attributes attr;
+    SDW4Attributes attr;
     std::memcpy(&attr, attributes, sizeof(attr));
     const bool codebook = attr.kind == SDW4Kind::W4A8;
     if ((!codebook && attr.kind != SDW4Kind::W4A4) || (attr.has_bias != 0 && attr.has_bias != 1) ||
@@ -348,7 +311,7 @@ static bool w4_infer(ggml_tensor* const* inputs, int count, const void* attribut
 
 ggml_ext_op_t sd_ggml_w4_op() {
     static const auto op = [] {
-        const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.linear_w4_convrot.v1", sizeof(W4Attributes), w4_infer};
+        const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.linear_w4_convrot.v1", sizeof(SDW4Attributes), w4_infer};
         auto result                 = ggml_ext_register_op(&desc);
         GGML_ASSERT(result);
         return result;
@@ -359,7 +322,7 @@ ggml_ext_op_t sd_ggml_w4_op() {
 bool sd_ggml_w4_describe(const ggml_tensor* node, SDW4Operation* op) {
     if (!node || ggml_ext_op_get(node) != sd_ggml_w4_op() || !ggml_is_contiguous(node))
         return false;
-    W4Attributes attr;
+    SDW4Attributes attr;
     std::memcpy(&attr, ggml_ext_op_attributes(node), sizeof(attr));
     const bool codebook = attr.kind == SDW4Kind::W4A8;
     const int count     = (codebook ? 5 : 3) + (attr.has_bias != 0);
@@ -372,30 +335,6 @@ bool sd_ggml_w4_describe(const ggml_tensor* node, SDW4Operation* op) {
     return true;
 }
 
-ggml_tensor* ggml_mul_mat_w4_convrot(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* ws, ggml_tensor* sc, ggml_tensor* sr, ggml_tensor* bias, SDW4Kind kind, int group) {
-    if (!w || !x || group < 0 || (group && !valid_group(group)))
-        return nullptr;
-    const int quant_group = group ? group : 1;
-    if (x->type == GGML_TYPE_F32)
-        x = ggml_quantize_i8_convrot(ctx, x, quant_group);
-    int packed_group;
-    if (!sd_ggml_convrot_describe(x, &packed_group) || packed_group != quant_group)
-        return nullptr;
-    const auto* logical    = x->src[0];
-    W4Attributes attr      = {{logical->ne[1], logical->ne[2], logical->ne[3]}, kind, bias ? 1 : 0};
-    ggml_tensor* inputs[6] = {w, x, ws};
-    int count              = 3;
-    if (kind == SDW4Kind::W4A8) {
-        inputs[count++] = sc;
-        inputs[count++] = sr;
-    } else if (sc || sr) {
-        return nullptr;
-    }
-    if (bias)
-        inputs[count++] = bias;
-    return ggml_ext_op_create(ctx, sd_ggml_w4_op(), inputs, count, &attr);
-}
-
 static bool sage_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
     if (count != 3 || !attributes)
         return false;
@@ -443,12 +382,6 @@ bool sd_ggml_sage_attn_describe(const ggml_tensor* node, SDSageAttentionParams*
     return true;
 }
 
-ggml_tensor* ggml_sage_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, ggml_sage_attn_mode mode) {
-    ggml_tensor* inputs[]              = {q, k, v};
-    const SDSageAttentionParams params = {scale, mode};
-    return ggml_ext_op_create(ctx, sd_ggml_sage_attn_op(), inputs, 3, &params);
-}
-
 static bool sol_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
     if (count != 3 || !attributes)
         return false;
@@ -495,13 +428,10 @@ bool sd_ggml_sol_attn_describe(const ggml_tensor* node, SDSolAttentionParams* pa
     return true;
 }
 
-ggml_tensor* ggml_sol_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, float tau) {
-    ggml_tensor* inputs[]             = {q, k, v};
-    const SDSolAttentionParams params = {scale, tau};
-    return ggml_ext_op_create(ctx, sd_ggml_sol_attn_op(), inputs, 3, &params);
-}
-
-static bool window_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
+bool sd_ggml_window_attn_infer(ggml_tensor* const* inputs,
+                               int count,
+                               const void* attributes,
+                               ggml_ext_tensor_desc* out) {
     if (count != 3 || !attributes)
         return false;
     SDWindowAttentionParams p;
@@ -531,7 +461,7 @@ static bool window_attn_infer(ggml_tensor* const* inputs, int count, const void*
 
 ggml_ext_op_t sd_ggml_window_attn_op() {
     static const auto op = [] {
-        const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.window_attn.v1", sizeof(SDWindowAttentionParams), window_attn_infer};
+        const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.window_attn.v1", sizeof(SDWindowAttentionParams), sd_ggml_window_attn_infer};
         auto result                 = ggml_ext_register_op(&desc);
         GGML_ASSERT(result);
         return result;
@@ -544,22 +474,12 @@ bool sd_ggml_window_attn_describe(const ggml_tensor* node, SDWindowAttentionPara
         return false;
     const auto* attributes = ggml_ext_op_attributes(node);
     ggml_ext_tensor_desc out;
-    if (!window_attn_infer(node->src, 3, attributes, &out) || node->type != out.type || std::memcmp(node->ne, out.ne, sizeof(out.ne)))
+    if (!sd_ggml_window_attn_infer(node->src, 3, attributes, &out) || node->type != out.type || std::memcmp(node->ne, out.ne, sizeof(out.ne)))
         return false;
     std::memcpy(params, attributes, sizeof(*params));
     return true;
 }
 
-ggml_tensor* sd_ggml_window_attn(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, const SDWindowAttentionParams& params) {
-    ggml_tensor* inputs[] = {q, k, v};
-    ggml_ext_tensor_desc out;
-    if (!backend || !window_attn_infer(inputs, 3, &params, &out))
-        return nullptr;
-    sd_ggml_register_extensions(backend);
-    auto result = ggml_ext_op_create(ctx, sd_ggml_window_attn_op(), inputs, 3, &params);
-    return ggml_backend_supports_op(backend, result) ? result : nullptr;
-}
-
 void sd_ggml_register_extensions(ggml_backend_t backend) {
     if (!backend) {
         return;
@@ -583,13 +503,6 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
     sd_cpu::register_extensions(device);
 }
 
-ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved) {
-    sd_ggml_register_extensions(backend);
-    ggml_tensor* inputs[] = {x, pe};
-    int32_t attributes    = interleaved ? 1 : 0;
-    return ggml_ext_op_create(ctx, sd_ggml_rope_op(), inputs, 2, &attributes);
-}
-
 #else
 void sd_ggml_register_extensions(ggml_backend_t) {}
 #endif

src/core/ggml_extensions.h

@@ -9,18 +9,24 @@ void sd_ggml_register_extensions(ggml_backend_t backend);
 #include "ggml-ext.h"
 
 ggml_ext_op_t sd_ggml_modulate_op();
+
 bool sd_ggml_modulate_describe(const ggml_tensor* node);
+
 size_t sd_ggml_modulate_workspace_size(const ggml_tensor* node);
 
 ggml_ext_op_t sd_ggml_rope_op();
+
 bool sd_ggml_rope_describe(const ggml_tensor* node);
-// Apply precomputed rotary matrices, with shapes in ne[] order:
-// a: F32 [head_dim, heads, tokens, batch], pe: F32 [2, 2, rot_dim/2, tokens].
-// pe stores [[cos, -sin], [sin, cos]] for each pair and token, shared across heads/batches.
-// Returns contiguous F32 [head_dim, tokens, heads*batch]. Input strides are preserved.
-// Interleaved pairs are (2*i, 2*i+1); otherwise pairs are (i, i+rot_dim/2).
-// Channels from rot_dim to head_dim are copied unchanged.
-ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved);
+
+struct SDI8Attributes {
+    int64_t shape[3];
+    int32_t group_size;
+    int32_t flags;
+};
+enum SDI8Flags { SD_I8_PACKED       = 1,
+                 SD_I8_INPUT_SCALE  = 2,
+                 SD_I8_WEIGHT_SCALE = 4,
+                 SD_I8_BIAS         = 8 };
 
 struct SDI8Operation {
     const ggml_tensor* weight;
@@ -32,16 +38,25 @@ struct SDI8Operation {
     bool packed;
 };
 
+bool sd_ggml_valid_convrot_group(int group);
+
 ggml_ext_op_t sd_ggml_i8_op();
+
 ggml_ext_op_t sd_ggml_convrot_op();
+
 bool sd_ggml_i8_describe(const ggml_tensor* node, SDI8Operation* operation);
+
 bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group_size);
-ggml_tensor* ggml_mul_mat_i8_tensorwise(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* weight_scale, ggml_tensor* bias, int group_size, ggml_tensor* input_scale = nullptr);
-ggml_tensor* ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* input, int group_size);
 
 enum class SDW4Kind : int32_t { W4A4 = 1,
                                 W4A8 = 2 };
 
+struct SDW4Attributes {
+    int64_t shape[3];
+    SDW4Kind kind;
+    int32_t has_bias;
+};
+
 struct SDW4Operation {
     const ggml_tensor* weight;
     const ggml_tensor* input;
@@ -53,10 +68,8 @@ struct SDW4Operation {
 };
 
 ggml_ext_op_t sd_ggml_w4_op();
+
 bool sd_ggml_w4_describe(const ggml_tensor* node, SDW4Operation* operation);
-// Packed I8 [K/2, N], low nibble first. W4A8 uses a 16-entry F32 codebook
-// and F8_E4M3 [K/16, N] relative scales. F32 inputs get an explicit quantize node.
-ggml_tensor* ggml_mul_mat_w4_convrot(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* weight_scale, ggml_tensor* s_channel, ggml_tensor* s_rel, ggml_tensor* bias, SDW4Kind kind, int group_size);
 
 enum ggml_sage_attn_mode {
     GGML_SAGE_ATTN_AUTO = 0,
@@ -70,10 +83,8 @@ struct SDSageAttentionParams {
 };
 
 ggml_ext_op_t sd_ggml_sage_attn_op();
+
 bool sd_ggml_sage_attn_describe(const ggml_tensor* node, SDSageAttentionParams* params);
-// Q/K: contiguous F32 [D, tokens, heads, batch], V: contiguous F16 in the same layout.
-// Output: F32 [D, query heads, query tokens, batch]. CUDA-only, unmasked inference.
-ggml_tensor* ggml_sage_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, ggml_sage_attn_mode mode);
 
 struct SDSolAttentionParams {
     float scale;
@@ -86,16 +97,19 @@ struct SDWindowAttentionParams {
     int32_t preserve_boundary;
 };
 
+bool sd_ggml_window_attn_infer(ggml_tensor* const* inputs,
+                               int count,
+                               const void* attributes,
+                               ggml_ext_tensor_desc* out);
+
 ggml_ext_op_t sd_ggml_window_attn_op();
+
 bool sd_ggml_window_attn_describe(const ggml_tensor* node, SDWindowAttentionParams* params);
-// Contiguous F32 Q [D*heads, W*H], K/V [D*kv_heads, W*H*planes].
-// Returns Q's layout, or nullptr when the backend has no compatible kernel.
-ggml_tensor* sd_ggml_window_attn(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, const SDWindowAttentionParams& params);
 
 ggml_ext_op_t sd_ggml_sol_attn_op();
+
 bool sd_ggml_sol_attn_describe(const ggml_tensor* node, SDSolAttentionParams* params);
-// Noncausal Sol-Attn: contiguous F32 Q/K/V [128, tokens, heads, batch], output [128, heads, tokens, batch].
-ggml_tensor* ggml_sol_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, float tau);
+
 #endif
 
 #endif

src/core/ggml_runner.cpp

@@ -13,22 +13,6 @@
 
 using namespace sd;
 
-ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
-                                    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,
-                                    bool* used_flash_attn) {
-    if (ctx->attn_scale > 0.f) {
-        kv_scale = ctx->attn_scale;
-    }
-    return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, ctx->sage_attn_enabled, used_flash_attn);
-}
-
 void GGMLRunner::alloc_params_ctx() {
     ggml_init_params params;
     params.mem_size   = static_cast<size_t>(MAX_PARAMS_TENSOR_NUM * ggml_tensor_overhead());

src/core/ggml_runner.h

@@ -23,6 +23,8 @@
 #define MAX_PARAMS_TENSOR_NUM 32768
 #define MAX_GRAPH_SIZE 327680
 
+struct GGMLRunnerContext;
+
 struct WeightAdapter {
     struct ForwardParams {
         enum class op_type_t {
@@ -46,22 +48,26 @@ struct WeightAdapter {
             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,
+
+    virtual ggml_tensor* patch_weight(GGMLRunnerContext* runner_ctx,
+                                      ggml_tensor* weight,
+                                      const std::string& weight_name) = 0;
+
+    virtual ggml_tensor* forward_with_lora(GGMLRunnerContext* runner_ctx,
                                            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,
+                                           ForwardParams forward_params) = 0;
+
+    virtual ggml_tensor* add_lora_to_output(GGMLRunnerContext* runner_ctx,
                                             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;
+                                            ForwardParams forward_params) = 0;
+
+    virtual size_t get_extra_graph_size() = 0;
 };
 
 struct GGMLRunnerContext {
@@ -124,17 +130,6 @@ struct GGMLRunnerContext {
     }
 };
 
-ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
-                                    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.f,
-                                    bool* used_flash_attn = nullptr);
-
 struct GGMLRunner {
 private:
     std::map<ggml_backend_t, size_t> logged_compute_bytes_;
@@ -142,11 +137,15 @@ private:
     ggml_status last_compute_status_ = GGML_STATUS_SUCCESS;
 
     sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, const std::unordered_set<const ggml_tensor*>* external_outputs = nullptr);
+
     std::vector<DeviceMemoryRequest> memory_requests(const std::vector<sd::BackendBufferSize>& sizes,
                                                      size_t pending_cache_bytes) const;
+
     bool fits(const std::vector<DeviceMemoryRequest>& requests,
               const std::vector<ggml_tensor*>& params) const;
+
     bool execute_segment(ggml_cgraph* graph, int n_threads, const std::function<void()>& on_before_compute);
+
     std::optional<sd::Tensor<float>> execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs);
 
 protected:

src/core/probe/flash_attn.cpp

@@ -9,6 +9,7 @@
 
 #include "core/ggml_extend.h"
 #include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
 #include "core/probe/backend_probe.h"
 #include "core/util.h"
 
@@ -41,8 +42,11 @@ static void build_flash_attention_probe(BackendProbeGraph& graph, ggml_backend_t
     auto k              = ggml_dup_tensor(graph.ctx, q);
     auto v              = ggml_dup_tensor(graph.ctx, q);
     graph.inputs        = {q, k, v};
+    GGMLRunnerContext runner_ctx;
+    runner_ctx.ggml_ctx = graph.ctx;
+    runner_ctx.backend  = backend;
     bool used_flash     = false;
-    auto output         = ggml_ext_attention_ext(graph.ctx, backend, q, k, v, heads, nullptr, false, flash, 1.f, false, &used_flash);
+    auto output         = ggml_ext_attention_ext(&runner_ctx, q, k, v, heads, nullptr, false, flash, 1.f, &used_flash);
     graph.output        = flash && !used_flash ? nullptr : output;
 }
 

src/model/adapter/lora.hpp

@@ -218,7 +218,8 @@ struct LoraModel : public GGMLRunner {
         lora_tensors = std::move(new_lora_tensors);
     }
 
-    ggml_tensor* get_lora_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
+    ggml_tensor* get_lora_weight_diff(const std::string& model_tensor_name, GGMLRunnerContext* runner_ctx) {
+        auto ctx            = runner_ctx->ggml_ctx;
         ggml_tensor* updown = nullptr;
         int index           = 0;
         while (true) {
@@ -241,21 +242,21 @@ struct LoraModel : public GGMLRunner {
 
             auto iter = lora_tensors.find(lora_up_name);
             if (iter != lora_tensors.end()) {
-                lora_up = ggml_ext_cast_f32(ctx, backend, iter->second);
+                lora_up = ggml_ext_cast_f32(runner_ctx, iter->second);
                 if (swapped_lora_up_tensors.count(lora_up_name)) {
-                    auto halves = ggml_ext_chunk(ctx, lora_up, 2, 1);
+                    auto halves = ggml_ext_chunk(runner_ctx, lora_up, 2, 1);
                     lora_up     = ggml_concat(ctx, halves[1], halves[0], 1);
                 }
             }
 
             iter = lora_tensors.find(lora_mid_name);
             if (iter != lora_tensors.end()) {
-                lora_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
+                lora_mid = ggml_ext_cast_f32(runner_ctx, iter->second);
             }
 
             iter = lora_tensors.find(lora_down_name);
             if (iter != lora_tensors.end()) {
-                lora_down = ggml_ext_cast_f32(ctx, backend, iter->second);
+                lora_down = ggml_ext_cast_f32(runner_ctx, iter->second);
             }
 
             if (lora_up == nullptr || lora_down == nullptr) {
@@ -288,8 +289,8 @@ struct LoraModel : public GGMLRunner {
             }
             scale_value *= multiplier;
 
-            auto curr_updown = ggml_ext_merge_lora(ctx, lora_down, lora_up, lora_mid);
-            curr_updown      = ggml_ext_scale(ctx, curr_updown, scale_value, true);
+            auto curr_updown = ggml_ext_merge_lora(runner_ctx, lora_down, lora_up, lora_mid);
+            curr_updown      = ggml_ext_scale(runner_ctx, curr_updown, scale_value, true);
 
             if (updown == nullptr) {
                 updown = curr_updown;
@@ -302,7 +303,8 @@ struct LoraModel : public GGMLRunner {
         return updown;
     }
 
-    ggml_tensor* get_raw_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
+    ggml_tensor* get_raw_weight_diff(const std::string& model_tensor_name, GGMLRunnerContext* runner_ctx) {
+        auto ctx            = runner_ctx->ggml_ctx;
         ggml_tensor* updown = nullptr;
         int index           = 0;
         while (true) {
@@ -319,7 +321,7 @@ struct LoraModel : public GGMLRunner {
 
             auto iter = lora_tensors.find(diff_name);
             if (iter != lora_tensors.end()) {
-                curr_updown = ggml_ext_cast_f32(ctx, backend, iter->second);
+                curr_updown = ggml_ext_cast_f32(runner_ctx, iter->second);
             } else {
                 break;
             }
@@ -329,7 +331,7 @@ struct LoraModel : public GGMLRunner {
             float scale_value = 1.0f;
             scale_value *= multiplier;
 
-            curr_updown = ggml_ext_scale(ctx, curr_updown, scale_value, true);
+            curr_updown = ggml_ext_scale(runner_ctx, curr_updown, scale_value, true);
 
             if (updown == nullptr) {
                 updown = curr_updown;
@@ -342,7 +344,8 @@ struct LoraModel : public GGMLRunner {
         return updown;
     }
 
-    ggml_tensor* get_loha_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
+    ggml_tensor* get_loha_weight_diff(const std::string& model_tensor_name, GGMLRunnerContext* runner_ctx) {
+        auto ctx            = runner_ctx->ggml_ctx;
         ggml_tensor* updown = nullptr;
         int index           = 0;
         while (true) {
@@ -370,17 +373,17 @@ struct LoraModel : public GGMLRunner {
 
             auto iter = lora_tensors.find(hada_1_down_name);
             if (iter != lora_tensors.end()) {
-                hada_1_down = ggml_ext_cast_f32(ctx, backend, iter->second);
+                hada_1_down = ggml_ext_cast_f32(runner_ctx, iter->second);
             }
 
             iter = lora_tensors.find(hada_1_up_name);
             if (iter != lora_tensors.end()) {
-                hada_1_up = ggml_ext_cast_f32(ctx, backend, iter->second);
+                hada_1_up = ggml_ext_cast_f32(runner_ctx, iter->second);
             }
 
             iter = lora_tensors.find(hada_1_mid_name);
             if (iter != lora_tensors.end()) {
-                hada_1_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
+                hada_1_mid = ggml_ext_cast_f32(runner_ctx, iter->second);
                 if (hada_1_up != nullptr) {
                     hada_1_up = ggml_cont(ctx, ggml_transpose(ctx, hada_1_up));
                 }
@@ -388,17 +391,17 @@ struct LoraModel : public GGMLRunner {
 
             iter = lora_tensors.find(hada_2_down_name);
             if (iter != lora_tensors.end()) {
-                hada_2_down = ggml_ext_cast_f32(ctx, backend, iter->second);
+                hada_2_down = ggml_ext_cast_f32(runner_ctx, iter->second);
             }
 
             iter = lora_tensors.find(hada_2_up_name);
             if (iter != lora_tensors.end()) {
-                hada_2_up = ggml_ext_cast_f32(ctx, backend, iter->second);
+                hada_2_up = ggml_ext_cast_f32(runner_ctx, iter->second);
             }
 
             iter = lora_tensors.find(hada_2_mid_name);
             if (iter != lora_tensors.end()) {
-                hada_2_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
+                hada_2_mid = ggml_ext_cast_f32(runner_ctx, iter->second);
                 if (hada_2_up != nullptr) {
                     hada_2_up = ggml_cont(ctx, ggml_transpose(ctx, hada_2_up));
                 }
@@ -435,10 +438,10 @@ struct LoraModel : public GGMLRunner {
             }
             scale_value *= multiplier;
 
-            ggml_tensor* updown_1 = ggml_ext_merge_lora(ctx, hada_1_down, hada_1_up, hada_1_mid);
-            ggml_tensor* updown_2 = ggml_ext_merge_lora(ctx, hada_2_down, hada_2_up, hada_2_mid);
+            ggml_tensor* updown_1 = ggml_ext_merge_lora(runner_ctx, hada_1_down, hada_1_up, hada_1_mid);
+            ggml_tensor* updown_2 = ggml_ext_merge_lora(runner_ctx, hada_2_down, hada_2_up, hada_2_mid);
             auto curr_updown      = ggml_mul_inplace(ctx, updown_1, updown_2);
-            curr_updown           = ggml_ext_scale(ctx, curr_updown, scale_value, true);
+            curr_updown           = ggml_ext_scale(runner_ctx, curr_updown, scale_value, true);
             if (updown == nullptr) {
                 updown = curr_updown;
             } else {
@@ -449,7 +452,8 @@ struct LoraModel : public GGMLRunner {
         return updown;
     }
 
-    ggml_tensor* get_lokr_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
+    ggml_tensor* get_lokr_weight_diff(const std::string& model_tensor_name, GGMLRunnerContext* runner_ctx) {
+        auto ctx            = runner_ctx->ggml_ctx;
         ggml_tensor* updown = nullptr;
         int index           = 0;
         while (true) {
@@ -476,24 +480,24 @@ struct LoraModel : public GGMLRunner {
 
             auto iter = lora_tensors.find(lokr_w1_name);
             if (iter != lora_tensors.end()) {
-                lokr_w1 = ggml_ext_cast_f32(ctx, backend, iter->second);
+                lokr_w1 = ggml_ext_cast_f32(runner_ctx, iter->second);
             }
 
             iter = lora_tensors.find(lokr_w2_name);
             if (iter != lora_tensors.end()) {
-                lokr_w2 = ggml_ext_cast_f32(ctx, backend, iter->second);
+                lokr_w2 = ggml_ext_cast_f32(runner_ctx, iter->second);
             }
 
             int64_t rank = 1;
             if (lokr_w1 == nullptr) {
                 iter = lora_tensors.find(lokr_w1_a_name);
                 if (iter != lora_tensors.end()) {
-                    lokr_w1_a = ggml_ext_cast_f32(ctx, backend, iter->second);
+                    lokr_w1_a = ggml_ext_cast_f32(runner_ctx, iter->second);
                 }
 
                 iter = lora_tensors.find(lokr_w1_b_name);
                 if (iter != lora_tensors.end()) {
-                    lokr_w1_b = ggml_ext_cast_f32(ctx, backend, iter->second);
+                    lokr_w1_b = ggml_ext_cast_f32(runner_ctx, iter->second);
                 }
 
                 if (lokr_w1_a == nullptr || lokr_w1_b == nullptr) {
@@ -502,18 +506,18 @@ struct LoraModel : public GGMLRunner {
 
                 rank = lokr_w1_b->ne[ggml_n_dims(lokr_w1_b) - 1];
 
-                lokr_w1 = ggml_ext_merge_lora(ctx, lokr_w1_b, lokr_w1_a);
+                lokr_w1 = ggml_ext_merge_lora(runner_ctx, lokr_w1_b, lokr_w1_a);
             }
 
             if (lokr_w2 == nullptr) {
                 iter = lora_tensors.find(lokr_w2_a_name);
                 if (iter != lora_tensors.end()) {
-                    lokr_w2_a = ggml_ext_cast_f32(ctx, backend, iter->second);
+                    lokr_w2_a = ggml_ext_cast_f32(runner_ctx, iter->second);
                 }
 
                 iter = lora_tensors.find(lokr_w2_b_name);
                 if (iter != lora_tensors.end()) {
-                    lokr_w2_b = ggml_ext_cast_f32(ctx, backend, iter->second);
+                    lokr_w2_b = ggml_ext_cast_f32(runner_ctx, iter->second);
                 }
 
                 if (lokr_w2_a == nullptr || lokr_w2_b == nullptr) {
@@ -522,7 +526,7 @@ struct LoraModel : public GGMLRunner {
 
                 rank = lokr_w2_b->ne[ggml_n_dims(lokr_w2_b) - 1];
 
-                lokr_w2 = ggml_ext_merge_lora(ctx, lokr_w2_b, lokr_w2_a);
+                lokr_w2 = ggml_ext_merge_lora(runner_ctx, lokr_w2_b, lokr_w2_a);
             }
 
             if (!lokr_w1_a) {
@@ -553,8 +557,8 @@ struct LoraModel : public GGMLRunner {
 
             scale_value *= multiplier;
 
-            auto curr_updown = ggml_ext_kronecker(ctx, lokr_w1, lokr_w2);
-            curr_updown      = ggml_ext_scale(ctx, curr_updown, scale_value, true);
+            auto curr_updown = ggml_ext_kronecker(runner_ctx, lokr_w1, lokr_w2);
+            curr_updown      = ggml_ext_scale(runner_ctx, curr_updown, scale_value, true);
 
             if (updown == nullptr) {
                 updown = curr_updown;
@@ -566,29 +570,33 @@ struct LoraModel : public GGMLRunner {
         return updown;
     }
 
-    ggml_tensor* get_weight_diff(const std::string& model_tensor_name, ggml_backend_t backend, ggml_context* ctx, ggml_tensor* model_tensor, bool with_lora_and_lokr = true) {
+    ggml_tensor* get_weight_diff(const std::string& model_tensor_name,
+                                 GGMLRunnerContext* runner_ctx,
+                                 ggml_tensor* model_tensor,
+                                 bool with_lora_and_lokr = true) {
+        auto ctx = runner_ctx->ggml_ctx;
         // lora
         ggml_tensor* diff = nullptr;
         if (with_lora_and_lokr) {
-            diff = get_lora_weight_diff(model_tensor_name, ctx, backend);
+            diff = get_lora_weight_diff(model_tensor_name, runner_ctx);
         }
         // diff
         if (diff == nullptr) {
-            diff = get_raw_weight_diff(model_tensor_name, ctx, backend);
+            diff = get_raw_weight_diff(model_tensor_name, runner_ctx);
         }
         // loha
         if (diff == nullptr) {
-            diff = get_loha_weight_diff(model_tensor_name, ctx, backend);
+            diff = get_loha_weight_diff(model_tensor_name, runner_ctx);
         }
         // lokr
         if (diff == nullptr && with_lora_and_lokr) {
-            diff = get_lokr_weight_diff(model_tensor_name, ctx, backend);
+            diff = get_lokr_weight_diff(model_tensor_name, runner_ctx);
         }
         if (diff != nullptr) {
             if (ggml_nelements(diff) < ggml_nelements(model_tensor)) {
                 if (ggml_n_dims(diff) == 2 && ggml_n_dims(model_tensor) == 2 && diff->ne[0] == model_tensor->ne[0]) {
                     LOG_WARN("pad for %s", model_tensor_name.c_str());
-                    auto pad_tensor = ggml_ext_zeros(ctx, diff->ne[0], model_tensor->ne[1] - diff->ne[1], 1, 1);
+                    auto pad_tensor = ggml_ext_zeros(runner_ctx, diff->ne[0], model_tensor->ne[1] - diff->ne[1], 1, 1);
                     diff            = ggml_concat(ctx, diff, pad_tensor, 1);
                 }
             }
@@ -619,12 +627,12 @@ struct LoraModel : public GGMLRunner {
         return diff;
     }
 
-    ggml_tensor* get_out_diff(ggml_context* ctx,
-                              ggml_backend_t backend,
+    ggml_tensor* get_out_diff(GGMLRunnerContext* runner_ctx,
                               ggml_tensor* x,
                               ggml_tensor* model_weight,
                               WeightAdapter::ForwardParams forward_params,
                               const std::string& model_tensor_name) {
+        auto ctx              = runner_ctx->ggml_ctx;
         ggml_tensor* out_diff = nullptr;
         int index             = 0;
 
@@ -712,7 +720,7 @@ struct LoraModel : public GGMLRunner {
                 }
                 scale_value *= multiplier;
 
-                auto curr_out_diff = ggml_ext_lokr_forward(ctx, backend, x, lokr_w1, lokr_w1_a, lokr_w1_b, lokr_w2, lokr_w2_a, lokr_w2_b, is_conv2d, forward_params.conv2d, scale_value);
+                auto curr_out_diff = ggml_ext_lokr_forward(runner_ctx, x, lokr_w1, lokr_w1_a, lokr_w1_b, lokr_w2, lokr_w2_a, lokr_w2_b, is_conv2d, forward_params.conv2d, scale_value);
                 if (out_diff == nullptr) {
                     out_diff = curr_out_diff;
                 } else {
@@ -761,7 +769,7 @@ struct LoraModel : public GGMLRunner {
             if (iter != lora_tensors.end()) {
                 lora_up = iter->second;
                 if (swapped_lora_up_tensors.count(lora_up_name)) {
-                    auto halves = ggml_ext_chunk(ctx, lora_up, 2, 1);
+                    auto halves = ggml_ext_chunk(runner_ctx, lora_up, 2, 1);
                     lora_up     = ggml_concat(ctx, halves[1], halves[0], 1);
                 }
                 if (is_conv2d && lora_up->type != GGML_TYPE_F16) {
@@ -850,13 +858,13 @@ struct LoraModel : public GGMLRunner {
 
             ggml_tensor* lx;
             if (!is_conv2d) {
-                lx = ggml_ext_linear(ctx, x, lora_down, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale);
+                lx = ggml_ext_linear(runner_ctx, x, lora_down, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale);
                 if (lora_mid) {
-                    lx = ggml_ext_linear(ctx, lx, lora_mid, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale);
+                    lx = ggml_ext_linear(runner_ctx, lx, lora_mid, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale);
                 }
-                lx = ggml_ext_linear(ctx, lx, lora_up, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale);
+                lx = ggml_ext_linear(runner_ctx, lx, lora_up, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale);
             } else {  // OP_CONV2D
-                lx = ggml_ext_conv_2d(ctx,
+                lx = ggml_ext_conv_2d(runner_ctx,
                                       x,
                                       lora_down,
                                       nullptr,
@@ -871,7 +879,7 @@ struct LoraModel : public GGMLRunner {
                                       forward_params.conv2d.circular_y,
                                       forward_params.conv2d.scale);
                 if (lora_mid) {
-                    lx = ggml_ext_conv_2d(ctx,
+                    lx = ggml_ext_conv_2d(runner_ctx,
                                           lx,
                                           lora_mid,
                                           nullptr,
@@ -886,7 +894,7 @@ struct LoraModel : public GGMLRunner {
                                           forward_params.conv2d.circular_y,
                                           forward_params.conv2d.scale);
                 }
-                lx = ggml_ext_conv_2d(ctx,
+                lx = ggml_ext_conv_2d(runner_ctx,
                                       lx,
                                       lora_up,
                                       nullptr,
@@ -902,7 +910,7 @@ struct LoraModel : public GGMLRunner {
                                       forward_params.conv2d.scale);
             }
 
-            auto curr_out_diff = ggml_ext_scale(ctx, lx, scale_value, true);
+            auto curr_out_diff = ggml_ext_scale(runner_ctx, lx, scale_value, true);
             if (out_diff == nullptr) {
                 out_diff = curr_out_diff;
             } else {
@@ -951,6 +959,7 @@ struct LoraModel : public GGMLRunner {
         size_t lora_graph_size = LORA_GRAPH_BASE_SIZE + lora_tensors.size() * 10;
         ggml_cgraph* gf        = ggml_new_graph_custom(compute_ctx, lora_graph_size, false);
 
+        auto runner_ctx = get_context();
         preprocess_lora_tensors(model_tensor_names);
 
         original_tensor_to_final_tensor.clear();
@@ -961,7 +970,7 @@ struct LoraModel : public GGMLRunner {
             ggml_tensor* model_tensor     = it.second;
 
             // lora
-            ggml_tensor* diff = get_weight_diff(model_tensor_name, runtime_backend, compute_ctx, model_tensor);
+            ggml_tensor* diff = get_weight_diff(model_tensor_name, &runner_ctx, model_tensor);
             if (diff == nullptr) {
                 continue;
             }
@@ -974,7 +983,7 @@ struct LoraModel : public GGMLRunner {
 
             ggml_tensor* final_tensor;
             if (model_tensor->type != GGML_TYPE_F32 && model_tensor->type != GGML_TYPE_F16) {
-                final_tensor = ggml_ext_cast_f32(compute_ctx, runtime_backend, model_tensor);
+                final_tensor = ggml_ext_cast_f32(&runner_ctx, model_tensor);
                 final_tensor = ggml_add_inplace(compute_ctx, final_tensor, diff);
                 final_tensor = ggml_cpy(compute_ctx, final_tensor, model_tensor);
             } else {
@@ -1061,41 +1070,47 @@ public:
         : lora_models(lora_models) {
     }
 
-    ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name, bool with_lora_and_lokr) {
+    ggml_tensor* patch_weight(GGMLRunnerContext* runner_ctx,
+                              ggml_tensor* weight,
+                              const std::string& weight_name,
+                              bool with_lora_and_lokr) {
+        auto ctx = runner_ctx->ggml_ctx;
         for (auto& lora_model : lora_models) {
-            ggml_tensor* diff = lora_model->get_weight_diff(weight_name, backend, ctx, weight, with_lora_and_lokr);
+            ggml_tensor* diff = lora_model->get_weight_diff(weight_name, runner_ctx, weight, with_lora_and_lokr);
             if (diff == nullptr) {
                 continue;
             }
 
             if (weight->type != GGML_TYPE_F32 && weight->type != GGML_TYPE_F16) {
-                weight = ggml_ext_cast_f32(ctx, backend, weight);
+                weight = ggml_ext_cast_f32(runner_ctx, weight);
             }
             weight = ggml_add(ctx, weight, diff);
         }
         return weight;
     }
 
-    ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name) override {
-        return patch_weight(ctx, backend, weight, weight_name, true);
+    ggml_tensor* patch_weight(GGMLRunnerContext* runner_ctx,
+                              ggml_tensor* weight,
+                              const std::string& weight_name) override {
+        return patch_weight(runner_ctx, weight, weight_name, true);
     }
 
-    ggml_tensor* forward_with_lora(ggml_context* ctx,
-                                   ggml_backend_t backend,
+    ggml_tensor* forward_with_lora(GGMLRunnerContext* runner_ctx,
                                    ggml_tensor* x,
                                    ggml_tensor* w,
                                    ggml_tensor* b,
                                    const std::string& prefix,
                                    WeightAdapter::ForwardParams forward_params) override {
-        w = patch_weight(ctx, backend, w, prefix + "weight", false);
+        auto ctx = runner_ctx->ggml_ctx;
+        w        = patch_weight(runner_ctx, w, prefix + "weight", false);
         if (b) {
-            b = patch_weight(ctx, backend, b, prefix + "bias", false);
+            b = patch_weight(runner_ctx, b, prefix + "bias", false);
         }
         ggml_tensor* out;
         if (forward_params.op_type == ForwardParams::op_type_t::OP_LINEAR) {
-            out = ggml_ext_linear(ctx, x, w, b, forward_params.linear.force_prec_f32, forward_params.linear.scale);
+            out = ggml_ext_linear(runner_ctx, x, w, b, forward_params.linear.force_prec_f32, forward_params.linear.scale);
         } else {  // OP_CONV2D
-            out = ggml_ext_conv_2d(ctx,
+            out = ggml_ext_conv_2d(runner_ctx,
                                    x,
                                    w,
                                    b,
@@ -1111,7 +1126,7 @@ public:
                                    forward_params.conv2d.scale);
         }
         for (auto& lora_model : lora_models) {
-            ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, w, forward_params, prefix + "weight");
+            ggml_tensor* out_diff = lora_model->get_out_diff(runner_ctx, x, w, forward_params, prefix + "weight");
             if (out_diff == nullptr) {
                 continue;
             }
@@ -1120,13 +1135,13 @@ public:
         return out;
     }
 
-    ggml_tensor* add_lora_to_output(ggml_context* ctx,
-                                    ggml_backend_t backend,
+    ggml_tensor* add_lora_to_output(GGMLRunnerContext* runner_ctx,
                                     ggml_tensor* x,
                                     ggml_tensor* w,
                                     ggml_tensor* output,
                                     const std::string& prefix,
                                     WeightAdapter::ForwardParams forward_params) override {
+        auto ctx = runner_ctx->ggml_ctx;
         // Packed INT4 stores half the logical input width. Adapter shape checks
         // and reshapes need the logical weight dimensions, without decoding it.
         ggml_tensor logical_weight = *w;
@@ -1138,10 +1153,10 @@ public:
             w = &logical_weight;
         }
         for (auto& lora_model : lora_models) {
-            ggml_tensor* weight_diff = lora_model->get_weight_diff(prefix + "weight", backend, ctx, w, false);
+            ggml_tensor* weight_diff = lora_model->get_weight_diff(prefix + "weight", runner_ctx, w, false);
             if (weight_diff != nullptr) {
                 GGML_ASSERT(forward_params.op_type == ForwardParams::op_type_t::OP_LINEAR);
-                ggml_tensor* out_diff = ggml_ext_linear(ctx,
+                ggml_tensor* out_diff = ggml_ext_linear(runner_ctx,
                                                         x,
                                                         weight_diff,
                                                         nullptr,
@@ -1150,7 +1165,7 @@ public:
                 output                = ggml_add_inplace(ctx, output, out_diff);
             }
 
-            ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, w, forward_params, prefix + "weight");
+            ggml_tensor* out_diff = lora_model->get_out_diff(runner_ctx, x, w, forward_params, prefix + "weight");
             if (out_diff != nullptr) {
                 output = ggml_add_inplace(ctx, output, out_diff);
             }

src/model/adapter/lora_ops.cpp

@@ -4,11 +4,46 @@
 
 #include "core/ggml_extend.h"
 #include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
 
-ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx,
+ggml_tensor* ggml_ext_mul_n_mode(GGMLRunnerContext* runner_ctx, ggml_tensor* a, ggml_tensor* b, int mode) {
+    auto ctx = runner_ctx->ggml_ctx;
+    // 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(GGMLRunnerContext* runner_ctx, ggml_tensor* a, ggml_tensor* b) {
+    auto ctx = runner_ctx->ggml_ctx;
+    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_merge_lora(GGMLRunnerContext* runner_ctx,
                                  ggml_tensor* lora_down,
                                  ggml_tensor* lora_up,
                                  ggml_tensor* lora_mid) {
+    auto ctx = runner_ctx->ggml_ctx;
     ggml_tensor* updown;
     // flat lora tensors to multiply it
     // ggml_n_dims drops the output dimension when a linear layer has one output.
@@ -32,25 +67,25 @@ ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx,
         // 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_ext_mul_n_mode(runner_ctx, ggml_ext_mul_n_mode(runner_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_tensor* ggml_ext_lokr_forward(GGMLRunnerContext* runner_ctx,
+                                   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) {
+    auto ctx     = runner_ctx->ggml_ctx;
+    auto backend = runner_ctx->backend;
     GGML_ASSERT((w1 != nullptr || (w1a != nullptr && w1b != nullptr)));
     GGML_ASSERT((w2 != nullptr || (w2a != nullptr && w2b != nullptr)));
 
@@ -73,7 +108,6 @@ ggml_tensor* ggml_ext_lokr_forward(
 
         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;
@@ -121,14 +155,14 @@ ggml_tensor* ggml_ext_lokr_forward(
 
         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);
+        return ggml_ext_scale(runner_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,
+            hb = ggml_ext_conv_2d(runner_ctx, h_split, w2, nullptr,
                                   conv_params.s0,
                                   conv_params.s1,
                                   conv_params.p0,
@@ -154,7 +188,7 @@ ggml_tensor* ggml_ext_lokr_forward(
                 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,
+            ggml_tensor* ha = ggml_ext_conv_2d(runner_ctx, h_split, a, nullptr,
                                                conv_params.s0,
                                                conv_params.s1,
                                                conv_params.p0,
@@ -167,8 +201,7 @@ ggml_tensor* ggml_ext_lokr_forward(
                                                conv_params.scale);
 
             // not supporting lora_mid here
-            hb = ggml_ext_conv_2d(ctx,
-                                  ha,
+            hb = ggml_ext_conv_2d(runner_ctx, ha,
                                   b,
                                   nullptr,
                                   1,
@@ -204,6 +237,6 @@ ggml_tensor* ggml_ext_lokr_forward(
         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);
+        return ggml_ext_scale(runner_ctx, out, scale);
     }
 }

src/model/adapter/lora_ops.h

@@ -3,23 +3,32 @@
 
 #include "core/ggml_runner.h"
 
-ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx,
+// 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(GGMLRunnerContext* runner_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(GGMLRunnerContext* runner_ctx, ggml_tensor* a, ggml_tensor* b);
+
+ggml_tensor* ggml_ext_merge_lora(GGMLRunnerContext* runner_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);
+ggml_tensor* ggml_ext_lokr_forward(GGMLRunnerContext* runner_ctx,
+                                   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

@@ -36,7 +36,7 @@ public:
         x = layer_norm->forward(ctx, x);
         // x = ggml_add(ctx, ggml_mul_mat(ctx, fc1_w, x),  fc1_b);
         x = fc1->forward(ctx, x);
-        x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+        x = ggml_ext_gelu(ctx, x, true);
         x = fc2->forward(ctx, x);
         // x = ggml_add(ctx, ggml_mul_mat(ctx, fc2_w, x),  fc2_b);
         if (use_residue)
@@ -132,8 +132,8 @@ public:
         k             = reshape_tensor(ctx->ggml_ctx, k, heads);
         v             = reshape_tensor(ctx->ggml_ctx, v, heads);
         scale         = 1.f / sqrt(sqrt((float)dim_head));
-        k             = ggml_ext_scale(ctx->ggml_ctx, k, scale, true);
-        q             = ggml_ext_scale(ctx->ggml_ctx, q, scale, true);
+        k             = ggml_ext_scale(ctx, k, scale, true);
+        q             = ggml_ext_scale(ctx, q, scale, true);
         // auto weight = ggml_mul_mat(ctx, q, k);
         auto weight = ggml_mul_mat(ctx->ggml_ctx, k, q);  // NOTE order of mul is opposite to pytorch
 

src/model/audio/wav2vec2.hpp

@@ -67,7 +67,7 @@ namespace Wav2Vec2 {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             auto conv = std::dynamic_pointer_cast<Conv1d>(blocks["conv"]);
             x         = conv->forward(ctx, x);
-            return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x));
+            return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx, x));
         }
     };
 
@@ -85,7 +85,7 @@ namespace Wav2Vec2 {
             x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
             x = layer_norm->forward(ctx, x);
             x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
-            return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x));
+            return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx, x));
         }
     };
 
@@ -103,7 +103,7 @@ namespace Wav2Vec2 {
             x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], 1, x->ne[1], x->ne[2]);
             x = layer_norm->forward(ctx, x);
             x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[2], x->ne[3]);
-            return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x));
+            return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx, x));
         }
     };
 
@@ -198,7 +198,7 @@ namespace Wav2Vec2 {
             auto w = weight(ctx);
             auto b = params.count("conv.bias") > 0 ? params["conv.bias"] : nullptr;
             x      = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
-            x      = ggml_ext_conv_1d(ctx->ggml_ctx, x, w, b, 1, kernel_size_ / 2, 1, groups_, true);
+            x      = ggml_ext_conv_1d(ctx, x, w, b, 1, kernel_size_ / 2, 1, groups_, true);
             // Apply GELU out of place before cropping to keep graph buffer reuse safe.
             x = ggml_gelu_erf(ctx->ggml_ctx, x);
             x = ggml_view_3d(ctx->ggml_ctx, x, x->ne[0] - 1, x->ne[1], x->ne[2], x->nb[1], x->nb[2], 0);
@@ -216,7 +216,7 @@ namespace Wav2Vec2 {
             auto intermediate_dense = std::dynamic_pointer_cast<Linear>(blocks["intermediate_dense"]);
             auto output_dense       = std::dynamic_pointer_cast<Linear>(blocks["output_dense"]);
             x                       = intermediate_dense->forward(ctx, x);
-            x                       = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+            x                       = ggml_ext_gelu(ctx, x, true);
             x                       = output_dense->forward(ctx, x);
             return x;
         }

src/model/common/block.hpp

@@ -32,7 +32,7 @@ public:
         if (vae_downsample) {
             auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
 
-            x = ggml_ext_pad(ctx->ggml_ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
+            x = ggml_ext_pad(ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
             x = conv->forward(ctx, x);
         } else {
             auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["op"]);
@@ -220,7 +220,7 @@ public:
         auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
 
         x = proj->forward(ctx, x);
-        x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+        x = ggml_ext_gelu(ctx, x, true);
         return x;
     }
 };
@@ -250,7 +250,7 @@ public:
         auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
 
         x = fc1->forward(ctx, x);
-        x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+        x = ggml_ext_gelu(ctx, x, true);
         x = fc2->forward(ctx, x);
         return x;
     }
@@ -293,9 +293,7 @@ public:
 
         auto net_0 = std::dynamic_pointer_cast<UnaryBlock>(blocks["net.0"]);
         auto net_2 = std::dynamic_pointer_cast<Linear>(blocks["net.2"]);
-        if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-            net_2->set_force_prec_f32(true);
-        }
+        net_2->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
 
         x = net_0->forward(ctx, x);  // [ne3, ne2, ne1, inner_dim]
         x = net_2->forward(ctx, x);  // [ne3, ne2, ne1, dim_out]
@@ -597,7 +595,7 @@ public:
         x = ggml_add(ctx->ggml_ctx, x, attn->forward(ctx, ln_1->forward(ctx, x)));
 
         auto h = mlp_c_fc->forward(ctx, ln_2->forward(ctx, x));
-        h      = ggml_ext_gelu_quick(ctx->ggml_ctx, h);
+        h      = ggml_ext_gelu_quick(ctx, h);
         h      = mlp_c_proj->forward(ctx, h);
         return ggml_add(ctx->ggml_ctx, x, h);
     }
@@ -614,7 +612,7 @@ protected:
     ggml_tensor* get_alpha(GGMLRunnerContext* ctx) {
         // image_only_indicator is always tensor([0.]) and since mix_factor.shape is [1,]
         // so learned_with_images is same as learned
-        auto mix_factor = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["mix_factor"]);
+        auto mix_factor = ggml_ext_cast_f32(ctx, params["mix_factor"]);
         return ggml_sigmoid(ctx->ggml_ctx, mix_factor);
     }
 

src/model/common/ggml_block.hpp

@@ -147,6 +147,7 @@ protected:
     bool bias;
     bool force_f32;
     bool force_prec_f32;
+    uint32_t precision_flags    = 0;
     bool has_weight_scale       = false;
     int w4_convrot_kind         = 0;
     bool int8_convrot           = false;
@@ -215,6 +216,10 @@ public:
         scale = scale_;
     }
 
+    void set_precision_flags(uint32_t flags) {
+        precision_flags = flags;
+    }
+
     void set_force_prec_f32(bool force_prec_f32_) {
         force_prec_f32 = force_prec_f32_;
     }
@@ -225,72 +230,34 @@ public:
     }
 
     ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        const bool force_prec_f32 = ggml_ext_force_prec_f32(ctx, this->force_prec_f32, precision_flags);
         ggml_tensor* w            = params["weight"];
         const float scale         = ctx->linear_scale > 0.f ? ctx->linear_scale : this->scale;
         ggml_tensor* weight_scale = has_weight_scale ? params["weight_scale"] : nullptr;
-#ifndef SD_USE_UPSTREAM_GGML
-        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_prec_set_acc(fp8_matmul, GGML_PREC_F32);
-                }
-                supports_fp8_matmul = ggml_backend_supports_op(ctx->backend, fp8_matmul) && !sd_backend_is_cpu(ctx->backend);
-            }
-            if (!supports_fp8_matmul) {
-                w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_BF16);
-            }
-        }
-#endif
-        ggml_tensor* b = nullptr;
+        w                         = ggml_ext_prepare_linear_weight(ctx, w, x, force_prec_f32);
+        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);
-            }
+            x                       = ggml_ext_prepare_i8_input(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");
+                b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias");
             }
-#ifndef SD_USE_UPSTREAM_GGML
-            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;
-                }
-            }
-#endif
+
 #ifndef SD_USE_UPSTREAM_GGML
             if (w4_convrot_kind) {
-                if (scale != 1.f)
-                    x = ggml_ext_scale(ctx->ggml_ctx, x, scale);
-                out = ggml_mul_mat_w4_convrot(ctx->ggml_ctx, w, x,
-                                              w4_convrot_kind == 2 ? params["weight_codebook"] : weight_scale,
-                                              w4_convrot_kind == 2 ? params["weight_s_channel"] : nullptr,
-                                              w4_convrot_kind == 2 ? params["weight_s_rel"] : nullptr,
-                                              scale == 1.f ? b : nullptr, static_cast<SDW4Kind>(w4_convrot_kind),
-                                              int8_convrot ? int8_convrot_group_size : 0);
-                GGML_ASSERT(out);
-                if (scale != 1.f) {
-                    out = ggml_ext_scale(ctx->ggml_ctx, out, 1.f / scale);
-                    if (b)
-                        out = ggml_add_inplace(ctx->ggml_ctx, out, b);
-                }
+                out = ggml_ext_linear_w4_convrot(ctx, x, w,
+                                                 w4_convrot_kind == 2 ? params["weight_codebook"] : weight_scale,
+                                                 w4_convrot_kind == 2 ? params["weight_s_channel"] : nullptr,
+                                                 w4_convrot_kind == 2 ? params["weight_s_rel"] : nullptr,
+                                                 b, static_cast<SDW4Kind>(w4_convrot_kind), int8_convrot ? int8_convrot_group_size : 0, scale);
             } else
 #endif
-                out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx,
+                out = ggml_ext_linear_i8_tensorwise(ctx,
                                                     x,
                                                     w,
                                                     weight_scale,
@@ -302,33 +269,21 @@ public:
                 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);
+                out                                  = ctx->weight_adapter->add_lora_to_output(ctx, 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_ext_linear(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);
+                out                                  = ctx->weight_adapter->add_lora_to_output(ctx, x, w, out, prefix, forward_params);
                 if (b != nullptr) {
-                    b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+                    b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias");
                 }
             }
             if (b != nullptr) {
@@ -341,9 +296,9 @@ public:
             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);
+            out                                  = ctx->weight_adapter->forward_with_lora(ctx, x, w, linear_bias, prefix, forward_params);
         } else {
-            out = ggml_ext_linear(ctx->ggml_ctx, x, w, linear_bias, force_prec_f32, scale);
+            out = ggml_ext_linear(ctx, x, w, linear_bias, force_prec_f32, scale);
         }
         return out;
     }
@@ -376,42 +331,6 @@ protected:
     bool is_int8_tensorwise     = false;
     int int8_convrot_group_size = 0;
 
-    static void get_rows_i8(ggml_tensor* dst, int ith, int nth, void* userdata) {
-        const auto* embedding     = static_cast<const Embedding*>(userdata);
-        const int group_size      = embedding->int8_convrot_group_size;
-        const auto* weight        = dst->src[0];
-        const auto* input_ids     = static_cast<const int32_t*>(dst->src[1]->data);
-        const auto* scales        = static_cast<const float*>(dst->src[2]->data);
-        const bool scalar_scale   = ggml_nelements(dst->src[2]) == 1;
-        const float normalization = group_size > 0 ? 1.f / std::sqrt(static_cast<float>(group_size)) : 1.f;
-        for (int64_t row = ith; row < dst->ne[1]; row += nth) {
-            const int32_t token = input_ids[row];
-            GGML_ASSERT(token >= 0 && token < weight->ne[1]);
-            const auto* src   = reinterpret_cast<const int8_t*>(static_cast<const char*>(weight->data) + token * weight->nb[1]);
-            auto* out         = reinterpret_cast<float*>(static_cast<char*>(dst->data) + row * dst->nb[1]);
-            const float scale = scales[scalar_scale ? 0 : token] * normalization;
-            for (int64_t i = 0; i < weight->ne[0]; ++i) {
-                out[i] = static_cast<float>(src[i]) * scale;
-            }
-            // The regular Hadamard rotation is symmetric and its own inverse.
-            for (int stride = 1; stride < group_size; stride *= 4) {
-                for (int64_t base = 0; base < weight->ne[0]; base += 4 * stride) {
-                    for (int j = 0; j < stride; ++j) {
-                        float* values      = out + base + j;
-                        const float a      = values[0];
-                        const float b      = values[stride];
-                        const float c      = values[2 * stride];
-                        const float d      = values[3 * stride];
-                        values[0]          = a + b + c - d;
-                        values[stride]     = a + b - c + d;
-                        values[2 * stride] = a - b + c + d;
-                        values[3 * stride] = -a + b + c + d;
-                    }
-                }
-            }
-        }
-    }
-
     void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override {
         auto weight_storage     = tensor_storage_map.find(prefix + "weight");
         is_int8_tensorwise      = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise;
@@ -451,30 +370,9 @@ public:
           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]);
-
-        ggml_tensor* embedding;
-        if (is_int8_tensorwise) {
-            GGML_ASSERT(input_ids->type == GGML_TYPE_I32);
-            ggml_tensor* args[] = {weight, input_ids, params["weight_scale"]};
-            embedding           = ggml_custom_4d(ctx->ggml_ctx, GGML_TYPE_F32, embedding_dim, input_ids->ne[0], 1, 1,
-                                                 args, 3, get_rows_i8, GGML_N_TASKS_MAX, this);
-        } else {
-            input_ids = ggml_reshape_3d(ctx->ggml_ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]);
-            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;
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input_ids) override {
+        return ggml_ext_embedding(ctx, params["weight"], input_ids,
+                                  is_int8_tensorwise ? params["weight_scale"] : nullptr, int8_convrot_group_size);
     }
 };
 
@@ -528,7 +426,7 @@ public:
 
     ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         GGML_ASSERT(x->ne[1] == in_channels);
-        return ggml_ext_conv_1d(ctx->ggml_ctx, x, params["weight"], bias ? params["bias"] : nullptr,
+        return ggml_ext_conv_1d(ctx, x, params["weight"], bias ? params["bias"] : nullptr,
                                 stride, padding, dilation, groups, force_prec_f32);
     }
 };
@@ -598,9 +496,9 @@ public:
             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 ctx->weight_adapter->forward_with_lora(ctx, x, w, b, prefix, forward_params);
         }
-        return ggml_ext_conv_2d(ctx->ggml_ctx,
+        return ggml_ext_conv_2d(ctx,
                                 x,
                                 w,
                                 b,
@@ -687,9 +585,9 @@ public:
                 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 ctx->weight_adapter->forward_with_lora(ctx, x, w, b, prefix, forward_params);
             }
-            return ggml_ext_conv_2d(ctx->ggml_ctx, x, w, b,
+            return ggml_ext_conv_2d(ctx, x, w, b,
                                     stride.second, stride.first,
                                     padding.second, padding.first,
                                     dilation.second, dilation.first,
@@ -699,24 +597,15 @@ public:
                                     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;
+        if (!ctx->weight_adapter || (groups == in_channels && groups == out_channels)) {
+            return ggml_ext_conv_2d_grouped(ctx, x, w, b, groups,
+                                            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);
         }
 
         int64_t ic_g = in_channels / groups;
@@ -743,33 +632,22 @@ public:
                 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);
-            }
+            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, x_i, w_i, b_i, prefix, forward_params);
         }
 
-        ggml_tensor* out = ggml_ext_vec_concat(ctx->ggml_ctx, out_slices, 2);
+        ggml_tensor* out = ggml_ext_vec_concat(ctx, out_slices, 2);
 
         return out;
     }
@@ -823,7 +701,7 @@ public:
         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");
+            w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight");
             if (w->type != GGML_TYPE_F16) {
                 w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_F16);
             }
@@ -831,10 +709,10 @@ public:
         if (bias) {
             b = params["bias"];
             if (ctx->weight_adapter) {
-                b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+                b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias");
             }
         }
-        return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
+        return ggml_ext_conv_3d(ctx, 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),
@@ -879,16 +757,16 @@ public:
         if (elementwise_affine) {
             w = params["weight"];
             if (ctx->weight_adapter) {
-                w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+                w = ctx->weight_adapter->patch_weight(ctx, 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");
+                    b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias");
                 }
             }
         }
-        return ggml_ext_layer_norm(ctx->ggml_ctx, x, w, b, eps);
+        return ggml_ext_layer_norm(ctx, x, w, b, eps);
     }
 };
 
@@ -927,11 +805,11 @@ public:
             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");
+                w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight");
+                b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias");
             }
         }
-        return ggml_ext_group_norm(ctx->ggml_ctx, x, w, b, num_groups, eps);
+        return ggml_ext_group_norm(ctx, x, w, b, num_groups, eps);
     }
 };
 
@@ -971,7 +849,7 @@ public:
         }
         ggml_tensor* w = params["weight"];
         if (ctx->weight_adapter) {
-            w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
+            w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight");
         }
         x = ggml_rms_norm(ctx->ggml_ctx, x, eps);
         // The affine multiply may run on a different backend after normalization falls back.
@@ -1032,7 +910,7 @@ public:
         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);
+            auto qkv_vec = ggml_ext_split_qkv(ctx, qkv);
             q            = qkv_vec[0];
             k            = qkv_vec[1];
             v            = qkv_vec[2];

src/model/common/rope.hpp

@@ -1071,58 +1071,7 @@ namespace Rope {
                                               ggml_tensor* x,
                                               ggml_tensor* pe,
                                               bool rope_interleaved = true) {
-        auto ctx = runner_ctx->ggml_ctx;
-#ifndef SD_USE_UPSTREAM_GGML
-        auto backend = runner_ctx->backend;
-        if (backend != nullptr && x->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32 &&
-            x->ne[0] % 2 == 0 && pe->ne[0] == 2 && pe->ne[1] == 2 &&
-            pe->ne[2] == x->ne[0] / 2 && pe->ne[3] == x->ne[2]) {
-            auto out = sd_ggml_rope_apply(ctx, backend, x, pe, rope_interleaved);
-            if (out && ggml_backend_supports_op(backend, out)) {
-                return out;
-            }
-        }
-#endif
-        // x: [N, L, n_head, d_head]
-        // pe: [L, d_head/2, 2, 2], [[cos, -sin], [sin, cos]]
-        int64_t d_head = x->ne[0];
-        int64_t n_head = x->ne[1];
-        int64_t L      = x->ne[2];
-        int64_t N      = x->ne[3];
-        x              = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N, n_head, L, d_head]
-        if (rope_interleaved) {
-            x = ggml_reshape_4d(ctx, x, 2, d_head / 2, L, n_head * N);  // [N * n_head, L, d_head/2, 2]
-            x = ggml_cont(ctx, ggml_permute(ctx, x, 3, 0, 1, 2));       // [2, N * n_head, L, d_head/2]
-        } else {
-            auto x_0 = ggml_view_3d(ctx, x, d_head / 2, L, n_head * N, x->nb[1], x->nb[2], 0);
-            auto x_1 = ggml_view_3d(ctx, x, d_head / 2, L, n_head * N, x->nb[1], x->nb[2], (d_head / 2) * x->nb[0]);
-
-            auto pe_2 = ggml_cont(ctx, ggml_permute(ctx, pe, 2, 3, 0, 1));
-            auto cos  = ggml_view_2d(ctx, pe_2, d_head / 2, L, pe_2->nb[1], 0);
-            auto sin  = ggml_view_2d(ctx, pe_2, d_head / 2, L, pe_2->nb[1], pe_2->nb[3]);
-
-            auto part_0 = ggml_sub(ctx, ggml_mul(ctx, x_0, cos), ggml_mul(ctx, x_1, sin));
-            auto part_1 = ggml_add(ctx, ggml_mul(ctx, x_0, sin), ggml_mul(ctx, x_1, cos));
-            return ggml_concat(ctx, part_0, part_1, 0);
-        }
-
-        int64_t offset = x->nb[2] * x->ne[2];
-        auto x_0       = ggml_view_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2], x->nb[1], x->nb[2], offset * 0);  // [N * n_head, L, d_head/2]
-        auto x_1       = ggml_view_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2], x->nb[1], x->nb[2], offset * 1);  // [N * n_head, L, d_head/2]
-        x_0            = ggml_reshape_4d(ctx, x_0, 1, x_0->ne[0], x_0->ne[1], x_0->ne[2]);                    // [N * n_head, L, d_head/2, 1]
-        x_1            = ggml_reshape_4d(ctx, x_1, 1, x_1->ne[0], x_1->ne[1], x_1->ne[2]);                    // [N * n_head, L, d_head/2, 1]
-        auto temp_x    = ggml_new_tensor_4d(ctx, x_0->type, 2, x_0->ne[1], x_0->ne[2], x_0->ne[3]);
-        x_0            = ggml_repeat(ctx, x_0, temp_x);  // [N * n_head, L, d_head/2, 2]
-        x_1            = ggml_repeat(ctx, x_1, temp_x);  // [N * n_head, L, d_head/2, 2]
-
-        pe        = ggml_cont(ctx, ggml_permute(ctx, pe, 3, 0, 1, 2));  // [2, L, d_head/2, 2]
-        offset    = pe->nb[2] * pe->ne[2];
-        auto pe_0 = ggml_view_3d(ctx, pe, pe->ne[0], pe->ne[1], pe->ne[2], pe->nb[1], pe->nb[2], offset * 0);  // [L, d_head/2, 2]
-        auto pe_1 = ggml_view_3d(ctx, pe, pe->ne[0], pe->ne[1], pe->ne[2], pe->nb[1], pe->nb[2], offset * 1);  // [L, d_head/2, 2]
-
-        auto x_out = ggml_add_inplace(ctx, ggml_mul(ctx, x_0, pe_0), ggml_mul(ctx, x_1, pe_1));  // [N * n_head, L, d_head/2, 2]
-        x_out      = ggml_reshape_3d(ctx, x_out, d_head, L, n_head * N);                         // [N*n_head, L, d_head]
-        return x_out;
+        return ggml_ext_rope_apply(runner_ctx, x, pe, rope_interleaved);
     }
 
     __STATIC_INLINE__ ggml_tensor* attention(GGMLRunnerContext* ctx,

src/model/common/vsr.hpp

@@ -177,7 +177,7 @@ namespace VSR {
                     params.window_height[i] = windows[i].first;
                     params.window_width[i]  = windows[i].second;
                 }
-                if (auto out = sd_ggml_window_attn(c, ctx->backend, q, k, v, params))
+                if (auto out = ggml_ext_window_attn(ctx, q, k, v, params))
                     return out;
             }
 #endif
@@ -192,7 +192,7 @@ namespace VSR {
                 auto tk   = ggml_get_rows(c, k, ki);
                 auto tv   = ggml_get_rows(c, v, ki);
                 auto mask = inputs.input(ctx, tile.mask);
-                auto out  = ggml_ext_attention_ext(c, ctx->backend, tq, tk, tv, heads, mask, false, ctx->flash_attn_enabled);
+                auto out  = ggml_ext_attention_ext(ctx, tq, tk, tv, heads, mask, false, ctx->flash_attn_enabled);
                 out       = ggml_reshape_3d(c, out, out->ne[0], tile.width, tile.height);
                 row       = row ? ggml_concat(c, row, out, 1) : out;
                 row_width += tile.width;

src/model/diffusion/anima.hpp

@@ -123,7 +123,7 @@ namespace Anima {
                 emb = ggml_add(ctx->ggml_ctx, emb, temb);
             }
 
-            auto emb_chunks = ggml_ext_chunk(ctx->ggml_ctx, emb, 3, 0);
+            auto emb_chunks = ggml_ext_chunk(ctx, emb, 3, 0);
             auto shift      = emb_chunks[0];
             auto scale      = emb_chunks[1];
             auto gate       = emb_chunks[2];
@@ -164,7 +164,7 @@ namespace Anima {
                 emb          = ggml_add(ctx->ggml_ctx, emb, temb_2c);
             }
 
-            auto emb_chunks = ggml_ext_chunk(ctx->ggml_ctx, emb, 2, 0);
+            auto emb_chunks = ggml_ext_chunk(ctx, emb, 2, 0);
             auto shift      = emb_chunks[0];
             auto scale      = emb_chunks[1];
 
@@ -275,7 +275,7 @@ namespace Anima {
             auto layer2 = std::dynamic_pointer_cast<Linear>(blocks["layer2"]);
 
             x = layer1->forward(ctx, x);
-            x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+            x = ggml_ext_gelu(ctx, x, true);
             x = layer2->forward(ctx, x);
             return x;
         }
@@ -293,7 +293,7 @@ namespace Anima {
             auto layer2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
 
             x = layer0->forward(ctx, x);
-            x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+            x = ggml_ext_gelu(ctx, x, true);
             x = layer2->forward(ctx, x);
             return x;
         }
@@ -500,14 +500,14 @@ namespace Anima {
             int64_t W = x->ne[0];
             int64_t H = x->ne[1];
 
-            auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, x->ne[0], x->ne[1], 1, x->ne[3]);
+            auto padding_mask = ggml_ext_zeros(ctx, x->ne[0], x->ne[1], 1, x->ne[3]);
             x                 = ggml_concat(ctx->ggml_ctx, x, padding_mask, 2);  // [N, C + 1, H, W]
 
             x               = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size);  // [N, h*w, (C+1)*ph*pw]
             int64_t img_len = x->ne[1];
             if (ref_latents.size() > 0) {
                 for (ggml_tensor* ref : ref_latents) {
-                    auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, ref->ne[0], ref->ne[1], 1, ref->ne[3]);
+                    auto padding_mask = ggml_ext_zeros(ctx, ref->ne[0], ref->ne[1], 1, ref->ne[3]);
                     ref               = ggml_concat(ctx->ggml_ctx, ref, padding_mask, 2);  // [N, C + 1, H, W]
                     ref               = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size);
                     x                 = ggml_concat(ctx->ggml_ctx, x, ref, 1);
@@ -515,7 +515,7 @@ namespace Anima {
             }
             x = x_embedder->forward(ctx, x);
 
-            auto timestep_proj     = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.hidden_size));
+            auto timestep_proj     = ggml_ext_timestep_embedding(ctx, timestep, static_cast<int>(config.hidden_size));
             auto temb              = t_embedder->forward(ctx, timestep_proj);
             auto embedded_timestep = t_embedding_norm->forward(ctx, timestep_proj);
 
@@ -530,14 +530,14 @@ namespace Anima {
                     adapted_context = ggml_mul(ctx->ggml_ctx, adapted_context, w);
                 }
                 if (adapted_context->ne[1] < 512) {
-                    auto pad_ctx    = ggml_ext_zeros(ctx->ggml_ctx,
+                    auto pad_ctx    = ggml_ext_zeros(ctx,
                                                      adapted_context->ne[0],
                                                      512 - adapted_context->ne[1],
                                                      adapted_context->ne[2],
                                                      1);
                     adapted_context = ggml_concat(ctx->ggml_ctx, adapted_context, pad_ctx, 1);
                 } else if (adapted_context->ne[1] > 512) {
-                    adapted_context = ggml_ext_slice(ctx->ggml_ctx, adapted_context, 1, 0, 512);
+                    adapted_context = ggml_ext_slice(ctx, adapted_context, 1, 0, 512);
                 }
                 encoder_hidden_states = adapted_context;
             }
@@ -552,11 +552,11 @@ namespace Anima {
                 x          = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe);
                 sd::ggml_graph_cut::mark_graph_cut(x, "anima.blocks." + std::to_string(i), "x");
             }
-            x = ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, img_len);
+            x = ggml_ext_slice(ctx, x, 1, 0, img_len);
 
             x = final_layer->forward(ctx, x, embedded_timestep, temb);  // [N, h*w, ph*pw*C]
 
-            x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, config.patch_size, config.patch_size, false);  // [N, C, H, W]
+            x = DiT::unpatchify_and_crop(ctx, x, H, W, config.patch_size, config.patch_size, false);  // [N, C, H, W]
 
             return x;
         }

src/model/diffusion/animatediff.hpp

@@ -121,13 +121,13 @@ namespace AnimateDiff {
             auto residual = x;
             auto h        = norm->forward(ctx, x);
 
-            h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1));
+            h = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1));
             h = ggml_reshape_3d(ctx->ggml_ctx, h, C, num_frames, W * H);
             h = proj_in->forward(ctx, h);
             h = tb0->forward(ctx, h);
             h = proj_out->forward(ctx, h);
             h = ggml_reshape_4d(ctx->ggml_ctx, h, C, num_frames, W, H);
-            h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1));
+            h = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1));
 
             return ggml_add(ctx->ggml_ctx, h, residual);
         }

src/model/diffusion/boogu.hpp

@@ -161,7 +161,7 @@ namespace Boogu {
             auto caption_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["caption_embedder.0"]);
             auto caption_embedder_1 = std::dynamic_pointer_cast<Linear>(blocks["caption_embedder.1"]);
 
-            auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(frequency_embedding_size), 10000, timestep_scale);
+            auto timestep_proj = ggml_ext_timestep_embedding(ctx, timestep, static_cast<int>(frequency_embedding_size), 10000, timestep_scale);
             auto time_embed    = timestep_embedder->forward(ctx, timestep_proj);
             auto caption_embed = caption_embedder_1->forward(ctx, caption_embedder_0->forward(ctx, text_hidden_states));
             return {time_embed, caption_embed};
@@ -179,7 +179,7 @@ namespace Boogu {
             auto norm   = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
 
             emb       = linear->forward(ctx, ggml_silu(ctx->ggml_ctx, emb));
-            auto mods = ggml_ext_chunk(ctx->ggml_ctx, emb, 4, 0);
+            auto mods = ggml_ext_chunk(ctx, emb, 4, 0);
 
             auto scale_msa = mods[0];
             auto gate_msa  = mods[1];
@@ -204,9 +204,7 @@ namespace Boogu {
             auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
             auto linear_3 = std::dynamic_pointer_cast<Linear>(blocks["linear_3"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-                linear_2->set_force_prec_f32(true);
-            }
+            linear_2->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
 
             auto h1 = linear_1->forward(ctx, x);
             auto h2 = linear_3->forward(ctx, x);
@@ -264,9 +262,7 @@ namespace Boogu {
             auto norm_k   = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
             auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-                to_out_0->set_force_prec_f32(true);
-            }
+            to_out_0->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
 
             int64_t N  = hidden_states->ne[2];
             int64_t Lq = hidden_states->ne[1];
@@ -388,9 +384,7 @@ namespace Boogu {
             auto instruct_out  = std::dynamic_pointer_cast<Linear>(blocks["processor.instruct_out"]);
             auto img_out       = std::dynamic_pointer_cast<Linear>(blocks["processor.img_out"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-                to_out_0->set_force_prec_f32(true);
-            }
+            to_out_0->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
 
             int64_t N          = img_hidden_states->ne[2];
             int64_t L_img      = img_hidden_states->ne[1];
@@ -417,8 +411,8 @@ namespace Boogu {
             k      = norm_k->forward(ctx, k);
 
             auto hidden_states = Rope::attention(ctx, q, k, v, rotary_emb, attention_mask);
-            auto instruct_attn = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, L_instruct);
-            auto img_attn      = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, L_instruct, L_instruct + L_img);
+            auto instruct_attn = ggml_ext_slice(ctx, hidden_states, 1, 0, L_instruct);
+            auto img_attn      = ggml_ext_slice(ctx, hidden_states, 1, L_instruct, L_instruct + L_img);
 
             instruct_attn = instruct_out->forward(ctx, instruct_attn);
             img_attn      = img_out->forward(ctx, img_attn);
@@ -504,8 +498,8 @@ namespace Boogu {
             auto instruct_shift_mlp = std::get<1>(instruct_norm2_out_vec);
 
             auto joint_attn_out    = img_instruct_attn->forward(ctx, img_norm1_out, instruct_norm1_out, joint_rotary_emb);
-            auto instruct_attn_out = ggml_ext_slice(ctx->ggml_ctx, joint_attn_out, 1, 0, L_instruct);
-            auto img_attn_out      = ggml_ext_slice(ctx->ggml_ctx, joint_attn_out, 1, L_instruct, joint_attn_out->ne[1]);
+            auto instruct_attn_out = ggml_ext_slice(ctx, joint_attn_out, 1, 0, L_instruct);
+            auto img_attn_out      = ggml_ext_slice(ctx, joint_attn_out, 1, L_instruct, joint_attn_out->ne[1]);
 
             auto img_self_attn_out = img_self_attn->forward(ctx, img_norm3_out, img_norm3_out, img_rotary_emb);
 
@@ -634,7 +628,7 @@ namespace Boogu {
             auto time_caption_embed = std::dynamic_pointer_cast<LuminaCombinedTimestepCaptionEmbedding>(blocks["time_caption_embed"]);
             auto norm_out           = std::dynamic_pointer_cast<LuminaLayerNormContinuous>(blocks["norm_out"]);
 
-            auto timestep = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx->ggml_ctx, timesteps), timesteps);
+            auto timestep = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx, timesteps), timesteps);
             auto embeds   = time_caption_embed->forward(ctx, timestep, context);
             auto temb     = embeds.first;
             auto txt      = embeds.second;
@@ -647,8 +641,8 @@ namespace Boogu {
             int64_t txt_len = txt->ne[1];
 
             GGML_ASSERT(pe->ne[3] == txt_len + ref_len + img_len);
-            auto txt_pe   = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt_len);
-            auto noise_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len + ref_len, txt_len + ref_len + img_len);
+            auto txt_pe   = ggml_ext_slice(ctx, pe, 3, 0, txt_len);
+            auto noise_pe = ggml_ext_slice(ctx, pe, 3, txt_len + ref_len, txt_len + ref_len + img_len);
 
             for (int i = 0; i < config.num_refiner_layers; i++) {
                 auto block = std::dynamic_pointer_cast<BooguImageTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
@@ -664,7 +658,7 @@ namespace Boogu {
 
             ggml_tensor* combined_img = img;
             if (ref_img != nullptr) {
-                auto ref_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len, txt_len + ref_len);
+                auto ref_pe = ggml_ext_slice(ctx, pe, 3, txt_len, txt_len + ref_len);
                 for (int i = 0; i < config.num_refiner_layers; i++) {
                     auto block = std::dynamic_pointer_cast<BooguImageTransformerBlock>(blocks["ref_image_refiner." + std::to_string(i)]);
                     ref_img    = block->forward(ctx, ref_img, ref_pe, temb);
@@ -673,7 +667,7 @@ namespace Boogu {
                 combined_img = ggml_concat(ctx->ggml_ctx, ref_img, img, 1);
             }
 
-            auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len, txt_len + combined_img->ne[1]);
+            auto img_pe = ggml_ext_slice(ctx, pe, 3, txt_len, txt_len + combined_img->ne[1]);
             for (int i = 0; i < config.num_double_stream_layers; i++) {
                 auto block   = std::dynamic_pointer_cast<BooguImageDoubleStreamBlock>(blocks["double_stream_layers." + std::to_string(i)]);
                 auto result  = block->forward(ctx, combined_img, txt, pe, img_pe, temb);
@@ -691,9 +685,9 @@ namespace Boogu {
             }
 
             hidden_states = norm_out->forward(ctx, hidden_states, temb);
-            hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, hidden_states->ne[1] - img_len, hidden_states->ne[1]);
-            hidden_states = DiT::unpatchify_and_crop(ctx->ggml_ctx, hidden_states, H, W, config.patch_size, config.patch_size, false);
-            hidden_states = ggml_ext_scale(ctx->ggml_ctx, hidden_states, -1.f);
+            hidden_states = ggml_ext_slice(ctx, hidden_states, 1, hidden_states->ne[1] - img_len, hidden_states->ne[1]);
+            hidden_states = DiT::unpatchify_and_crop(ctx, hidden_states, H, W, config.patch_size, config.patch_size, false);
+            hidden_states = ggml_ext_scale(ctx, hidden_states, -1.f);
             return hidden_states;
         }
     };

src/model/diffusion/control.hpp

@@ -230,7 +230,7 @@ public:
         auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["control_add_embedding.linear_2"]);
 
         // [N*num_control_types, CONTROL_TYPE_EMBED_DIM]
-        auto e = ggml_ext_timestep_embedding(ctx->ggml_ctx, control_types, ControlNetUnionConfig::CONTROL_TYPE_EMBED_DIM);
+        auto e = ggml_ext_timestep_embedding(ctx, control_types, ControlNetUnionConfig::CONTROL_TYPE_EMBED_DIM);
         e      = ggml_reshape_2d(ctx->ggml_ctx, e, ControlNetUnionConfig::CONTROL_TYPE_EMBED_DIM * config.num_control_types, batch);
         e      = linear_1->forward(ctx, e);
         e      = ggml_silu_inplace(ctx->ggml_ctx, e);
@@ -238,8 +238,9 @@ public:
     }
 
     // [N, C, H, W] => [N, C]
-    static ggml_tensor* spatial_mean(ggml_context* ctx, ggml_tensor* x) {
-        auto flat = ggml_reshape_3d(ctx, ggml_ext_cont(ctx, x), x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
+    static ggml_tensor* spatial_mean(GGMLRunnerContext* runner_ctx, ggml_tensor* x) {
+        auto ctx  = runner_ctx->ggml_ctx;
+        auto flat = ggml_reshape_3d(ctx, ggml_ext_cont(runner_ctx, x), x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
         return ggml_reshape_2d(ctx, ggml_mean(ctx, flat), x->ne[2], x->ne[3]);
     }
 
@@ -255,11 +256,11 @@ public:
         auto task_embedding = params["task_embedding"];  // [num_control_types, C]
         auto task           = ggml_view_2d(c, task_embedding, channels, 1, task_embedding->nb[1], control_type_index * task_embedding->nb[1]);
 
-        auto hint_feat = ggml_add(c, spatial_mean(c, hint), task);  // [N, C]
+        auto hint_feat = ggml_add(c, spatial_mean(ctx, hint), task);  // [N, C]
         if (hint_feat->ne[1] != batch) {
             hint_feat = ggml_repeat(c, hint_feat, ggml_new_tensor_2d(c, GGML_TYPE_F32, channels, batch));
         }
-        auto sample_feat = spatial_mean(c, sample);  // [N, C]
+        auto sample_feat = spatial_mean(ctx, sample);  // [N, C]
 
         // as in the reference (nn.MultiheadAttention without batch_first) the attention is over
         // the batch, and the 2 features (hint, sample) are its batch: [2, N, C]
@@ -273,7 +274,7 @@ public:
         }
 
         auto projection = std::dynamic_pointer_cast<Linear>(blocks["spatial_ch_projs"]);
-        auto alpha      = ggml_reshape_2d(c, ggml_ext_slice(c, x, 2, 0, 1), channels, batch);
+        auto alpha      = ggml_reshape_2d(c, ggml_ext_slice(ctx, x, 2, 0, 1), channels, batch);
         alpha           = projection->forward(ctx, alpha);
         alpha           = ggml_reshape_4d(c, alpha, 1, 1, channels, batch);
 
@@ -346,7 +347,7 @@ public:
 
         auto middle_block_out = std::dynamic_pointer_cast<Conv2d>(blocks["middle_block_out.0"]);
 
-        auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, model_channels);  // [N, model_channels]
+        auto t_emb = ggml_ext_timestep_embedding(ctx, timesteps, model_channels);  // [N, model_channels]
 
         auto emb = time_embed_0->forward(ctx, t_emb);
         emb      = ggml_silu_inplace(ctx->ggml_ctx, emb);

src/model/diffusion/dit.hpp

@@ -5,11 +5,12 @@
 #include "core/ggml_runner.h"
 
 namespace DiT {
-    inline ggml_tensor* patchify(ggml_context* ctx,
+    inline ggml_tensor* patchify(GGMLRunnerContext* runner_ctx,
                                  ggml_tensor* x,
                                  int pw,
                                  int ph,
                                  bool patch_last = true) {
+        auto ctx = runner_ctx->ggml_ctx;
         // x: [N, C, H, W]
         // return: [N, h*w, C*ph*pw] if patch_last else [N, h*w, ph*pw*C]
         int64_t N = x->ne[3];
@@ -28,19 +29,20 @@ namespace DiT {
             x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));  // [N, h*w, C, ph*pw]
             x = ggml_reshape_3d(ctx, x, pw * ph * C, w * h, N);    // [N, h*w, C*ph*pw]
         } else {
-            x = ggml_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));  // [N, h*w, C, ph*pw]
-            x = ggml_reshape_3d(ctx, x, C * pw * ph, w * h, N);              // [N, h*w, ph*pw*C]
+            x = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, x, 2, 0, 1, 3));  // [N, h*w, C, ph*pw]
+            x = ggml_reshape_3d(ctx, x, C * pw * ph, w * h, N);                     // [N, h*w, ph*pw*C]
         }
         return x;
     }
 
-    inline ggml_tensor* unpatchify(ggml_context* ctx,
+    inline ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx,
                                    ggml_tensor* x,
                                    int64_t h,
                                    int64_t w,
                                    int ph,
                                    int pw,
                                    bool patch_last = true) {
+        auto ctx = runner_ctx->ggml_ctx;
         // x: [N, h*w, C*ph*pw] if patch_last else [N, h*w, ph*pw*C]
         // return: [N, C, H, W]
         int64_t N = x->ne[2];
@@ -74,7 +76,7 @@ namespace DiT {
 
         int pad_h = (ph - H % ph) % ph;
         int pad_w = (pw - W % pw) % pw;
-        x         = ggml_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
+        x         = ggml_ext_pad(ctx, x, pad_w, pad_h, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
         return x;
     }
 
@@ -84,11 +86,11 @@ namespace DiT {
                                          int pw,
                                          bool patch_last = true) {
         x = pad_to_patch_size(ctx, x, ph, pw);
-        x = patchify(ctx->ggml_ctx, x, ph, pw, patch_last);
+        x = patchify(ctx, x, ph, pw, patch_last);
         return x;
     }
 
-    inline ggml_tensor* unpatchify_and_crop(ggml_context* ctx,
+    inline ggml_tensor* unpatchify_and_crop(GGMLRunnerContext* runner_ctx,
                                             ggml_tensor* x,
                                             int64_t H,
                                             int64_t W,
@@ -99,19 +101,20 @@ namespace DiT {
         int pad_w = (pw - W % pw) % pw;
         int64_t h = ((H + pad_h) / ph);
         int64_t w = ((W + pad_w) / pw);
-        x         = unpatchify(ctx, x, h, w, ph, pw, patch_last);  // [N, C, H + pad_h, W + pad_w]
-        x         = ggml_ext_slice(ctx, x, 1, 0, H);               // [N, C, H, W + pad_w]
-        x         = ggml_ext_slice(ctx, x, 0, 0, W);               // [N, C, H, W]
+        x         = unpatchify(runner_ctx, x, h, w, ph, pw, patch_last);  // [N, C, H + pad_h, W + pad_w]
+        x         = ggml_ext_slice(runner_ctx, x, 1, 0, H);               // [N, C, H, W + pad_w]
+        x         = ggml_ext_slice(runner_ctx, x, 0, 0, W);               // [N, C, H, W]
         return x;
     }
 
-    inline ggml_tensor* patchify_3d(ggml_context* ctx,
+    inline ggml_tensor* patchify_3d(GGMLRunnerContext* runner_ctx,
                                     ggml_tensor* x,
                                     int pt,
                                     int ph,
                                     int pw,
                                     int64_t N       = 1,
                                     bool patch_last = true) {
+        auto ctx = runner_ctx->ggml_ctx;
         // x: [N*C, T, H, W]
         // return: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C]
         int64_t C     = x->ne[3] / N;
@@ -125,24 +128,24 @@ namespace DiT {
         GGML_ASSERT(C * N == x->ne[3]);
         GGML_ASSERT(t_len * pt == T && h_len * ph == H && w_len * pw == W);
 
-        x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N);  // [N*C*t_len, pt, h_len*ph, w_len*pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len, h_len*ph, pt, w_len*pw]
-        x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N);  // [N*C*t_len*h_len, ph, pt, w_len*pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, pt, ph, w_len*pw]
-        x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N);  // [N*C*t_len*h_len, pt*ph, w_len, pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, w_len, pt*ph, pw]
-        x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N);  // [N, C, t_len*h_len*w_len, pt*ph*pw]
+        x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N);            // [N*C*t_len, pt, h_len*ph, w_len*pw]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len, h_len*ph, pt, w_len*pw]
+        x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N);            // [N*C*t_len*h_len, ph, pt, w_len*pw]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len*h_len, pt, ph, w_len*pw]
+        x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N);            // [N*C*t_len*h_len, pt*ph, w_len, pw]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len*h_len, w_len, pt*ph, pw]
+        x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N);            // [N, C, t_len*h_len*w_len, pt*ph*pw]
         if (patch_last) {
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));  // [N, t_len*h_len*w_len, C, pt*ph*pw]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N, t_len*h_len*w_len, C, pt*ph*pw]
         } else {
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));  // [N, t_len*h_len*w_len, pt*ph*pw, C]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 2, 0, 1, 3));  // [N, t_len*h_len*w_len, pt*ph*pw, C]
         }
         // [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C]
         x = ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1);
         return x;
     }
 
-    inline ggml_tensor* unpatchify_3d(ggml_context* ctx,
+    inline ggml_tensor* unpatchify_3d(GGMLRunnerContext* runner_ctx,
                                       ggml_tensor* x,
                                       int64_t t_len,
                                       int64_t h_len,
@@ -151,6 +154,7 @@ namespace DiT {
                                       int ph,
                                       int pw,
                                       bool patch_last = true) {
+        auto ctx = runner_ctx->ggml_ctx;
         // x: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, pt*ph*pw*C]
         // return: [N*C, t_len*pt, h_len*ph, w_len*pw]
         int64_t N = x->ne[2];
@@ -159,20 +163,20 @@ namespace DiT {
         GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
 
         if (patch_last) {
-            x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, C, pt*ph*pw]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
+            x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N);            // [N, t_len*h_len*w_len, C, pt*ph*pw]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N, C, t_len*h_len*w_len, pt*ph*pw]
         } else {
-            x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, pt*ph*pw, C]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
+            x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);            // [N, t_len*h_len*w_len, pt*ph*pw, C]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 1, 2, 0, 3));  // [N, C, t_len*h_len*w_len, pt*ph*pw]
         }
 
-        x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);  // [N*C*t_len*h_len, w_len, pt*ph, pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, pt*ph, w_len, pw]
-        x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);  // [N*C*t_len*h_len, pt, ph, w_len*pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, ph, pt, w_len*pw]
-        x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N);  // [N*C*t_len, h_len*ph, pt, w_len*pw]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len, pt, h_len*ph, w_len*pw]
-        x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N);  // [N*C, t_len*pt, h_len*ph, w_len*pw]
+        x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);            // [N*C*t_len*h_len, w_len, pt*ph, pw]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len*h_len, pt*ph, w_len, pw]
+        x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);            // [N*C*t_len*h_len, pt, ph, w_len*pw]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len*h_len, ph, pt, w_len*pw]
+        x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N);            // [N*C*t_len, h_len*ph, pt, w_len*pw]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len, pt, h_len*ph, w_len*pw]
+        x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N);            // [N*C, t_len*pt, h_len*ph, w_len*pw]
         return x;
     }
 }  // namespace DiT

src/model/diffusion/ernie_image.hpp

@@ -85,11 +85,12 @@ namespace ErnieImage {
         }
     };
 
-    __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(ggml_context* ctx,
+    __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(GGMLRunnerContext* runner_ctx,
                                                               ggml_tensor* timesteps,
                                                               int dim,
                                                               int max_period = 10000) {
-        auto emb       = ggml_ext_timestep_embedding(ctx, timesteps, dim, max_period, 1.0f);
+        auto ctx       = runner_ctx->ggml_ctx;
+        auto emb       = ggml_ext_timestep_embedding(runner_ctx, timesteps, dim, max_period, 1.0f);
         int64_t half   = dim / 2;
         auto cos_part  = ggml_view_2d(ctx, emb, half, emb->ne[1], emb->nb[1], 0);
         auto sin_part  = ggml_view_2d(ctx, emb, half, emb->ne[1], emb->nb[1], half * emb->nb[0]);
@@ -97,7 +98,8 @@ namespace ErnieImage {
         return sin_first;
     }
 
-    __STATIC_INLINE__ ggml_tensor* apply_rotary_emb(ggml_context* ctx, ggml_tensor* x, ggml_tensor* pe) {
+    __STATIC_INLINE__ ggml_tensor* apply_rotary_emb(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* pe) {
+        auto ctx = runner_ctx->ggml_ctx;
         // x: [N, S, heads, head_dim]
         // pe: [2, S, 1, head_dim], stored as ggml [head_dim, 1, S, 2].
         int64_t head_dim = x->ne[0];
@@ -110,8 +112,8 @@ namespace ErnieImage {
         GGML_ASSERT(pe->ne[1] == 1 && pe->ne[2] == S && pe->ne[3] == 2);
 
         x           = ggml_cont(ctx, x);
-        auto x_rot  = ggml_ext_slice(ctx, x, 0, 0, rot_dim, false);
-        auto x_pass = rot_dim < head_dim ? ggml_ext_slice(ctx, x, 0, rot_dim, head_dim, false) : nullptr;
+        auto x_rot  = ggml_ext_slice(runner_ctx, x, 0, 0, rot_dim, false);
+        auto x_pass = rot_dim < head_dim ? ggml_ext_slice(runner_ctx, x, 0, rot_dim, head_dim, false) : nullptr;
 
         int64_t half = rot_dim / 2;
         auto x1      = ggml_view_4d(ctx, x_rot, half, heads, S, N, x_rot->nb[1], x_rot->nb[2], x_rot->nb[3], 0);
@@ -120,8 +122,8 @@ namespace ErnieImage {
         x2           = ggml_cont(ctx, x2);
         auto rotated = ggml_concat(ctx, ggml_neg(ctx, x2), x1, 0);
 
-        auto cos_emb = ggml_ext_slice(ctx, pe, 3, 0, 1, false);
-        auto sin_emb = ggml_ext_slice(ctx, pe, 3, 1, 2, false);
+        auto cos_emb = ggml_ext_slice(runner_ctx, pe, 3, 0, 1, false);
+        auto sin_emb = ggml_ext_slice(runner_ctx, pe, 3, 1, 2, false);
 
         auto out = ggml_add(ctx, ggml_mul(ctx, x_rot, cos_emb), ggml_mul(ctx, rotated, sin_emb));
         if (x_pass != nullptr) {
@@ -175,8 +177,8 @@ namespace ErnieImage {
             q = norm_q->forward(ctx, q);
             k = norm_k->forward(ctx, k);
 
-            q = apply_rotary_emb(ctx->ggml_ctx, q, pe);
-            k = apply_rotary_emb(ctx->ggml_ctx, k, pe);
+            q = apply_rotary_emb(ctx, q, pe);
+            k = apply_rotary_emb(ctx, k, pe);
 
             q = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, q, 0, 2, 1, 3));  // [N, heads, S, head_dim]
             q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]);
@@ -269,7 +271,7 @@ namespace ErnieImage {
             auto norm   = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
             auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
 
-            auto mods  = ggml_ext_chunk(ctx->ggml_ctx, linear->forward(ctx, conditioning), 2, 0);
+            auto mods  = ggml_ext_chunk(ctx, linear->forward(ctx, conditioning), 2, 0);
             auto scale = mods[0];
             auto shift = mods[1];
 
@@ -336,9 +338,9 @@ namespace ErnieImage {
             auto final_norm      = std::dynamic_pointer_cast<ErnieImageAdaLNContinuous>(blocks["final_norm"]);
             auto final_linear    = std::dynamic_pointer_cast<Linear>(blocks["final_linear"]);
 
-            auto img = x_embedder_proj->forward(ctx, x);                                                  // [N, hidden_size, Hp, Wp]
-            img      = ggml_reshape_3d(ctx->ggml_ctx, img, img->ne[0] * img->ne[1], img->ne[2], N);       // [N, hidden_size, image_tokens]
-            img      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));  // [N, image_tokens, hidden_size]
+            auto img = x_embedder_proj->forward(ctx, x);                                             // [N, hidden_size, Hp, Wp]
+            img      = ggml_reshape_3d(ctx->ggml_ctx, img, img->ne[0] * img->ne[1], img->ne[2], N);  // [N, hidden_size, image_tokens]
+            img      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, img, 1, 0, 2, 3));       // [N, image_tokens, hidden_size]
 
             auto txt       = context;
             auto text_proj = std::dynamic_pointer_cast<Linear>(blocks["text_proj"]);
@@ -348,13 +350,13 @@ namespace ErnieImage {
 
             auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1);  // [N, image_tokens + text_tokens, hidden_size]
 
-            auto sample = timestep_embedding_sin_cos(ctx->ggml_ctx, timestep, static_cast<int>(config.hidden_size));
+            auto sample = timestep_embedding_sin_cos(ctx, timestep, static_cast<int>(config.hidden_size));
             auto c      = time_embedding->forward(ctx, sample);  // [N, hidden_size]
 
             auto mod_params = adaLN_mod->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, 6 * hidden_size]
             sd::ggml_graph_cut::mark_graph_cut(hidden_states, "ernie_image.prelude", "hidden_states");
             // sd::ggml_graph_cut::mark_graph_cut(mod_params, "ernie_image.prelude", "mod_params");
-            auto chunks = ggml_ext_chunk(ctx->ggml_ctx, mod_params, 6, 0);
+            auto chunks = ggml_ext_chunk(ctx, mod_params, 6, 0);
             std::vector<ggml_tensor*> temb;
             temb.reserve(6);
             for (auto chunk : chunks) {
@@ -368,10 +370,10 @@ namespace ErnieImage {
             }
 
             hidden_states = final_norm->forward(ctx, hidden_states, c);
-            hidden_states = final_linear->forward(ctx, hidden_states);                  // [N, image_tokens, p*p*out_channels]
-            auto patches  = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img);  // [N, image_tokens, hidden_size]
+            hidden_states = final_linear->forward(ctx, hidden_states);        // [N, image_tokens, p*p*out_channels]
+            auto patches  = ggml_ext_slice(ctx, hidden_states, 1, 0, n_img);  // [N, image_tokens, hidden_size]
 
-            auto out = DiT::unpatchify(ctx->ggml_ctx,
+            auto out = DiT::unpatchify(ctx,
                                        patches,
                                        Hp,
                                        Wp,

src/model/diffusion/fibo2.hpp

@@ -108,14 +108,14 @@ namespace Fibo2 {
             };
             ggml_tensor *q, *k, *v;
             if (fused_qkv) {
-                auto parts = ggml_ext_chunk(c, project("qkv", x), 3, 0);
+                auto parts = ggml_ext_chunk(ctx, project("qkv", x), 3, 0);
                 q          = parts[0];
                 k          = parts[1];
                 v          = parts[2];
             } else {
                 q = project("to_q", x);
                 if (fused_kv) {
-                    auto parts = ggml_ext_chunk(c, project("to_kv", kv), 2, 0);
+                    auto parts = ggml_ext_chunk(ctx, project("to_kv", kv), 2, 0);
                     k          = parts[0];
                     v          = parts[1];
                 } else {
@@ -158,9 +158,7 @@ namespace Fibo2 {
                 return ZImage::FeedForward::forward(ctx, x);
             }
             auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
-            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-                w2->set_force_prec_f32(true);
-            }
+            w2->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
             x = std::dynamic_pointer_cast<Linear>(blocks["w13"])->forward(ctx, x);
             return w2->forward(ctx, ggml_swiglu(ctx->ggml_ctx, x));
         }
@@ -206,7 +204,7 @@ namespace Fibo2 {
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* text, int gist_count) {
-            auto x = ggml_ext_slice(ctx->ggml_ctx, params["latent_queries"], 1, 0, gist_count);
+            auto x = ggml_ext_slice(ctx, params["latent_queries"], 1, 0, gist_count);
             for (int i = 0; i < config.perceiver_depth; ++i) {
                 x = std::dynamic_pointer_cast<PerceiverLayer>(blocks["layers." + std::to_string(i)])->forward(ctx, x, text);
             }
@@ -244,8 +242,8 @@ namespace Fibo2 {
                 return std::dynamic_pointer_cast<RMSNorm>(blocks[name])->forward(ctx, h);
             };
             auto modulation = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.0"]);
-            auto mn         = ggml_ext_chunk(c, modulation->forward(ctx, noisy), 4, 0);
-            auto mc         = ggml_ext_chunk(c, modulation->forward(ctx, clean), 4, 0);
+            auto mn         = ggml_ext_chunk(ctx, modulation->forward(ctx, noisy), 4, 0);
+            auto mc         = ggml_ext_chunk(ctx, modulation->forward(ctx, clean), 4, 0);
             auto select     = [&](int index, bool gate) {
                 auto n      = gate ? ggml_tanh(c, mn[index]) : mn[index];
                 auto k      = gate ? ggml_tanh(c, mc[index]) : mc[index];
@@ -311,12 +309,12 @@ namespace Fibo2 {
             auto embedder = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
             auto noisy    = embedder->forward(ctx, t);
             auto clean    = embedder->forward(ctx, ggml_scale(c, t, 0.f));
-            auto image_pe = ggml_ext_slice(c, pe, 3, gist_count, pe->ne[3]);
+            auto image_pe = ggml_ext_slice(ctx, pe, 3, gist_count, pe->ne[3]);
             for (int i = 0; i < config.refiner_depth; ++i) {
                 x = std::dynamic_pointer_cast<TransformerBlock>(blocks["noise_refiner." + std::to_string(i)])->forward(ctx, x, image_pe, noisy, clean, 0, image_count);
                 sd::ggml_graph_cut::mark_graph_cut(x, "fibo2.noise_refiner." + std::to_string(i), "x");
             }
-            auto bundle      = [&](int i) { return ggml_ext_slice(c, text, 0, i * config.text_dim, (i + 1) * config.text_dim); };
+            auto bundle      = [&](int i) { return ggml_ext_slice(ctx, text, 0, i * config.text_dim, (i + 1) * config.text_dim); };
             auto gist        = std::dynamic_pointer_cast<Perceiver>(blocks["perceiver"])->forward(ctx, bundle(0), gist_count);
             x                = ggml_concat(c, gist, x, 1);
             int bundle_index = 1;
@@ -326,7 +324,7 @@ namespace Fibo2 {
                 x            = block->forward(ctx, x, pe, noisy, clean, gist_count, image_count, context);
                 sd::ggml_graph_cut::mark_graph_cut(x, "fibo2.layers." + std::to_string(i), "x");
             }
-            x = ggml_ext_slice(c, x, 1, gist_count, gist_count + image_count);
+            x = ggml_ext_slice(ctx, x, 1, gist_count, gist_count + image_count);
             x = std::dynamic_pointer_cast<ZImage::FinalLayer>(blocks["final_layer"])->forward(ctx, x, noisy);
             x = ggml_cont(c, ggml_permute(c, x, 1, 0, 2, 3));
             return ggml_reshape_4d(c, x, width, height, config.channels, 1);

src/model/diffusion/flux.hpp

@@ -335,9 +335,9 @@ namespace Flux {
 
             x = mlp_0->forward(ctx, x);
             if (use_mlp_silu_act) {
-                x = ggml_ext_silu_act(ctx->ggml_ctx, x);
+                x = ggml_ext_silu_act(ctx, x);
             } else {
-                x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+                x = ggml_ext_gelu(ctx, x, true);
             }
             x = mlp_2->forward(ctx, x);
             return x;
@@ -686,11 +686,11 @@ namespace Flux {
 
             auto mlp = ggml_view_3d(ctx->ggml_ctx, qkv_mlp, mlp_hidden_dim * mlp_mult_factor, qkv_mlp->ne[1], qkv_mlp->ne[2], qkv_mlp->nb[1], qkv_mlp->nb[2], hidden_size * 3 * qkv_mlp->nb[0]);
             if (use_yak_mlp) {
-                mlp = ggml_ext_silu_act(ctx->ggml_ctx, mlp, false);
+                mlp = ggml_ext_silu_act(ctx, mlp, false);
             } else if (use_mlp_silu_act) {
-                mlp = ggml_ext_silu_act(ctx->ggml_ctx, mlp);
+                mlp = ggml_ext_silu_act(ctx, mlp);
             } else {
-                mlp = ggml_ext_gelu(ctx->ggml_ctx, mlp, true);
+                mlp = ggml_ext_gelu(ctx, mlp, true);
             }
             auto attn_mlp = ggml_concat(ctx->ggml_ctx, attn, mlp, 0);  // [N, n_token, hidden_size + mlp_hidden_dim]
             auto output   = linear2->forward(ctx, attn_mlp);           // [N, n_token, hidden_size]
@@ -745,7 +745,7 @@ namespace Flux {
                 auto adaLN_modulation_1 = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
 
                 auto m     = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, 2 * hidden_size]
-                auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, 2, 0);
+                auto m_vec = ggml_ext_chunk(ctx, m, 2, 0);
                 shift      = m_vec[0];
                 scale      = m_vec[1];
             }
@@ -832,16 +832,16 @@ namespace Flux {
             int64_t hidden_size_x = x->ne[0];
 
             auto mlp_params = param_generator->forward(ctx, s);
-            auto fc_params  = ggml_ext_chunk(ctx->ggml_ctx, mlp_params, 3, 0);
+            auto fc_params  = ggml_ext_chunk(ctx, mlp_params, 3, 0);
             auto fc1_gate   = ggml_reshape_3d(ctx->ggml_ctx, fc_params[0], hidden_size_x * mlp_ratio, hidden_size_x, batch_size);
             auto fc1_value  = ggml_reshape_3d(ctx->ggml_ctx, fc_params[1], hidden_size_x * mlp_ratio, hidden_size_x, batch_size);
             auto fc2        = ggml_reshape_3d(ctx->ggml_ctx, fc_params[2], hidden_size_x, mlp_ratio * hidden_size_x, batch_size);
 
-            fc1_gate  = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, fc1_gate, 1, 0, 2, 3));  // [batch_size, hidden_size_x*mlp_ratio, hidden_size_x]
+            fc1_gate  = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, fc1_gate, 1, 0, 2, 3));  // [batch_size, hidden_size_x*mlp_ratio, hidden_size_x]
             fc1_gate  = ggml_l2_norm(ctx->ggml_ctx, fc1_gate, 1e-12f);
-            fc1_value = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, fc1_value, 1, 0, 2, 3));  // [batch_size, hidden_size_x*mlp_ratio, hidden_size_x]
+            fc1_value = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, fc1_value, 1, 0, 2, 3));  // [batch_size, hidden_size_x*mlp_ratio, hidden_size_x]
             fc1_value = ggml_l2_norm(ctx->ggml_ctx, fc1_value, 1e-12f);
-            fc2       = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, fc2, 1, 0, 2, 3));  // [batch_size, hidden_size_x, hidden_size_x*mlp_ratio]
+            fc2       = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, fc2, 1, 0, 2, 3));  // [batch_size, hidden_size_x, hidden_size_x*mlp_ratio]
             fc2       = ggml_l2_norm(ctx->ggml_ctx, fc2, 1e-12f);
 
             auto res_x = x;
@@ -892,9 +892,9 @@ namespace Flux {
             auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
             auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
 
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));  // [N, H, W, C]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));  // [N, H, W, C]
             x = norm->forward(ctx, x);
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));  // [N, C, H, W]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));  // [N, C, H, W]
             x = conv->forward(ctx, x);
 
             return x;
@@ -1028,13 +1028,13 @@ namespace Flux {
             if (config.is_chroma) {
                 int64_t mod_index_length = 344;
                 auto approx              = std::dynamic_pointer_cast<ChromaApproximator>(blocks["distilled_guidance_layer"]);
-                auto distill_timestep    = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 16, 10000, 1000.f);
-                auto distill_guidance    = ggml_ext_timestep_embedding(ctx->ggml_ctx, guidance, 16, 10000, 1000.f);
+                auto distill_timestep    = ggml_ext_timestep_embedding(ctx, timesteps, 16, 10000, 1000.f);
+                auto distill_guidance    = ggml_ext_timestep_embedding(ctx, guidance, 16, 10000, 1000.f);
 
                 // auto mod_index_arange  = ggml_arange(ctx, 0, (float)mod_index_length, 1);
                 // ggml_arange tot working on a lot of backends, precomputing it on CPU instead
                 GGML_ASSERT(mod_index_arange != nullptr);
-                auto modulation_index = ggml_ext_timestep_embedding(ctx->ggml_ctx, mod_index_arange, 32, 10000, 1000.f);  // [1, 344, 32]
+                auto modulation_index = ggml_ext_timestep_embedding(ctx, mod_index_arange, 32, 10000, 1000.f);  // [1, 344, 32]
 
                 // Batch broadcast (will it ever be useful)
                 modulation_index = ggml_repeat(ctx->ggml_ctx, modulation_index, ggml_new_tensor_3d(ctx->ggml_ctx, GGML_TYPE_F32, modulation_index->ne[0], modulation_index->ne[1], img->ne[2]));  // [N, 344, 32]
@@ -1057,12 +1057,12 @@ namespace Flux {
                 vec                  = dual_time_embed->forward(ctx, timestep_sem, timestep_tex);
             } else {
                 auto time_in = std::dynamic_pointer_cast<MLPEmbedder>(blocks["time_in"]);
-                vec          = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1000.f));
+                vec          = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx, timesteps, 256, 10000, 1000.f));
                 if (config.guidance_embed) {
                     GGML_ASSERT(guidance != nullptr);
                     auto guidance_in = std::dynamic_pointer_cast<MLPEmbedder>(blocks["guidance_in"]);
                     // bf16 and fp16 result is different
-                    auto g_in = ggml_ext_timestep_embedding(ctx->ggml_ctx, guidance, 256, 10000, 1000.f);
+                    auto g_in = ggml_ext_timestep_embedding(ctx, guidance, 256, 10000, 1000.f);
                     vec       = ggml_add(ctx->ggml_ctx, vec, guidance_in->forward(ctx, g_in));
                 }
 
@@ -1223,7 +1223,7 @@ namespace Flux {
 
             img = img_in_patch->forward(ctx, img);                                                       // [N, hidden_size, H/patch_size, W/patch_size]
             img = ggml_reshape_3d(ctx->ggml_ctx, img, img->ne[0] * img->ne[1], img->ne[2], img->ne[3]);  // [N, hidden_size, H/patch_size*W/patch_size]
-            img = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));      // [N, H/patch_size*W/patch_size, hidden_size]
+            img = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, img, 1, 0, 2, 3));                // [N, H/patch_size*W/patch_size, hidden_size]
 
             auto out = forward_orig(ctx, img, context, timestep, y, guidance, pe, mod_index_arange, skip_layers,
                                     pulid_id, pulid_id_weight);  // [N, n_img_token, hidden_size]
@@ -1232,14 +1232,14 @@ namespace Flux {
             auto nerf_image_embedder   = std::dynamic_pointer_cast<NerfEmbedder>(blocks["nerf_image_embedder"]);
             auto nerf_final_layer_conv = std::dynamic_pointer_cast<NerfFinalLayerConv>(blocks["nerf_final_layer_conv"]);
 
-            auto nerf_pixels    = DiT::patchify(ctx->ggml_ctx, orig_img, patch_size, patch_size);  // [N, num_patches, C * patch_size * patch_size]
+            auto nerf_pixels    = DiT::patchify(ctx, orig_img, patch_size, patch_size);  // [N, num_patches, C * patch_size * patch_size]
             int64_t num_patches = nerf_pixels->ne[1];
             nerf_pixels         = ggml_reshape_3d(ctx->ggml_ctx,
                                                   nerf_pixels,
                                                   nerf_pixels->ne[0] / C,
                                                   C,
-                                                  nerf_pixels->ne[1] * nerf_pixels->ne[2]);                                  // [N*num_patches, C, patch_size*patch_size]
-            nerf_pixels         = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, nerf_pixels, 1, 0, 2, 3));  // [N*num_patches, patch_size*patch_size, C]
+                                                  nerf_pixels->ne[1] * nerf_pixels->ne[2]);                        // [N*num_patches, C, patch_size*patch_size]
+            nerf_pixels         = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, nerf_pixels, 1, 0, 2, 3));  // [N*num_patches, patch_size*patch_size, C]
 
             auto nerf_hidden = ggml_reshape_2d(ctx->ggml_ctx, out, out->ne[0], out->ne[1] * out->ne[2]);  // [N*num_patches, hidden_size]
             auto img_dct     = nerf_image_embedder->forward(ctx, nerf_pixels, dct);                       // [N*num_patches, patch_size*patch_size, nerf_hidden_size]
@@ -1250,9 +1250,9 @@ namespace Flux {
                 img_dct = block->forward(ctx, img_dct, nerf_hidden);
             }
 
-            img_dct = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img_dct, 1, 0, 2, 3));                                 // [N*num_patches, nerf_hidden_size, patch_size*patch_size]
+            img_dct = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, img_dct, 1, 0, 2, 3));                                           // [N*num_patches, nerf_hidden_size, patch_size*patch_size]
             img_dct = ggml_reshape_3d(ctx->ggml_ctx, img_dct, img_dct->ne[0] * img_dct->ne[1], num_patches, img_dct->ne[2] / num_patches);  // [N, num_patches, nerf_hidden_size*patch_size*patch_size]
-            img_dct = DiT::unpatchify(ctx->ggml_ctx, img_dct, (H + pad_h) / patch_size, (W + pad_w) / patch_size, patch_size, patch_size);  // [N, nerf_hidden_size, H, W]
+            img_dct = DiT::unpatchify(ctx, img_dct, (H + pad_h) / patch_size, (W + pad_w) / patch_size, patch_size, patch_size);            // [N, nerf_hidden_size, H, W]
 
             out = nerf_final_layer_conv->forward(ctx, img_dct);  // [N, C, H, W]
 
@@ -1331,7 +1331,7 @@ namespace Flux {
                 out = ggml_cont(ctx->ggml_ctx, out);
             }
 
-            out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size);  // [N, C, H, W]
+            out = DiT::unpatchify_and_crop(ctx, out, H, W, patch_size, patch_size);  // [N, C, H, W]
             return out;
         }
 

src/model/diffusion/flux3.hpp

@@ -95,7 +95,7 @@ namespace Flux3 {
             auto k           = norm->key_norm(ctx, split(linear(ctx, "k_proj", x_mod)));
             auto v           = split(linear(ctx, "v_proj", x_mod));
             auto attn        = Rope::attention(ctx, q, k, v, pe, nullptr);
-            auto mlp         = ggml_ext_silu_act(gc, linear(ctx, "mlp_in", x_mod));
+            auto mlp         = ggml_ext_silu_act(ctx, linear(ctx, "mlp_in", x_mod));
             auto update      = ggml_add(gc, linear(ctx, "attn_out", attn), linear(ctx, "mlp_out", mlp));
             ggml_tensor* out = nullptr;
             offset           = 0;
@@ -140,16 +140,16 @@ namespace Flux3 {
             auto gc             = ctx->ggml_ctx;
             const int64_t width = image->ne[0], height = image->ne[1];
             auto time_in   = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["time_in"]);
-            auto image_vec = time_in->forward(ctx, ggml_ext_timestep_embedding(gc, timestep, 256, 10000, 1000.f));
-            auto clean_vec = time_in->forward(ctx, ggml_ext_timestep_embedding(gc, zero, 256, 10000, 1000.f));
+            auto image_vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx, timestep, 256, 10000, 1000.f));
+            auto clean_vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx, zero, 256, 10000, 1000.f));
             if (config.flow_map) {
                 auto ref_time = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["ref_time_in"]);
-                auto ref_vec  = ref_time->forward(ctx, ggml_ext_timestep_embedding(gc, target, 256, 10000, 1000.f));
+                auto ref_vec  = ref_time->forward(ctx, ggml_ext_timestep_embedding(ctx, target, 256, 10000, 1000.f));
                 image_vec     = ggml_add(gc, image_vec, ref_vec);
                 clean_vec     = ggml_add(gc, clean_vec, ref_vec);
             }
             std::vector<std::string> modes{"image"};
-            std::vector<ggml_tensor*> inputs{DiT::patchify(gc, image, 1, 1)};
+            std::vector<ggml_tensor*> inputs{DiT::patchify(ctx, image, 1, 1)};
             if (reference) {
                 modes.push_back("image_cond");
                 inputs.push_back(reference);
@@ -184,7 +184,7 @@ namespace Flux3 {
                 sd::ggml_graph_cut::mark_graph_cut(joint, "flux3.single_blocks." + std::to_string(i), "joint");
             }
             auto out = std::dynamic_pointer_cast<Flux::LastLayer>(blocks["final_layer.image"])->forward(ctx, token_slice(gc, joint, 0, lengths[0]), image_vec);
-            return DiT::unpatchify(gc, out, height, width, 1, 1);
+            return DiT::unpatchify(ctx, out, height, width, 1, 1);
         }
     };
 
@@ -215,9 +215,10 @@ namespace Flux3 {
             GGML_ASSERT(params.x && params.context && params.timesteps);
             GGML_ASSERT(!config.flow_map || extra->target_timestep);
             auto get_graph = [&]() {
-                auto gf   = new_graph_custom(FLUX_GRAPH_SIZE);
-                auto x    = make_input(*params.x);
-                auto text = make_input(*params.context);
+                auto gf         = new_graph_custom(FLUX_GRAPH_SIZE);
+                auto runner_ctx = get_context();
+                auto x          = make_input(*params.x);
+                auto text       = make_input(*params.context);
                 GGML_ASSERT(x->ne[3] == 1 && text->ne[2] == 1);
                 std::vector<ggml_tensor*> refs;
                 ggml_tensor* references = nullptr;
@@ -225,13 +226,13 @@ namespace Flux3 {
                     for (const auto& ref : *params.ref_latents) {
                         auto r = make_input(ref);
                         refs.push_back(r);
-                        r          = DiT::patchify(compute_ctx, r, 1, 1);
+                        r          = DiT::patchify(&runner_ctx, r, 1, 1);
                         references = references ? ggml_concat(compute_ctx, references, r, 1) : r;
                     }
                 }
                 sd::PECacheKey pe_key("flux3.pe");
                 pe_key.add_shape(x).add_shape(text).add_shapes(refs).add(circular_x_enabled).add(circular_y_enabled);
-                auto pe         = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                auto pe  = make_input(pe_cache_.get_or_create(pe_key, [&]() {
                     Rope::Embedding embedding;
                     auto append_image = [&](int64_t w, int64_t h, float index) {
                         embedding.positions.append_image(static_cast<int>(h), static_cast<int>(w), 1, 0.5f, 0.5f);
@@ -254,9 +255,8 @@ namespace Flux3 {
                     const int64_t tokens = static_cast<int64_t>(values.size() / 256);
                     return sd::Tensor<float>({2, 2, 64, tokens}, std::move(values));
                 }));
-                auto runner_ctx = get_context();
-                auto out        = model.forward(&runner_ctx, x, text, references, make_input(*params.timesteps), make_input(zero),
-                                                config.flow_map ? make_input(*extra->target_timestep) : nullptr, pe);
+                auto out = model.forward(&runner_ctx, x, text, references, make_input(*params.timesteps), make_input(zero),
+                                         config.flow_map ? make_input(*extra->target_timestep) : nullptr, pe);
                 ggml_build_forward_expand(gf, out);
                 return gf;
             };

src/model/diffusion/hidream_o1.hpp

@@ -197,7 +197,7 @@ namespace HiDreamO1 {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
             auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
             auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
-            auto emb   = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10000, 1000.0f);
+            auto emb   = ggml_ext_timestep_embedding(ctx, t, frequency_embedding_size, 10000, 1000.0f);
             emb        = mlp_0->forward(ctx, emb);
             emb        = ggml_silu_inplace(ctx->ggml_ctx, emb);
             emb        = mlp_2->forward(ctx, emb);
@@ -292,7 +292,6 @@ namespace HiDreamO1 {
 
         ggml_tensor* encode_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) {
             return LLM::LLMRunner::encode_image_common(this,
-                                                       compute_ctx,
                                                        runner_ctx,
                                                        image,
                                                        config.llm.vision,
@@ -398,7 +397,7 @@ namespace HiDreamO1 {
             auto t_emb          = t_embedder1->forward(&runner_ctx, timestep);
             int64_t txt_seq_len = txt->ne[1];
             if (txt_seq_len > 1) {
-                auto prefix = ggml_ext_slice(compute_ctx, txt, 1, 0, txt_seq_len - 1);
+                auto prefix = ggml_ext_slice(&runner_ctx, txt, 1, 0, txt_seq_len - 1);
                 txt         = ggml_concat(compute_ctx, prefix, ggml_reshape_3d(compute_ctx, t_emb, t_emb->ne[0], 1, 1), 1);
             } else {
                 txt = ggml_reshape_3d(compute_ctx, t_emb, t_emb->ne[0], 1, 1);
@@ -425,8 +424,8 @@ namespace HiDreamO1 {
                 }
                 x_pred_start = first_vinput;
             }
-            auto x_pred = ggml_ext_slice(compute_ctx, x_pred_all, 1, x_pred_start, x_pred_start + target_tokens);
-            x_pred      = DiT::unpatchify_and_crop(compute_ctx, x_pred, x->ne[1], x->ne[0], PATCH_SIZE, PATCH_SIZE);
+            auto x_pred = ggml_ext_slice(&runner_ctx, x_pred_all, 1, x_pred_start, x_pred_start + target_tokens);
+            x_pred      = DiT::unpatchify_and_crop(&runner_ctx, x_pred, x->ne[1], x->ne[0], PATCH_SIZE, PATCH_SIZE);
 
             float sigma = 1.0f - timestep_tensor.values()[0];
             sigma       = std::max(1e-6f, sigma);

src/model/diffusion/hunyuan.hpp

@@ -50,7 +50,7 @@ namespace Hunyuan {
 
             // self attn
             auto qkv     = self_attn_qkv->forward(ctx, norm1->forward(ctx, txt));
-            auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
+            auto qkv_vec = ggml_ext_split_qkv(ctx, qkv);
             auto q       = qkv_vec[0];
             auto k       = qkv_vec[1];
             auto v       = qkv_vec[2];
@@ -60,7 +60,7 @@ namespace Hunyuan {
 
             // adaLN_modulation
             auto emb  = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, t_emb));
-            auto mods = ggml_ext_chunk(ctx->ggml_ctx, emb, 2, 0);
+            auto mods = ggml_ext_chunk(ctx, emb, 2, 0);
 
             txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, attn_out, mods[0]));
 
@@ -122,11 +122,11 @@ namespace Hunyuan {
             auto c_embedder               = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["c_embedder"]);
             auto individual_token_refiner = std::dynamic_pointer_cast<IndividualTokenRefiner>(blocks["individual_token_refiner"]);
 
-            auto t_emb = t_embedder->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f));
+            auto t_emb = t_embedder->forward(ctx, ggml_ext_timestep_embedding(ctx, timestep, 256, 10000, 1.f));
 
-            auto h                  = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, txt, 1, 0, 2, 3));
+            auto h                  = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, txt, 1, 0, 2, 3));
             auto pooled_projections = ggml_scale(ctx->ggml_ctx, ggml_sum_rows(ctx->ggml_ctx, h), 1.f / txt->ne[1]);
-            pooled_projections      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, pooled_projections, 1, 0, 2, 3));
+            pooled_projections      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, pooled_projections, 1, 0, 2, 3));
             auto c_emb              = c_embedder->forward(ctx, pooled_projections);
 
             t_emb = ggml_add(ctx->ggml_ctx, t_emb, c_emb);
@@ -151,9 +151,9 @@ namespace Hunyuan {
             auto fc3       = std::dynamic_pointer_cast<Linear>(blocks["fc3"]);
 
             x = fc1->forward(ctx, layernorm->forward(ctx, x));
-            x = ggml_ext_gelu(ctx->ggml_ctx, x);
+            x = ggml_ext_gelu(ctx, x);
             x = fc2->forward(ctx, x);
-            x = ggml_ext_gelu(ctx->ggml_ctx, x);
+            x = ggml_ext_gelu(ctx, x);
             return fc3->forward(ctx, x);
         }
     };
@@ -364,11 +364,12 @@ namespace Hunyuan {
             return x;
         }
 
-        ggml_tensor* unpatchify(struct ggml_context* ctx,
+        ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx,
                                 ggml_tensor* x,
                                 int64_t t_len,
                                 int64_t h_len,
                                 int64_t w_len) {
+            auto ctx = runner_ctx->ggml_ctx;
             // x: [N, t_len*h_len*w_len, C*pt*ph*pw]
             // return: [N*C, t_len*pt, h_len*ph, w_len*pw]
             int64_t N  = x->ne[3];
@@ -379,15 +380,15 @@ namespace Hunyuan {
 
             GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
 
-            x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, pt*ph*pw, C]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
-            x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);  // [N*C*t_len*h_len, w_len, pt*ph, pw]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, pt*ph, w_len, pw]
-            x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);  // [N*C*t_len*h_len, pt, ph, w_len*pw]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, ph, pt, w_len*pw]
-            x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N);  // [N*C*t_len, h_len*ph, pt, w_len*pw]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len, pt, h_len*ph, w_len*pw]
-            x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N);  // [N*C, t_len*pt, h_len*ph, w_len*pw]
+            x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);            // [N, t_len*h_len*w_len, pt*ph*pw, C]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 1, 2, 0, 3));  // [N, C, t_len*h_len*w_len, pt*ph*pw]
+            x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);            // [N*C*t_len*h_len, w_len, pt*ph, pw]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len*h_len, pt*ph, w_len, pw]
+            x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);            // [N*C*t_len*h_len, pt, ph, w_len*pw]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len*h_len, ph, pt, w_len*pw]
+            x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N);            // [N*C*t_len, h_len*ph, pt, w_len*pw]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len, pt, h_len*ph, w_len*pw]
+            x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N);            // [N*C, t_len*pt, h_len*ph, w_len*pw]
             return x;
         }
 
@@ -431,10 +432,10 @@ namespace Hunyuan {
 
             img      = img_in->forward(ctx, img);                     // [N*C, t_len*h_len*w_len, hidden_size]
             txt      = txt_in->forward(ctx, txt, timestep, nullptr);  // [N, n_txt_token, hidden_size]
-            auto vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f));
+            auto vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx, timestep, 256, 10000, 1.f));
             if (config.use_meanflow && timestep_r != nullptr) {
                 auto time_r_in = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["time_r_in"]);
-                auto vec_r     = time_r_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_r, 256, 10000, 1000.f));
+                auto vec_r     = time_r_in->forward(ctx, ggml_ext_timestep_embedding(ctx, timestep_r, 256, 10000, 1000.f));
                 vec            = ggml_add(ctx->ggml_ctx, vec, vec_r);
                 if (!config.use_meanflow_sum) {
                     vec = ggml_scale(ctx->ggml_ctx, vec, 0.5f);
@@ -446,7 +447,7 @@ namespace Hunyuan {
             }
             if (config.guidance_embed && guidance != nullptr) {
                 auto guidance_in  = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["guidance_in"]);
-                auto guidance_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, guidance, 256, 10000, 1.f);
+                auto guidance_emb = ggml_ext_timestep_embedding(ctx, guidance, 256, 10000, 1.f);
                 vec               = ggml_add(ctx->ggml_ctx, vec, guidance_in->forward(ctx, guidance_emb));
             }
 
@@ -529,12 +530,12 @@ namespace Hunyuan {
 
             auto out = forward_orig(ctx, x, context, timestep, pe, guidance, y, txt_byt5, clip_fea, timestep_r, N);
 
-            out = unpatchify(ctx->ggml_ctx, out, t_len, h_len, w_len);  // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w]
+            out = unpatchify(ctx, out, t_len, h_len, w_len);  // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w]
 
             // slice
-            out = ggml_ext_slice(ctx->ggml_ctx, out, 2, 0, T);  // [N*C, T, H + pad_h, W + pad_w]
-            out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H);  // [N*C, T, H, W + pad_w]
-            out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W);  // [N*C, T, H, W]
+            out = ggml_ext_slice(ctx, out, 2, 0, T);  // [N*C, T, H + pad_h, W + pad_w]
+            out = ggml_ext_slice(ctx, out, 1, 0, H);  // [N*C, T, H, W + pad_w]
+            out = ggml_ext_slice(ctx, out, 0, 0, W);  // [N*C, T, H, W]
 
             return out;
         }

src/model/diffusion/ideogram4.hpp

@@ -73,12 +73,13 @@ namespace Ideogram4 {
         }
     };
 
-    __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(ggml_context* ctx,
+    __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(GGMLRunnerContext* runner_ctx,
                                                               ggml_tensor* timesteps,
                                                               int dim) {
+        auto ctx = runner_ctx->ggml_ctx;
         GGML_ASSERT(dim % 2 == 0);
-        auto embedding = ggml_ext_timestep_embedding(ctx, timesteps, dim, TIMESTEP_MAX_PERIOD, 10.f);
-        auto chunks    = ggml_ext_chunk(ctx, embedding, 2, 0);
+        auto embedding = ggml_ext_timestep_embedding(runner_ctx, timesteps, dim, TIMESTEP_MAX_PERIOD, 10.f);
+        auto chunks    = ggml_ext_chunk(runner_ctx, embedding, 2, 0);
         return ggml_concat(ctx, chunks[1], chunks[0], 0);
     }
 
@@ -108,7 +109,10 @@ namespace Ideogram4 {
         return ggml_ext_modulate(ctx, x, nullptr, scale);
     }
 
-    __STATIC_INLINE__ ggml_tensor* patchify(ggml_context* ctx, ggml_tensor* x, const Ideogram4Config& config) {
+    __STATIC_INLINE__ ggml_tensor* patchify(GGMLRunnerContext* runner_ctx,
+                                            ggml_tensor* x,
+                                            const Ideogram4Config& config) {
+        auto ctx = runner_ctx->ggml_ctx;
         // x: [N, 128, H, W] with channel order [ae, ph, pw].
         // return: [N, H*W, 128] with token channel order [ph, pw, ae].
         const int64_t W = x->ne[0];
@@ -126,11 +130,12 @@ namespace Ideogram4 {
         return x;
     }
 
-    __STATIC_INLINE__ ggml_tensor* unpatchify(ggml_context* ctx,
+    __STATIC_INLINE__ ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx,
                                               ggml_tensor* x,
                                               int64_t H,
                                               int64_t W,
                                               const Ideogram4Config& config) {
+        auto ctx        = runner_ctx->ggml_ctx;
         const int64_t C = x->ne[0];
         const int64_t N = x->ne[2];
 
@@ -217,7 +222,7 @@ namespace Ideogram4 {
             auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["o"]);
 
             auto qkv     = qkv_proj->forward(ctx, x);
-            auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
+            auto qkv_vec = ggml_ext_split_qkv(ctx, qkv);
             auto q       = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, n_token, N);
             auto k       = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, n_token, N);
             auto v       = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, n_token, N);
@@ -278,7 +283,7 @@ namespace Ideogram4 {
             auto adaln_modulation = std::dynamic_pointer_cast<Linear>(blocks["adaln_modulation"]);
 
             auto mod       = adaln_modulation->forward(ctx, adaln_input);
-            auto mods      = ggml_ext_chunk(ctx->ggml_ctx, mod, 4, 0);
+            auto mods      = ggml_ext_chunk(ctx, mod, 4, 0);
             auto scale_msa = mods[0];
             auto gate_msa  = to_token_modulation(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, mods[1]));
             auto scale_mlp = mods[2];
@@ -315,7 +320,7 @@ namespace Ideogram4 {
             auto mlp_in  = std::dynamic_pointer_cast<Linear>(blocks["mlp_in"]);
             auto mlp_out = std::dynamic_pointer_cast<Linear>(blocks["mlp_out"]);
 
-            x = timestep_embedding_sin_cos(ctx->ggml_ctx, x, static_cast<int>(dim));
+            x = timestep_embedding_sin_cos(ctx, x, static_cast<int>(dim));
             x = ggml_silu(ctx->ggml_ctx, mlp_in->forward(ctx, x));
             x = mlp_out->forward(ctx, x);
             return x;
@@ -383,7 +388,7 @@ namespace Ideogram4 {
             auto embed_image_indicator = std::dynamic_pointer_cast<Embedding>(blocks["embed_image_indicator"]);
             auto final_layer           = std::dynamic_pointer_cast<Ideogram4FinalLayer>(blocks["final_layer"]);
 
-            auto img = patchify(ctx->ggml_ctx, x, config);
+            auto img = patchify(ctx, x, config);
             img      = input_proj->forward(ctx, img);
 
             ggml_tensor* h      = img;
@@ -413,11 +418,11 @@ namespace Ideogram4 {
 
             h = final_layer->forward(ctx, h, adaln_input);
             if (context_len > 0) {
-                h = ggml_ext_slice(ctx->ggml_ctx, h, 1, context_len, h->ne[1]);
+                h = ggml_ext_slice(ctx, h, 1, context_len, h->ne[1]);
             }
 
-            h = unpatchify(ctx->ggml_ctx, h, H, W, config);
-            h = ggml_ext_scale(ctx->ggml_ctx, h, -1.f);
+            h = unpatchify(ctx, h, H, W, config);
+            h = ggml_ext_scale(ctx, h, -1.f);
             return h;
         }
     };

src/model/diffusion/iris.hpp

@@ -92,7 +92,7 @@ namespace Iris {
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe = nullptr, ggml_tensor* mask = nullptr) {
             auto c   = ctx->ggml_ctx;
-            auto qkv = split_qkv(c, std::dynamic_pointer_cast<Linear>(blocks["qkv"])->forward(ctx, x));
+            auto qkv = ggml_ext_split_qkv(ctx, std::dynamic_pointer_cast<Linear>(blocks["qkv"])->forward(ctx, x));
             auto q   = ggml_reshape_4d(c, qkv[0], dim / heads, heads, x->ne[1], x->ne[2]);
             auto k   = ggml_reshape_4d(c, qkv[1], dim / heads, heads, x->ne[1], x->ne[2]);
             auto v   = ggml_reshape_4d(c, qkv[2], dim / heads, heads, x->ne[1], x->ne[2]);
@@ -247,24 +247,24 @@ namespace Iris {
         std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* y, ggml_tensor* img_mod, ggml_tensor* txt_mod, ggml_tensor* pe) {
             auto c     = ctx->ggml_ctx;
             int64_t nt = y->ne[1];
-            auto mx    = ggml_ext_chunk(c, ggml_add(c, img_mod, params[dual ? "adaln_img.bias" : "adaln.bias"]), 6, 0);
+            auto mx    = ggml_ext_chunk(ctx, ggml_add(c, img_mod, params[dual ? "adaln_img.bias" : "adaln.bias"]), 6, 0);
             if (!dual) {
                 auto tokens = ggml_concat(c, y, x, 1);
                 auto h      = Pid::apply_adaln(c, norm(ctx, "norm1", tokens), mx[0], mx[1]);
                 auto qkv    = project(ctx, h, "");
                 auto out    = Rope::attention(ctx, qkv[0], qkv[1], qkv[2], pe, nullptr);
                 tokens      = finish(ctx, tokens, h, out, mx, "");
-                return {ggml_ext_slice(c, tokens, 1, nt, tokens->ne[1]), ggml_ext_slice(c, tokens, 1, 0, nt)};
+                return {ggml_ext_slice(ctx, tokens, 1, nt, tokens->ne[1]), ggml_ext_slice(ctx, tokens, 1, 0, nt)};
             }
-            auto my  = ggml_ext_chunk(c, ggml_add(c, txt_mod, params["adaln_txt.bias"]), 6, 0);
+            auto my  = ggml_ext_chunk(ctx, ggml_add(c, txt_mod, params["adaln_txt.bias"]), 6, 0);
             auto hx  = Pid::apply_adaln(c, norm(ctx, "norm_x1", x), mx[0], mx[1]);
             auto hy  = Pid::apply_adaln(c, norm(ctx, "norm_y1", y), my[0], my[1]);
             auto qx  = project(ctx, hx, "_x");
             auto qy  = project(ctx, hy, "_y");
             auto out = Rope::attention(ctx, ggml_concat(c, qy[0], qx[0], 2), ggml_concat(c, qy[1], qx[1], 2),
                                        ggml_concat(c, qy[2], qx[2], 2), pe, nullptr);
-            x        = finish(ctx, x, hx, ggml_ext_slice(c, out, 1, nt, out->ne[1]), mx, "_x");
-            y        = finish(ctx, y, hy, ggml_ext_slice(c, out, 1, 0, nt), my, "_y");
+            x        = finish(ctx, x, hx, ggml_ext_slice(ctx, out, 1, nt, out->ne[1]), mx, "_x");
+            y        = finish(ctx, y, hy, ggml_ext_slice(ctx, out, 1, 0, nt), my, "_y");
             return {x, y};
         }
     };
@@ -294,7 +294,7 @@ namespace Iris {
             auto batch   = cond->ne[2];
             auto mods    = std::dynamic_pointer_cast<Linear>(blocks["adaln"])->forward(ctx, cond);
             mods         = ggml_reshape_3d(c, mods, 4 * d, p2, patches * batch);
-            auto m       = ggml_ext_chunk(c, mods, 4, 0);
+            auto m       = ggml_ext_chunk(ctx, mods, 4, 0);
             auto h       = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"])->forward(ctx, x);
             h            = ggml_reshape_3d(c, h, d * p2, patches, batch);
             h            = std::dynamic_pointer_cast<Linear>(blocks["compress"])->forward(ctx, h);
@@ -337,7 +337,7 @@ namespace Iris {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image, ggml_tensor* t, ggml_tensor* text, ggml_tensor* mask, ggml_tensor* joint_pe, ggml_tensor* pixel_pe, ggml_tensor* pixel_pos) {
             auto c    = ctx->ggml_ctx;
             int p     = static_cast<int>(config.patch_size);
-            auto x    = DiT::patchify(c, image, p, p, true);
+            auto x    = DiT::patchify(ctx, image, p, p, true);
             x         = std::dynamic_pointer_cast<Pid::PatchTokenEmbedder>(blocks["s_embedder"])->forward(ctx, x);
             auto te   = std::dynamic_pointer_cast<Pid::PixelDiTTimestepEmbedder>(blocks["t_embedder"])->forward(ctx, t);
             te        = ggml_reshape_3d(c, te, config.hidden_size, 1, image->ne[3]);
@@ -360,7 +360,7 @@ namespace Iris {
             }
             x = std::dynamic_pointer_cast<Pid::FinalLayer>(blocks["final_layer"])->forward(ctx, x);
             x = ggml_reshape_3d(c, x, 3 * p * p, cond->ne[1], image->ne[3]);
-            return DiT::unpatchify(c, x, image->ne[1] / p, image->ne[0] / p, p, p, false);
+            return DiT::unpatchify(ctx, x, image->ne[1] / p, image->ne[0] / p, p, p, false);
         }
     };
 

src/model/diffusion/krea2.hpp

@@ -185,9 +185,9 @@ namespace Krea2 {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             ggml_tensor* scale = params["scale"];
             if (ctx->weight_adapter) {
-                scale = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, scale, prefix + "scale.weight");
+                scale = ctx->weight_adapter->patch_weight(ctx, scale, prefix + "scale.weight");
             }
-            scale = ggml_add(ctx->ggml_ctx, scale, ggml_ext_ones(ctx->ggml_ctx, scale->ne[0], 1, 1, 1));
+            scale = ggml_add(ctx->ggml_ctx, scale, ggml_ext_ones(ctx, scale->ne[0], 1, 1, 1));
             x     = ggml_rms_norm(ctx->ggml_ctx, x, eps);
             x     = ggml_mul_inplace(ctx->ggml_ctx, x, scale);
             return x;
@@ -273,9 +273,7 @@ namespace Krea2 {
             auto knorm = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["qknorm.knorm"]);
             auto wo    = std::dynamic_pointer_cast<Linear>(blocks["wo"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-                wo->set_force_prec_f32(true);
-            }
+            wo->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
 
             int64_t L = x->ne[1];
             int64_t N = x->ne[2];
@@ -316,11 +314,11 @@ namespace Krea2 {
         std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
             auto lin = params["lin"];
             if (ctx->weight_adapter) {
-                lin = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, lin, prefix + "lin.weight");
+                lin = ctx->weight_adapter->patch_weight(ctx, lin, prefix + "lin.weight");
             }
             lin      = ggml_repeat(ctx->ggml_ctx, lin, vec);
             auto out = ggml_add(ctx->ggml_ctx, vec, lin);
-            return ggml_ext_chunk(ctx->ggml_ctx, out, 6, 0);
+            return ggml_ext_chunk(ctx, out, 6, 0);
         }
     };
 
@@ -342,10 +340,10 @@ namespace Krea2 {
         std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
             auto lin = params["lin"];
             if (ctx->weight_adapter) {
-                lin = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, lin, prefix + "lin.weight");
+                lin = ctx->weight_adapter->patch_weight(ctx, lin, prefix + "lin.weight");
             }
             auto out = ggml_add(ctx->ggml_ctx, lin, vec);
-            return ggml_ext_chunk(ctx->ggml_ctx, out, 2, 1);
+            return ggml_ext_chunk(ctx, out, 2, 1);
         }
     };
 
@@ -541,7 +539,7 @@ namespace Krea2 {
             auto linear_0 = std::dynamic_pointer_cast<Linear>(blocks["0"]);
             auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
             x             = linear_0->forward(ctx, x);
-            x             = ggml_ext_gelu(ctx->ggml_ctx, x, false);
+            x             = ggml_ext_gelu(ctx, x, false);
             x             = linear_2->forward(ctx, x);
             return x;
         }
@@ -555,7 +553,7 @@ namespace Krea2 {
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["1"]);
-            x             = ggml_ext_gelu(ctx->ggml_ctx, x, false);
+            x             = ggml_ext_gelu(ctx, x, false);
             x             = linear_1->forward(ctx, x);
             return x;
         }
@@ -575,7 +573,7 @@ namespace Krea2 {
             auto linear_3 = std::dynamic_pointer_cast<Linear>(blocks["3"]);
             x             = norm->forward(ctx, x);
             x             = linear_1->forward(ctx, x);
-            x             = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+            x             = ggml_ext_gelu(ctx, x, true);
             x             = linear_3->forward(ctx, x);
             return x;
         }
@@ -656,7 +654,7 @@ namespace Krea2 {
             int64_t ref_len = img->ne[1] - img_len;
             img             = first->forward(ctx, img);
 
-            auto t    = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.timestep_dim), 10000, 1000.f);
+            auto t    = ggml_ext_timestep_embedding(ctx, timestep, static_cast<int>(config.timestep_dim), 10000, 1000.f);
             t         = tmlp->forward(ctx, t);
             t         = ggml_reshape_3d(ctx->ggml_ctx, t, t->ne[0], 1, t->ne[1]);
             auto tvec = tproj->forward(ctx, t);
@@ -665,7 +663,7 @@ namespace Krea2 {
             if (ref_latents.size() > 0 && zero_timestep_refs) {
                 // "index_timestep_zero" mode: use timestep = 0 for ref latents
                 auto timestep_0 = ggml_scale(ctx->ggml_ctx, timestep, 0.0f);
-                auto t_0        = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_0, static_cast<int>(config.timestep_dim), 10000, 1000.f);
+                auto t_0        = ggml_ext_timestep_embedding(ctx, timestep_0, static_cast<int>(config.timestep_dim), 10000, 1000.f);
                 t_0             = tmlp->forward(ctx, t_0);
                 t_0             = ggml_reshape_3d(ctx->ggml_ctx, t_0, t_0->ne[0], 1, t_0->ne[1]);
                 tvec_0          = tproj->forward(ctx, t_0);
@@ -683,9 +681,9 @@ namespace Krea2 {
                 sd::ggml_graph_cut::mark_graph_cut(hidden_states, "krea2.blocks." + std::to_string(i), "hidden_states");
             }
 
-            hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, txt_len, txt_len + img_len);
+            hidden_states = ggml_ext_slice(ctx, hidden_states, 1, txt_len, txt_len + img_len);
             hidden_states = last->forward(ctx, hidden_states, t);
-            hidden_states = DiT::unpatchify_and_crop(ctx->ggml_ctx, hidden_states, H, W, config.patch_size, config.patch_size, true);
+            hidden_states = DiT::unpatchify_and_crop(ctx, hidden_states, H, W, config.patch_size, config.patch_size, true);
             return hidden_states;
         }
     };

src/model/diffusion/lens.hpp

@@ -87,7 +87,7 @@ namespace Lens {
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timesteps) {
             auto timestep_embedder = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["timestep_embedder"]);
-            auto timesteps_proj    = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1000.f);
+            auto timesteps_proj    = ggml_ext_timestep_embedding(ctx, timesteps, 256, 10000, 1000.f);
             return timestep_embedder->forward(ctx, timesteps_proj);
         }
     };
@@ -151,8 +151,8 @@ namespace Lens {
             int64_t N       = img->ne[2];
             int64_t inner   = dim_head * num_heads;
 
-            auto img_qkv_vec = split_qkv(ctx->ggml_ctx, img_qkv->forward(ctx, img));
-            auto txt_qkv_vec = split_qkv(ctx->ggml_ctx, txt_qkv->forward(ctx, txt));
+            auto img_qkv_vec = ggml_ext_split_qkv(ctx, img_qkv->forward(ctx, img));
+            auto txt_qkv_vec = ggml_ext_split_qkv(ctx, txt_qkv->forward(ctx, txt));
 
             auto img_q = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[0], dim_head, num_heads, n_img, N);
             auto img_k = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[1], dim_head, num_heads, n_img, N);
@@ -232,9 +232,9 @@ namespace Lens {
             auto temb = ggml_silu(ctx->ggml_ctx, t_emb);
 
             auto img_mod_params = img_mod_1->forward(ctx, temb);
-            auto img_mod_vec    = ggml_ext_chunk(ctx->ggml_ctx, img_mod_params, 6, 0);
+            auto img_mod_vec    = ggml_ext_chunk(ctx, img_mod_params, 6, 0);
             auto txt_mod_params = txt_mod_1->forward(ctx, temb);
-            auto txt_mod_vec    = ggml_ext_chunk(ctx->ggml_ctx, txt_mod_params, 6, 0);
+            auto txt_mod_vec    = ggml_ext_chunk(ctx, txt_mod_params, 6, 0);
 
             auto img_normed    = img_norm1->forward(ctx, img);
             auto img_modulated = Flux::modulate(ctx, img_normed, img_mod_vec[0], img_mod_vec[1]);
@@ -268,7 +268,7 @@ namespace Lens {
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* conditioning) {
             auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
-            auto mods   = ggml_ext_chunk(ctx->ggml_ctx, linear->forward(ctx, ggml_silu(ctx->ggml_ctx, conditioning)), 2, 0);
+            auto mods   = ggml_ext_chunk(ctx, linear->forward(ctx, ggml_silu(ctx->ggml_ctx, conditioning)), 2, 0);
             auto scale  = mods[0];
             auto shift  = mods[1];
             x           = ggml_norm(ctx->ggml_ctx, x, eps);
@@ -319,10 +319,10 @@ namespace Lens {
             auto t_emb = time_text_embed->forward(ctx, timestep);
 
             auto img = ggml_reshape_3d(ctx->ggml_ctx, x, W * H, C, N);
-            img      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));
+            img      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, img, 1, 0, 2, 3));
             img      = img_in->forward(ctx, img);
 
-            std::vector<ggml_tensor*> txt_chunks = ggml_ext_chunk(ctx->ggml_ctx, context, config.selected_layer_count, 0);
+            std::vector<ggml_tensor*> txt_chunks = ggml_ext_chunk(ctx, context, config.selected_layer_count, 0);
             ggml_tensor* txt                     = nullptr;
             for (int i = 0; i < config.selected_layer_count; ++i) {
                 auto txt_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm." + std::to_string(i)]);
@@ -346,7 +346,7 @@ namespace Lens {
             img = norm_out->forward(ctx, img, t_emb);
             img = proj_out->forward(ctx, img);
 
-            auto out = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));
+            auto out = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, img, 1, 0, 2, 3));
             out      = ggml_reshape_4d(ctx->ggml_ctx, out, W, H, config.patch_size * config.patch_size * config.out_channels, N);
             return out;
         }

src/model/diffusion/lingbot_video.hpp

@@ -365,7 +365,7 @@ namespace LingBotVideo {
             }
             selected_mask = ggml_reshape_2d(gctx, selected_mask, num_experts, n_token_total);
 
-            ggml_tensor* excluded_group_mask = ggml_sub(gctx, selected_mask, ggml_ext_ones_like(gctx, selected_mask));
+            ggml_tensor* excluded_group_mask = ggml_sub(gctx, selected_mask, ggml_ext_ones_like(ctx, selected_mask));
             ggml_tensor* masked_scores       = ggml_add(gctx, choice_scores, ggml_scale(gctx, excluded_group_mask, 1.0e9f));
             return ggml_argsort_top_k(gctx, masked_scores, static_cast<int>(num_experts_per_tok));
         }
@@ -476,7 +476,7 @@ namespace LingBotVideo {
 
             auto table = ggml_reshape_3d(ctx->ggml_ctx, params["scale_shift_table"], hidden_size * 6, 1, 1);
             auto mod   = ggml_add(ctx->ggml_ctx, temb6, table);
-            auto mods  = ggml_ext_chunk(ctx->ggml_ctx, mod, 6, 0);
+            auto mods  = ggml_ext_chunk(ctx, mod, 6, 0);
 
             auto shift_msa = mods[0];
             auto scale_msa = mods[1];
@@ -556,13 +556,13 @@ namespace LingBotVideo {
             int64_t w_len = W / config.patch_w;
             int64_t n_img = t_len * h_len * w_len;
 
-            auto img = DiT::patchify_3d(ctx->ggml_ctx, x, config.patch_t, config.patch_h, config.patch_w, N, false);
+            auto img = DiT::patchify_3d(ctx, x, config.patch_t, config.patch_h, config.patch_w, N, false);
             img      = patch_embedder->forward(ctx, img);
 
             auto txt           = text_embedder->forward(ctx, context);
             auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1);
 
-            auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx,
+            auto timestep_proj = ggml_ext_timestep_embedding(ctx,
                                                              timestep,
                                                              static_cast<int>(config.freq_dim),
                                                              10000,
@@ -582,16 +582,16 @@ namespace LingBotVideo {
                 sd::ggml_graph_cut::mark_graph_cut(hidden_states, "lingbot_video.blocks." + std::to_string(i), "hidden_states");
             }
 
-            auto final_mods = ggml_ext_chunk(ctx->ggml_ctx,
+            auto final_mods = ggml_ext_chunk(ctx,
                                              norm_out_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, temb_tokens)),
                                              2,
                                              0);
             hidden_states   = norm_out->forward(ctx, hidden_states);
             hidden_states   = Flux::modulate(ctx, hidden_states, final_mods[0], final_mods[1], true);
             hidden_states   = proj_out->forward(ctx, hidden_states);
-            hidden_states   = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img);
+            hidden_states   = ggml_ext_slice(ctx, hidden_states, 1, 0, n_img);
 
-            auto out = DiT::unpatchify_3d(ctx->ggml_ctx,
+            auto out = DiT::unpatchify_3d(ctx,
                                           hidden_states,
                                           t_len,
                                           h_len,

src/model/diffusion/llada_image.hpp

@@ -217,8 +217,8 @@ namespace LLaDAImage {
 
             GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]);
 
-            auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]);
-            auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]);
+            auto txt_pe = ggml_ext_slice(ctx, pe, 3, 0, txt->ne[1]);
+            auto img_pe = ggml_ext_slice(ctx, pe, 3, txt->ne[1], pe->ne[3]);
 
             for (int i = 0; i < config.num_refiner_layers; i++) {
                 auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
@@ -246,7 +246,7 @@ namespace LLaDAImage {
 
             txt_img = final_layer->forward(ctx, txt_img, t_emb);
 
-            return ggml_ext_slice(ctx->ggml_ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token);
+            return ggml_ext_slice(ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token);
         }
 
         ggml_tensor* pad_stream(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pad_token) {
@@ -308,8 +308,8 @@ namespace LLaDAImage {
 
             GGML_ASSERT(cap_len * 2 + img_len * 2 + sig_len == pe->ne[3]);
 
-            auto cap_pe = ggml_ext_slice(gctx, pe, 3, 0, cap_len * 2);
-            auto img_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2, cap_len * 2 + img_len * 2);
+            auto cap_pe = ggml_ext_slice(ctx, pe, 3, 0, cap_len * 2);
+            auto img_pe = ggml_ext_slice(ctx, pe, 3, cap_len * 2, cap_len * 2 + img_len * 2);
 
             auto img_adaln = ggml_concat(gctx, per_token(t_clean, img_len), per_token(t_noisy, img_len), 1);
 
@@ -322,7 +322,7 @@ namespace LLaDAImage {
                 img        = block->forward(ctx, img, img_pe, nullptr, img_adaln);
             }
             if (sig != nullptr) {
-                auto sig_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2 + img_len * 2, pe->ne[3]);
+                auto sig_pe = ggml_ext_slice(ctx, pe, 3, cap_len * 2 + img_len * 2, pe->ne[3]);
                 for (int i = 0; i < config.num_refiner_layers; i++) {
                     auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["sigvq_refiner." + std::to_string(i)]);
                     sig        = block->forward(ctx, sig, sig_pe, nullptr, nullptr);
@@ -350,7 +350,7 @@ namespace LLaDAImage {
             // The stream is padded to SEQ_MULTI_OF, so drop the pad tokens: they are not part of
             // the latent grid that unpatchify reconstructs.
             int64_t target_start = cap_len * 2 + img_len;
-            return ggml_ext_slice(gctx, seq, 1, target_start, target_start + n_img_token);
+            return ggml_ext_slice(ctx, seq, 1, target_start, target_start + n_img_token);
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx,
@@ -372,10 +372,10 @@ namespace LLaDAImage {
 
             auto out = forward_core(ctx, img, timestep, context, pe);
 
-            out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size, false);
+            out = DiT::unpatchify_and_crop(ctx, out, H, W, patch_size, patch_size, false);
 
             // The reference pipeline negates the model output before the scheduler step.
-            return ggml_ext_scale(ctx->ggml_ctx, out, -1.f);
+            return ggml_ext_scale(ctx, out, -1.f);
         }
     };
 
@@ -488,8 +488,8 @@ namespace LLaDAImage {
             auto src    = DiT::pad_and_patchify(&runner_ctx, source, config.patch_size, config.patch_size, false);
 
             auto out = llada_image.forward_editing(&runner_ctx, target, timesteps, context, semantic, src, pe);
-            out      = DiT::unpatchify_and_crop(runner_ctx.ggml_ctx, out, H, W, config.patch_size, config.patch_size, false);
-            out      = ggml_ext_scale(runner_ctx.ggml_ctx, out, -1.f);
+            out      = DiT::unpatchify_and_crop(&runner_ctx, out, H, W, config.patch_size, config.patch_size, false);
+            out      = ggml_ext_scale(&runner_ctx, out, -1.f);
 
             ggml_build_forward_expand(gf, out);
             return gf;

src/model/diffusion/ltxv.hpp

@@ -581,7 +581,7 @@ namespace LTXV {
             auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
             auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
 
-            auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, frequency_embedding_size);
+            auto t_emb = ggml_ext_timestep_embedding(ctx, timestep, frequency_embedding_size);
             t_emb      = linear_1->forward(ctx, t_emb);
             t_emb      = ggml_silu_inplace(ctx->ggml_ctx, t_emb);
             t_emb      = linear_2->forward(ctx, t_emb);
@@ -631,7 +631,7 @@ namespace LTXV {
             auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
 
             caption = linear_1->forward(ctx, caption);
-            caption = ggml_ext_gelu(ctx->ggml_ctx, caption, true);
+            caption = ggml_ext_gelu(ctx, caption, true);
             caption = linear_2->forward(ctx, caption);
             return caption;
         }
@@ -650,7 +650,7 @@ namespace LTXV {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* caption) {
             auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
             caption       = ggml_rms_norm(ctx->ggml_ctx, caption, 1e-6f);
-            caption       = ggml_ext_scale(ctx->ggml_ctx, caption, std::sqrt(static_cast<float>(hidden_size) / static_cast<float>(in_features)));
+            caption       = ggml_ext_scale(ctx, caption, std::sqrt(static_cast<float>(hidden_size) / static_cast<float>(in_features)));
             return linear_1->forward(ctx, caption);
         }
     };
@@ -724,7 +724,7 @@ namespace LTXV {
                 auto to_gate_logits = std::dynamic_pointer_cast<Linear>(blocks["to_gate_logits"]);
                 auto gate_logits    = to_gate_logits->forward(ctx, x);
                 auto gates          = ggml_sigmoid(ctx->ggml_ctx, gate_logits);
-                gates               = ggml_ext_scale(ctx->ggml_ctx, gates, 2.0f, true);
+                gates               = ggml_ext_scale(ctx, gates, 2.0f, true);
                 gates               = ggml_reshape_4d(ctx->ggml_ctx, gates, 1, heads, gate_logits->ne[1], gate_logits->ne[2]);
 
                 auto out4 = ggml_reshape_4d(ctx->ggml_ctx, out, dim_head, heads, out->ne[1], out->ne[2]);
@@ -783,7 +783,7 @@ namespace LTXV {
             s             = ggml_repeat(ctx->ggml_ctx, s, e);
             t             = ggml_repeat(ctx->ggml_ctx, t, e);
             auto out      = ggml_add(ctx->ggml_ctx, s, t);
-            return ggml_ext_chunk(ctx->ggml_ctx, out, static_cast<int>(coeff), 1);
+            return ggml_ext_chunk(ctx, out, static_cast<int>(coeff), 1);
         }
 
         std::vector<ggml_tensor*> get_prompt_scale_shift_values(GGMLRunnerContext* ctx,
@@ -797,7 +797,7 @@ namespace LTXV {
             s        = ggml_repeat(ctx->ggml_ctx, s, e);
             t        = ggml_repeat(ctx->ggml_ctx, t, e);
             auto out = ggml_add(ctx->ggml_ctx, s, t);
-            return ggml_ext_chunk(ctx->ggml_ctx, out, 2, 1);
+            return ggml_ext_chunk(ctx, out, 2, 1);
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx,
@@ -945,7 +945,7 @@ namespace LTXV {
             for (int64_t i = 1; i < duplications; i++) {
                 regs_full = ggml_concat(ctx->ggml_ctx, regs_full, regs, 1);
             }
-            regs_full = ggml_ext_slice(ctx->ggml_ctx, regs_full, 1, seq_len, seq_len + total_to_keep);
+            regs_full = ggml_ext_slice(ctx, regs_full, 1, seq_len, seq_len + total_to_keep);
             return ggml_concat(ctx->ggml_ctx, hidden_states, regs_full, 1);
         }
 
@@ -1182,7 +1182,7 @@ namespace LTXV {
             t           = ggml_repeat(ctx->ggml_ctx, t, e);
             s           = ggml_repeat(ctx->ggml_ctx, s, e);
             auto out    = ggml_add(ctx->ggml_ctx, s, t);
-            auto chunks = ggml_ext_chunk(ctx->ggml_ctx, out, static_cast<int>(coeff), 1);
+            auto chunks = ggml_ext_chunk(ctx, out, static_cast<int>(coeff), 1);
             return std::vector<ggml_tensor*>(chunks.begin() + start, chunks.begin() + start + count);
         }
 
@@ -1287,27 +1287,27 @@ namespace LTXV {
                 auto ax_norm3 = rms_norm(ctx->ggml_ctx, ax);
 
                 if (run_a2v) {
-                    auto a2v_audio_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_audio"], 1, 0, 4);
-                    auto a2v_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
+                    auto a2v_audio_table = ggml_ext_slice(ctx, params["scale_shift_table_a2v_ca_audio"], 1, 0, 4);
+                    auto a2v_video_table = ggml_ext_slice(ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
                     auto a2v_audio       = get_ada_values(ctx, a2v_audio_table, a_cross_scale_shift_timestep, a_dim, 4);
                     auto a2v_video       = get_ada_values(ctx, a2v_video_table, v_cross_scale_shift_timestep, v_dim, 4);
                     auto vx_scaled       = LTXV::modulate(ctx, vx_norm3, a2v_video[1], a2v_video[0]);
                     auto ax_scaled       = LTXV::modulate(ctx, ax_norm3, a2v_audio[1], a2v_audio[0]);
                     auto a2v_out         = audio_to_video_attn->forward(ctx, vx_scaled, ax_scaled, nullptr, v_cross_pe, a_cross_pe);
-                    auto a2v_gate_table  = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 4, 5);
+                    auto a2v_gate_table  = ggml_ext_slice(ctx, params["scale_shift_table_a2v_ca_video"], 1, 4, 5);
                     auto a2v_gate        = get_ada_values(ctx, a2v_gate_table, v_cross_gate_timestep, v_dim, 1)[0];
                     vx                   = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, a2v_out, a2v_gate));
                 }
 
                 if (run_v2a) {
-                    auto v2a_audio_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_audio"], 1, 0, 4);
-                    auto v2a_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
+                    auto v2a_audio_table = ggml_ext_slice(ctx, params["scale_shift_table_a2v_ca_audio"], 1, 0, 4);
+                    auto v2a_video_table = ggml_ext_slice(ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
                     auto v2a_audio       = get_ada_values(ctx, v2a_audio_table, a_cross_scale_shift_timestep, a_dim, 4);
                     auto v2a_video       = get_ada_values(ctx, v2a_video_table, v_cross_scale_shift_timestep, v_dim, 4);
                     auto ax_scaled       = LTXV::modulate(ctx, ax_norm3, v2a_audio[3], v2a_audio[2]);
                     auto vx_scaled       = LTXV::modulate(ctx, vx_norm3, v2a_video[3], v2a_video[2]);
                     auto v2a_out         = video_to_audio_attn->forward(ctx, ax_scaled, vx_scaled, nullptr, a_cross_pe, v_cross_pe);
-                    auto v2a_gate_table  = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_audio"], 1, 4, 5);
+                    auto v2a_gate_table  = ggml_ext_slice(ctx, params["scale_shift_table_a2v_ca_audio"], 1, 4, 5);
                     auto v2a_gate        = get_ada_values(ctx, v2a_gate_table, a_cross_gate_timestep, a_dim, 1)[0];
                     ax                   = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, v2a_out, v2a_gate));
                 }
@@ -1427,7 +1427,7 @@ namespace LTXV {
 
         ggml_tensor* patchify_video(GGMLRunnerContext* ctx, ggml_tensor* x, int64_t n) {
             x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], x->ne[3] / n, n);
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3));
             return x;
         }
 
@@ -1436,14 +1436,14 @@ namespace LTXV {
                                       int64_t width,
                                       int64_t height,
                                       int64_t frames) {
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3));
             x = ggml_reshape_4d(ctx->ggml_ctx, x, width, height, frames, x->ne[1] * x->ne[2]);
             return x;
         }
 
         ggml_tensor* patchify_audio(GGMLRunnerContext* ctx, ggml_tensor* ax) {
             // ax: [b, c, t, f]
-            ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ax, 0, 2, 1, 3));  // [b, t, c, f]
+            ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, ax, 0, 2, 1, 3));            // [b, t, c, f]
             ax = ggml_reshape_3d(ctx->ggml_ctx, ax, ax->ne[0] * ax->ne[1], ax->ne[2], ax->ne[3]);  // [b, t, c*f]
             return ax;
         }
@@ -1462,7 +1462,7 @@ namespace LTXV {
                 return nullptr;
             }
             ax = ggml_reshape_4d(ctx->ggml_ctx, ax, config.audio_frequency_bins, config.num_audio_channels, audio_length, ax->ne[2]);  // [b, t, c, f]
-            ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ax, 0, 2, 1, 3));                                      // [b, c, t, f]
+            ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, ax, 0, 2, 1, 3));                                                // [b, c, t, f]
             return ax;
         }
 
@@ -1483,10 +1483,10 @@ namespace LTXV {
                 context->ne[1] < 1024;
 
             if (is_fully_processed_context) {
-                auto v_context         = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.cross_attention_dim);
+                auto v_context         = ggml_ext_slice(ctx, context, 0, 0, config.cross_attention_dim);
                 ggml_tensor* a_context = nullptr;
                 if (process_audio_context) {
-                    a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim);
+                    a_context = ggml_ext_slice(ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim);
                 }
                 return {v_context, a_context};
             }
@@ -1494,14 +1494,14 @@ namespace LTXV {
             ggml_tensor* v_context = context;
             ggml_tensor* a_context = process_audio_context ? context : nullptr;
             if (is_unprocessed_dual_context) {
-                v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.cross_attention_dim);
+                v_context = ggml_ext_slice(ctx, context, 0, 0, config.cross_attention_dim);
                 if (process_audio_context) {
-                    a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim);
+                    a_context = ggml_ext_slice(ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim);
                 }
             } else if (context->ne[0] == config.caption_channels * 2) {
-                v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.caption_channels);
+                v_context = ggml_ext_slice(ctx, context, 0, 0, config.caption_channels);
                 if (process_audio_context) {
-                    a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.caption_channels, config.caption_channels * 2);
+                    a_context = ggml_ext_slice(ctx, context, 0, config.caption_channels, config.caption_channels * 2);
                 }
             }
 
@@ -1604,7 +1604,7 @@ namespace LTXV {
             auto t    = ggml_repeat(ctx->ggml_ctx, ggml_reshape_3d(ctx->ggml_ctx, embedded_timestep, dim, 1, embedded_timestep->ne[1]), temp);
             auto s    = ggml_repeat(ctx->ggml_ctx, ggml_reshape_3d(ctx->ggml_ctx, table, dim, 2, 1), temp);
             auto out  = ggml_add(ctx->ggml_ctx, s, t);
-            return ggml_ext_chunk(ctx->ggml_ctx, out, 2, 1);
+            return ggml_ext_chunk(ctx, out, 2, 1);
         }
 
         std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
@@ -1653,13 +1653,13 @@ namespace LTXV {
                 a_context = ggml_cont(ctx->ggml_ctx, a_context);
             }
 
-            auto v_timestep_scaled = ggml_ext_scale(ctx->ggml_ctx, timestep, config.timestep_scale_multiplier);
+            auto v_timestep_scaled = ggml_ext_scale(ctx, timestep, config.timestep_scale_multiplier);
             auto v_pair            = adaln_single->forward(ctx, v_timestep_scaled);
             auto v_timestep_mod    = v_pair.first;
             auto v_embedded_time   = v_pair.second;
 
             ggml_tensor* effective_audio_timestep = audio_timestep != nullptr ? audio_timestep : timestep;
-            auto a_timestep_scaled                = ggml_ext_scale(ctx->ggml_ctx, effective_audio_timestep, config.timestep_scale_multiplier);
+            auto a_timestep_scaled                = ggml_ext_scale(ctx, effective_audio_timestep, config.timestep_scale_multiplier);
             auto a_pair                           = audio_adaln_single->forward(ctx, a_timestep_scaled);
             auto a_timestep_mod                   = a_pair.first;
             auto a_embedded_time                  = a_pair.second;
@@ -1680,13 +1680,13 @@ namespace LTXV {
                 std::dynamic_pointer_cast<AdaLayerNormSingle>(blocks["av_ca_video_scale_shift_adaln_single"])->forward(ctx, av_ca_video_timestep).first;
             auto av_ca_a2v_gate_noise_timestep =
                 std::dynamic_pointer_cast<AdaLayerNormSingle>(blocks["av_ca_a2v_gate_adaln_single"])
-                    ->forward(ctx, ggml_ext_scale(ctx->ggml_ctx, av_ca_video_timestep, av_ca_factor))
+                    ->forward(ctx, ggml_ext_scale(ctx, av_ca_video_timestep, av_ca_factor))
                     .first;
             auto av_ca_audio_scale_shift_timestep =
                 std::dynamic_pointer_cast<AdaLayerNormSingle>(blocks["av_ca_audio_scale_shift_adaln_single"])->forward(ctx, av_ca_audio_timestep).first;
             auto av_ca_v2a_gate_noise_timestep =
                 std::dynamic_pointer_cast<AdaLayerNormSingle>(blocks["av_ca_v2a_gate_adaln_single"])
-                    ->forward(ctx, ggml_ext_scale(ctx->ggml_ctx, av_ca_audio_timestep, av_ca_factor))
+                    ->forward(ctx, ggml_ext_scale(ctx, av_ca_audio_timestep, av_ca_factor))
                     .first;
 
             sd::ggml_graph_cut::mark_graph_cut(vx, "ltxav.prelude", "vx");
@@ -1797,9 +1797,8 @@ namespace LTXV {
             return {vx, ax};
         }
 
-        ggml_tensor* merge_av_latents(ggml_context* ctx,
-                                      ggml_tensor* vx,
-                                      ggml_tensor* ax) const {
+        ggml_tensor* merge_av_latents(GGMLRunnerContext* runner_ctx, ggml_tensor* vx, ggml_tensor* ax) const {
+            auto ctx = runner_ctx->ggml_ctx;
             if (ax == nullptr || ggml_nelements(ax) == 0 || ax->ne[1] == 0) {
                 return vx;
             }
@@ -1815,7 +1814,7 @@ namespace LTXV {
             ax = ggml_cont(ctx, ax);
             ax = ggml_reshape_4d(ctx, ax, audio_values, 1, 1, 1);
             if (pad_values > 0) {
-                ax = ggml_ext_pad(ctx, ax, static_cast<int>(pad_values), 0, 0, 0);
+                ax = ggml_ext_pad(runner_ctx, ax, static_cast<int>(pad_values), 0, 0, 0);
             }
             int64_t extra_channels = padded_len / divisor;
             ax                     = ggml_reshape_4d(ctx, ax, width, height, frames, extra_channels);
@@ -2008,7 +2007,7 @@ namespace LTXV {
                                             audio_cross_pe,
                                             video_connector_pe,
                                             audio_connector_pe);
-            auto out        = merge_av_latents(compute_ctx, out_pair.first, out_pair.second);
+            auto out        = merge_av_latents(&runner_ctx, out_pair.first, out_pair.second);
             ggml_build_forward_expand(gf, out);
             return gf;
         }

src/model/diffusion/ming_image.hpp

@@ -50,8 +50,8 @@ namespace MingImage {
             txt                 = ggml_concat(gctx, txt, direct, 1);
             auto t              = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"])->forward(ctx, timestep);
             const int64_t n_txt = txt->ne[1], n_img = img->ne[1];
-            auto txt_pe = ggml_ext_slice(gctx, pe, 3, 0, n_txt);
-            auto img_pe = ggml_ext_slice(gctx, pe, 3, n_txt, n_txt + n_img);
+            auto txt_pe = ggml_ext_slice(ctx, pe, 3, 0, n_txt);
+            auto img_pe = ggml_ext_slice(ctx, pe, 3, n_txt, n_txt + n_img);
             for (int64_t i = 0; i < config.num_refiner_layers; ++i) {
                 txt = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)])->forward(ctx, txt, txt_pe);
                 sd::ggml_graph_cut::mark_graph_cut(txt, "ming_image.context_refiner." + std::to_string(i), "txt");
@@ -65,9 +65,9 @@ namespace MingImage {
                 combined = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["layers." + std::to_string(i)])->forward(ctx, combined, pe, nullptr, t);
                 sd::ggml_graph_cut::mark_graph_cut(combined, "ming_image.layers." + std::to_string(i), "combined");
             }
-            img = ggml_ext_slice(gctx, combined, 1, n_txt, n_txt + n_img);
+            img = ggml_ext_slice(ctx, combined, 1, n_txt, n_txt + n_img);
             img = std::dynamic_pointer_cast<ZImage::FinalLayer>(blocks["final_layer"])->forward(ctx, img, t);
-            img = DiT::unpatchify_and_crop(gctx, img, height, width, config.patch_size, config.patch_size, false);
+            img = DiT::unpatchify_and_crop(ctx, img, height, width, config.patch_size, config.patch_size, false);
             return ggml_scale(gctx, img, -1.f);
         }
     };

src/model/diffusion/minimax_h3.hpp

@@ -263,13 +263,13 @@ namespace MiniMaxH3 {
 
             std::vector<ggml_tensor*> chunks;
             for (int64_t start = 0; start < x->ne[1]; start += chunk_size) {
-                auto input = ggml_ext_slice(ctx->ggml_ctx, x, 1, start, std::min(start + chunk_size, x->ne[1]), false);
+                auto input = ggml_ext_slice(ctx, x, 1, start, std::min(start + chunk_size, x->ne[1]), false);
                 if (!ggml_is_contiguous(input)) {
                     input = ggml_cont(ctx->ggml_ctx, input);
                 }
                 chunks.push_back(forward_chunk(input));
             }
-            return ggml_ext_vec_concat(ctx->ggml_ctx, chunks, 1);
+            return ggml_ext_vec_concat(ctx, chunks, 1);
         }
     };
 
@@ -278,29 +278,8 @@ namespace MiniMaxH3 {
         return ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2] * x->ne[3]);
     }
 
-    static ggml_tensor* apply_partial_rope(GGMLRunnerContext* runner_ctx,
-                                           ggml_tensor* x,
-                                           ggml_tensor* pe) {
-        auto ctx        = runner_ctx->ggml_ctx;
-        int64_t rot_dim = pe->ne[2] * 2;
-        GGML_ASSERT(rot_dim <= x->ne[0]);
-#ifndef SD_USE_UPSTREAM_GGML
-        if (runner_ctx->backend != nullptr && x->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32) {
-            auto out = sd_ggml_rope_apply(ctx, runner_ctx->backend, x, pe, false);
-            if (out && ggml_backend_supports_op(runner_ctx->backend, out)) {
-                return out;
-            }
-        }
-#endif
-        auto rotated = Rope::apply_rope(runner_ctx,
-                                        ggml_ext_slice(ctx, x, 0, 0, rot_dim),
-                                        pe,
-                                        false);
-        if (rot_dim == x->ne[0]) {
-            return rotated;
-        }
-        auto tail = attention_layout(ctx, ggml_ext_slice(ctx, x, 0, rot_dim, x->ne[0]));
-        return ggml_concat(ctx, rotated, tail, 0);
+    static ggml_tensor* apply_partial_rope(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* pe) {
+        return ggml_ext_rope_apply(runner_ctx, x, pe, false);
     }
 
     struct Attention : public GGMLBlock {
@@ -381,7 +360,7 @@ namespace MiniMaxH3 {
             for (size_t i = 0; i < inputs.size(); ++i) {
                 auto input = inputs[i];
                 for (int64_t start = 0; start < input->ne[1]; start += chunk_size) {
-                    auto part = ggml_ext_slice(ctx->ggml_ctx, input, 1, start, std::min(start + chunk_size, input->ne[1]), false);
+                    auto part = ggml_ext_slice(ctx, input, 1, start, std::min(start + chunk_size, input->ne[1]), false);
                     if (!ggml_is_contiguous(part)) {
                         part = ggml_cont(ctx->ggml_ctx, part);
                     }
@@ -397,7 +376,7 @@ namespace MiniMaxH3 {
             }
             std::vector<ggml_tensor*> result;
             for (auto& parts : chunks) {
-                result.push_back(ggml_ext_vec_concat(ctx->ggml_ctx, parts, 2));
+                result.push_back(ggml_ext_vec_concat(ctx, parts, 2));
             }
             return result;
         }
@@ -420,13 +399,13 @@ namespace MiniMaxH3 {
                 for (int64_t start = 0; start < heads; start += head_chunk) {
                     std::vector<ggml_tensor*> part;
                     for (auto tensor : qkv) {
-                        part.push_back(ggml_ext_slice(ctx->ggml_ctx, tensor, 1, start, std::min(start + head_chunk, heads), false));
+                        part.push_back(ggml_ext_slice(ctx, tensor, 1, start, std::min(start + head_chunk, heads), false));
                     }
                     outputs.push_back(attend_heads(ctx, part, pe));
                     // Keep each head group's temporaries out of the concatenation tree's live set.
                     ctx->expand_graph(outputs.back());
                 }
-                return ggml_ext_vec_concat(ctx->ggml_ctx, outputs, 0);
+                return ggml_ext_vec_concat(ctx, outputs, 0);
             }
             return attend_heads(ctx, qkv, pe);
         }
@@ -562,19 +541,20 @@ namespace MiniMaxH3 {
         int32_t source_index = -1;
     };
 
-    static std::vector<ggml_tensor*> modulation_row(ggml_context* ctx,
+    static std::vector<ggml_tensor*> modulation_row(GGMLRunnerContext* runner_ctx,
                                                     ggml_tensor* projection,
                                                     int64_t hidden_size,
                                                     int expand,
                                                     int modalities,
                                                     int64_t row) {
+        auto ctx              = runner_ctx->ggml_ctx;
         int64_t timestep_rows = projection->ne[1];
         auto reshaped         = ggml_reshape_2d(ctx,
                                                 projection,
                                                 hidden_size * expand,
                                                 timestep_rows * modalities);
-        auto selected         = ggml_ext_slice(ctx, reshaped, 1, row, row + 1);
-        return ggml_ext_chunk(ctx, selected, expand, 0);
+        auto selected         = ggml_ext_slice(runner_ctx, reshaped, 1, row, row + 1);
+        return ggml_ext_chunk(runner_ctx, selected, expand, 0);
     }
 
     struct TransformerBlock : public GGMLBlock {
@@ -616,9 +596,9 @@ namespace MiniMaxH3 {
             std::vector<std::vector<ggml_tensor*>> rows;
             std::vector<ggml_tensor*> inputs;
             for (const auto& segment : segments) {
-                rows.push_back(modulation_row(ctx->ggml_ctx, mods, config.hidden_size,
+                rows.push_back(modulation_row(ctx, mods, config.hidden_size,
                                               6, 3, segment.modulation_row));
-                inputs.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 1, segment.start, segment.end, false));
+                inputs.push_back(ggml_ext_slice(ctx, x, 1, segment.start, segment.end, false));
             }
             auto modulate = [&](ggml_tensor* value, const std::vector<ggml_tensor*>& row, int shift, int scale) {
                 return ggml_ext_modulate(ctx, value, row[shift], row[scale]);
@@ -629,7 +609,7 @@ namespace MiniMaxH3 {
             auto attended = attn->attend(ctx, qkv, pe);
             std::vector<ggml_tensor*> projections;
             for (const auto& segment : segments) {
-                auto part = ggml_ext_slice(ctx->ggml_ctx, attended, 1, segment.start, segment.end, false);
+                auto part = ggml_ext_slice(ctx, attended, 1, segment.start, segment.end, false);
                 projections.push_back(attn->project_output(ctx, part));
             }
             ggml_tensor* output = nullptr;
@@ -672,8 +652,8 @@ namespace MiniMaxH3 {
             auto audio_out = std::dynamic_pointer_cast<Linear>(blocks["audio_out"]);
             auto mods      = adaln->forward(ctx, t_emb);
             auto apply     = [&](const TokenModulationSpan& segment) {
-                auto row   = modulation_row(ctx->ggml_ctx, mods, config.hidden_size, 2, 1, segment.modulation_row);
-                auto value = norm->forward(ctx, ggml_ext_slice(ctx->ggml_ctx, x, 1, segment.start, segment.end));
+                auto row   = modulation_row(ctx, mods, config.hidden_size, 2, 1, segment.modulation_row);
+                auto value = norm->forward(ctx, ggml_ext_slice(ctx, x, 1, segment.start, segment.end));
                 return ggml_add(ctx->ggml_ctx,
                                 ggml_add(ctx->ggml_ctx, value, ggml_mul(ctx->ggml_ctx, value, row[1])),
                                 row[0]);
@@ -774,7 +754,7 @@ namespace MiniMaxH3 {
                                                   1);
             ggml_tensor* angles = nullptr;
             for (int axis = 0; axis < 3; ++axis) {
-                auto pos          = ggml_ext_slice(ctx->ggml_ctx, position_ids, 0, axis, axis + 1);
+                auto pos          = ggml_ext_slice(ctx, position_ids, 0, axis, axis + 1);
                 auto expanded_inv = ggml_repeat_4d(ctx->ggml_ctx,
                                                    inv,
                                                    inv->ne[0],
@@ -816,7 +796,7 @@ namespace MiniMaxH3 {
             ggml_tensor* video_rows = nullptr;
             int64_t video_offset    = 0;
             for (auto condition : condition_videos) {
-                auto rows = DiT::patchify_3d(ctx->ggml_ctx,
+                auto rows = DiT::patchify_3d(ctx,
                                              condition,
                                              config.patch_t,
                                              config.patch_h,
@@ -827,7 +807,7 @@ namespace MiniMaxH3 {
                 video_offset += rows->ne[1];
                 video_rows = video_rows == nullptr ? rows : ggml_concat(ctx->ggml_ctx, video_rows, rows, 1);
             }
-            auto target_video_rows                         = DiT::patchify_3d(ctx->ggml_ctx,
+            auto target_video_rows                         = DiT::patchify_3d(ctx,
                                                                               video,
                                                                               config.patch_t,
                                                                               config.patch_h,
@@ -840,7 +820,7 @@ namespace MiniMaxH3 {
             auto video_embeds                              = video_proj->forward(ctx, video_rows);
 
             auto pack_audio_rows = [&](ggml_tensor* value) {
-                value = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, value, 2, 0, 1, 3));
+                value = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, value, 2, 0, 1, 3));
                 return ggml_reshape_3d(ctx->ggml_ctx,
                                        value,
                                        value->ne[0],
@@ -873,20 +853,20 @@ namespace MiniMaxH3 {
                     GGML_ASSERT(sequence.source_index >= 0 &&
                                 sequence.source_index < static_cast<int32_t>(condition_video_ranges.size()));
                     auto range = condition_video_ranges[static_cast<size_t>(sequence.source_index)];
-                    append(ggml_ext_slice(ctx->ggml_ctx, video_embeds, 1, range.first, range.second));
+                    append(ggml_ext_slice(ctx, video_embeds, 1, range.first, range.second));
                 } else if (sequence.kind == SequenceKind::CONDITION_AUDIO) {
                     GGML_ASSERT(sequence.source_index >= 0 &&
                                 sequence.source_index < static_cast<int32_t>(condition_audio_ranges.size()));
                     auto range = condition_audio_ranges[static_cast<size_t>(sequence.source_index)];
-                    append(ggml_ext_slice(ctx->ggml_ctx, audio_embeds, 1, range.first, range.second));
+                    append(ggml_ext_slice(ctx, audio_embeds, 1, range.first, range.second));
                 } else if (sequence.kind == SequenceKind::TARGET_AUDIO) {
-                    append(ggml_ext_slice(ctx->ggml_ctx,
+                    append(ggml_ext_slice(ctx,
                                           audio_embeds,
                                           1,
                                           target_audio_range.first,
                                           target_audio_range.second));
                 } else {
-                    append(ggml_ext_slice(ctx->ggml_ctx,
+                    append(ggml_ext_slice(ctx,
                                           video_embeds,
                                           1,
                                           target_video_range.first,
@@ -912,7 +892,7 @@ namespace MiniMaxH3 {
 
             auto final_layer = std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer"]);
             auto output      = final_layer->forward(ctx, h, t_emb, video_segment, audio_segment);
-            auto video_out   = DiT::unpatchify_3d(ctx->ggml_ctx,
+            auto video_out   = DiT::unpatchify_3d(ctx,
                                                   output.first,
                                                   video->ne[2] / config.patch_t,
                                                   video->ne[1] / config.patch_h,
@@ -927,9 +907,9 @@ namespace MiniMaxH3 {
                                                audio->ne[1] / 2,
                                                2,
                                                audio->ne[2]);
-            audio_out        = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, audio_out, 1, 2, 0, 3));
-            video_out        = ggml_ext_scale(ctx->ggml_ctx, video_out, -1.f);
-            audio_out        = ggml_ext_scale(ctx->ggml_ctx, audio_out, -1.f);
+            audio_out        = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, audio_out, 1, 2, 0, 3));
+            video_out        = ggml_ext_scale(ctx, video_out, -1.f);
+            audio_out        = ggml_ext_scale(ctx, audio_out, -1.f);
             return {video_out, audio_out};
         }
     };
@@ -1221,15 +1201,14 @@ namespace MiniMaxH3 {
             model.get_param_tensors(tensors, prefix);
         }
 
-        ggml_tensor* merge_av_latents(ggml_context* ctx,
-                                      ggml_tensor* video,
-                                      ggml_tensor* audio) const {
+        ggml_tensor* merge_av_latents(GGMLRunnerContext* runner_ctx, ggml_tensor* video, ggml_tensor* audio) const {
+            auto ctx        = runner_ctx->ggml_ctx;
             int64_t divisor = video->ne[0] * video->ne[1] * video->ne[2];
             int64_t values  = ggml_nelements(audio);
             int64_t padding = (divisor - values % divisor) % divisor;
             audio           = ggml_reshape_4d(ctx, ggml_cont(ctx, audio), values, 1, 1, 1);
             if (padding > 0) {
-                audio = ggml_ext_pad(ctx, audio, static_cast<int>(padding), 0, 0, 0);
+                audio = ggml_ext_pad(runner_ctx, audio, static_cast<int>(padding), 0, 0, 0);
             }
             audio = ggml_reshape_4d(ctx,
                                     audio,
@@ -1302,21 +1281,22 @@ namespace MiniMaxH3 {
             }
             // The sampler carries c_a = (sigma_v / sigma_a) * x_a so the packed
             // latent follows one sigma schedule. Restore x_a for the H3 network.
-            auto audio  = ggml_ext_scale(compute_ctx, audio_carrier, sigma_a / sigma_v);
-            auto layout = build_layout(context_tensor.shape()[1],
-                                       packed.shape()[2],
-                                       packed.shape()[1],
-                                       packed.shape()[0],
-                                       audio_length,
-                                       condition_videos,
-                                       condition_audios,
-                                       keyframe_indices,
-                                       reference_blocks,
-                                       text_tags,
-                                       t_v,
-                                       t_a,
-                                       r_v,
-                                       r_a);
+            auto audio_ctx = get_context();
+            auto audio     = ggml_ext_scale(&audio_ctx, audio_carrier, sigma_a / sigma_v);
+            auto layout    = build_layout(context_tensor.shape()[1],
+                                          packed.shape()[2],
+                                          packed.shape()[1],
+                                          packed.shape()[0],
+                                          audio_length,
+                                          condition_videos,
+                                          condition_audios,
+                                          keyframe_indices,
+                                          reference_blocks,
+                                          text_tags,
+                                          t_v,
+                                          t_a,
+                                          r_v,
+                                          r_a);
 
             position_input_cache = sd::Tensor<float>(
                 {3, static_cast<int64_t>(layout.positions.size() / 3)},
@@ -1393,11 +1373,11 @@ namespace MiniMaxH3 {
                                             layout.audio_segment);
             // Convert the model's audio velocity to d(c_a) / d(sigma_v).
             output.second = ggml_add(compute_ctx,
-                                     ggml_ext_scale(compute_ctx, audio, 1.f - audio_scale),
-                                     ggml_ext_scale(compute_ctx,
+                                     ggml_ext_scale(&runner_ctx, audio, 1.f - audio_scale),
+                                     ggml_ext_scale(&runner_ctx,
                                                     output.second,
                                                     1.f + (audio_scale - 1.f) * sigma_a));
-            auto merged   = merge_av_latents(compute_ctx, output.first, output.second);
+            auto merged   = merge_av_latents(&runner_ctx, output.first, output.second);
             ggml_build_forward_expand(graph, merged);
             return graph;
         }

src/model/diffusion/minit2i.hpp

@@ -257,14 +257,18 @@ namespace MiniT2I {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
             auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
             auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
-            auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10000, 1.0f);
+            auto t_emb = ggml_ext_timestep_embedding(ctx, t, frequency_embedding_size, 10000, 1.0f);
             t_emb      = mlp_0->forward(ctx, t_emb);
             t_emb      = ggml_silu_inplace(ctx->ggml_ctx, t_emb);
             return mlp_2->forward(ctx, t_emb);
         }
     };
 
-    inline std::vector<ggml_tensor*> split_qkv(ggml_context* ctx, ggml_tensor* qkv, int64_t num_heads, int64_t head_dim) {
+    inline std::vector<ggml_tensor*> split_qkv(GGMLRunnerContext* runner_ctx,
+                                               ggml_tensor* qkv,
+                                               int64_t num_heads,
+                                               int64_t head_dim) {
+        auto ctx  = runner_ctx->ggml_ctx;
         int64_t N = qkv->ne[2];
         int64_t L = qkv->ne[1];
         auto q    = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N,
@@ -301,7 +305,7 @@ namespace MiniT2I {
             auto q_norm    = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
             auto k_norm    = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
 
-            auto qkv = split_qkv(ctx->ggml_ctx, qkv_proj->forward(ctx, norm1->forward(ctx, txt)), num_heads, head_dim);
+            auto qkv = split_qkv(ctx, qkv_proj->forward(ctx, norm1->forward(ctx, txt)), num_heads, head_dim);
             auto q   = q_norm->forward(ctx, qkv[0]);
             auto k   = k_norm->forward(ctx, qkv[1]);
             auto v   = qkv[2];
@@ -353,16 +357,16 @@ namespace MiniT2I {
             int64_t li = img->ne[1];
             int64_t lt = txt->ne[1];
 
-            auto img_qkv = split_qkv(ctx->ggml_ctx, img_qkv_p->forward(ctx, img_norm1->forward(ctx, img)), num_heads, head_dim);
-            auto txt_qkv = split_qkv(ctx->ggml_ctx, txt_qkv_p->forward(ctx, txt_norm1->forward(ctx, txt)), num_heads, head_dim);
+            auto img_qkv = split_qkv(ctx, img_qkv_p->forward(ctx, img_norm1->forward(ctx, img)), num_heads, head_dim);
+            auto txt_qkv = split_qkv(ctx, txt_qkv_p->forward(ctx, txt_norm1->forward(ctx, txt)), num_heads, head_dim);
 
             auto q = ggml_concat(ctx->ggml_ctx, q_norm->forward(ctx, txt_qkv[0]), q_norm->forward(ctx, img_qkv[0]), 2);
             auto k = ggml_concat(ctx->ggml_ctx, k_norm->forward(ctx, txt_qkv[1]), k_norm->forward(ctx, img_qkv[1]), 2);
             auto v = ggml_concat(ctx->ggml_ctx, txt_qkv[2], img_qkv[2], 2);
 
             auto out     = Rope::attention(ctx, q, k, v, pe, nullptr, 1.0f, false);
-            auto out_txt = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, lt);
-            auto out_img = ggml_ext_slice(ctx->ggml_ctx, out, 1, lt, lt + li);
+            auto out_txt = ggml_ext_slice(ctx, out, 1, 0, lt);
+            auto out_img = ggml_ext_slice(ctx, out, 1, lt, lt + li);
 
             img = ggml_add(ctx->ggml_ctx, img, img_proj->forward(ctx, out_img));
             txt = ggml_add(ctx->ggml_ctx, txt, txt_proj->forward(ctx, out_txt));
@@ -416,7 +420,7 @@ namespace MiniT2I {
             mask            = ggml_reshape_3d(ctx->ggml_ctx, mask, 1, mask->ne[0], mask->ne[1]);
             mask            = ggml_repeat(ctx->ggml_ctx, mask, context);
             auto keep       = ggml_mul(ctx->ggml_ctx, context, mask);
-            auto inv        = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx->ggml_ctx, mask), mask);
+            auto inv        = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx, mask), mask);
             auto mask_token = ggml_repeat(ctx->ggml_ctx, params["mask_token"], context);
             return ggml_add(ctx->ggml_ctx, keep, ggml_mul(ctx->ggml_ctx, mask_token, inv));
         }
@@ -468,8 +472,8 @@ namespace MiniT2I {
             }
             auto combined = ggml_concat(ctx->ggml_ctx, txt, x, 1);
             auto out      = final_layer->forward(ctx, combined);
-            auto img_out  = ggml_ext_slice(ctx->ggml_ctx, out, 1, txt->ne[1], txt->ne[1] + x->ne[1]);
-            return DiT::unpatchify(ctx->ggml_ctx, img_out, hp, wp, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size), false);
+            auto img_out  = ggml_ext_slice(ctx, out, 1, txt->ne[1], txt->ne[1] + x->ne[1]);
+            return DiT::unpatchify(ctx, img_out, hp, wp, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size), false);
         }
     };
 

src/model/diffusion/mmdit.hpp

@@ -272,7 +272,7 @@ public:
         auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
         auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
 
-        auto t_freq = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size);  // [N, frequency_embedding_size]
+        auto t_freq = ggml_ext_timestep_embedding(ctx, t, frequency_embedding_size);  // [N, frequency_embedding_size]
 
         auto t_emb = mlp_0->forward(ctx, t_freq);
         t_emb      = ggml_silu_inplace(ctx->ggml_ctx, t_emb);
@@ -334,7 +334,7 @@ public:
         auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
 
         auto qkv         = qkv_proj->forward(ctx, x);
-        auto qkv_vec     = split_qkv(ctx->ggml_ctx, qkv);
+        auto qkv_vec     = ggml_ext_split_qkv(ctx, qkv);
         int64_t head_dim = qkv_vec[0]->ne[0] / num_heads;
         auto q           = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, qkv_vec[0]->ne[1], qkv_vec[0]->ne[2]);  // [N, n_token, n_head, d_head]
         auto k           = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, qkv_vec[1]->ne[1], qkv_vec[1]->ne[2]);  // [N, n_token, n_head, d_head]
@@ -439,7 +439,7 @@ public:
 
         int n_mods = 9;
         auto m     = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, n_mods * hidden_size]
-        auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, n_mods, 0);
+        auto m_vec = ggml_ext_chunk(ctx, m, n_mods, 0);
 
         auto shift_msa  = m_vec[0];  // [N, hidden_size]
         auto scale_msa  = m_vec[1];  // [N, hidden_size]
@@ -476,7 +476,7 @@ public:
             n_mods = 2;
         }
         auto m     = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, n_mods * hidden_size]
-        auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, n_mods, 0);
+        auto m_vec = ggml_ext_chunk(ctx, m, n_mods, 0);
 
         auto shift_msa = m_vec[0];  // [N, hidden_size]
         auto scale_msa = m_vec[1];  // [N, hidden_size]
@@ -750,7 +750,7 @@ public:
         auto adaLN_modulation_1 = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
 
         auto m     = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, 2 * hidden_size]
-        auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, 2, 0);
+        auto m_vec = ggml_ext_chunk(ctx, m, 2, 0);
         auto shift = m_vec[0];  // [N, hidden_size]
         auto scale = m_vec[1];  // [N, hidden_size]
 
@@ -926,7 +926,7 @@ public:
 
         x = forward_core_with_concat(ctx, x, c, context, skip_layers);  // (N, H*W, patch_size ** 2 * out_channels)
 
-        x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, config.patch_size, config.patch_size, /*patch_last*/ false);  // [N, C, H, W]
+        x = DiT::unpatchify_and_crop(ctx, x, H, W, config.patch_size, config.patch_size, /*patch_last*/ false);  // [N, C, H, W]
 
         return x;
     }

src/model/diffusion/moss_image.hpp

@@ -164,9 +164,9 @@ namespace MOSSImage {
             auto c            = ctx->ggml_ctx;
             auto linear       = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, h); };
             int64_t image_len = image->ne[0] * image->ne[1];
-            auto x            = linear("video_embedder", DiT::patchify(c, image, 1, 1));
+            auto x            = linear("video_embedder", DiT::patchify(ctx, image, 1, 1));
             text              = linear("text_embedder.proj", text);
-            auto text_pe      = ggml_ext_slice(c, pe, 3, image_len, image_len + text->ne[1]);
+            auto text_pe      = ggml_ext_slice(ctx, pe, 3, image_len, image_len + text->ne[1]);
             for (int i = 0; i < 2; ++i) {
                 text = std::dynamic_pointer_cast<TextBlock>(blocks["text_embedder.blocks." + std::to_string(i)])->forward(ctx, text, text_pe);
             }
@@ -175,7 +175,7 @@ namespace MOSSImage {
             auto head = linear("video_head_modulation", time);
             std::array<ggml_tensor*, 6> parts;
             for (int i = 0; i < 6; ++i) {
-                parts[i] = ggml_ext_slice(c, mod, 0, i * config.dim, (i + 1) * config.dim);
+                parts[i] = ggml_ext_slice(ctx, mod, 0, i * config.dim, (i + 1) * config.dim);
                 if (ctx->graph != nullptr) {
                     ggml_build_forward_expand(ctx->graph, parts[i]);
                 }
@@ -187,13 +187,13 @@ namespace MOSSImage {
                     sd::ggml_graph_cut::mark_graph_cut(x, "moss_image.layers." + std::to_string(i), "x");
                 }
             }
-            x          = ggml_ext_slice(c, x, 1, 0, image_len);
+            x          = ggml_ext_slice(ctx, x, 1, 0, image_len);
             x          = std::dynamic_pointer_cast<RMSNorm>(blocks["final_norm_video"])->forward(ctx, x);
-            auto shift = ggml_ext_slice(c, head, 0, 0, config.dim);
-            auto scale = ggml_ext_slice(c, head, 0, config.dim, 2 * config.dim);
+            auto shift = ggml_ext_slice(ctx, head, 0, 0, config.dim);
+            auto scale = ggml_ext_slice(ctx, head, 0, config.dim, 2 * config.dim);
             x          = ggml_ext_modulate(ctx, x, shift, scale);
             x          = linear("final_linear_video", x);
-            return DiT::unpatchify(c, x, image->ne[1], image->ne[0], 1, 1);
+            return DiT::unpatchify(ctx, x, image->ne[1], image->ne[0], 1, 1);
         }
     };
 

src/model/diffusion/pid.hpp

@@ -185,14 +185,15 @@ namespace Pid {
         return ggml_add(ctx, ggml_add(ctx, x, ggml_mul(ctx, x, scale)), shift);
     }
 
-    inline ggml_tensor* replicate_pad_2d(ggml_context* ctx, ggml_tensor* x) {
-        auto left  = ggml_ext_slice(ctx, x, 0, 0, 1);
-        auto right = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]);
+    inline ggml_tensor* replicate_pad_2d(GGMLRunnerContext* runner_ctx, ggml_tensor* x) {
+        auto ctx   = runner_ctx->ggml_ctx;
+        auto left  = ggml_ext_slice(runner_ctx, x, 0, 0, 1);
+        auto right = ggml_ext_slice(runner_ctx, x, 0, x->ne[0] - 1, x->ne[0]);
         x          = ggml_concat(ctx, left, x, 0);
         x          = ggml_concat(ctx, x, right, 0);
 
-        auto top    = ggml_ext_slice(ctx, x, 1, 0, 1);
-        auto bottom = ggml_ext_slice(ctx, x, 1, x->ne[1] - 1, x->ne[1]);
+        auto top    = ggml_ext_slice(runner_ctx, x, 1, 0, 1);
+        auto bottom = ggml_ext_slice(runner_ctx, x, 1, x->ne[1] - 1, x->ne[1]);
         x           = ggml_concat(ctx, top, x, 1);
         return ggml_concat(ctx, x, bottom, 1);
     }
@@ -235,7 +236,7 @@ namespace Pid {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
             auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
             auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
-            auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10);
+            auto t_emb = ggml_ext_timestep_embedding(ctx, t, frequency_embedding_size, 10);
             t_emb      = mlp_0->forward(ctx, t_emb);
             t_emb      = ggml_silu_inplace(ctx->ggml_ctx, t_emb);
             return mlp_2->forward(ctx, t_emb);
@@ -292,7 +293,7 @@ namespace Pid {
             auto proj     = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
 
             auto qkv         = qkv_proj->forward(ctx, x);
-            auto qkv_vec     = split_qkv(ctx->ggml_ctx, qkv);
+            auto qkv_vec     = ggml_ext_split_qkv(ctx, qkv);
             int64_t L        = x->ne[1];
             int64_t N        = x->ne[2];
             int64_t head_dim = dim / num_heads;
@@ -342,14 +343,14 @@ namespace Pid {
             int64_t N        = x->ne[2];
             int64_t head_dim = dim / num_heads;
 
-            auto qkv_x = split_qkv(ctx->ggml_ctx, qkv_x_proj->forward(ctx, x));
+            auto qkv_x = ggml_ext_split_qkv(ctx, qkv_x_proj->forward(ctx, x));
             auto qx    = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[0], head_dim, num_heads, Nx, N);
             auto kx    = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[1], head_dim, num_heads, Nx, N);
             auto vx    = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[2], head_dim, num_heads, Nx, N);
             qx         = q_norm_x->forward(ctx, qx);
             kx         = k_norm_x->forward(ctx, kx);
 
-            auto qkv_y = split_qkv(ctx->ggml_ctx, qkv_y_proj->forward(ctx, y));
+            auto qkv_y = ggml_ext_split_qkv(ctx, qkv_y_proj->forward(ctx, y));
             auto qy    = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[0], head_dim, num_heads, Ny, N);
             auto ky    = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[1], head_dim, num_heads, Ny, N);
             auto vy    = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[2], head_dim, num_heads, Ny, N);
@@ -362,8 +363,8 @@ namespace Pid {
             auto pos_joint = ggml_concat(ctx->ggml_ctx, pos_txt, pos_img, 3);
             auto out       = Rope::attention(ctx, q_joint, k_joint, v_joint, pos_joint, nullptr, 1.0f, true);
 
-            auto out_y = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, Ny);
-            auto out_x = ggml_ext_slice(ctx->ggml_ctx, out, 1, Ny, Ny + Nx);
+            auto out_y = ggml_ext_slice(ctx, out, 1, 0, Ny);
+            auto out_x = ggml_ext_slice(ctx, out, 1, Ny, Ny + Nx);
             return {proj_x->forward(ctx, out_x), proj_y->forward(ctx, out_y)};
         }
     };
@@ -400,8 +401,8 @@ namespace Pid {
             auto ada_img = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation_img.0"]);
             auto ada_txt = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation_txt.0"]);
 
-            auto mx = ggml_ext_chunk(ctx->ggml_ctx, ada_img->forward(ctx, c), 6, 0);
-            auto my = ggml_ext_chunk(ctx->ggml_ctx, ada_txt->forward(ctx, c), 6, 0);
+            auto mx = ggml_ext_chunk(ctx, ada_img->forward(ctx, c), 6, 0);
+            auto my = ggml_ext_chunk(ctx, ada_txt->forward(ctx, c), 6, 0);
 
             auto x_norm   = apply_adaln(ctx->ggml_ctx, norm_x1->forward(ctx, x), mx[0], mx[1]);
             auto y_norm   = apply_adaln(ctx->ggml_ctx, norm_y1->forward(ctx, y), my[0], my[1]);
@@ -438,13 +439,13 @@ namespace Pid {
             int64_t L  = (W / patch_size) * (H / patch_size);
             int64_t P2 = patch_size * patch_size;
 
-            auto x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, inputs, 2, 0, 1, 3));
+            auto x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, inputs, 2, 0, 1, 3));
             x      = ggml_reshape_3d(ctx->ggml_ctx, x, in_channels, W * H, B);
             x      = proj->forward(ctx, x);
             x      = ggml_add(ctx->ggml_ctx, x, pos_full);
             x      = ggml_reshape_4d(ctx->ggml_ctx, x, hidden_size_output, W, H, B);
-            x      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
-            x      = DiT::patchify(ctx->ggml_ctx, x, static_cast<int>(patch_size), static_cast<int>(patch_size), false);
+            x      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));
+            x      = DiT::patchify(ctx, x, static_cast<int>(patch_size), static_cast<int>(patch_size), false);
             x      = ggml_reshape_3d(ctx->ggml_ctx, x, hidden_size_output, P2, L * B);
             return x;
         }
@@ -500,7 +501,7 @@ namespace Pid {
 
             auto ada_params = ada->forward(ctx, s_cond);
             ada_params      = ggml_reshape_3d(ctx->ggml_ctx, ada_params, 6 * pixel_dim, P2, BL);
-            auto mod        = ggml_ext_chunk(ctx->ggml_ctx, ada_params, 6, 0);
+            auto mod        = ggml_ext_chunk(ctx, ada_params, 6, 0);
 
             auto x_norm    = apply_adaln(ctx->ggml_ctx, norm1->forward(ctx, x), mod[0], mod[1]);
             auto x_flat    = ggml_reshape_2d(ctx->ggml_ctx, x_norm, P2 * pixel_dim, BL);
@@ -564,12 +565,12 @@ namespace Pid {
             auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["block.5"]);
             auto h     = ggml_silu_inplace(ctx->ggml_ctx, norm1->forward(ctx, x));
             if (replicate_padding) {
-                h = replicate_pad_2d(ctx->ggml_ctx, h);
+                h = replicate_pad_2d(ctx, h);
             }
             h = conv1->forward(ctx, h);
             h = ggml_silu_inplace(ctx->ggml_ctx, norm2->forward(ctx, h));
             if (replicate_padding) {
-                h = replicate_pad_2d(ctx->ggml_ctx, h);
+                h = replicate_pad_2d(ctx, h);
             }
             h = conv2->forward(ctx, h);
             return ggml_add(ctx->ggml_ctx, x, h);
@@ -628,13 +629,13 @@ namespace Pid {
             if (unpatchify_factor > 1) {
                 int64_t latent_h = lq_latent->ne[1];
                 int64_t latent_w = lq_latent->ne[0];
-                lq_latent        = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, lq_latent, 2, 0, 1, 3));
+                lq_latent        = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, lq_latent, 2, 0, 1, 3));
                 lq_latent        = ggml_reshape_3d(ctx->ggml_ctx,
                                                    lq_latent,
                                                    lq_latent->ne[0],
                                                    lq_latent->ne[1] * lq_latent->ne[2],
                                                    lq_latent->ne[3]);
-                lq_latent        = DiT::unpatchify(ctx->ggml_ctx,
+                lq_latent        = DiT::unpatchify(ctx,
                                                    lq_latent,
                                                    latent_h,
                                                    latent_w,
@@ -658,12 +659,12 @@ namespace Pid {
             }
 
             if (config.lq_replicate_padding) {
-                lq_latent = replicate_pad_2d(ctx->ggml_ctx, lq_latent);
+                lq_latent = replicate_pad_2d(ctx, lq_latent);
             }
             auto feat = conv0->forward(ctx, lq_latent);
             feat      = ggml_silu_inplace(ctx->ggml_ctx, feat);
             if (config.lq_replicate_padding) {
-                feat = replicate_pad_2d(ctx->ggml_ctx, feat);
+                feat = replicate_pad_2d(ctx, feat);
             }
             feat = conv2->forward(ctx, feat);
             for (int i = 0; i < config.lq_num_res_blocks; ++i) {
@@ -674,7 +675,7 @@ namespace Pid {
             int64_t B   = feat->ne[3];
             int64_t C   = feat->ne[2];
             int64_t L   = target_pH * target_pW;
-            auto tokens = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, feat, 2, 0, 1, 3));
+            auto tokens = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, feat, 2, 0, 1, 3));
             tokens      = ggml_reshape_3d(ctx->ggml_ctx, tokens, C, L, B);
 
             int num_outputs = static_cast<int>((config.patch_depth + config.lq_interval - 1) / config.lq_interval);
@@ -754,15 +755,15 @@ namespace Pid {
             int64_t L      = Hs * Ws;
             int64_t P2     = config.patch_size * config.patch_size;
 
-            auto x_patches = DiT::patchify(ctx->ggml_ctx, x, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size), true);
+            auto x_patches = DiT::patchify(ctx, x, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size), true);
             auto t_emb     = t_embedder->forward(ctx, timesteps);
             auto condition = ggml_silu(ctx->ggml_ctx, t_emb);
 
             GGML_ASSERT(context != nullptr);
             int64_t Ltxt = std::min<int64_t>(context->ne[1], config.txt_max_length);
-            auto y       = ggml_ext_slice(ctx->ggml_ctx, context, 1, 0, Ltxt);
+            auto y       = ggml_ext_slice(ctx, context, 1, 0, Ltxt);
             auto y_emb   = y_embedder->forward(ctx, y);
-            auto y_pos   = ggml_ext_slice(ctx->ggml_ctx, params["y_pos_embedding"], 1, 0, Ltxt);
+            auto y_pos   = ggml_ext_slice(ctx, params["y_pos_embedding"], 1, 0, Ltxt);
             y_emb        = ggml_add(ctx->ggml_ctx, y_emb, y_pos);
 
             std::vector<ggml_tensor*> lq_features = lq_proj->forward(ctx, lq_latent, Hs, Ws);
@@ -809,15 +810,15 @@ namespace Pid {
 
             pixels   = final_layer->forward(ctx, pixels);
             pixels   = ggml_reshape_3d(ctx->ggml_ctx, pixels, config.in_channels * P2, L, B);
-            auto out = DiT::unpatchify(ctx->ggml_ctx,
+            auto out = DiT::unpatchify(ctx,
                                        pixels,
                                        Hs,
                                        Ws,
                                        static_cast<int>(config.patch_size),
                                        static_cast<int>(config.patch_size),
                                        false);
-            out      = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H_orig);
-            out      = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W_orig);
+            out      = ggml_ext_slice(ctx, out, 1, 0, H_orig);
+            out      = ggml_ext_slice(ctx, out, 0, 0, W_orig);
             return out;
         }
     };

src/model/diffusion/pixart.hpp

@@ -140,11 +140,11 @@ namespace PixArt {
             std::vector<ggml_tensor*> qkv;
             if (self_attention) {
                 auto projected = std::dynamic_pointer_cast<Linear>(blocks["qkv"])->forward(ctx, x);
-                qkv            = ggml_ext_chunk(ctx->ggml_ctx, projected, 3, 0);
+                qkv            = ggml_ext_chunk(ctx, projected, 3, 0);
             } else {
                 auto q     = std::dynamic_pointer_cast<Linear>(blocks["q_linear"])->forward(ctx, x);
                 auto kv    = std::dynamic_pointer_cast<Linear>(blocks["kv_linear"])->forward(ctx, context);
-                auto parts = ggml_ext_chunk(ctx->ggml_ctx, kv, 2, 0);
+                auto parts = ggml_ext_chunk(ctx, kv, 2, 0);
                 qkv        = {q, parts[0], parts[1]};
             }
             auto out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, mask, false, ctx->flash_attn_enabled);
@@ -186,7 +186,7 @@ namespace PixArt {
             }
             table   = ggml_reshape_3d(ctx->ggml_ctx, table, dim, 6, 1);
             auto m  = ggml_add(ctx->ggml_ctx, ggml_reshape_3d(ctx->ggml_ctx, mod, dim, 6, N), table);
-            auto mv = ggml_ext_chunk(ctx->ggml_ctx, ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, m), dim * 6, N), 6, 0);
+            auto mv = ggml_ext_chunk(ctx, ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx, m), dim * 6, N), 6, 0);
 
             auto attn1 = std::dynamic_pointer_cast<PixArtAttention>(blocks["attn"]);
             auto attn2 = std::dynamic_pointer_cast<PixArtAttention>(blocks["cross_attn"]);
@@ -204,7 +204,7 @@ namespace PixArt {
             x = ggml_add(ctx->ggml_ctx, x, attn2->forward(ctx, x, context, context_mask));
             h = modulate(ctx, norm(ctx->ggml_ctx, x), mv[3], mv[4]);
             h = proj->forward(ctx, h);
-            h = ggml_ext_gelu(ctx->ggml_ctx, h, true);
+            h = ggml_ext_gelu(ctx, h, true);
             h = fc2->forward(ctx, h);
             return gate(h, mv[5]);
         }
@@ -256,18 +256,18 @@ namespace PixArt {
             int64_t hp = H / p;
 
             auto h = std::dynamic_pointer_cast<Conv2d>(blocks["x_embedder.proj"])->forward(ctx, x);  // [N, hidden, hp, wp]
-            h      = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 1, 2, 0, 3));       // [N, hp, wp, hidden] -> [N, hp*wp, hidden]
+            h      = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, h, 1, 2, 0, 3));                 // [N, hp, wp, hidden] -> [N, hp*wp, hidden]
             h      = ggml_reshape_3d(ctx->ggml_ctx, h, config.hidden_size, wp * hp, N);              // [N, hp*wp, hidden]
             h      = ggml_add(ctx->ggml_ctx, h, pos_embed);
 
-            auto t   = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, ADALN_EMBED_DIM, 10000);
+            auto t   = ggml_ext_timestep_embedding(ctx, timesteps, ADALN_EMBED_DIM, 10000);
             auto emb = std::dynamic_pointer_cast<PixArtTimestepEmbedding>(blocks["t_embedder"])->forward(ctx, t);
 
             auto mod = std::dynamic_pointer_cast<Linear>(blocks["t_block.1"])
                            ->forward(ctx, ggml_silu(ctx->ggml_ctx, emb));  // [N, 6 * hidden]
 
             auto ctx_emb = std::dynamic_pointer_cast<Linear>(blocks["y_embedder.y_proj.fc1"])->forward(ctx, context);
-            ctx_emb      = ggml_ext_gelu(ctx->ggml_ctx, ctx_emb, true);
+            ctx_emb      = ggml_ext_gelu(ctx, ctx_emb, true);
             ctx_emb      = std::dynamic_pointer_cast<Linear>(blocks["y_embedder.y_proj.fc2"])->forward(ctx, ctx_emb);
 
             for (int i = 0; i < config.num_layers; ++i) {
@@ -284,13 +284,13 @@ namespace PixArt {
             auto ss    = ggml_add(ctx->ggml_ctx,
                                   ggml_reshape_3d(ctx->ggml_ctx, tail_table, config.hidden_size, 2, 1),
                                   ggml_reshape_3d(ctx->ggml_ctx, emb, config.hidden_size, 1, N));  // [2, hidden, N]
-            auto parts = ggml_ext_chunk(ctx->ggml_ctx,
-                                        ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, ss), config.hidden_size * 2, N),
+            auto parts = ggml_ext_chunk(ctx,
+                                        ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx, ss), config.hidden_size * 2, N),
                                         2, 0);
             h          = std::dynamic_pointer_cast<LayerNorm>(blocks["final_layer.norm_final"])->forward(ctx, h);
             h          = modulate(ctx, h, parts[0], parts[1]);
             h          = std::dynamic_pointer_cast<Linear>(blocks["final_layer.linear"])->forward(ctx, h);  // [N, hp*wp, p*p*out_ch]
-            h          = DiT::unpatchify(ctx->ggml_ctx, h, hp, wp, static_cast<int>(p), static_cast<int>(p), false);
+            h          = DiT::unpatchify(ctx, h, hp, wp, static_cast<int>(p), static_cast<int>(p), false);
             return h;  // [N, out_channels, H, W]
         }
     };
@@ -369,7 +369,7 @@ namespace PixArt {
             auto runner_ctx  = get_context();
             ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pos, context_mask);
             // learn_sigma: keep the noise prediction half of the output channels
-            out = ggml_ext_slice(compute_ctx, out, 2, 0, config.in_channels);
+            out = ggml_ext_slice(&runner_ctx, out, 2, 0, config.in_channels);
             ggml_build_forward_expand(gf, out);
             return gf;
         }

src/model/diffusion/qwen_image.hpp

@@ -113,7 +113,7 @@ namespace Qwen {
             // return: [N, embedding_dim]
             auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
 
-            auto timesteps_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1.f);
+            auto timesteps_proj = ggml_ext_timestep_embedding(ctx, timesteps, 256, 10000, 1.f);
             auto timesteps_emb  = timestep_embedder->forward(ctx, timesteps_proj);
             if (use_additional_t_cond) {
                 GGML_ASSERT(addition_t_cond != nullptr);
@@ -186,9 +186,7 @@ namespace Qwen {
             auto to_v     = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
             auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-                to_out_0->set_force_prec_f32(true);
-            }
+            to_out_0->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
 
             auto norm_added_q = std::dynamic_pointer_cast<UnaryBlock>(blocks["norm_added_q"]);
             auto norm_added_k = std::dynamic_pointer_cast<UnaryBlock>(blocks["norm_added_k"]);
@@ -289,14 +287,17 @@ namespace Qwen {
                                                                                eps));
         }
 
-        std::vector<ggml_tensor*> get_mod_params_vec(ggml_context* ctx, ggml_tensor* mod_params, ggml_tensor* index = nullptr) {
+        std::vector<ggml_tensor*> get_mod_params_vec(GGMLRunnerContext* runner_ctx,
+                                                     ggml_tensor* mod_params,
+                                                     ggml_tensor* index = nullptr) {
+            auto ctx = runner_ctx->ggml_ctx;
             // index: [N, n_img_token]
             // mod_params: [N, hidden_size * 12]
             if (index == nullptr) {
-                return ggml_ext_chunk(ctx, mod_params, 6, 0);
+                return ggml_ext_chunk(runner_ctx, mod_params, 6, 0);
             }
             mod_params          = ggml_reshape_1d(ctx, mod_params, ggml_nelements(mod_params));
-            auto mod_params_vec = ggml_ext_chunk(ctx, mod_params, 12, 0);
+            auto mod_params_vec = ggml_ext_chunk(runner_ctx, mod_params, 12, 0);
             index               = ggml_reshape_3d(ctx, index, 1, index->ne[0], index->ne[1]);                                      // [N, n_img_token, 1]
             index               = ggml_repeat_4d(ctx, index, mod_params_vec[0]->ne[0], index->ne[1], index->ne[2], index->ne[3]);  // [N, n_img_token, hidden_size]
             std::vector<ggml_tensor*> mod_results;
@@ -339,15 +340,15 @@ namespace Qwen {
 
             auto img_mod_params    = ggml_silu(ctx->ggml_ctx, t_emb);
             img_mod_params         = img_mod_1->forward(ctx, img_mod_params);
-            auto img_mod_param_vec = get_mod_params_vec(ctx->ggml_ctx, img_mod_params, modulate_index);
+            auto img_mod_param_vec = get_mod_params_vec(ctx, img_mod_params, modulate_index);
 
             if (zero_cond_t) {
-                t_emb = ggml_ext_chunk(ctx->ggml_ctx, t_emb, 2, 1)[0];
+                t_emb = ggml_ext_chunk(ctx, t_emb, 2, 1)[0];
             }
 
             auto txt_mod_params    = ggml_silu(ctx->ggml_ctx, t_emb);
             txt_mod_params         = txt_mod_1->forward(ctx, txt_mod_params);
-            auto txt_mod_param_vec = get_mod_params_vec(ctx->ggml_ctx, txt_mod_params);
+            auto txt_mod_param_vec = get_mod_params_vec(ctx, txt_mod_params);
 
             auto img_normed    = img_norm1->forward(ctx, img);
             auto img_modulated = Flux::modulate(ctx, img_normed, img_mod_param_vec[0], img_mod_param_vec[1], modulate_index != nullptr);
@@ -402,7 +403,7 @@ namespace Qwen {
             auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
 
             auto emb   = linear->forward(ctx, ggml_silu(ctx->ggml_ctx, c));
-            auto mods  = ggml_ext_chunk(ctx->ggml_ctx, emb, 2, 0);
+            auto mods  = ggml_ext_chunk(ctx, emb, 2, 0);
             auto scale = mods[0];
             auto shift = mods[1];
 
@@ -461,7 +462,7 @@ namespace Qwen {
 
             auto t_emb = time_text_embed->forward(ctx, timestep, addition_t_cond);
             if (config.zero_cond_t) {
-                auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), addition_t_cond);
+                auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx, timestep), addition_t_cond);
                 t_emb        = ggml_concat(ctx->ggml_ctx, t_emb, t_emb_0, 1);
             }
             auto img = img_in->forward(ctx, x);
@@ -482,7 +483,7 @@ namespace Qwen {
             }
 
             if (config.zero_cond_t) {
-                t_emb = ggml_ext_chunk(ctx->ggml_ctx, t_emb, 2, 1)[0];
+                t_emb = ggml_ext_chunk(ctx, t_emb, 2, 1)[0];
             }
 
             img = norm_out->forward(ctx, img, t_emb);
@@ -519,12 +520,12 @@ namespace Qwen {
             auto patchify_input = [&](ggml_tensor* input) -> ggml_tensor* {
                 input = DiT::pad_to_patch_size(ctx, input, config.patch_size, config.patch_size);
                 if (!has_time_axis) {
-                    return DiT::patchify(ctx->ggml_ctx, input, config.patch_size, config.patch_size);
+                    return DiT::patchify(ctx, input, config.patch_size, config.patch_size);
                 }
                 if (input->ne[3] == 1) {
                     input = ggml_reshape_4d(ctx->ggml_ctx, input, input->ne[0], input->ne[1], 1, input->ne[2]);
                 }
-                return DiT::patchify_3d(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
+                return DiT::patchify_3d(ctx, input, 1, config.patch_size, config.patch_size, N);
             };
 
             auto img           = patchify_input(x);
@@ -550,11 +551,11 @@ namespace Qwen {
                 int pad_w = (config.patch_size - W % config.patch_size) % config.patch_size;
                 int h_len = static_cast<int>((H + pad_h) / config.patch_size);
                 int w_len = static_cast<int>((W + pad_w) / config.patch_size);
-                out       = DiT::unpatchify_3d(ctx->ggml_ctx, out, T, h_len, w_len, 1, config.patch_size, config.patch_size);
-                out       = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H);  // [N*C, T, H, W + pad_w]
-                out       = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W);  // [N*C, T, H, W]
+                out       = DiT::unpatchify_3d(ctx, out, T, h_len, w_len, 1, config.patch_size, config.patch_size);
+                out       = ggml_ext_slice(ctx, out, 1, 0, H);  // [N*C, T, H, W + pad_w]
+                out       = ggml_ext_slice(ctx, out, 0, 0, W);  // [N*C, T, H, W]
             } else {
-                out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size);  // [N, C, H, W]
+                out = DiT::unpatchify_and_crop(ctx, out, H, W, config.patch_size, config.patch_size);  // [N, C, H, W]
             }
 
             return out;

src/model/diffusion/qwen_image_2_1.hpp

@@ -146,7 +146,7 @@ namespace Qwen {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             auto weight = params["weight"];
             if (ctx->weight_adapter) {
-                weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight");
+                weight = ctx->weight_adapter->patch_weight(ctx, weight, prefix + "weight");
             }
             weight = ggml_scale_bias(ctx->ggml_ctx, weight, 1.f, 1.f);
             return ggml_mul(ctx->ggml_ctx, ggml_rms_norm(ctx->ggml_ctx, x, eps), weight);
@@ -164,7 +164,7 @@ namespace Qwen {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
             x = std::dynamic_pointer_cast<QwenImage21ZeroCenterRMSNorm>(blocks["text_norm"])->forward(ctx, x);
             x = std::dynamic_pointer_cast<Linear>(blocks["in_layer"])->forward(ctx, x);
-            x = ggml_ext_gelu(ctx->ggml_ctx, x);
+            x = ggml_ext_gelu(ctx, x);
             return std::dynamic_pointer_cast<Linear>(blocks["out_layer"])->forward(ctx, x);
         }
     };
@@ -196,7 +196,7 @@ namespace Qwen {
                 // prefix before its full-sequence K/V can accumulate across layers.
                 ctx->expand_graph(q);
                 auto persist = [&](ggml_tensor* tensor, int axis, const char* name) {
-                    auto part = ggml_ext_slice(ctx->ggml_ctx, tensor, axis, 0, cache.prefix_length);
+                    auto part = ggml_ext_slice(ctx, tensor, axis, 0, cache.prefix_length);
                     // Pack the contiguous data into wider rows so quantization blocks
                     // can exceed head_dim without padding or changing element order.
                     part      = ggml_reshape_2d(ctx->ggml_ctx, part, x->ne[0], cache.prefix_length);
@@ -236,17 +236,15 @@ namespace Qwen {
             } else {
                 for (size_t i = 0; i < segments.size(); ++i) {
                     const auto& segment = segments[i];
-                    auto sq             = ggml_ext_slice(ctx->ggml_ctx, q, 1, segment.start, segment.end);
-                    auto sk             = ggml_ext_slice(ctx->ggml_ctx, k, 1, 0, segment.end);
-                    auto sv             = ggml_ext_slice(ctx->ggml_ctx, v, 2, 0, segment.end);
+                    auto sq             = ggml_ext_slice(ctx, q, 1, segment.start, segment.end);
+                    auto sk             = ggml_ext_slice(ctx, k, 1, 0, segment.end);
+                    auto sv             = ggml_ext_slice(ctx, v, 2, 0, segment.end);
                     auto out            = attend(sq, sk, sv, masks[i]);
                     result              = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1);
                 }
             }
             auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
-            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-                to_out->set_force_prec_f32(true);
-            }
+            to_out->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
             return to_out->forward(ctx, result);
         }
     };
@@ -268,7 +266,7 @@ namespace Qwen {
 
         static ggml_tensor* modulate(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* params, int64_t prefix_length, bool gate = false) {
             auto ctx   = runner_ctx->ggml_ctx;
-            auto rows  = ggml_ext_chunk(ctx, params, 2, 1);
+            auto rows  = ggml_ext_chunk(runner_ctx, params, 2, 1);
             auto apply = [&](ggml_tensor* part, ggml_tensor* row) {
                 row = gate ? ggml_tanh(ctx, row) : ggml_scale_bias(ctx, row, 1.f, 1.f);
                 runner_ctx->expand_graph(row);
@@ -277,8 +275,8 @@ namespace Qwen {
             if (prefix_length == 0) {
                 return apply(x, rows[0]);
             }
-            auto target = apply(ggml_ext_slice(ctx, x, 1, prefix_length, x->ne[1]), rows[0]);
-            auto prefix = apply(ggml_ext_slice(ctx, x, 1, 0, prefix_length), rows[1]);
+            auto target = apply(ggml_ext_slice(runner_ctx, x, 1, prefix_length, x->ne[1]), rows[0]);
+            auto prefix = apply(ggml_ext_slice(runner_ctx, x, 1, 0, prefix_length), rows[1]);
             return ggml_concat(ctx, prefix, target, 1);
         }
 
@@ -323,27 +321,27 @@ namespace Qwen {
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector<ggml_tensor*>& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
-            auto time = ggml_concat(ctx->ggml_ctx, timestep, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), 0);
+            auto time = ggml_concat(ctx->ggml_ctx, timestep, ggml_ext_zeros_like(ctx, timestep), 0);
             // Runtime flow timesteps already use the [0, 1000] scale.
-            time               = ggml_ext_timestep_embedding(ctx->ggml_ctx, time, 256, 10000, 1.f);
+            time               = ggml_ext_timestep_embedding(ctx, time, 256, 10000, 1.f);
             time               = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["time_text_embed.timestep_embedder"])->forward(ctx, time);
             time               = ggml_silu(ctx->ggml_ctx, time);
             auto modulation    = std::dynamic_pointer_cast<Linear>(blocks["modulation.1"])->forward(ctx, time);
-            auto mod           = ggml_ext_chunk(ctx->ggml_ctx, modulation, 4, 0);
+            auto mod           = ggml_ext_chunk(ctx, modulation, 4, 0);
             auto img_in        = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
             ggml_tensor* joint = nullptr;
             if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) {
-                joint = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, x, 1, 1));
+                joint = img_in->forward(ctx, DiT::patchify(ctx, x, 1, 1));
             } else {
                 auto text = std::dynamic_pointer_cast<QwenImage21TextProjection>(blocks["txt_in"])->forward(ctx, context);
                 for (const auto& segment : layout.segments) {
                     ggml_tensor* h;
                     if (segment.image_index < 0) {
-                        h = ggml_ext_slice(ctx->ggml_ctx, text, 1, segment.context_start,
+                        h = ggml_ext_slice(ctx, text, 1, segment.context_start,
                                            segment.context_start + segment.end - segment.start);
                     } else {
                         auto image = segment.image_index == static_cast<int>(refs.size()) ? x : refs[segment.image_index];
-                        h          = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1));
+                        h          = img_in->forward(ctx, DiT::patchify(ctx, image, 1, 1));
                     }
                     joint = joint == nullptr ? h : ggml_concat(ctx->ggml_ctx, joint, h, 1);
                 }
@@ -359,13 +357,13 @@ namespace Qwen {
                 sd::ggml_graph_cut::mark_graph_cut(joint, layer_cache.cut_group, "joint");
             }
             if (cache.mode != QwenImage21PrefixCache::Mode::REUSE) {
-                joint = ggml_ext_slice(ctx->ggml_ctx, joint, 1, layout.prefix_length, joint->ne[1]);
+                joint = ggml_ext_slice(ctx, joint, 1, layout.prefix_length, joint->ne[1]);
             }
-            auto scale = std::dynamic_pointer_cast<Linear>(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx->ggml_ctx, time, 2, 1)[0]);
+            auto scale = std::dynamic_pointer_cast<Linear>(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx, time, 2, 1)[0]);
             joint      = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out.norm"])->forward(ctx, joint);
             joint      = ggml_mul(ctx->ggml_ctx, joint, ggml_scale_bias(ctx->ggml_ctx, scale, 1.f, 1.f));
             joint      = std::dynamic_pointer_cast<Linear>(blocks["proj_out"])->forward(ctx, joint);
-            return DiT::unpatchify_and_crop(ctx->ggml_ctx, joint, x->ne[1], x->ne[0], 1, 1);
+            return DiT::unpatchify_and_crop(ctx, joint, x->ne[1], x->ne[0], 1, 1);
         }
     };
 

src/model/diffusion/relay_vsr.hpp

@@ -124,7 +124,7 @@ namespace RelayVSR {
                 k   = VSR::apply_rope(ctx, k, key_rope, heads_);
                 out = local->forward(ctx, q, k, v, heads_);
             } else {
-                out = ggml_ext_attention_ext(c, ctx->backend, q, k, v, heads_, nullptr, false, ctx->flash_attn_enabled);
+                out = ggml_ext_attention_ext(ctx, q, k, v, heads_, nullptr, false, ctx->flash_attn_enabled);
             }
             return to_out_0->forward(ctx, out);
         }
@@ -157,7 +157,7 @@ namespace RelayVSR {
             auto ffn_net_2      = std::dynamic_pointer_cast<Linear>(blocks["ffn.net.2"]);
             auto c              = ctx->ggml_ctx;
             auto modulation     = ggml_add(c, timestep_proj, params["scale_shift_table"]);
-            auto pieces         = ggml_ext_chunk(c, modulation, 6, 1);
+            auto pieces         = ggml_ext_chunk(ctx, modulation, 6, 1);
             auto attn1          = std::dynamic_pointer_cast<WanAttention>(blocks["attn1"]);
             auto attn2          = std::dynamic_pointer_cast<WanAttention>(blocks["attn2"]);
             auto y              = attn1->forward(ctx, ggml_ext_modulate(ctx, norm1->forward(ctx, x), pieces[0], pieces[1]), nullptr, qpe, kpe, &local, name + ".attn1");
@@ -165,7 +165,7 @@ namespace RelayVSR {
             x                   = ggml_add(c, x, attn2->forward(ctx, norm2->forward(ctx, x), text, nullptr, nullptr, nullptr, name + ".attn2"));
             y                   = ggml_ext_modulate(ctx, norm3->forward(ctx, x), pieces[3], pieces[4]);
             y                   = ffn_net_0_proj->forward(ctx, y);
-            y                   = ffn_net_2->forward(ctx, ggml_ext_gelu(c, y));
+            y                   = ffn_net_2->forward(ctx, ggml_ext_gelu(ctx, y));
             return ggml_add(c, x, ggml_mul(c, y, pieces[5]));
         }
     };
@@ -213,7 +213,7 @@ namespace RelayVSR {
             auto e0   = condition_embedder_time_proj->forward(ctx, ggml_silu(c, e));
             e0        = ggml_reshape_3d(c, e0, config_.dim, 6, 1);
             auto text = condition_embedder_text_embedder_linear_1->forward(ctx, params["prompt_embeds"]);
-            text      = condition_embedder_text_embedder_linear_2->forward(ctx, ggml_ext_gelu(c, text));
+            text      = condition_embedder_text_embedder_linear_2->forward(ctx, ggml_ext_gelu(ctx, text));
             for (int i = 0; i < config_.num_layers; ++i) {
                 auto name = "blocks." + std::to_string(i);
                 x         = std::dynamic_pointer_cast<WanTransformerBlock>(blocks[name])->forward(ctx, x, text, e0, qpe, kpe, local, name);
@@ -221,7 +221,7 @@ namespace RelayVSR {
             }
             e               = ggml_reshape_3d(c, e, config_.dim, 1, 1);
             auto modulation = ggml_add(c, params["scale_shift_table"], e);
-            auto pieces     = ggml_ext_chunk(c, modulation, 2, 1);
+            auto pieces     = ggml_ext_chunk(ctx, modulation, 2, 1);
             x               = norm_out->forward(ctx, x);
             x               = ggml_ext_modulate(ctx, x, pieces[0], pieces[1]);
             x               = proj_out->forward(ctx, x);

src/model/diffusion/seedvr2.hpp

@@ -139,7 +139,7 @@ namespace SeedVR2 {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             auto up     = std::dynamic_pointer_cast<Linear>(blocks["proj_in"])->forward(ctx, x);
             auto hidden = swiglu_ ? ggml_swiglu_split(ctx->ggml_ctx, std::dynamic_pointer_cast<Linear>(blocks["proj_in_gate"])->forward(ctx, x), up)
-                                  : ggml_ext_gelu(ctx->ggml_ctx, up);
+                                  : ggml_ext_gelu(ctx, up);
             return std::dynamic_pointer_cast<Linear>(blocks["proj_out"])->forward(ctx, hidden);
         }
     };
@@ -200,7 +200,7 @@ namespace SeedVR2 {
             auto vpe  = rotation(ctx, inputs.input(ctx, plan.positions));
             auto tpe  = rotation(ctx, inputs.input(ctx, plan.text_positions));
             auto qkv  = [&](ggml_tensor* x, bool text, ggml_tensor* pe) {
-                auto parts = ggml_ext_chunk(c, x, 3, 0);
+                auto parts = ggml_ext_chunk(ctx, x, 3, 0);
                 for (int i = 0; i < 2; ++i) {
                     auto h   = ggml_reshape_3d(c, parts[i], 128, heads, x->ne[1]);
                     h        = std::dynamic_pointer_cast<RMSNorm>(blocks[std::string(i == 0 ? "attn.norm_q." : "attn.norm_k.") + branch(text)])->forward(ctx, h);
@@ -216,15 +216,15 @@ namespace SeedVR2 {
             for (int length : plan.lengths) {
                 std::vector<ggml_tensor*> joined;
                 for (int i = 0; i < 3; ++i)
-                    joined.push_back(ggml_concat(c, ggml_ext_slice(c, vp[i], 1, start, start + length), tp[i], 1));
-                auto out = ggml_ext_attention_ext(c, ctx->backend, joined[0], joined[1], joined[2], heads, nullptr, false, ctx->flash_attn_enabled);
-                auto vo  = ggml_ext_slice(c, out, 1, 0, length);
-                auto to  = ggml_ext_slice(c, out, 1, length, length + txt->ne[1]);
+                    joined.push_back(ggml_concat(c, ggml_ext_slice(ctx, vp[i], 1, start, start + length), tp[i], 1));
+                auto out = ggml_ext_attention_ext(ctx, joined[0], joined[1], joined[2], heads, nullptr, false, ctx->flash_attn_enabled);
+                auto vo  = ggml_ext_slice(ctx, out, 1, 0, length);
+                auto to  = ggml_ext_slice(ctx, out, 1, length, length + txt->ne[1]);
                 video_windows.push_back(vo);
                 tout = tout ? ggml_add(c, tout, to) : to;
                 start += length;
             }
-            auto vout = ggml_ext_vec_concat(c, video_windows, 1);
+            auto vout = ggml_ext_vec_concat(ctx, video_windows, 1);
             vout      = ggml_get_rows(c, vout, inputs.input(ctx, plan.reverse));
             tout      = ggml_scale(c, tout, 1.f / plan.lengths.size());
             vid       = ggml_add(c, vid, ada(linear("attn.proj_out", false, vout), false, "attn", true));

src/model/diffusion/sefi_image.hpp

@@ -77,8 +77,8 @@ namespace SefiImage {
             auto semantic_embedder = std::dynamic_pointer_cast<SefiTimestepEmbedding>(blocks["semantic_embedder"]);
             auto texture_embedder  = std::dynamic_pointer_cast<SefiTimestepEmbedding>(blocks["texture_embedder"]);
 
-            auto sem_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_sem, (int)timestep_guidance_in_dim, 10000, 1.f);
-            auto tex_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_tex, (int)timestep_guidance_in_dim, 10000, 1.f);
+            auto sem_proj = ggml_ext_timestep_embedding(ctx, timestep_sem, (int)timestep_guidance_in_dim, 10000, 1.f);
+            auto tex_proj = ggml_ext_timestep_embedding(ctx, timestep_tex, (int)timestep_guidance_in_dim, 10000, 1.f);
             auto sem_emb  = semantic_embedder->forward(ctx, sem_proj);
             auto tex_emb  = texture_embedder->forward(ctx, tex_proj);
             return ggml_concat(ctx->ggml_ctx, sem_emb, tex_emb, 0);

src/model/diffusion/sensenova_u1.h

@@ -143,7 +143,7 @@ namespace SenseNovaU1 {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timesteps) {
             auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
             auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
-            auto x     = ggml_ext_timestep_embedding(ctx->ggml_ctx,
+            auto x     = ggml_ext_timestep_embedding(ctx,
                                                      timesteps,
                                                      static_cast<int>(frequency_embedding_size),
                                                      10000,
@@ -166,8 +166,8 @@ namespace SenseNovaU1 {
         // head axis used by the RoPE kernel: [hidden, 1, tokens, batch].
         x                  = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], 1, x->ne[1], x->ne[2]);
         const int64_t half = x->ne[0] / 2;
-        auto x_part        = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, half);
-        auto y_part        = ggml_ext_slice(ctx->ggml_ctx, x, 0, half, x->ne[0]);
+        auto x_part        = ggml_ext_slice(ctx, x, 0, 0, half);
+        auto y_part        = ggml_ext_slice(ctx, x, 0, half, x->ne[0]);
         x_part             = ggml_rope_ext(ctx->ggml_ctx,
                                            x_part,
                                            position_x,
@@ -255,10 +255,10 @@ namespace SenseNovaU1 {
         const int64_t h = x->ne[1];
         const int64_t w = x->ne[0];
         GGML_ASSERT(x->ne[2] % (upscale_factor * upscale_factor) == 0);
-        x = ggml_ext_cont(ctx->ggml_ctx,
-                          ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
+        x = ggml_ext_cont(ctx,
+                          ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));
         x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]);
-        return DiT::unpatchify(ctx->ggml_ctx, x, h, w, upscale_factor, upscale_factor, true);
+        return DiT::unpatchify(ctx, x, h, w, upscale_factor, upscale_factor, true);
     }
 
     struct PixelDecoder : public GGMLBlock {
@@ -352,13 +352,13 @@ namespace SenseNovaU1 {
             const int64_t spatial_dim  = config.head_dim - temporal_dim;
             const int64_t axis_dim     = spatial_dim / 2;
 
-            auto temporal = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, temporal_dim);
-            auto spatial  = ggml_ext_slice(ctx->ggml_ctx, x, 0, temporal_dim, config.head_dim);
+            auto temporal = ggml_ext_slice(ctx, x, 0, 0, temporal_dim);
+            auto spatial  = ggml_ext_slice(ctx, x, 0, temporal_dim, config.head_dim);
             temporal      = std::dynamic_pointer_cast<LLM::LLMRMSNorm>(blocks[norm_name])->forward(ctx, temporal);
             spatial       = std::dynamic_pointer_cast<LLM::LLMRMSNorm>(blocks[norm_hw_name])->forward(ctx, spatial);
 
-            auto height = ggml_ext_slice(ctx->ggml_ctx, spatial, 0, 0, axis_dim);
-            auto width  = ggml_ext_slice(ctx->ggml_ctx, spatial, 0, axis_dim, spatial_dim);
+            auto height = ggml_ext_slice(ctx, spatial, 0, 0, axis_dim);
+            auto width  = ggml_ext_slice(ctx, spatial, 0, axis_dim, spatial_dim);
             temporal    = apply_axis_rope(ctx,
                                           temporal,
                                           position_t,
@@ -448,14 +448,13 @@ namespace SenseNovaU1 {
             }
 
             q = ggml_cont(ctx->ggml_ctx,
-                          ggml_ext_torch_permute(ctx->ggml_ctx, q, 0, 2, 1, 3));
+                          ggml_ext_torch_permute(ctx, q, 0, 2, 1, 3));
             q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]);
             k = ggml_cont(ctx->ggml_ctx,
-                          ggml_ext_torch_permute(ctx->ggml_ctx, k, 0, 2, 1, 3));
+                          ggml_ext_torch_permute(ctx, k, 0, 2, 1, 3));
             k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);
 
-            auto out = ggml_ext_attention_ext(ctx->ggml_ctx,
-                                              ctx->backend,
+            auto out = ggml_ext_attention_ext(ctx,
                                               q,
                                               k,
                                               v,

src/model/diffusion/unet.hpp

@@ -247,7 +247,7 @@ public:
 
         auto num_frames = ggml_arange(ctx->ggml_ctx, 0.f, static_cast<float>(timesteps), 1.f);
         // since b is 1, no need to do repeat
-        auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, num_frames, static_cast<int>(in_channels), max_time_embed_period);  // [N, in_channels]
+        auto t_emb = ggml_ext_timestep_embedding(ctx, num_frames, static_cast<int>(in_channels), max_time_embed_period);  // [N, in_channels]
 
         auto emb = time_pos_embed_0->forward(ctx, t_emb);
         emb      = ggml_silu_inplace(ctx->ggml_ctx, emb);
@@ -572,7 +572,7 @@ public:
         auto out_0 = std::dynamic_pointer_cast<GroupNorm32>(blocks["out.0"]);
         auto out_2 = std::dynamic_pointer_cast<Conv2d>(blocks["out.2"]);
 
-        auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, model_channels);  // [N, model_channels]
+        auto t_emb = ggml_ext_timestep_embedding(ctx, timesteps, model_channels);  // [N, model_channels]
 
         auto emb = time_embed_0->forward(ctx, t_emb);
         emb      = ggml_silu_inplace(ctx->ggml_ctx, emb);
@@ -685,7 +685,7 @@ public:
         sd::ggml_graph_cut::mark_graph_cut(h, "unet.middle_block", "h");
         if (controls.size() > 0) {
             // not in place, it segfaults the Vulkan backend with SDXL ControlNets (seen on RADV)
-            auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[controls.size() - 1], control_strength, false);
+            auto cs = ggml_ext_scale(ctx, controls[controls.size() - 1], control_strength, false);
             h       = ggml_add(ctx->ggml_ctx, h, cs);  // middle control
         }
         int control_offset = static_cast<int>(controls.size() - 2);
@@ -699,7 +699,7 @@ public:
 
                 if (controls.size() > 0) {
                     // not in place, see the middle control
-                    auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[control_offset], control_strength, false);
+                    auto cs = ggml_ext_scale(ctx, controls[control_offset], control_strength, false);
                     h_skip  = ggml_add(ctx->ggml_ctx, h_skip, cs);  // control net condition
                     control_offset--;
                 }

src/model/diffusion/wan.hpp

@@ -362,7 +362,7 @@ namespace WAN {
 
             auto modulation = params["modulation"];
             e               = ggml_add(ctx->ggml_ctx, e, modulation);  // [N, 6, dim] or [N, T, 6, dim]
-            auto es         = ggml_ext_chunk(ctx->ggml_ctx, e, 6, 1);  // ([N, 1, dim], ...) or [N, T, 1, dim]
+            auto es         = ggml_ext_chunk(ctx, e, 6, 1);            // ([N, 1, dim], ...) or [N, T, 1, dim]
 
             auto norm1      = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
             auto self_attn  = std::dynamic_pointer_cast<WanSelfAttention>(blocks["self_attn"]);
@@ -391,7 +391,7 @@ namespace WAN {
             y = modulate_add(ctx->ggml_ctx, y, es[3]);
 
             y = ffn_0->forward(ctx, y);
-            y = ggml_ext_gelu(ctx->ggml_ctx, y, true);
+            y = ggml_ext_gelu(ctx, y, true);
             y = ffn_2->forward(ctx, y);
 
             x = ggml_add(ctx->ggml_ctx, x, modulate_mul(ctx->ggml_ctx, y, es[5]));
@@ -484,7 +484,7 @@ namespace WAN {
             e               = ggml_repeat_4d(ctx->ggml_ctx, e, e->ne[0], 2, e->ne[2], e->ne[3]);   // [N, 2, dim] or [N, T, 2, dim]
 
             e       = ggml_add(ctx->ggml_ctx, e, modulation);  // [N, 2, dim] or [N, T, 2, dim]
-            auto es = ggml_ext_chunk(ctx->ggml_ctx, e, 2, 1);  // ([N, 1, dim], ...) or ([N, T, 1, dim], ...)
+            auto es = ggml_ext_chunk(ctx, e, 2, 1);            // ([N, 1, dim], ...) or ([N, T, 1, dim], ...)
 
             auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
             auto head = std::dynamic_pointer_cast<Linear>(blocks["head"]);
@@ -525,8 +525,8 @@ namespace WAN {
             if (flf_pos_embed_token_number > 0) {
                 auto emb_pos = params["emb_pos"];
 
-                auto a = ggml_ext_slice(ctx->ggml_ctx, image_embeds, 1, 0, emb_pos->ne[1]);
-                auto b = ggml_ext_slice(ctx->ggml_ctx, emb_pos, 1, 0, image_embeds->ne[1]);
+                auto a = ggml_ext_slice(ctx, image_embeds, 1, 0, emb_pos->ne[1]);
+                auto b = ggml_ext_slice(ctx, emb_pos, 1, 0, image_embeds->ne[1]);
 
                 image_embeds = ggml_add(ctx->ggml_ctx, a, b);
             }
@@ -538,7 +538,7 @@ namespace WAN {
 
             auto x = proj_0->forward(ctx, image_embeds);
             x      = proj_1->forward(ctx, x);
-            x      = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+            x      = ggml_ext_gelu(ctx, x, true);
             x      = proj_3->forward(ctx, x);
             x      = proj_4->forward(ctx, x);
 
@@ -639,15 +639,16 @@ namespace WAN {
             int pad_t = (std::get<0>(config.patch_size) - T % std::get<0>(config.patch_size)) % std::get<0>(config.patch_size);
             int pad_h = (std::get<1>(config.patch_size) - H % std::get<1>(config.patch_size)) % std::get<1>(config.patch_size);
             int pad_w = (std::get<2>(config.patch_size) - W % std::get<2>(config.patch_size)) % std::get<2>(config.patch_size);
-            ggml_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, pad_t, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
+            ggml_ext_pad(ctx, x, pad_w, pad_h, pad_t, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
             return x;
         }
 
-        ggml_tensor* unpatchify(ggml_context* ctx,
+        ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx,
                                 ggml_tensor* x,
                                 int64_t t_len,
                                 int64_t h_len,
                                 int64_t w_len) {
+            auto ctx = runner_ctx->ggml_ctx;
             // x: [N, t_len*h_len*w_len, pt*ph*pw*C]
             // return: [N*C, t_len*pt, h_len*ph, w_len*pw]
             int64_t N  = x->ne[3];
@@ -658,15 +659,15 @@ namespace WAN {
 
             GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
 
-            x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, pt*ph*pw, C]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
-            x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);  // [N*C*t_len*h_len, w_len, pt*ph, pw]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, pt*ph, w_len, pw]
-            x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);  // [N*C*t_len*h_len, pt, ph, w_len*pw]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, ph, pt, w_len*pw]
-            x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N);  // [N*C*t_len, h_len*ph, pt, w_len*pw]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len, pt, h_len*ph, w_len*pw]
-            x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N);  // [N*C, t_len*pt, h_len*ph, w_len*pw]
+            x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);            // [N, t_len*h_len*w_len, pt*ph*pw, C]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 1, 2, 0, 3));  // [N, C, t_len*h_len*w_len, pt*ph*pw]
+            x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);            // [N*C*t_len*h_len, w_len, pt*ph, pw]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len*h_len, pt*ph, w_len, pw]
+            x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);            // [N*C*t_len*h_len, pt, ph, w_len*pw]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len*h_len, ph, pt, w_len*pw]
+            x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N);            // [N*C*t_len, h_len*ph, pt, w_len*pw]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [N*C*t_len, pt, h_len*ph, w_len*pw]
+            x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N);            // [N*C, t_len*pt, h_len*ph, w_len*pw]
             return x;
         }
 
@@ -707,7 +708,7 @@ namespace WAN {
             // patch_embedding
             x = patch_embedding->forward(ctx, x);                                                    // [N*dim, t_len, h_len, w_len]
             x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], x->ne[3] / N, N);  // [N, dim, t_len*h_len*w_len]
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));  // [N, t_len*h_len*w_len, dim]
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3));                      // [N, t_len*h_len*w_len, dim]
 
             ggml_tensor* audio_local  = nullptr;
             ggml_tensor* audio_global = nullptr;
@@ -725,26 +726,26 @@ namespace WAN {
 
                 // video tokens get cond_mask[0], reference tokens cond_mask[1]
                 auto cond_mask = params["trainable_cond_mask.weight"];
-                auto cm0       = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, cond_mask, 1, 0, 1), config.dim, 1, 1);
+                auto cm0       = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx, cond_mask, 1, 0, 1), config.dim, 1, 1);
                 x              = ggml_add(ctx->ggml_ctx, x, cm0);
 
                 if (reference_latent != nullptr) {
                     t_ref_len = reference_latent->ne[2];
                     auto ref  = patch_embedding->forward(ctx, reference_latent);
                     ref       = ggml_reshape_3d(ctx->ggml_ctx, ref, ref->ne[0] * ref->ne[1] * ref->ne[2], ref->ne[3] / N, N);
-                    ref       = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ref, 1, 0, 2, 3));  // [N, t_ref*h_len*w_len, dim]
-                    auto cm1  = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, cond_mask, 1, 1, 2), config.dim, 1, 1);
+                    ref       = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, ref, 1, 0, 2, 3));  // [N, t_ref*h_len*w_len, dim]
+                    auto cm1  = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx, cond_mask, 1, 1, 2), config.dim, 1, 1);
                     ref       = ggml_add(ctx->ggml_ctx, ref, cm1);
                     x         = ggml_concat(ctx->ggml_ctx, x, ref, 1);
 
                     // Reference tokens use timestep 0.
                     GGML_ASSERT(timestep->ne[0] == T);
-                    timestep = ggml_ext_pad(ctx->ggml_ctx, timestep, (int)t_ref_len, 0, 0, 0);
+                    timestep = ggml_ext_pad(ctx, timestep, (int)t_ref_len, 0, 0, 0);
                 }
             }
 
             // time_embedding
-            auto e = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, config.freq_dim);
+            auto e = ggml_ext_timestep_embedding(ctx, timestep, config.freq_dim);
             e      = time_embedding_0->forward(ctx, e);
             e      = ggml_silu_inplace(ctx->ggml_ctx, e);
             e      = time_embedding_2->forward(ctx, e);  // [N, dim] or [N, T, dim]
@@ -755,7 +756,7 @@ namespace WAN {
             e0      = ggml_reshape_4d(ctx->ggml_ctx, e0, e0->ne[0] / 6, 6, e0->ne[1], e0->ne[2]);  //  [N, 6, dim] or [N, T, 6, dim]
 
             context = text_embedding_0->forward(ctx, context);
-            context = ggml_ext_gelu(ctx->ggml_ctx, context);
+            context = ggml_ext_gelu(ctx, context);
             context = text_embedding_2->forward(ctx, context);  // [N, context_txt_len, dim]
 
             int64_t context_img_len = 0;
@@ -775,7 +776,7 @@ namespace WAN {
 
                 c = vace_patch_embedding->forward(ctx, vace_context);                                    // [N*dim, t_len, h_len, w_len]
                 c = ggml_reshape_3d(ctx->ggml_ctx, c, c->ne[0] * c->ne[1] * c->ne[2], c->ne[3] / N, N);  // [N, dim, t_len*h_len*w_len]
-                c = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, c, 1, 0, 2, 3));  // [N, t_len*h_len*w_len, dim]
+                c = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, c, 1, 0, 2, 3));                      // [N, t_len*h_len*w_len, dim]
             }
             sd::ggml_graph_cut::mark_graph_cut(x, "wan.prelude", "x");
             // sd::ggml_graph_cut::mark_graph_cut(e, "wan.prelude", "e");
@@ -806,7 +807,7 @@ namespace WAN {
                     auto result = vace_block->forward(ctx, c, x_orig, e0, pe, context, context_img_len);
                     auto c_skip = result.first;
                     c           = result.second;
-                    c_skip      = ggml_ext_scale(ctx->ggml_ctx, c_skip, vace_strength);
+                    c_skip      = ggml_ext_scale(ctx, c_skip, vace_strength);
                     x           = ggml_add(ctx->ggml_ctx, x, c_skip);
                 }
 
@@ -870,16 +871,16 @@ namespace WAN {
 
             if (reference_latent != nullptr) {
                 // Exclude reference tokens from the generated video.
-                out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, t_len * h_len * w_len);
+                out = ggml_ext_slice(ctx, out, 1, 0, t_len * h_len * w_len);
             }
 
-            out = unpatchify(ctx->ggml_ctx, out, t_len, h_len, w_len);  // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w]
+            out = unpatchify(ctx, out, t_len, h_len, w_len);  // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w]
 
             // slice
 
-            out = ggml_ext_slice(ctx->ggml_ctx, out, 2, 0, T);  // [N*C, T, H + pad_h, W + pad_w]
-            out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H);  // [N*C, T, H, W + pad_w]
-            out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W);  // [N*C, T, H, W]
+            out = ggml_ext_slice(ctx, out, 2, 0, T);  // [N*C, T, H + pad_h, W + pad_w]
+            out = ggml_ext_slice(ctx, out, 1, 0, H);  // [N*C, T, H, W + pad_w]
+            out = ggml_ext_slice(ctx, out, 0, 0, W);  // [N*C, T, H, W]
 
             return out;
         }

src/model/diffusion/wan_audio.hpp

@@ -27,7 +27,7 @@ namespace WAN {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             // Replicate the first sample for causal left padding.
             if (kernel_size_ > 1) {
-                auto first = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, 1);
+                auto first = ggml_ext_slice(ctx, x, 0, 0, 1);
                 for (int i = 0; i < kernel_size_ - 1; i++) {
                     x = ggml_concat(ctx->ggml_ctx, first, x, 0);
                 }
@@ -57,7 +57,7 @@ namespace WAN {
             x = std::dynamic_pointer_cast<LayerNorm>(blocks[norm_key])->forward(ctx, x);
             x = ggml_silu(ctx->ggml_ctx, x);
             if (to_conv_layout) {
-                x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+                x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
             }
             return x;
         }
@@ -87,11 +87,11 @@ namespace WAN {
             auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
             std::vector<ggml_tensor*> tokens;
             // Each token group is normalized independently over channels.
-            for (auto& group : ggml_ext_chunk(ctx->ggml_ctx, local, num_token_, 1)) {
+            for (auto& group : ggml_ext_chunk(ctx, local, num_token_, 1)) {
                 ggml_tensor* s = ggml_permute(ctx->ggml_ctx, group, 1, 0, 2, 3);
                 s              = norm1->forward(ctx, s);
                 s              = ggml_silu(ctx->ggml_ctx, s);
-                s              = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, s, 1, 0, 2, 3));
+                s              = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, s, 1, 0, 2, 3));
                 s              = conv_norm_silu(ctx, s, "conv2", "norm2", true);
                 s              = conv_norm_silu(ctx, s, "conv3", "norm3", false);
                 tokens.push_back(ggml_reshape_3d(ctx->ggml_ctx, s, s->ne[0], 1, s->ne[1]));
@@ -99,7 +99,7 @@ namespace WAN {
             auto padding = ggml_reshape_3d(ctx->ggml_ctx, params["padding_tokens"], hidden_dim_, 1, 1);
             padding      = ggml_repeat(ctx->ggml_ctx, padding, tokens[0]);
             tokens.push_back(padding);
-            ggml_tensor* local_out = ggml_ext_vec_concat(ctx->ggml_ctx, tokens, 1);
+            ggml_tensor* local_out = ggml_ext_vec_concat(ctx, tokens, 1);
 
             if (!need_global_) {
                 return {local_out, nullptr};
@@ -141,10 +141,10 @@ namespace WAN {
             auto x       = ggml_mul(ctx->ggml_ctx, features, ggml_reshape_3d(ctx->ggml_ctx, weights, 1, 1, num_layers_));
             x            = ggml_div(ctx->ggml_ctx, x, ggml_sum(ctx->ggml_ctx, weights));
             // Move the layer axis to ggml dimension 0 for reduction.
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));
             x = ggml_sum_rows(ctx->ggml_ctx, x);
             x = ggml_reshape_2d(ctx->ggml_ctx, x, x->ne[1], x->ne[2]);
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3));
             return std::dynamic_pointer_cast<WanMotionEncoder>(blocks["encoder"])->forward(ctx, x);
         }
     };
@@ -187,12 +187,12 @@ namespace WAN {
 
             auto temb  = ggml_silu(ctx->ggml_ctx, audio_global);
             temb       = adain_linear->forward(ctx, temb);
-            auto shift = ggml_ext_slice(ctx->ggml_ctx, temb, 0, 0, dim_);
-            auto scale = ggml_ext_slice(ctx->ggml_ctx, temb, 0, dim_, dim_ * 2);
+            auto shift = ggml_ext_slice(ctx, temb, 0, 0, dim_);
+            auto scale = ggml_ext_slice(ctx, temb, 0, dim_, dim_ * 2);
             shift      = ggml_reshape_3d(ctx->ggml_ctx, shift, dim_, 1, T);
             scale      = ggml_reshape_3d(ctx->ggml_ctx, scale, dim_, 1, T);
 
-            auto x_vid = ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, seq_len);
+            auto x_vid = ggml_ext_slice(ctx, x, 1, 0, seq_len);
             auto h     = ggml_reshape_3d(ctx->ggml_ctx, x_vid, dim_, n_tok, T);
             h          = adain_norm->forward(ctx, h);
             h          = ggml_add(ctx->ggml_ctx, h, ggml_mul(ctx->ggml_ctx, h, scale));
@@ -203,7 +203,7 @@ namespace WAN {
 
             auto x_head = ggml_add(ctx->ggml_ctx, x_vid, res);
             if (seq_len < n_token) {
-                auto x_tail = ggml_ext_slice(ctx->ggml_ctx, x, 1, seq_len, n_token);
+                auto x_tail = ggml_ext_slice(ctx, x, 1, seq_len, n_token);
                 return ggml_concat(ctx->ggml_ctx, x_head, x_tail, 1);
             }
             return x_head;

src/model/diffusion/z_image.hpp

@@ -179,12 +179,10 @@ namespace ZImage {
                 auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
                 auto v_proj = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
 
-                if (sd_backend_is(ctx->backend, "ROCm")) {
-                    out_proj->set_force_prec_f32(true);
-                    q_proj->set_force_prec_f32(true);
-                    k_proj->set_force_prec_f32(true);
-                    v_proj->set_force_prec_f32(true);
-                }
+                out_proj->set_precision_flags(GGML_PRECISION_F32_ROCM);
+                q_proj->set_precision_flags(GGML_PRECISION_F32_ROCM);
+                k_proj->set_precision_flags(GGML_PRECISION_F32_ROCM);
+                v_proj->set_precision_flags(GGML_PRECISION_F32_ROCM);
 
                 auto q = ggml_reshape_4d(ctx->ggml_ctx, q_proj->forward(ctx, x), head_dim, num_heads, n_token, N);
                 auto k = ggml_reshape_4d(ctx->ggml_ctx, k_proj->forward(ctx, x), head_dim, num_kv_heads, n_token, N);
@@ -201,10 +199,8 @@ namespace ZImage {
 
             auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
 
-            if (sd_backend_is(ctx->backend, "ROCm")) {
-                out_proj->set_force_prec_f32(true);
-                qkv_proj->set_force_prec_f32(true);
-            }
+            out_proj->set_precision_flags(GGML_PRECISION_F32_ROCM);
+            qkv_proj->set_precision_flags(GGML_PRECISION_F32_ROCM);
 
             auto qkv = qkv_proj->forward(ctx, x);                                                                            // [N, n_token, (num_heads + num_kv_heads*2)*head_dim]
             qkv      = ggml_reshape_4d(ctx->ggml_ctx, qkv, head_dim, num_heads + num_kv_heads * 2, qkv->ne[1], qkv->ne[2]);  // [N, n_token, num_heads + num_kv_heads*2, head_dim]
@@ -281,9 +277,7 @@ namespace ZImage {
             auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
             auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
 
-            if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
-                w2->set_force_prec_f32(true);
-            }
+            w2->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM);
 
             auto x1 = w1->forward(ctx, x);
             auto x3 = w3->forward(ctx, x);
@@ -353,7 +347,7 @@ namespace ZImage {
                 auto adaLN_modulation_0 = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.0"]);
 
                 auto m         = adaLN_modulation_0->forward(ctx, adaln_input);  // [N, 4 * hidden_size]
-                auto mods      = ggml_ext_chunk(ctx->ggml_ctx, m, 4, 0);
+                auto mods      = ggml_ext_chunk(ctx, m, 4, 0);
                 auto scale_msa = mods[0];
                 auto gate_msa  = mods[1];
                 auto scale_mlp = mods[2];
@@ -457,12 +451,10 @@ namespace ZImage {
             int64_t n_token = c->ne[1];
             int64_t N       = c->ne[2];
 
-            if (sd_backend_is(ctx->backend, "ROCm")) {
-                out_proj->set_force_prec_f32(true);
-                to_q->set_force_prec_f32(true);
-                to_k->set_force_prec_f32(true);
-                to_v->set_force_prec_f32(true);
-            }
+            out_proj->set_precision_flags(GGML_PRECISION_F32_ROCM);
+            to_q->set_precision_flags(GGML_PRECISION_F32_ROCM);
+            to_k->set_precision_flags(GGML_PRECISION_F32_ROCM);
+            to_v->set_precision_flags(GGML_PRECISION_F32_ROCM);
 
             auto q = ggml_reshape_4d(ctx->ggml_ctx, to_q->forward(ctx, c), head_dim, num_heads, n_token, N);
             auto k = ggml_reshape_4d(ctx->ggml_ctx, to_k->forward(ctx, c), head_dim, num_kv_heads, n_token, N);
@@ -497,7 +489,7 @@ namespace ZImage {
             }
 
             auto m         = adaLN_modulation_0->forward(ctx, t_emb);
-            auto mods      = ggml_ext_chunk(ctx->ggml_ctx, m, 4, 0);
+            auto mods      = ggml_ext_chunk(ctx, m, 4, 0);
             auto scale_msa = mods[0];
             auto gate_msa  = mods[1];
             auto scale_mlp = mods[2];
@@ -713,7 +705,7 @@ namespace ZImage {
 
                 auto out = block->forward(ctx, c, img_pe, img, t_emb);
                 c        = out.second;
-                ref_hints.push_back(ggml_ext_scale(ctx->ggml_ctx, out.first, strength));
+                ref_hints.push_back(ggml_ext_scale(ctx, out.first, strength));
                 sd::ggml_graph_cut::mark_graph_cut(c, "z_image.control.noise_refiner." + std::to_string(i), "control");
             }
 
@@ -757,7 +749,7 @@ namespace ZImage {
 
                 auto out     = block->forward(ctx, ctrl_unified, pe, unified_c, t_emb);
                 ctrl_unified = out.second;
-                main_hints.push_back(ggml_ext_scale(ctx->ggml_ctx, out.first, strength));
+                main_hints.push_back(ggml_ext_scale(ctx, out.first, strength));
                 sd::ggml_graph_cut::mark_graph_cut(ctrl_unified, "z_image.control.layers." + std::to_string(i), "control");
             }
 
@@ -807,8 +799,8 @@ namespace ZImage {
 
             GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]);
 
-            auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]);
-            auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]);
+            auto txt_pe = ggml_ext_slice(ctx, pe, 3, 0, txt->ne[1]);
+            auto img_pe = ggml_ext_slice(ctx, pe, 3, txt->ne[1], pe->ne[3]);
 
             std::vector<ggml_tensor*> ref_hints;
             ggml_tensor* control_ctx = nullptr;
@@ -858,7 +850,7 @@ namespace ZImage {
 
             txt_img = final_layer->forward(ctx, txt_img, t_emb);  // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, ph*pw*C]
 
-            img = ggml_ext_slice(ctx->ggml_ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token);  // [N, n_img_token, ph*pw*C]
+            img = ggml_ext_slice(ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token);  // [N, n_img_token, ph*pw*C]
 
             return img;
         }
@@ -906,10 +898,10 @@ namespace ZImage {
 
             auto out = forward_core(ctx, img, timestep, context, pe, control, control_strength);
 
-            out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, n_img_token);                              // [N, n_img_token, ph*pw*C]
-            out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size, false);  // [N, C, H, W]
+            out = ggml_ext_slice(ctx, out, 1, 0, n_img_token);                              // [N, n_img_token, ph*pw*C]
+            out = DiT::unpatchify_and_crop(ctx, out, H, W, patch_size, patch_size, false);  // [N, C, H, W]
 
-            out = ggml_ext_scale(ctx->ggml_ctx, out, -1.f);
+            out = ggml_ext_scale(ctx, out, -1.f);
 
             return out;
         }

src/model/diffusion/z_image_l2p.hpp

@@ -134,7 +134,7 @@ namespace ZImageL2P {
             const int64_t n_img = w * h;
             const int64_t n_txt = context->ne[1];
 
-            auto img = DiT::patchify(gctx, x, config.patch_size, config.patch_size, false);
+            auto img = DiT::patchify(ctx, x, config.patch_size, config.patch_size, false);
             img      = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"])->forward(ctx, img);
             auto txt = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"])->forward(ctx, context);
             txt      = std::dynamic_pointer_cast<Linear>(blocks["cap_embedder.1"])->forward(ctx, txt);
@@ -155,8 +155,8 @@ namespace ZImageL2P {
             }
             GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]);
 
-            auto txt_pe = ggml_ext_slice(gctx, pe, 3, 0, txt->ne[1]);
-            auto img_pe = ggml_ext_slice(gctx, pe, 3, txt->ne[1], pe->ne[3]);
+            auto txt_pe = ggml_ext_slice(ctx, pe, 3, 0, txt->ne[1]);
+            auto img_pe = ggml_ext_slice(ctx, pe, 3, txt->ne[1], pe->ne[3]);
             for (int64_t i = 0; i < config.num_refiner_layers; ++i) {
                 txt = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)])->forward(ctx, txt, txt_pe);
                 sd::ggml_graph_cut::mark_graph_cut(txt, "z_image_l2p.context_refiner." + std::to_string(i), "txt");
@@ -175,11 +175,11 @@ namespace ZImageL2P {
             // The local decoder consumes the raw last-layer hidden states: there is no final
             // norm or adaLN modulation in front of it, unlike Z-Image's final_layer.
             const int64_t img_start = n_txt + n_txt_pad;
-            auto feat               = ggml_ext_slice(gctx, txt_img, 1, img_start, img_start + n_img);                           // [N, h*w, hidden_size]
+            auto feat               = ggml_ext_slice(ctx, txt_img, 1, img_start, img_start + n_img);                            // [N, h*w, hidden_size]
             feat                    = ggml_reshape_4d(gctx, feat, config.hidden_size, w, h, N);                                 // [N, h, w, hidden_size]
             feat                    = ggml_cont(gctx, ggml_permute(gctx, feat, 2, 0, 1, 3));                                    // [N, hidden_size, h, w]
             auto out                = std::dynamic_pointer_cast<LocalDecoder>(blocks["local_decoder"])->forward(ctx, x, feat);  // [N, C, H, W]
-            return ggml_ext_scale(gctx, out, -1.f);
+            return ggml_ext_scale(ctx, out, -1.f);
         }
     };
 

src/model/te/clip.hpp

@@ -35,9 +35,9 @@ public:
 
         x = fc1->forward(ctx, x);
         if (use_gelu) {
-            x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+            x = ggml_ext_gelu(ctx, x, true);
         } else {
-            x = ggml_ext_gelu_quick(ctx->ggml_ctx, x, true);
+            x = ggml_ext_gelu_quick(ctx, x, true);
         }
         x = fc2->forward(ctx, x);
         return x;
@@ -228,10 +228,10 @@ public:
         // concat(patch_embedding, class_embedding) + position_embedding
         ggml_tensor* patch_embedding;
         int64_t N       = pixel_values->ne[3];
-        patch_embedding = ggml_ext_conv_2d(ctx->ggml_ctx, pixel_values, patch_embed_weight, nullptr, patch_size, patch_size);  // [N, embed_dim, image_size // pacht_size, image_size // pacht_size]
-        patch_embedding = ggml_reshape_3d(ctx->ggml_ctx, patch_embedding, num_patches, embed_dim, N);                          // [N, embed_dim, num_patches]
-        patch_embedding = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, patch_embedding, 1, 0, 2, 3));                  // [N, num_patches, embed_dim]
-        patch_embedding = ggml_reshape_4d(ctx->ggml_ctx, patch_embedding, 1, embed_dim, num_patches, N);                       // [N, num_patches, embed_dim, 1]
+        patch_embedding = ggml_ext_conv_2d(ctx, pixel_values, patch_embed_weight, nullptr, patch_size, patch_size);  // [N, embed_dim, image_size // pacht_size, image_size // pacht_size]
+        patch_embedding = ggml_reshape_3d(ctx->ggml_ctx, patch_embedding, num_patches, embed_dim, N);                // [N, embed_dim, num_patches]
+        patch_embedding = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, patch_embedding, 1, 0, 2, 3));        // [N, num_patches, embed_dim]
+        patch_embedding = ggml_reshape_4d(ctx->ggml_ctx, patch_embedding, 1, embed_dim, num_patches, N);             // [N, num_patches, embed_dim, 1]
 
         ggml_tensor* class_embedding = ggml_new_tensor_2d(ctx->ggml_ctx, GGML_TYPE_F32, embed_dim, N);
         class_embedding              = ggml_repeat(ctx->ggml_ctx, class_embed_weight, class_embedding);      // [N, embed_dim]
@@ -325,7 +325,7 @@ public:
             auto text_projection = params["text_projection"];
             ggml_tensor* pooled  = ggml_view_1d(ctx->ggml_ctx, x, hidden_size, x->nb[1] * max_token_idx);
             if (text_projection != nullptr) {
-                pooled = ggml_ext_linear(ctx->ggml_ctx, pooled, text_projection, nullptr);
+                pooled = ggml_ext_linear(ctx, pooled, text_projection, nullptr);
             } else {
                 LOG_VERBOSE("identity projection");
             }
@@ -426,7 +426,7 @@ public:
         if (transpose_weight) {
             w = ggml_cont(ctx->ggml_ctx, ggml_transpose(ctx->ggml_ctx, w));
         }
-        return ggml_ext_linear(ctx->ggml_ctx, x, w, nullptr);
+        return ggml_ext_linear(ctx, x, w, nullptr);
     }
 };
 

src/model/te/llada_image_te.hpp

@@ -78,15 +78,15 @@ namespace LLaDAImageTE {
             auto b = params["in_proj_bias"];
 
             auto slice_w = [&](int64_t index) {
-                return ggml_ext_slice(gctx, w, 1, index * hidden_size, (index + 1) * hidden_size);
+                return ggml_ext_slice(ctx, w, 1, index * hidden_size, (index + 1) * hidden_size);
             };
             auto slice_b = [&](int64_t index) {
-                return ggml_ext_slice(gctx, b, 0, index * hidden_size, (index + 1) * hidden_size);
+                return ggml_ext_slice(ctx, b, 0, index * hidden_size, (index + 1) * hidden_size);
             };
 
-            auto q = ggml_ext_linear(gctx, query, slice_w(0), slice_b(0));
-            auto k = ggml_ext_linear(gctx, context, slice_w(1), slice_b(1));
-            auto v = ggml_ext_linear(gctx, context, slice_w(2), slice_b(2));
+            auto q = ggml_ext_linear(ctx, query, slice_w(0), slice_b(0));
+            auto k = ggml_ext_linear(ctx, context, slice_w(1), slice_b(1));
+            auto v = ggml_ext_linear(ctx, context, slice_w(2), slice_b(2));
 
             auto x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask);  // [N, num_queries, hidden_size]
             return out_proj->forward(ctx, x);
@@ -297,9 +297,9 @@ namespace LLaDAImageTE {
 
             int64_t hidden_size = num_heads * head_dim;
             auto qkv            = qkv_proj->forward(ctx, x);
-            auto q              = ggml_ext_slice(gctx, qkv, 0, 0, hidden_size);
-            auto k              = ggml_ext_slice(gctx, qkv, 0, hidden_size, hidden_size * 2);
-            auto v              = ggml_ext_slice(gctx, qkv, 0, hidden_size * 2, hidden_size * 3);
+            auto q              = ggml_ext_slice(ctx, qkv, 0, 0, hidden_size);
+            auto k              = ggml_ext_slice(ctx, qkv, 0, hidden_size, hidden_size * 2);
+            auto v              = ggml_ext_slice(ctx, qkv, 0, hidden_size * 2, hidden_size * 3);
 
             auto out = ggml_ext_attention_ext(ctx, q, k, v, num_heads);
             return out_proj->forward(ctx, out);

src/model/te/llm.hpp

@@ -485,7 +485,7 @@ namespace LLM {
         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");
+                w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight");
             }
             x           = ggml_rms_norm(ctx->ggml_ctx, x, eps);
             auto scaled = ggml_mul(ctx->ggml_ctx, x, w);
@@ -941,7 +941,7 @@ namespace LLM {
             // An embed can sit flush against the previous one or at the very start/end of the
             // sequence, leaving no text tokens to splice around it.
             if (txt_token_end > txt_token_start) {
-                auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
+                auto txt_embed = ggml_ext_slice(ctx, raw_x, 1, txt_token_start, txt_token_end);
                 if (input_embed == nullptr) {
                     input_embed = txt_embed;
                 } else {
@@ -960,7 +960,7 @@ namespace LLM {
         txt_token_end   = raw_x->ne[1];
 
         if (txt_token_end > txt_token_start) {
-            auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
+            auto final_txt_embed = ggml_ext_slice(ctx, raw_x, 1, txt_token_start, txt_token_end);
             input_embed          = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1);
         }
         GGML_ASSERT(raw_x->ne[1] + removed_tokens == input_embed->ne[1]);
@@ -982,16 +982,16 @@ namespace LLM {
             GGML_ASSERT(index >= token_start);
             GGML_ASSERT(index + image_embed->ne[1] <= raw_x->ne[1]);
             if (index > token_start) {
-                auto text_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, token_start, index);
+                auto text_embed = ggml_ext_slice(ctx, raw_x, 1, token_start, index);
                 output          = output == nullptr ? text_embed : ggml_concat(ctx->ggml_ctx, output, text_embed, 1);
             }
-            auto visual_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, index, index + image_embed->ne[1]);
+            auto visual_embed = ggml_ext_slice(ctx, raw_x, 1, index, index + image_embed->ne[1]);
             visual_embed      = ggml_add(ctx->ggml_ctx, visual_embed, image_embed);
             output            = output == nullptr ? visual_embed : ggml_concat(ctx->ggml_ctx, output, visual_embed, 1);
             token_start       = index + image_embed->ne[1];
         }
         if (token_start < raw_x->ne[1]) {
-            auto text_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, token_start, raw_x->ne[1]);
+            auto text_embed = ggml_ext_slice(ctx, raw_x, 1, token_start, raw_x->ne[1]);
             output          = output == nullptr ? text_embed : ggml_concat(ctx->ggml_ctx, output, text_embed, 1);
         }
         GGML_ASSERT(output != nullptr && output->ne[1] == raw_x->ne[1]);
@@ -1020,7 +1020,7 @@ namespace LLM {
                 auto linear_fc1 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc1"]);
                 auto linear_fc2 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc2"]);
                 x               = linear_fc1->forward(ctx, x);
-                x               = ggml_ext_gelu(ctx->ggml_ctx, x);
+                x               = ggml_ext_gelu(ctx, x);
                 x               = linear_fc2->forward(ctx, x);
             } else {
                 auto gate_proj = std::dynamic_pointer_cast<Linear>(blocks["gate_proj"]);
@@ -1105,11 +1105,11 @@ namespace LLM {
                 auto proj_0 = std::dynamic_pointer_cast<Conv2d>(blocks["proj.0"]);
                 auto proj_1 = std::dynamic_pointer_cast<Conv2d>(blocks["proj.1"]);
 
-                auto x0 = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                auto x0 = ggml_ext_slice(ctx, x, 2, 0, 1);
                 x0      = ggml_reshape_4d(ctx->ggml_ctx, x0, x0->ne[0], x0->ne[1], in_channels, x0->ne[3] / in_channels);
                 x0      = proj_0->forward(ctx, x0);
 
-                auto x1 = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, 2);
+                auto x1 = ggml_ext_slice(ctx, x, 2, 1, 2);
                 x1      = ggml_reshape_4d(ctx->ggml_ctx, x1, x1->ne[0], x1->ne[1], in_channels, x1->ne[3] / in_channels);
                 x1      = proj_1->forward(ctx, x1);
 
@@ -1173,7 +1173,7 @@ namespace LLM {
             x = ln_q->forward(ctx, x);
             x = ggml_reshape_2d(ctx->ggml_ctx, x, hidden_size, ggml_nelements(x) / hidden_size);
             x = mlp_0->forward(ctx, x);
-            x = ggml_ext_gelu(ctx->ggml_ctx, x);
+            x = ggml_ext_gelu(ctx, x);
             x = mlp_2->forward(ctx, x);
             return x;
         }
@@ -1252,7 +1252,7 @@ namespace LLM {
             } else {
                 auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
                 auto qkv      = qkv_proj->forward(ctx, x);
-                qkv_vec       = split_qkv(ctx->ggml_ctx, qkv);
+                qkv_vec       = ggml_ext_split_qkv(ctx, qkv);
             }
 
             auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, qkv_vec[0]->ne[1], qkv_vec[0]->ne[2]);  // [N, n_token, n_head, d_head]
@@ -1487,12 +1487,12 @@ namespace LLM {
             for (int i = 0; i < kernel; ++i) {
                 auto w    = ggml_cont(c, ggml_view_2d(c, params["conv1d.weight"], 1, conv_dim, kernel * sizeof(float), i * sizeof(float)));
                 w         = ggml_reshape_1d(c, w, conv_dim);
-                auto part = ggml_mul(c, ggml_ext_slice(c, padded, 1, i, i + n), w);
+                auto part = ggml_mul(c, ggml_ext_slice(ctx, padded, 1, i, i + n), w);
                 mixed     = mixed ? ggml_add(c, mixed, part) : part;
             }
             mixed      = ggml_silu(c, mixed);
             auto heads = [&](int start, int count) {
-                auto h = ggml_ext_slice(c, mixed, 0, start, start + head_dim * count);
+                auto h = ggml_ext_slice(ctx, mixed, 0, start, start + head_dim * count);
                 return ggml_reshape_4d(c, h, head_dim, count, n, 1);
             };
             auto repeat_keys = [&](ggml_tensor* h) {
@@ -1514,7 +1514,7 @@ namespace LLM {
             auto a     = gguf_layout ? params["A"] : ggml_neg(c, ggml_exp(c, params["A_log"]));
             decay      = ggml_mul(c, decay, a);
             decay      = ggml_reshape_4d(c, decay, 1, value_heads, n, 1);
-            auto state = ggml_ext_zeros(c, head_dim, head_dim, value_heads, 1);
+            auto state = ggml_ext_zeros(ctx, head_dim, head_dim, value_heads, 1);
             auto out   = ggml_gated_delta_net(c, q, k, v, decay, beta, state, 1);
             out        = ggml_view_3d(c, out, head_dim, value_heads, n, head_dim * sizeof(float), head_dim * value_heads * sizeof(float), 0);
             out        = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["norm"])->forward(ctx, out);
@@ -1594,13 +1594,14 @@ namespace LLM {
         // ggml applies theta_i / freq_factors[i], so a huge factor collapses the angle to zero and
         // leaves that pair unrotated. This reproduces transformers' "proportional" RoPE, whose
         // inv_freq is zero-padded past `rope_pairs`, without reordering the head.
-        ggml_tensor* rope_freq_factors(ggml_context* ctx) const {
+        ggml_tensor* rope_freq_factors(GGMLRunnerContext* runner_ctx) const {
+            auto ctx  = runner_ctx->ggml_ctx;
             int pairs = head_dim / 2;
             if (rope_pairs >= pairs) {
                 return nullptr;
             }
-            auto rotated   = ggml_ext_ones(ctx, rope_pairs, 1, 1, 1);
-            auto unrotated = ggml_ext_full(ctx, 1e30f, pairs - rope_pairs, 1, 1, 1);
+            auto rotated   = ggml_ext_ones(runner_ctx, rope_pairs, 1, 1, 1);
+            auto unrotated = ggml_ext_full(runner_ctx, 1e30f, pairs - rope_pairs, 1, 1, 1);
             return ggml_concat(ctx, rotated, unrotated, 0);
         }
 
@@ -1624,9 +1625,9 @@ namespace LLM {
                 auto qkv      = qkv_proj->forward(ctx, x);  // [N, n_token, (num_heads + num_kv_heads*2)*head_dim]
                 int64_t q_len = num_heads * head_dim;
                 int64_t k_len = num_kv_heads * head_dim;
-                q             = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, 0, q_len);
-                k             = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, q_len, q_len + k_len);
-                v             = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, q_len + k_len, q_len + k_len * 2);
+                q             = ggml_ext_slice(ctx, qkv, 0, 0, q_len);
+                k             = ggml_ext_slice(ctx, qkv, 0, q_len, q_len + k_len);
+                v             = ggml_ext_slice(ctx, qkv, 0, q_len + k_len, q_len + k_len * 2);
             } else {
                 auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["q_proj"]);
                 auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
@@ -1640,9 +1641,9 @@ namespace LLM {
             ggml_tensor* output_gate = nullptr;
             if (arch == LLMArch::QWEN3_5) {
                 q           = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim * 2, num_heads, n_token, N);
-                output_gate = ggml_ext_slice(ctx->ggml_ctx, q, 0, head_dim, head_dim * 2);
+                output_gate = ggml_ext_slice(ctx, q, 0, head_dim, head_dim * 2);
                 output_gate = ggml_reshape_3d(ctx->ggml_ctx, output_gate, head_dim * num_heads, n_token, N);
-                q           = ggml_ext_slice(ctx->ggml_ctx, q, 0, 0, head_dim);
+                q           = ggml_ext_slice(ctx, q, 0, 0, head_dim);
             }
             q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, n_token, N);     // [N, n_token, num_heads, head_dim]
             k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_kv_heads, n_token, N);  // [N, n_token, num_kv_heads, head_dim]
@@ -1666,12 +1667,12 @@ namespace LLM {
                 auto rotate = [&](ggml_tensor* input) {
                     auto gctx    = ctx->ggml_ctx;
                     int64_t half = rope_cos->ne[0];
-                    auto first   = ggml_ext_slice(gctx, input, 0, 0, half);
-                    auto second  = ggml_ext_slice(gctx, input, 0, half, 2 * half);
+                    auto first   = ggml_ext_slice(ctx, input, 0, 0, half);
+                    auto second  = ggml_ext_slice(ctx, input, 0, half, 2 * half);
                     auto left    = ggml_sub(gctx, ggml_mul(gctx, first, rope_cos), ggml_mul(gctx, second, rope_sin));
                     auto right   = ggml_add(gctx, ggml_mul(gctx, second, rope_cos), ggml_mul(gctx, first, rope_sin));
                     auto rotated = ggml_concat(gctx, left, right, 0);
-                    return ggml_concat(gctx, rotated, ggml_ext_slice(gctx, input, 0, 2 * half, input->ne[0]), 0);
+                    return ggml_concat(gctx, rotated, ggml_ext_slice(ctx, input, 0, 2 * half, input->ne[0]), 0);
                 };
                 q = rotate(q);
                 k = rotate(k);
@@ -1746,7 +1747,7 @@ namespace LLM {
                                                  1.f);
             } else if (arch == LLMArch::GEMMA4_12B) {
                 float rope_theta  = (rope_index == 1 ? 10000.0f : 1000000.0f);
-                auto freq_factors = rope_freq_factors(ctx->ggml_ctx);
+                auto freq_factors = rope_freq_factors(ctx);
                 q                 = ggml_rope_ext(ctx->ggml_ctx,
                                                   q,
                                                   input_pos,
@@ -1844,18 +1845,18 @@ namespace LLM {
 
             if (unscaled_attention) {
                 // Gemma 4 attends with scaling=1.0; undo the helper's own 1/sqrt(head_dim).
-                q = ggml_ext_scale(ctx->ggml_ctx, q, std::sqrt(static_cast<float>(head_dim)));
+                q = ggml_ext_scale(ctx, q, std::sqrt(static_cast<float>(head_dim)));
             }
 
-            q = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, q, 0, 2, 1, 3));  // [N, num_heads, n_token, head_dim]
-            q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]);      // [N*num_heads, n_token, head_dim]
+            q = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, q, 0, 2, 1, 3));        // [N, num_heads, n_token, head_dim]
+            q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]);  // [N*num_heads, n_token, head_dim]
 
-            k = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, k, 0, 2, 1, 3));  // [N, num_kv_heads, n_token, head_dim]
-            k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);      // [N*num_kv_heads, n_token, head_dim]
+            k = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, k, 0, 2, 1, 3));        // [N, num_kv_heads, n_token, head_dim]
+            k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);  // [N*num_kv_heads, n_token, head_dim]
 
             if (arch == LLMArch::GPT_OSS_20B) {
                 GGML_ASSERT(N == 1);
-                auto v_attn = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, v, 1, 2, 0, 3));  // [N, kv_heads, head_dim, tokens]
+                auto v_attn = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, v, 1, 2, 0, 3));  // [N, kv_heads, head_dim, tokens]
                 v_attn      = ggml_reshape_3d(ctx->ggml_ctx, v_attn, n_token, head_dim, num_kv_heads * N);
 
                 auto kq = ggml_mul_mat(ctx->ggml_ctx, k, q);
@@ -1870,7 +1871,7 @@ namespace LLM {
                 auto kqv = ggml_mul_mat(ctx->ggml_ctx, v_attn, kq);
                 kqv      = ggml_reshape_4d(ctx->ggml_ctx, kqv, head_dim, n_token, num_heads, N);
                 kqv      = ggml_permute(ctx->ggml_ctx, kqv, 0, 2, 1, 3);
-                x        = ggml_ext_cont(ctx->ggml_ctx, kqv);
+                x        = ggml_ext_cont(ctx, kqv);
                 x        = ggml_reshape_3d(ctx->ggml_ctx, x, head_dim * num_heads, n_token, N);
             } else {
                 x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled);  // [N, n_token, hidden_size]
@@ -2059,7 +2060,7 @@ namespace LLM {
             std::vector<ggml_tensor*> intermediate_outputs;
 
             if (config.normalize_input) {
-                x = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(static_cast<float>(config.hidden_size)), true);
+                x = ggml_ext_scale(ctx, x, std::sqrt(static_cast<float>(config.hidden_size)), true);
             }
             if (return_all_hidden_states) {
                 intermediate_outputs.push_back(x);
@@ -2207,9 +2208,10 @@ namespace LLM {
         std::array<std::vector<int32_t>, 4> pos_embed_idx_data_;
         std::array<std::vector<float>, 4> pos_embed_weight_data_;
 
-        static ggml_tensor* process_image_common(ggml_context* ctx,
+        static ggml_tensor* process_image_common(GGMLRunnerContext* runner_ctx,
                                                  ggml_tensor* image,
                                                  const LLMVisionConfig& vision_params) {
+            auto ctx = runner_ctx->ggml_ctx;
             // image: [C, H, W]
             // return: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw], grid_t == 1
             int64_t C  = image->ne[2];
@@ -2222,16 +2224,16 @@ namespace LLM {
             int64_t pw = vision_params.patch_size;
 
             image = ggml_reshape_4d(ctx, image, pw, mw, (W / mw / pw), H * C);                               // [C*H, (W/mw/pw), mw, pw]
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1));                          // [mw, C*H, (W/mw/pw), pw]
+            image = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, image, 0, 2, 3, 1));                   // [mw, C*H, (W/mw/pw), pw]
             image = ggml_reshape_4d(ctx, image, pw * (W / mw / pw), H, C, mw);                               // [mw, C, H, (W/mw/pw)*pw]
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1));                          // [H, mw, C, (W/mw/pw)*pw]
+            image = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, image, 0, 2, 3, 1));                   // [H, mw, C, (W/mw/pw)*pw]
             image = ggml_reshape_4d(ctx, image, pw, (W / mw / pw) * C * mw, ph, mh * (H / mh / ph));         // [(H/mh/ph)*mh, ph, mw*C*(W/mw/pw), pw]
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3));                          // [(H/mh/ph)*mh, mw*C*(W/mw/pw), ph, pw]
+            image = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, image, 0, 2, 1, 3));                   // [(H/mh/ph)*mh, mw*C*(W/mw/pw), ph, pw]
             image = ggml_reshape_4d(ctx, image, pw * ph, (W / mw / pw), C, mw * mh * (H / mh / ph));         // [(H/mh/ph)*mh*mw, C, (W/mw/pw), ph*pw]
             image = ggml_concat(ctx, image, image, 0);                                                       // [(H/mh/ph)*mh*mw, C, (W/mw/pw), pt*ph*pw]
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3));                          // [(H/mh/ph)*mh*mw, (W/mw/pw), C, pt*ph*pw]
+            image = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, image, 0, 2, 1, 3));                   // [(H/mh/ph)*mh*mw, (W/mw/pw), C, pt*ph*pw]
             image = ggml_reshape_4d(ctx, image, pw * ph * pt * C, (W / mw / pw), mw * mh, (H / mh / ph));    // [(H/mh/ph), mh*mw, (W/mw/pw), C*pt*ph*pw]
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3));                          // [(H/mh/ph), (W/mw/pw), mh*mw, C*pt*ph*pw]
+            image = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, image, 0, 2, 1, 3));                   // [(H/mh/ph), (W/mw/pw), mh*mw, C*pt*ph*pw]
             image = ggml_reshape_2d(ctx, image, pw * ph * pt * C, mw * mh * (W / mw / pw) * (H / mh / ph));  // [(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw]
             return image;
         }
@@ -2301,7 +2303,6 @@ namespace LLM {
         }
 
         static ggml_tensor* encode_image_common(GGMLRunner* runner,
-                                                ggml_context* compute_ctx,
                                                 GGMLRunnerContext* runner_ctx,
                                                 ggml_tensor* image,
                                                 const LLMVisionConfig& vision_params,
@@ -2313,13 +2314,14 @@ namespace LLM {
                                                 std::array<std::vector<int32_t>, 4>& pos_embed_idx_data,
                                                 std::array<std::vector<float>, 4>& pos_embed_weight_data,
                                                 std::vector<ggml_tensor*>* output_tensors = nullptr) {
+            auto compute_ctx = runner_ctx->ggml_ctx;
             GGML_ASSERT(image->ne[1] % (vision_params.patch_size * vision_params.spatial_merge_size) == 0);
             GGML_ASSERT(image->ne[0] % (vision_params.patch_size * vision_params.spatial_merge_size) == 0);
 
             int grid_h = static_cast<int>(image->ne[1]) / vision_params.patch_size;
             int grid_w = static_cast<int>(image->ne[0]) / vision_params.patch_size;
 
-            auto pixel_values = process_image_common(compute_ctx, image, vision_params);
+            auto pixel_values = process_image_common(runner_ctx, image, vision_params);
             int head_dim      = static_cast<int>(vision_params.hidden_size / vision_params.num_heads);
 
             if (vision_params.arch == LLMVisionArch::QWEN3_VL) {
@@ -2678,8 +2680,8 @@ namespace LLM {
             return grid_t * grid_h * grid_w;
         }
 
-        ggml_tensor* process_image(ggml_context* ctx, ggml_tensor* image) {
-            return process_image_common(ctx, image, config.vision);
+        ggml_tensor* process_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) {
+            return process_image_common(runner_ctx, image, config.vision);
         }
 
         ggml_tensor* build_patch_pos_embeds(GGMLRunnerContext* runner_ctx,
@@ -2698,7 +2700,6 @@ namespace LLM {
 
         ggml_tensor* encode_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) {
             return encode_image_common(this,
-                                       compute_ctx,
                                        runner_ctx,
                                        image,
                                        config.vision,
@@ -2714,7 +2715,6 @@ namespace LLM {
         std::vector<ggml_tensor*> encode_image_outputs(GGMLRunnerContext* runner_ctx, ggml_tensor* image) {
             std::vector<ggml_tensor*> outputs;
             encode_image_common(this,
-                                compute_ctx,
                                 runner_ctx,
                                 image,
                                 config.vision,

src/model/te/ming_image_te.hpp

@@ -73,16 +73,16 @@ namespace MingImageTE {
             auto connector            = std::dynamic_pointer_cast<ConnectorModel>(blocks["connector"]);
             auto x                    = backbone->embed(ctx, ids);
             const int64_t query_start = prompt_length + 1;
-            auto before               = ggml_ext_slice(gctx, x, 1, 0, query_start);
-            auto after                = ggml_ext_slice(gctx, x, 1, query_start + config.num_queries, x->ne[1]);
+            auto before               = ggml_ext_slice(ctx, x, 1, 0, query_start);
+            auto after                = ggml_ext_slice(ctx, x, 1, query_start + config.num_queries, x->ne[1]);
             x                         = ggml_concat(gctx, ggml_concat(gctx, before, params["query_tokens_dict.16x16"], 1), after, 1);
             // HF hidden_states[20] includes final RMSNorm; sd.cpp selects it as num_layers + 1.
             x            = backbone->forward_embeds(ctx, x, nullptr, mask, {5, 12, 21}, {}, nullptr, false, image_mask, cos, sin);
-            auto direct  = ggml_ext_slice(gctx, x, 1, 0, prompt_length);
+            auto direct  = ggml_ext_slice(ctx, x, 1, 0, prompt_length);
             direct       = std::dynamic_pointer_cast<RMSNorm>(blocks["proj_directvlm.0"])->forward(ctx, direct);
             direct       = std::dynamic_pointer_cast<Linear>(blocks["proj_directvlm.1"])->forward(ctx, direct);
-            auto queries = ggml_ext_slice(gctx, x, 0, config.backbone.hidden_size * 2, config.backbone.hidden_size * 3);
-            queries      = ggml_ext_slice(gctx, queries, 1, query_start, query_start + config.num_queries);
+            auto queries = ggml_ext_slice(ctx, x, 0, config.backbone.hidden_size * 2, config.backbone.hidden_size * 3);
+            queries      = ggml_ext_slice(ctx, queries, 1, query_start, query_start + config.num_queries);
             queries      = std::dynamic_pointer_cast<Linear>(blocks["proj_in"])->forward(ctx, queries);
             queries      = connector->forward_embeds(ctx, queries, connector_positions, nullptr, {});
             queries      = std::dynamic_pointer_cast<Linear>(blocks["proj_out"])->forward(ctx, queries);

src/model/te/ph3_image_connector.hpp

@@ -28,11 +28,11 @@ struct PH3ImageTextConnector : public UnaryBlock {
         x         = ggml_div(ctx->ggml_ctx, ggml_sub(ctx->ggml_ctx, x, params["input_mean"]), params["input_std"]);
         x         = proj->forward(ctx, x);
         x         = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, x, params["output_std"]), params["output_mean"]);
-        auto sink = ggml_repeat(ctx->ggml_ctx, params["sink_embedding"], ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, 1));
+        auto sink = ggml_repeat(ctx->ggml_ctx, params["sink_embedding"], ggml_ext_slice(ctx, x, 1, 0, 1));
         if (x->ne[1] == 1) {
             return sink;
         }
-        return ggml_concat(ctx->ggml_ctx, sink, ggml_ext_slice(ctx->ggml_ctx, x, 1, 1, x->ne[1]), 1);
+        return ggml_concat(ctx->ggml_ctx, sink, ggml_ext_slice(ctx, x, 1, 1, x->ne[1]), 1);
     }
 };
 

src/model/te/t5.hpp

@@ -249,7 +249,7 @@ public:
             }
         }
 
-        k = ggml_ext_scale(ctx->ggml_ctx, k, ::sqrtf(static_cast<float>(d_head)), true);
+        k = ggml_ext_scale(ctx, k, ::sqrtf(static_cast<float>(d_head)), true);
 
         x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask);  // [N, n_token, d_head * n_head]
 

src/model/upscaler/esrgan.hpp

@@ -127,7 +127,7 @@ public:
         x_cat      = ggml_concat(ctx->ggml_ctx, x_cat, x4, 2);
         auto x5    = conv5->forward(ctx, x_cat);
 
-        x5 = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, x5, 0.2f), x);
+        x5 = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx, x5, 0.2f), x);
         return x5;
     }
 };
@@ -152,7 +152,7 @@ public:
         out      = rdb2->forward(ctx, out);
         out      = rdb3->forward(ctx, out);
 
-        out = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, out, 0.2f), x);
+        out = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx, out, 0.2f), x);
         return out;
     }
 };

src/model/upscaler/flash_vsr.hpp

@@ -146,7 +146,7 @@ namespace RelayVSR {
             auto y                         = attention->forward(ctx, attention_norm->forward(ctx, x), qpe, kpe, local, name, condition_slot, start_slot, end_slot, history);
             x                              = ggml_add(c, x, y);
             y                              = feed_forward_gate_up_proj->forward(ctx, feed_forward_norm->forward(ctx, x));
-            auto parts                     = ggml_ext_chunk(c, y, 2, 0);
+            auto parts                     = ggml_ext_chunk(ctx, y, 2, 0);
             y                              = ggml_mul(c, ggml_silu(c, parts[0]), parts[1]);
             return ggml_add(c, x, feed_forward_down->forward(ctx, y));
         }

src/model/upscaler/h3_latent_upscaler.hpp

@@ -195,8 +195,8 @@ namespace H3VUpsampler {
             ggml_tensor* emb_silu = ggml_silu(ctx->ggml_ctx, emb);
             ggml_tensor* emb_out  = emb_layers->forward(ctx, emb_silu);
             ggml_tensor* emb_1d   = ggml_reshape_1d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, emb_out), 2 * channels);
-            ggml_tensor* scale    = ggml_reshape_4d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, emb_1d, 0, 0, channels), 1, 1, 1, channels);
-            ggml_tensor* shift    = ggml_reshape_4d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, emb_1d, 0, channels, 2 * channels), 1, 1, 1, channels);
+            ggml_tensor* scale    = ggml_reshape_4d(ctx->ggml_ctx, ggml_ext_slice(ctx, emb_1d, 0, 0, channels), 1, 1, 1, channels);
+            ggml_tensor* shift    = ggml_reshape_4d(ctx->ggml_ctx, ggml_ext_slice(ctx, emb_1d, 0, channels, 2 * channels), 1, 1, 1, channels);
             ggml_tensor* gain     = ggml_scale_bias(ctx->ggml_ctx, scale, 1.f, 1.f);
 
             h = out_norm->forward(ctx, h);
@@ -236,7 +236,7 @@ namespace H3VUpsampler {
         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");
+                w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight");
             }
             ggml_tensor* w32 = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_F32);
             ggml_tensor* xa  = x;
@@ -250,7 +250,7 @@ namespace H3VUpsampler {
             // ggml_pad only pads the end; concatenate zeros for symmetric padding.
             ggml_tensor* xp = xa;
             if (pad > 0) {
-                ggml_tensor* zero = ggml_scale(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, xa, 2, 0, 1), 0.0f);
+                ggml_tensor* zero = ggml_scale(ctx->ggml_ctx, ggml_ext_slice(ctx, xa, 2, 0, 1), 0.0f);
                 ggml_tensor* lead = zero;
                 for (int i = 1; i < pad; ++i) {
                     lead = ggml_concat(ctx->ggml_ctx, lead, zero, 2);
@@ -261,15 +261,15 @@ namespace H3VUpsampler {
 
             ggml_tensor* acc = nullptr;
             for (int j = 0; j < kernel; ++j) {
-                ggml_tensor* xs   = ggml_ext_slice(ctx->ggml_ctx, xp, 2, j, j + T);
-                ggml_tensor* wj   = ggml_ext_slice(ctx->ggml_ctx, w32, 2, j, j + 1);  // [1, 1, 1, C]
+                ggml_tensor* xs   = ggml_ext_slice(ctx, xp, 2, j, j + T);
+                ggml_tensor* wj   = ggml_ext_slice(ctx, w32, 2, j, j + 1);  // [1, 1, 1, C]
                 ggml_tensor* term = ggml_mul(ctx->ggml_ctx, xs, wj);
                 acc               = acc == nullptr ? term : ggml_add(ctx->ggml_ctx, acc, term);
             }
             if (has_bias) {
                 ggml_tensor* b = params["bias"];
                 if (ctx->weight_adapter) {
-                    b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
+                    b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias");
                 }
                 acc = ggml_add(ctx->ggml_ctx, acc, ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, 1, channels));
             }

src/model/upscaler/ltx_latent_upscaler.hpp

@@ -158,15 +158,15 @@ namespace LTXVUpsampler {
             const int64_t W = x->ne[0];
             const int64_t H = x->ne[1];
             const int64_t T = x->ne[2];
-            x               = ggml_ext_cont(ctx->ggml_ctx, x);
+            x               = ggml_ext_cont(ctx, x);
             x               = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, num_channels, 1);
             x               = ggml_group_norm(ctx->ggml_ctx, x, num_groups, eps);
 
             ggml_tensor* weight = params["weight"];
             ggml_tensor* bias   = params["bias"];
             if (ctx->weight_adapter) {
-                weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight");
-                bias   = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, bias, prefix + "bias");
+                weight = ctx->weight_adapter->patch_weight(ctx, weight, prefix + "weight");
+                bias   = ctx->weight_adapter->patch_weight(ctx, bias, prefix + "bias");
             }
             weight = ggml_reshape_4d(ctx->ggml_ctx, weight, 1, 1, num_channels, 1);
             bias   = ggml_reshape_4d(ctx->ggml_ctx, bias, 1, 1, num_channels, 1);
@@ -218,9 +218,9 @@ namespace LTXVUpsampler {
             int64_t w = x->ne[0];
             GGML_ASSERT(x->ne[2] % (upscale_factor * upscale_factor) == 0);
             // x: [b*f, c*p1*p2, h, w] -> [b*f, c, h*p1, w*p2]
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));  // [b*f, h, w, c*p1*p2]
-            x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]);          // [b*f, h*w, c*p1*p2]
-            return DiT::unpatchify(ctx->ggml_ctx, x, h, w, upscale_factor, upscale_factor, true);
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));              // [b*f, h, w, c*p1*p2]
+            x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]);  // [b*f, h*w, c*p1*p2]
+            return DiT::unpatchify(ctx, x, h, w, upscale_factor, upscale_factor, true);
         }
     };
 
@@ -241,9 +241,9 @@ namespace LTXVUpsampler {
             const int64_t C = x->ne[3] / upscale_factor;
 
             // x: [b, c*p, f, h, w] -> [b, c, f*p, h, w]
-            x = ggml_ext_cont(ctx->ggml_ctx, x);
+            x = ggml_ext_cont(ctx, x);
             x = ggml_reshape_4d(ctx->ggml_ctx, x, W * H, F, upscale_factor, C);
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
+            x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
             return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, F * upscale_factor, C);
         }
     };
@@ -324,11 +324,11 @@ namespace LTXVUpsampler {
             auto blur_down     = std::dynamic_pointer_cast<BlurDownsample>(blocks["blur_down"]);
 
             // rearrange(x, "b c f h w -> (b f) c h w")
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));
             x = conv->forward(ctx, x);
             x = pixel_shuffle->forward(ctx, x);
             x = blur_down->forward(ctx, x);
-            return ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+            return ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));
         }
     };
 
@@ -403,18 +403,18 @@ namespace LTXVUpsampler {
             } else if (config.temporal_upsample) {
                 auto upsample_conv = std::dynamic_pointer_cast<Conv3d>(blocks["upsampler.0"]);
                 auto pixel_shuffle = std::dynamic_pointer_cast<TemporalPixelShuffleND>(blocks["upsampler.1"]);
-                x                  = upsample_conv->forward(ctx, x);                    // [b, c*2, f, h, w]
-                x                  = pixel_shuffle->forward(ctx, x);                    // [b, c, f*2, h, w]
-                x                  = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]);  // x[:, :, 1:, :, :]
+                x                  = upsample_conv->forward(ctx, x);          // [b, c*2, f, h, w]
+                x                  = pixel_shuffle->forward(ctx, x);          // [b, c, f*2, h, w]
+                x                  = ggml_ext_slice(ctx, x, 2, 1, x->ne[2]);  // x[:, :, 1:, :, :]
             } else {
                 auto upsample_conv = std::dynamic_pointer_cast<Conv2d>(blocks["upsampler.0"]);
                 auto pixel_shuffle = std::dynamic_pointer_cast<PixelShuffleND>(blocks["upsampler.1"]);
 
                 // rearrange(x, "b c f h w -> (b f) c h w"),
-                x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // [b*f, c, h, w]
-                x = upsample_conv->forward(ctx, x);                                                      // [b*f, c*4, h, w]
-                x = pixel_shuffle->forward(ctx, x);                                                      // [b*f, c, h*2, w*2]
-                x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // [b*c, f, h, w]
+                x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));  // [b*f, c, h, w]
+                x = upsample_conv->forward(ctx, x);                                  // [b*f, c*4, h, w]
+                x = pixel_shuffle->forward(ctx, x);                                  // [b*f, c, h*2, w*2]
+                x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));  // [b*c, f, h, w]
             }
             sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.spatial_up", "x");
 

src/model/upscaler/mosr.hpp

@@ -118,9 +118,9 @@ public:
 
     ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         // [N, C, H, W] -> [N, H, W, C] so the norm reduces over channels.
-        x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
-        x = ggml_ext_layer_norm(ctx->ggml_ctx, x, params["weight"], params["bias"], 1e-6f);
-        return ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
+        x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
+        x = ggml_ext_layer_norm(ctx, x, params["weight"], params["bias"], 1e-6f);
+        return ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
     }
 };
 
@@ -141,9 +141,8 @@ public:
 
     ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         const int pad = kernel_size / 2;
-        x             = ggml_conv_2d_dw_direct(ctx->ggml_ctx, params["weight"], ggml_ext_cont(ctx->ggml_ctx, x), 1, 1, pad, pad, 1, 1);
-        auto b        = ggml_reshape_4d(ctx->ggml_ctx, params["bias"], 1, 1, channels, 1);
-        return ggml_add_inplace(ctx->ggml_ctx, x, b);
+        return ggml_ext_conv_2d_grouped(ctx, ggml_ext_cont(ctx, x), params["weight"], params["bias"], channels,
+                                        1, 1, pad, pad, 1, 1, true);
     }
 };
 
@@ -179,10 +178,10 @@ public:
         auto i = channel_view(hidden, hidden - conv_ch);
         auto c = conv->forward(ctx, channel_view(2 * hidden - conv_ch, conv_ch));
         if (hidden > conv_ch) {
-            c = ggml_concat(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, i), c, 2);
+            c = ggml_concat(ctx->ggml_ctx, ggml_ext_cont(ctx, i), c, 2);
         }
 
-        h = ggml_mul(ctx->ggml_ctx, c, mosr_mish(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, g)));
+        h = ggml_mul(ctx->ggml_ctx, c, mosr_mish(ctx->ggml_ctx, ggml_ext_cont(ctx, g)));
         h = mosr_mish(ctx->ggml_ctx, fc2->forward(ctx, h));
         return ggml_add(ctx->ggml_ctx, h, ggml_scale_bias(ctx->ggml_ctx, x, 1.f, -0.5f));
     }

src/model/vae/auto_encoder_kl.hpp

@@ -216,7 +216,7 @@ protected:
     }
 
     ggml_tensor* get_alpha(GGMLRunnerContext* ctx) {
-        auto mix_factor = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["mix_factor"]);
+        auto mix_factor = ggml_ext_cast_f32(ctx, params["mix_factor"]);
         return ggml_sigmoid(ctx->ggml_ctx, mix_factor);
     }
 
@@ -601,11 +601,11 @@ public:
             int64_t H = h * p;
             int64_t W = w * p;
 
-            z = ggml_reshape_4d(ctx->ggml_ctx, z, w * h, p * p, C, N);                           // [N, C, p*p, h*w]
-            z = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, z, 1, 0, 2, 3));  // [N, C, h*w, p*p]
-            z = ggml_reshape_4d(ctx->ggml_ctx, z, p, p, w, h * C * N);                           // [N*C*h, w, p, p]
-            z = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, z, 0, 2, 1, 3));  // [N*C*h, p, w, p]
-            z = ggml_reshape_4d(ctx->ggml_ctx, z, W, H, C, N);                                   // [N, C, h*p, w*p]
+            z = ggml_reshape_4d(ctx->ggml_ctx, z, w * h, p * p, C, N);                 // [N, C, p*p, h*w]
+            z = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, z, 1, 0, 2, 3));  // [N, C, h*w, p*p]
+            z = ggml_reshape_4d(ctx->ggml_ctx, z, p, p, w, h * C * N);                 // [N*C*h, w, p, p]
+            z = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, z, 0, 2, 1, 3));  // [N*C*h, p, w, p]
+            z = ggml_reshape_4d(ctx->ggml_ctx, z, W, H, C, N);                         // [N, C, h*p, w*p]
         }
 
         if (use_quant) {
@@ -632,7 +632,7 @@ public:
             // sd::ggml_graph_cut::mark_graph_cut(z, "vae.encode.final", "z");
         }
         if (sd_version_uses_flux2_vae(version)) {
-            z = ggml_ext_chunk(ctx->ggml_ctx, z, 2, 2)[0];
+            z = ggml_ext_chunk(ctx, z, 2, 2)[0];
 
             // [N, C, H, W] -> [N, C*p*p, H/p, W/p]
             int64_t p = 2;

src/model/vae/hunyuan_vae.hpp

@@ -96,20 +96,20 @@ namespace Hunyuan {
                     auto previous = carry->take();
                     GGML_ASSERT(previous->ne[2] <= pad_t);
                     for (int64_t frame = 0; frame < previous->ne[2]; frame++) {
-                        temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, previous, 2, frame, frame + 1));
+                        temporal_frames.push_back(ggml_ext_slice(ctx, previous, 2, frame, frame + 1));
                     }
                     for (int64_t frame = previous->ne[2]; frame < pad_t; frame++) {
-                        temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1));
+                        temporal_frames.push_back(ggml_ext_slice(ctx, x, 2, 0, 1));
                     }
                 } else {
-                    auto first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                    auto first = ggml_ext_slice(ctx, x, 2, 0, 1);
                     for (int frame = 0; frame < pad_t; frame++) {
                         temporal_frames.push_back(first);
                     }
                 }
             }
             for (int64_t frame = 0; frame < x->ne[2]; frame++) {
-                temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 2, frame, frame + 1));
+                temporal_frames.push_back(ggml_ext_slice(ctx, x, 2, frame, frame + 1));
             }
 
             if (pad_t > 0 && carry != nullptr && carry->output != nullptr) {
@@ -127,13 +127,13 @@ namespace Hunyuan {
             }
             auto replicate_pad = [&](ggml_tensor* input, int dim, int left, int right) {
                 if (left > 0) {
-                    auto first = ggml_ext_slice(ctx->ggml_ctx, input, dim, 0, 1);
+                    auto first = ggml_ext_slice(ctx, input, dim, 0, 1);
                     for (int i = 0; i < left; i++) {
                         input = ggml_concat(ctx->ggml_ctx, first, input, dim);
                     }
                 }
                 if (right > 0) {
-                    auto last = ggml_ext_slice(ctx->ggml_ctx, input, dim, input->ne[dim] - 1, input->ne[dim]);
+                    auto last = ggml_ext_slice(ctx, input, dim, input->ne[dim] - 1, input->ne[dim]);
                     for (int i = 0; i < right; i++) {
                         input = ggml_concat(ctx->ggml_ctx, input, last, dim);
                     }
@@ -184,19 +184,19 @@ namespace Hunyuan {
             auto k = k_proj->forward(ctx, x);  // [b*c, t, h, w]
             auto v = v_proj->forward(ctx, x);  // [b*c, t, h, w]
 
-            q = ggml_reshape_3d(ctx->ggml_ctx, q, w * h * t, c, b);                                  // [b, c, t*h*w]
-            q = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, q, 1, 0, 2, 3));  // [b, t*h*w, c]
+            q = ggml_reshape_3d(ctx->ggml_ctx, q, w * h * t, c, b);              // [b, c, t*h*w]
+            q = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, q, 1, 0, 2, 3));  // [b, t*h*w, c]
 
-            k = ggml_reshape_3d(ctx->ggml_ctx, k, w * h * t, c, b);                                  // [b, c, t*h*w]
-            k = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, k, 1, 0, 2, 3));  // [b, t*h*w, c]
+            k = ggml_reshape_3d(ctx->ggml_ctx, k, w * h * t, c, b);              // [b, c, t*h*w]
+            k = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, k, 1, 0, 2, 3));  // [b, t*h*w, c]
 
-            v = ggml_reshape_3d(ctx->ggml_ctx, v, w * h * t, c, b);                                  // [b, c, t*h*w]
-            v = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3));  // [b, t*h*w, c]
+            v = ggml_reshape_3d(ctx->ggml_ctx, v, w * h * t, c, b);              // [b, c, t*h*w]
+            v = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, v, 1, 0, 2, 3));  // [b, t*h*w, c]
 
             x = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);  // [b, t*h*w, c]
 
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));  // [b, c, t*h*w]
-            x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, t, c * b);                         // [b*c, t, h, w]
+            x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));  // [b, c, t*h*w]
+            x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, t, c * b);               // [b*c, t, h, w]
 
             x = proj_out->forward(ctx, x);
 
@@ -301,13 +301,13 @@ namespace Hunyuan {
             int64_t W      = x->ne[0];
 
             x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, C, factor);                          // [factor, C, T*H, W]
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // [C, factor, T*H, W]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));            // [C, factor, T*H, W]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, factor_s, factor_s * factor_t * C);  // [C*factor_t*factor_s, factor_s, T*H, W]
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));  // [C*factor_t*factor_s, T*H, W, factor_s]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));            // [C*factor_t*factor_s, T*H, W, factor_s]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, H * T, factor_s, factor_t * C);  // [C*factor_t, factor_s, T*H, W*factor_s]
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));  // [C*factor_t, T*H, factor_s, W*factor_s]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));            // [C*factor_t, T*H, factor_s, W*factor_s]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W * factor_s * H, T, factor_t, C);  // [C, factor_t, T, H*factor_s*W*factor_s]
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));  // [C, T, factor_t, H*factor_s*W*factor_s]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));            // [C, T, factor_t, H*factor_s*W*factor_s]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, factor_s * H, factor_t * T, C);  // [C, T*factor_t, H*factor_s, W*factor_s]
 
             return x;
@@ -328,11 +328,11 @@ namespace Hunyuan {
 
             ggml_tensor* shortcut = nullptr;
             if (factor_t == 2 && !continuation) {
-                auto h_first = ggml_ext_slice(ctx->ggml_ctx, h, 2, 0, 1);                   // [B*factor*OC, 1, H, W]
-                h_first      = _pixel_shuffle_3d(ctx, h_first, 1, factor_s, B);             // [B*2*OC, 1, H*factor_s, W*factor_s]
-                h_first      = ggml_ext_slice(ctx->ggml_ctx, h_first, 3, 0, out_channels);  // [B*OC, 1, H*factor_s, W*factor_s]
+                auto h_first = ggml_ext_slice(ctx, h, 2, 0, 1);                   // [B*factor*OC, 1, H, W]
+                h_first      = _pixel_shuffle_3d(ctx, h_first, 1, factor_s, B);   // [B*2*OC, 1, H*factor_s, W*factor_s]
+                h_first      = ggml_ext_slice(ctx, h_first, 3, 0, out_channels);  // [B*OC, 1, H*factor_s, W*factor_s]
 
-                auto x_first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                auto x_first = ggml_ext_slice(ctx, x, 2, 0, 1);
                 x_first      = repeat_interleave_channels(ctx, x_first, repeats / 2, x->ne[0], x->ne[1], 1);
                 x_first      = _pixel_shuffle_3d(ctx, x_first, 1, factor_s, B);
 
@@ -340,12 +340,12 @@ namespace Hunyuan {
                     return ggml_add(ctx->ggml_ctx, h_first, x_first);
                 }
 
-                auto h_next = ggml_ext_slice(ctx->ggml_ctx, h, 2, 1, h->ne[2]);       // [B*factor*OC, T - 1, H, W]
+                auto h_next = ggml_ext_slice(ctx, h, 2, 1, h->ne[2]);                 // [B*factor*OC, T - 1, H, W]
                 h_next      = _pixel_shuffle_3d(ctx, h_next, factor_t, factor_s, B);  // [B*OC, (T - 1)*factor_t, H*factor_s, W*factor_s]
 
                 h = ggml_concat(ctx->ggml_ctx, h_first, h_next, 2);  // [B*OC, 1 + (T - 1)*factor_t, H*factor_s, W*factor_s]
 
-                auto x_next = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]);
+                auto x_next = ggml_ext_slice(ctx, x, 2, 1, x->ne[2]);
                 x_next      = repeat_interleave_channels(ctx, x_next, repeats, x->ne[0], x->ne[1], x->ne[2] - 1);
                 x_next      = _pixel_shuffle_3d(ctx, x_next, factor_t, factor_s, B);
 
@@ -374,13 +374,13 @@ namespace Hunyuan {
         int64_t factor = factor_t * factor_s * factor_s;
 
         x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W * factor_s * H, factor_t, T, C);
-        x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));
         x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, factor_s, H * T, factor_t * C);
-        x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));
         x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s, W, H * T, factor_s * factor_t * C);
-        x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));
         x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, factor, C);
-        x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));
         return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, C * factor);
     }
 
@@ -398,9 +398,9 @@ namespace Hunyuan {
         int64_t spatial = W * H * T;
         int64_t groups  = x->ne[3] / group_size;
         x               = ggml_reshape_3d(ctx->ggml_ctx, x, spatial, group_size, groups);
-        x               = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+        x               = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3));
         x               = ggml_sum_rows(ctx->ggml_ctx, x);
-        x               = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+        x               = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3));
         x               = ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, groups);
         return ggml_scale(ctx->ggml_ctx, x, 1.f / static_cast<float>(group_size));
     }
@@ -435,19 +435,19 @@ namespace Hunyuan {
             ggml_tensor* h_first = nullptr;
             ggml_tensor* x_first = nullptr;
             if (factor_t == 2) {
-                h_first = ggml_ext_slice(ctx->ggml_ctx, h, 2, 0, 1);
+                h_first = ggml_ext_slice(ctx, h, 2, 0, 1);
                 h_first = pixel_unshuffle_3d(ctx, h_first, 1, factor_s);
                 h_first = ggml_concat(ctx->ggml_ctx, h_first, h_first, 3);
 
-                x_first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                x_first = ggml_ext_slice(ctx, x, 2, 0, 1);
                 x_first = pixel_unshuffle_3d(ctx, x_first, 1, factor_s);
                 x_first = mean_channel_groups(ctx, x_first, group_size / 2);
 
                 if (x->ne[2] == 1) {
                     return ggml_add(ctx->ggml_ctx, h_first, x_first);
                 }
-                h = ggml_ext_slice(ctx->ggml_ctx, h, 2, 1, h->ne[2]);
-                x = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]);
+                h = ggml_ext_slice(ctx, h, 2, 1, h->ne[2]);
+                x = ggml_ext_slice(ctx, x, 2, 1, x->ne[2]);
             }
 
             GGML_ASSERT(h->ne[2] % factor_t == 0);
@@ -648,7 +648,7 @@ namespace Hunyuan {
             x             = ggml_silu_inplace(ctx->ggml_ctx, x);
             x             = conv_out->forward(ctx, x);
             x             = ggml_add(ctx->ggml_ctx, x, shortcut);
-            return ggml_ext_slice(ctx->ggml_ctx, x, 3, 0, z_channels);
+            return ggml_ext_slice(ctx, x, 3, 0, z_channels);
         }
     };
 
@@ -707,7 +707,7 @@ namespace Hunyuan {
             const int64_t frames = h->ne[2];
             for (int64_t start = 0; start < frames; start += HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE) {
                 const int64_t end = std::min(start + HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE, frames);
-                auto chunk        = ggml_ext_slice(ctx->ggml_ctx, h, 2, start, end);
+                auto chunk        = ggml_ext_slice(ctx, h, 2, start, end);
 
                 std::vector<ggml_tensor*> carry_output;
                 TemporalConvCarry carry{

src/model/vae/ltx_audio_vae.hpp

@@ -220,7 +220,7 @@ namespace LTXV {
 
         auto x = ggml_reshape_3d(ctx, waveform, time, 1, channels * batch);
         if (left_pad > 0) {
-            x = ggml_ext_pad_ext(ctx, runner_ctx->backend, x, static_cast<int>(left_pad), 0, 0, 0, 0, 0, 0, 0);
+            x = ggml_ext_pad_ext(runner_ctx, x, static_cast<int>(left_pad), 0, 0, 0, 0, 0, 0, 0);
         }
 
         auto frames = ggml_conv_1d(ctx, forward_basis, x, hop_length, 0, 1);
@@ -228,8 +228,8 @@ namespace LTXV {
         GGML_ASSERT(frames->ne[1] == stft_channels);
         GGML_ASSERT(frames->ne[2] == channels * batch);
 
-        auto real      = ggml_ext_slice(ctx, frames, 1, 0, n_freqs);
-        auto imag      = ggml_ext_slice(ctx, frames, 1, n_freqs, stft_channels);
+        auto real      = ggml_ext_slice(runner_ctx, frames, 1, 0, n_freqs);
+        auto imag      = ggml_ext_slice(runner_ctx, frames, 1, n_freqs, stft_channels);
         auto magnitude = ggml_sqrt(ctx,
                                    ggml_add(ctx,
                                             ggml_sqr(ctx, real),
@@ -265,42 +265,19 @@ namespace LTXV {
         return type == GGML_TYPE_BF16 ? GGML_TYPE_F16 : type;
     }
 
-    static ggml_tensor* repeat_with_vulkan_f32_workaround(ggml_backend_t backend,
-                                                          ggml_context* ctx,
-                                                          ggml_tensor* x,
-                                                          int64_t ne0,
-                                                          int64_t ne1,
-                                                          int64_t ne2,
-                                                          int64_t ne3) {
-        if (x->type != GGML_TYPE_F32 &&
-            (x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) &&
-            sd_backend_is(backend, "vulkan")) {
-            auto x_f32    = ggml_cast(ctx, x, GGML_TYPE_F32);
-            auto repeated = ggml_repeat_4d(ctx,
-                                           x_f32,
-                                           ne0,
-                                           ne1,
-                                           ne2,
-                                           ne3);
-            return ggml_cast(ctx, repeated, x->type);
-        }
-        return ggml_repeat_4d(ctx, x, ne0, ne1, ne2, ne3);
-    }
-
     static ggml_tensor* repeat_1d_value(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t count) {
-        auto ctx = runner_ctx->ggml_ctx;
         GGML_ASSERT(x->ne[0] == 1);
-        return repeat_with_vulkan_f32_workaround(runner_ctx->backend, ctx, x, count, x->ne[1], x->ne[2], x->ne[3]);
+        return ggml_ext_repeat_4d(runner_ctx, x, count, x->ne[1], x->ne[2], x->ne[3]);
     }
 
     static ggml_tensor* replicate_pad_1d(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t left, int64_t right) {
         auto ctx = runner_ctx->ggml_ctx;
         if (left > 0) {
-            auto first = ggml_ext_slice(ctx, x, 0, 0, 1);
+            auto first = ggml_ext_slice(runner_ctx, x, 0, 0, 1);
             x          = ggml_concat(ctx, repeat_1d_value(runner_ctx, first, left), x, 0);
         }
         if (right > 0) {
-            auto last = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]);
+            auto last = ggml_ext_slice(runner_ctx, x, 0, x->ne[0] - 1, x->ne[0]);
             x         = ggml_concat(ctx, x, repeat_1d_value(runner_ctx, last, right), 0);
         }
         return x;
@@ -314,7 +291,7 @@ namespace LTXV {
         } else if (ggml_n_dims(base) == 1) {
             base = ggml_reshape_4d(ctx, base, base->ne[0], 1, 1, 1);
         }
-        return repeat_with_vulkan_f32_workaround(runner_ctx->backend, ctx, base, base->ne[0], 1, channels, 1);
+        return ggml_ext_repeat_4d(runner_ctx, base, base->ne[0], 1, channels, 1);
     }
 
     static ggml_tensor* depthwise_conv1d(GGMLRunnerContext* runner_ctx,
@@ -329,7 +306,8 @@ namespace LTXV {
         return ggml_reshape_4d(ctx, out, out->ne[0], out->ne[1], 1, 1);
     }
 
-    static ggml_tensor* reverse_1d_filter(ggml_context* ctx, ggml_tensor* filter) {
+    static ggml_tensor* reverse_1d_filter(GGMLRunnerContext* runner_ctx, ggml_tensor* filter) {
+        auto ctx = runner_ctx->ggml_ctx;
         GGML_ASSERT(ctx != nullptr);
         GGML_ASSERT(filter != nullptr);
         GGML_ASSERT(filter->ne[1] == 1);
@@ -338,16 +316,17 @@ namespace LTXV {
 
         ggml_tensor* reversed = nullptr;
         for (int64_t k = filter->ne[0] - 1; k >= 0; --k) {
-            auto slice = ggml_ext_slice(ctx, filter, 0, k, k + 1);
+            auto slice = ggml_ext_slice(runner_ctx, filter, 0, k, k + 1);
             reversed   = reversed == nullptr ? slice : ggml_concat(ctx, reversed, slice, 0);
         }
         return reversed;
     }
 
-    static ggml_tensor* depthwise_conv_transpose1d(ggml_context* ctx,
+    static ggml_tensor* depthwise_conv_transpose1d(GGMLRunnerContext* runner_ctx,
                                                    ggml_tensor* x,
                                                    ggml_tensor* filter,
                                                    int stride) {
+        auto ctx = runner_ctx->ggml_ctx;
         GGML_ASSERT(x->ne[2] == 1 && x->ne[3] == 1);
         GGML_ASSERT(filter->ne[1] == 1);
         GGML_ASSERT(filter->ne[2] == 1 && filter->ne[3] == 1);
@@ -359,7 +338,7 @@ namespace LTXV {
 
         auto x_flat = ggml_reshape_3d(ctx, x, 1, time, channels);
         if (stride > 1) {
-            auto zero_unit = ggml_ext_scale(ctx, x_flat, 0.0f);
+            auto zero_unit = ggml_ext_scale(runner_ctx, x_flat, 0.0f);
             auto zero_tail = zero_unit;
             for (int i = 1; i < stride - 1; ++i) {
                 zero_tail = ggml_concat(ctx, zero_tail, zero_unit, 0);
@@ -368,16 +347,16 @@ namespace LTXV {
         }
         x_flat = ggml_reshape_3d(ctx, x_flat, time * stride, 1, channels);
 
-        auto reversed_filter = reverse_1d_filter(ctx, filter);
+        auto reversed_filter = reverse_1d_filter(runner_ctx, filter);
         auto out             = ggml_conv_1d(ctx, reversed_filter, x_flat, 1, static_cast<int>(kernel_size - 1), 1);
         if (out->ne[0] > out_time) {
-            out = ggml_ext_slice(ctx, out, 0, 0, out_time);
+            out = ggml_ext_slice(runner_ctx, out, 0, 0, out_time);
         }
         GGML_ASSERT(out->ne[0] == out_time);
         GGML_ASSERT(out->ne[1] == 1);
         GGML_ASSERT(out->ne[2] == channels);
 
-        out = ggml_ext_scale(ctx, out, static_cast<float>(stride));
+        out = ggml_ext_scale(runner_ctx, out, static_cast<float>(stride));
         return ggml_reshape_4d(ctx, out, out_time, channels, 1, 1);
     }
 
@@ -409,21 +388,22 @@ namespace LTXV {
 
         auto x = ggml_reshape_3d(ctx, waveform, time, channels * batch, 1);
         x      = replicate_pad_1d(runner_ctx, x, pad, pad);
-        x      = depthwise_conv_transpose1d(ctx, x, filter, ratio);
-        x      = ggml_ext_slice(ctx, x, 0, pad_left, x->ne[0] - pad_right);
+        x      = depthwise_conv_transpose1d(runner_ctx, x, filter, ratio);
+        x      = ggml_ext_slice(runner_ctx, x, 0, pad_left, x->ne[0] - pad_right);
         return ggml_reshape_3d(ctx, x, x->ne[0], channels, batch);
     }
 
-    static ggml_tensor* crop_waveform_samples(ggml_context* ctx,
+    static ggml_tensor* crop_waveform_samples(GGMLRunnerContext* runner_ctx,
                                               ggml_tensor* waveform,
                                               int64_t target_samples) {
+        auto ctx = runner_ctx->ggml_ctx;
         GGML_ASSERT(ctx != nullptr);
         GGML_ASSERT(waveform != nullptr);
         if (waveform->ne[0] == target_samples) {
             return waveform;
         }
         GGML_ASSERT(waveform->ne[0] > target_samples);
-        return ggml_ext_slice(ctx, waveform, 0, 0, target_samples);
+        return ggml_ext_slice(runner_ctx, waveform, 0, 0, target_samples);
     }
 
     struct PixelNorm2D : public UnaryBlock {
@@ -433,9 +413,9 @@ namespace LTXV {
             : eps(eps) {}
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
-            auto h = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
+            auto h = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));
             h      = ggml_rms_norm(ctx->ggml_ctx, h, eps);
-            h      = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, h, 1, 2, 0, 3));
+            h      = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, h, 1, 2, 0, 3));
             return h;
         }
     };
@@ -456,8 +436,7 @@ namespace LTXV {
             auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
             int pad_h = kernel_size.first - 1;
             int pad_w = kernel_size.second - 1;
-            x         = ggml_ext_pad_ext(ctx->ggml_ctx,
-                                         ctx->backend,
+            x         = ggml_ext_pad_ext(ctx,
                                          x,
                                          pad_w / 2,
                                          pad_w - pad_w / 2,
@@ -481,7 +460,7 @@ namespace LTXV {
             auto conv = std::dynamic_pointer_cast<HeightCausalConv2D>(blocks["conv"]);
             x         = ggml_upscale(ctx->ggml_ctx, x, 2, GGML_SCALE_MODE_NEAREST);
             x         = conv->forward(ctx, x);
-            return ggml_ext_slice(ctx->ggml_ctx, x, 1, 1, x->ne[1]);
+            return ggml_ext_slice(ctx, x, 1, 1, x->ne[1]);
         }
     };
 
@@ -606,7 +585,7 @@ namespace LTXV {
             GGML_ASSERT(dilation == 1);
             x = ggml_conv_transpose_1d(ctx->ggml_ctx, params["weight"], x, stride, 0, dilation);
             if (padding > 0) {
-                x = ggml_ext_slice(ctx->ggml_ctx, x, 0, padding, x->ne[0] - padding);
+                x = ggml_ext_slice(ctx, x, 0, padding, x->ne[0] - padding);
             }
             if (bias) {
                 auto b = ggml_reshape_4d(ctx->ggml_ctx, params["bias"], 1, params["bias"]->ne[0], 1, 1);
@@ -639,7 +618,7 @@ namespace LTXV {
             beta             = ggml_reshape_4d(ctx->ggml_ctx, beta, 1, beta->ne[0], 1, 1);
             auto oscillation = ggml_sin(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, x, alpha));
             oscillation      = ggml_mul(ctx->ggml_ctx, oscillation, oscillation);
-            auto eps_tensor  = ggml_ext_scale(ctx->ggml_ctx, ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1), eps);
+            auto eps_tensor  = ggml_ext_scale(ctx, ggml_ext_ones(ctx, 1, 1, 1, 1), eps);
             oscillation      = ggml_div(ctx->ggml_ctx, oscillation, ggml_add(ctx->ggml_ctx, beta, eps_tensor));
             return ggml_add(ctx->ggml_ctx, x, oscillation);
         }
@@ -675,8 +654,8 @@ namespace LTXV {
             int up_pad_right = up_pad * up_ratio + (up_kernel_size - up_ratio + 1) / 2;
 
             x = replicate_pad_1d(ctx, x, up_pad, up_pad);
-            x = depthwise_conv_transpose1d(ctx->ggml_ctx, x, up_filter, up_ratio);
-            x = ggml_ext_slice(ctx->ggml_ctx, x, 0, up_pad_left, x->ne[0] - up_pad_right);
+            x = depthwise_conv_transpose1d(ctx, x, up_filter, up_ratio);
+            x = ggml_ext_slice(ctx, x, 0, up_pad_left, x->ne[0] - up_pad_right);
 
             x = act->forward(ctx, x);
 
@@ -778,7 +757,7 @@ namespace LTXV {
 
         ggml_tensor* prepare_input(GGMLRunnerContext* ctx, ggml_tensor* mel) {
             mel       = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, mel, 1, 0, 2, 3));
-            auto mels = ggml_ext_chunk(ctx->ggml_ctx, mel, 2, 2);
+            auto mels = ggml_ext_chunk(ctx, mel, 2, 2);
             mel       = ggml_concat(ctx->ggml_ctx, mels[0], mels[1], 1);
             // mel = ggml_reshape_4d(ctx->ggml_ctx,
             //                       mel,
@@ -810,7 +789,7 @@ namespace LTXV {
                     auto block_out = resblock->forward(ctx, x);
                     sum            = sum == nullptr ? block_out : ggml_add(ctx->ggml_ctx, sum, block_out);
                 }
-                x = ggml_ext_scale(ctx->ggml_ctx, sum, 1.0f / static_cast<float>(resblock_kernel_sizes.size()));
+                x = ggml_ext_scale(ctx, sum, 1.0f / static_cast<float>(resblock_kernel_sizes.size()));
             }
 
             x = act_post->forward(ctx, x);
@@ -876,8 +855,8 @@ namespace LTXV {
                                          int target_freq) {
             int64_t time = std::min<int64_t>(decoded->ne[1], target_time);
             int64_t freq = std::min<int64_t>(decoded->ne[0], target_freq);
-            decoded      = ggml_ext_slice(ctx->ggml_ctx, decoded, 0, 0, freq);
-            decoded      = ggml_ext_slice(ctx->ggml_ctx, decoded, 1, 0, time);
+            decoded      = ggml_ext_slice(ctx, decoded, 0, 0, freq);
+            decoded      = ggml_ext_slice(ctx, decoded, 1, 0, time);
             return decoded;
         }
 
@@ -995,7 +974,7 @@ namespace LTXV {
                                        ggml_add(ctx->ggml_ctx, residual, skip),
                                        -1.0f,
                                        1.0f);
-                waveform  = crop_waveform_samples(ctx->ggml_ctx, waveform, out_time);
+                waveform  = crop_waveform_samples(ctx, waveform, out_time);
             }
 
             return waveform;

src/model/vae/ltx_vae.hpp

@@ -32,28 +32,26 @@ namespace LTXVAE {
         return ggml_reshape_4d(ctx, x, 1, 1, 1, ggml_nelements(x));
     }
 
-    static inline std::pair<ggml_tensor*, ggml_tensor*> get_shift_scale(ggml_context* ctx,
-                                                                        ggml_tensor* table,
-                                                                        ggml_tensor* timestep,
-                                                                        int64_t channels,
-                                                                        int parts) {
+    static inline std::pair<ggml_tensor*, ggml_tensor*> get_shift_scale(GGMLRunnerContext* runner_ctx, ggml_tensor* table, ggml_tensor* timestep, int64_t channels, int parts) {
+        auto ctx = runner_ctx->ggml_ctx;
         GGML_ASSERT(timestep != nullptr);
         GGML_ASSERT(ggml_nelements(timestep) == channels * parts);
 
         auto timestep_view = ggml_reshape_2d(ctx, timestep, channels, parts);
         auto values        = ggml_add(ctx, table, timestep_view);
-        auto chunks        = ggml_ext_chunk(ctx, values, parts, 1, false);
+        auto chunks        = ggml_ext_chunk(runner_ctx, values, parts, 1, false);
         auto shift         = reshape_channel_broadcast(ctx, ggml_cont(ctx, chunks[0]));
         auto scale         = reshape_channel_broadcast(ctx, ggml_cont(ctx, chunks[1]));
         return {shift, scale};
     }
 
-    static inline ggml_tensor* depth_to_space_3d(ggml_context* ctx,
+    static inline ggml_tensor* depth_to_space_3d(GGMLRunnerContext* runner_ctx,
                                                  ggml_tensor* x,
                                                  int64_t c,
                                                  int factor_t,
                                                  int factor_s,
                                                  bool drop_first_temporal_frame) {
+        auto ctx = runner_ctx->ggml_ctx;
         // x: [B*c*p1*p2*p3, T, H, W], B == 1, p2 == p3 == factor_s, p1 == factor_t
         // return: [B*c, T*p1, H*p2, W*p2]
         // Match: rearrange(x, "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)")
@@ -61,26 +59,24 @@ namespace LTXVAE {
         const int64_t H = x->ne[1];
         const int64_t W = x->ne[0];
 
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));        // [T, C, H, W]
-        x = ggml_reshape_4d(ctx, x, W, H, factor_s, factor_s * factor_t * c * T);  // [T*c*p1*p2, p3, H, W]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));        // [T*c*p1*p2, H, W, p3]
-        x = ggml_reshape_4d(ctx, x, factor_s * W, H, factor_s, factor_t * c * T);  // [T*c*p1, p2, H, W*p3]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));        // [T*c*p1, H, p2, W*p3]
-        x = ggml_reshape_4d(ctx, x, factor_s * W * factor_s * H, factor_t, c, T);  // [T, c, p1, H*p2*W*p3]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));        // [c, T, p1, H*p2*W*p3]
-        x = ggml_reshape_4d(ctx, x, factor_s * W, factor_s * H, factor_t * T, c);  // [T, c, T*p1, H*p2*W*p3]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 1, 3, 2));  // [T, C, H, W]
+        x = ggml_reshape_4d(ctx, x, W, H, factor_s, factor_s * factor_t * c * T);          // [T*c*p1*p2, p3, H, W]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 2, 0, 1, 3));  // [T*c*p1*p2, H, W, p3]
+        x = ggml_reshape_4d(ctx, x, factor_s * W, H, factor_s, factor_t * c * T);          // [T*c*p1, p2, H, W*p3]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [T*c*p1, H, p2, W*p3]
+        x = ggml_reshape_4d(ctx, x, factor_s * W * factor_s * H, factor_t, c, T);          // [T, c, p1, H*p2*W*p3]
+        x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 1, 3, 2));  // [c, T, p1, H*p2*W*p3]
+        x = ggml_reshape_4d(ctx, x, factor_s * W, factor_s * H, factor_t * T, c);          // [T, c, T*p1, H*p2*W*p3]
 
         if (drop_first_temporal_frame && factor_t > 1 && x->ne[2] > 0) {
-            x = ggml_ext_slice(ctx, x, 2, 1, x->ne[2]);
+            x = ggml_ext_slice(runner_ctx, x, 2, 1, x->ne[2]);
         }
 
         return x;
     }
 
-    static inline ggml_tensor* patchify(ggml_context* ctx,
-                                        ggml_tensor* x,
-                                        int patch_size) {
-        return WAN::WanVAE::patchify(ctx, x, patch_size, 1);
+    static inline ggml_tensor* patchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int patch_size) {
+        return WAN::WanVAE::patchify(runner_ctx, x, patch_size, 1);
     }
 
     class CausalConv3d : public GGMLBlock {
@@ -113,20 +109,20 @@ namespace LTXVAE {
             auto conv = std::dynamic_pointer_cast<Conv3d>(blocks["conv"]);
 
             if (causal) {
-                auto first_frame     = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                auto first_frame     = ggml_ext_slice(ctx, x, 2, 0, 1);
                 auto first_frame_pad = first_frame;
                 for (int i = 1; i < time_kernel_size - 1; i++) {
                     first_frame_pad = ggml_concat(ctx->ggml_ctx, first_frame_pad, first_frame, 2);
                 }
                 x = ggml_concat(ctx->ggml_ctx, first_frame_pad, x, 2);
             } else {
-                auto first_frame     = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                auto first_frame     = ggml_ext_slice(ctx, x, 2, 0, 1);
                 auto first_frame_pad = first_frame;
                 for (int i = 1; i < (time_kernel_size - 1) / 2; i++) {
                     first_frame_pad = ggml_concat(ctx->ggml_ctx, first_frame_pad, first_frame, 2);
                 }
 
-                auto last_frame     = ggml_ext_slice(ctx->ggml_ctx, x, 2, x->ne[2] - 1, x->ne[2]);
+                auto last_frame     = ggml_ext_slice(ctx, x, 2, x->ne[2] - 1, x->ne[2]);
                 auto last_frame_pad = last_frame;
                 for (int i = 1; i < (time_kernel_size - 1) / 2; i++) {
                     last_frame_pad = ggml_concat(ctx->ggml_ctx, last_frame_pad, last_frame, 2);
@@ -165,7 +161,7 @@ namespace LTXVAE {
                 GGML_ASSERT(start_idx >= 0);
                 GGML_ASSERT(end_idx > 0);
 
-                auto slice         = ggml_ext_slice(ctx->ggml_ctx, x, 2, start_idx, end_idx);
+                auto slice         = ggml_ext_slice(ctx, x, 2, start_idx, end_idx);
                 feat_map[feat_idx] = ggml_cont(ctx->ggml_ctx, slice);
             }
             feat_idx++;
@@ -175,7 +171,7 @@ namespace LTXVAE {
                 if (prev != nullptr) {
                     left_pad = prev;
                 } else {
-                    auto first_frame = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                    auto first_frame = ggml_ext_slice(ctx, x, 2, 0, 1);
                     left_pad         = first_frame;
                     for (int i = 1; i < pad; i++) {
                         left_pad = ggml_concat(ctx->ggml_ctx, left_pad, first_frame, 2);
@@ -185,7 +181,7 @@ namespace LTXVAE {
             }
 
             if (!causal && pad > 0) {
-                auto last_frame = ggml_ext_slice(ctx->ggml_ctx, x, 2, x->ne[2] - 1, x->ne[2]);
+                auto last_frame = ggml_ext_slice(ctx, x, 2, x->ne[2] - 1, x->ne[2]);
                 auto right_pad  = last_frame;
                 for (int i = 1; i < pad; i++) {
                     right_pad = ggml_concat(ctx->ggml_ctx, right_pad, last_frame, 2);
@@ -204,9 +200,9 @@ namespace LTXVAE {
             : eps(eps) {}
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
-            auto h = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 3, 0, 1, 2));
+            auto h = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 3, 0, 1, 2));
             h      = ggml_rms_norm(ctx->ggml_ctx, h, eps);
-            h      = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, h, 1, 2, 3, 0));
+            h      = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, h, 1, 2, 3, 0));
             return h;
         }
     };
@@ -267,7 +263,7 @@ namespace LTXVAE {
                 auto values = ggml_add(ctx->ggml_ctx,
                                        params["scale_shift_table"],
                                        ggml_reshape_2d(ctx->ggml_ctx, timestep, channels, 4));
-                auto chunks = ggml_ext_chunk(ctx->ggml_ctx, values, 4, 1, false);
+                auto chunks = ggml_ext_chunk(ctx, values, 4, 1, false);
                 shift1      = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[0]));
                 scale1      = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[1]));
                 shift2      = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[2]));
@@ -313,7 +309,7 @@ namespace LTXVAE {
                 auto values = ggml_add(ctx->ggml_ctx,
                                        params["scale_shift_table"],
                                        ggml_reshape_2d(ctx->ggml_ctx, timestep, channels, 4));
-                auto chunks = ggml_ext_chunk(ctx->ggml_ctx, values, 4, 1, false);
+                auto chunks = ggml_ext_chunk(ctx, values, 4, 1, false);
                 shift1      = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[0]));
                 scale1      = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[1]));
                 shift2      = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[2]));
@@ -430,7 +426,7 @@ namespace LTXVAE {
 
             ggml_tensor* x_in = nullptr;
             if (residual) {
-                x_in       = depth_to_space_3d(ctx->ggml_ctx, x, in_channels / (factor_t * factor_s * factor_s), factor_t, factor_s, factor_t > 1);
+                x_in       = depth_to_space_3d(ctx, x, in_channels / (factor_t * factor_s * factor_s), factor_t, factor_s, factor_t > 1);
                 int repeat = (factor_t * factor_s * factor_s) / out_channels_reduction_factor;
                 auto res   = x_in;
                 for (int i = 1; i < repeat; i++) {
@@ -440,7 +436,7 @@ namespace LTXVAE {
             }
 
             x = conv->forward(ctx, x, causal);
-            x = depth_to_space_3d(ctx->ggml_ctx, x, get_output_channels(), factor_t, factor_s, factor_t > 1);
+            x = depth_to_space_3d(ctx, x, get_output_channels(), factor_t, factor_s, factor_t > 1);
             if (residual) {
                 x = ggml_add(ctx->ggml_ctx, x, x_in);
             }
@@ -460,7 +456,7 @@ namespace LTXVAE {
 
             ggml_tensor* x_in = nullptr;
             if (residual) {
-                x_in       = depth_to_space_3d(ctx->ggml_ctx, x, in_channels / (factor_t * factor_s * factor_s), factor_t, factor_s, drop_first);
+                x_in       = depth_to_space_3d(ctx, x, in_channels / (factor_t * factor_s * factor_s), factor_t, factor_s, drop_first);
                 int repeat = (factor_t * factor_s * factor_s) / out_channels_reduction_factor;
                 auto res   = x_in;
                 for (int i = 1; i < repeat; i++) {
@@ -470,7 +466,7 @@ namespace LTXVAE {
             }
 
             x = conv->forward(ctx, x, feat_map, feat_idx, chunk_idx, causal, temporal_pad);
-            x = depth_to_space_3d(ctx->ggml_ctx, x, get_output_channels(), factor_t, factor_s, drop_first);
+            x = depth_to_space_3d(ctx, x, get_output_channels(), factor_t, factor_s, drop_first);
             if (residual) {
                 x = ggml_add(ctx->ggml_ctx, x, x_in);
             }
@@ -515,7 +511,7 @@ namespace LTXVAE {
             auto conv_downsample = std::dynamic_pointer_cast<WAN::AvgDown3D>(blocks["conv_downsample"]);
 
             if (factor_t > 1 && x->ne[2] > 0) {
-                auto first_frame     = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                auto first_frame     = ggml_ext_slice(ctx, x, 2, 0, 1);
                 auto first_frame_pad = first_frame;
                 for (int i = 1; i < factor_t - 1; ++i) {
                     first_frame_pad = ggml_concat(ctx->ggml_ctx, first_frame_pad, first_frame, 2);
@@ -860,7 +856,7 @@ namespace LTXVAE {
             x = ggml_silu_inplace(ctx->ggml_ctx, x);
             x = conv_out->forward(ctx, x, true);
 
-            auto last_channel = ggml_ext_slice(ctx->ggml_ctx, x, 3, x->ne[3] - 1, x->ne[3]);
+            auto last_channel = ggml_ext_slice(ctx, x, 3, x->ne[3] - 1, x->ne[3]);
             auto repeat_shape = ggml_new_tensor_4d(ctx->ggml_ctx, last_channel->type, last_channel->ne[0], last_channel->ne[1], last_channel->ne[2], latent_channels - 1);
             auto repeated     = ggml_repeat(ctx->ggml_ctx, last_channel, repeat_shape);
             return ggml_concat(ctx->ggml_ctx, x, repeated, 3);
@@ -959,7 +955,7 @@ namespace LTXVAE {
 
             ggml_tensor* scaled_timestep = timestep;
             if (timestep_conditioning) {
-                auto multiplier = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["timestep_scale_multiplier"]);
+                auto multiplier = ggml_ext_cast_f32(ctx, params["timestep_scale_multiplier"]);
                 scaled_timestep = ggml_mul(ctx->ggml_ctx, timestep, multiplier);
             }
 
@@ -981,7 +977,7 @@ namespace LTXVAE {
             if (timestep_conditioning) {
                 auto last_time_embedder = std::dynamic_pointer_cast<PixArtAlphaCombinedTimestepSizeEmbeddings>(blocks["last_time_embedder"]);
                 auto timestep_embed     = last_time_embedder->forward(ctx, scaled_timestep);
-                auto [shift, scale]     = get_shift_scale(ctx->ggml_ctx,
+                auto [shift, scale]     = get_shift_scale(ctx,
                                                           params["last_scale_shift_table"],
                                                           timestep_embed,
                                                           hidden_channels,
@@ -1010,7 +1006,7 @@ namespace LTXVAE {
 
             ggml_tensor* scaled_timestep = timestep;
             if (timestep_conditioning && timestep != nullptr) {
-                auto multiplier = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["timestep_scale_multiplier"]);
+                auto multiplier = ggml_ext_cast_f32(ctx, params["timestep_scale_multiplier"]);
                 scaled_timestep = ggml_mul(ctx->ggml_ctx, timestep, multiplier);
             }
 
@@ -1038,7 +1034,7 @@ namespace LTXVAE {
             if (timestep_conditioning) {
                 auto last_time_embedder = std::dynamic_pointer_cast<PixArtAlphaCombinedTimestepSizeEmbeddings>(blocks["last_time_embedder"]);
                 auto timestep_embed     = last_time_embedder->forward(ctx, scaled_timestep);
-                auto [shift, scale]     = get_shift_scale(ctx->ggml_ctx,
+                auto [shift, scale]     = get_shift_scale(ctx,
                                                           params["last_scale_shift_table"],
                                                           timestep_embed,
                                                           hidden_channels,
@@ -1092,7 +1088,7 @@ namespace LTXVAE {
             auto processor = std::dynamic_pointer_cast<PerChannelStatistics>(blocks["per_channel_statistics"]);
             auto latents   = processor->un_normalize(ctx, z);
             auto out       = decoder->forward(ctx, latents, timestep);
-            out            = WAN::WanVAE::unpatchify(ctx->ggml_ctx, out, patch_size, 1);
+            out            = WAN::WanVAE::unpatchify(ctx, out, patch_size, 1);
             return out;
         }
 
@@ -1112,7 +1108,7 @@ namespace LTXVAE {
             const int64_t T = z->ne[2];
             if (T <= 1) {
                 auto out = decoder->forward(ctx, latents, timestep);
-                return WAN::WanVAE::unpatchify(ctx->ggml_ctx, out, patch_size, 1);
+                return WAN::WanVAE::unpatchify(ctx, out, patch_size, 1);
             }
 
             // feat_map holds ggml_tensor* nodes (contiguous copies at each conv layer).
@@ -1145,20 +1141,20 @@ namespace LTXVAE {
 
                 int chunk_overlap = overlap;  // modified by forward_tiled_frame temporal inflation
 
-                auto z_chunk = ggml_ext_slice(ctx->ggml_ctx, latents, 2, i, end_i);
+                auto z_chunk = ggml_ext_slice(ctx, latents, 2, i, end_i);
 
                 auto out_chunk = decoder->forward_tiled_frame(ctx, z_chunk, timestep,
                                                               feat_map, feat_idx, i, chunk_overlap);
 
                 // discard overlap frames if it's not the final chunk
                 if (overlap > 0 && end_i < (int)T) {
-                    out_chunk = ggml_ext_slice(ctx->ggml_ctx, out_chunk, 2, 0, out_chunk->ne[2] - chunk_overlap);
+                    out_chunk = ggml_ext_slice(ctx, out_chunk, 2, 0, out_chunk->ne[2] - chunk_overlap);
                 }
 
                 out = (out == nullptr) ? out_chunk : ggml_concat(ctx->ggml_ctx, out, out_chunk, 2);
             }
 
-            return WAN::WanVAE::unpatchify(ctx->ggml_ctx, out, patch_size, 1);
+            return WAN::WanVAE::unpatchify(ctx, out, patch_size, 1);
         }
 
         ggml_tensor* decode_tiled_chunk(GGMLRunnerContext* ctx,
@@ -1177,9 +1173,9 @@ namespace LTXVAE {
             auto out_chunk    = decoder->forward_tiled_frame(ctx, latents, timestep,
                                                              feat_map, feat_idx, chunk_idx, chunk_overlap);
             if (chunk_overlap > 0) {
-                out_chunk = ggml_ext_slice(ctx->ggml_ctx, out_chunk, 2, 0, out_chunk->ne[2] - chunk_overlap);
+                out_chunk = ggml_ext_slice(ctx, out_chunk, 2, 0, out_chunk->ne[2] - chunk_overlap);
             }
-            return WAN::WanVAE::unpatchify(ctx->ggml_ctx, out_chunk, patch_size, 1);
+            return WAN::WanVAE::unpatchify(ctx, out_chunk, patch_size, 1);
         }
 
         ggml_tensor* encode(GGMLRunnerContext* ctx,
@@ -1188,9 +1184,9 @@ namespace LTXVAE {
             auto encoder   = std::dynamic_pointer_cast<Encoder>(blocks["encoder"]);
             auto processor = std::dynamic_pointer_cast<PerChannelStatistics>(blocks["per_channel_statistics"]);
 
-            x         = patchify(ctx->ggml_ctx, x, patch_size);
+            x         = patchify(ctx, x, patch_size);
             auto out  = encoder->forward(ctx, x);
-            auto mean = ggml_ext_chunk(ctx->ggml_ctx, out, 2, 3, false)[0];
+            auto mean = ggml_ext_chunk(ctx, out, 2, 3, false)[0];
             mean      = ggml_cont(ctx->ggml_ctx, mean);
             return processor->normalize(ctx, mean);
         }

src/model/vae/mage_vae.hpp

@@ -32,23 +32,24 @@ namespace MageVAE {
             ggml_tensor* weight = affine ? params["weight"] : nullptr;
             ggml_tensor* bias   = affine ? params["bias"] : nullptr;
             if (affine && ctx->weight_adapter) {
-                weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight");
-                bias   = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, bias, prefix + "bias");
+                weight = ctx->weight_adapter->patch_weight(ctx, weight, prefix + "weight");
+                bias   = ctx->weight_adapter->patch_weight(ctx, bias, prefix + "bias");
             }
             // [N, C, H, W] -> [N, H, W, C] so layer norm reduces over channels.
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
-            x = ggml_ext_layer_norm(ctx->ggml_ctx, x, weight, bias, 1e-6f);
-            return ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
+            x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
+            x = ggml_ext_layer_norm(ctx, x, weight, bias, 1e-6f);
+            return ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
         }
     };
 
-    inline ggml_tensor* modulate_2d(ggml_context* ctx,
+    inline ggml_tensor* modulate_2d(GGMLRunnerContext* runner_ctx,
                                     ggml_tensor* x,
                                     ggml_tensor* shift,
                                     ggml_tensor* scale) {
-        shift = ggml_reshape_4d(ctx, shift, 1, 1, shift->ne[0], shift->ne[1]);
-        scale = ggml_reshape_4d(ctx, scale, 1, 1, scale->ne[0], scale->ne[1]);
-        return ggml_add(ctx, ggml_mul(ctx, x, ggml_add(ctx, scale, ggml_ext_ones(ctx, 1, 1, 1, 1))), shift);
+        auto ctx = runner_ctx->ggml_ctx;
+        shift    = ggml_reshape_4d(ctx, shift, 1, 1, shift->ne[0], shift->ne[1]);
+        scale    = ggml_reshape_4d(ctx, scale, 1, 1, scale->ne[0], scale->ne[1]);
+        return ggml_add(ctx, ggml_mul(ctx, x, ggml_add(ctx, scale, ggml_ext_ones(runner_ctx, 1, 1, 1, 1))), shift);
     }
 
     inline ggml_tensor* channel_attention(GGMLRunnerContext* ctx,
@@ -70,7 +71,7 @@ namespace MageVAE {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timestep) {
             auto linear_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
             auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
-            auto x        = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f);
+            auto x        = ggml_ext_timestep_embedding(ctx, timestep, 256, 10000, 1.f);
             x             = linear_0->forward(ctx, x);
             x             = ggml_silu_inplace(ctx->ggml_ctx, x);
             return linear_2->forward(ctx, x);
@@ -139,9 +140,9 @@ namespace MageVAE {
             auto ada   = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
 
             auto params = ada->forward(ctx, ggml_silu(ctx->ggml_ctx, condition));
-            auto chunks = ggml_ext_chunk(ctx->ggml_ctx, params, 6, 0);
+            auto chunks = ggml_ext_chunk(ctx, params, 6, 0);
             auto x      = norm1->forward(ctx, input);
-            x           = modulate_2d(ctx->ggml_ctx, x, chunks[0], chunks[1]);
+            x           = modulate_2d(ctx, x, chunks[0], chunks[1]);
             x           = conv1->forward(ctx, x);
             x           = conv2->forward(ctx, x);
             x           = ggml_gelu(ctx->ggml_ctx, x);
@@ -151,7 +152,7 @@ namespace MageVAE {
             x           = ggml_add(ctx->ggml_ctx, input, ggml_mul(ctx->ggml_ctx, x, gate_1));
 
             auto h      = norm2->forward(ctx, x);
-            h           = modulate_2d(ctx->ggml_ctx, h, chunks[3], chunks[4]);
+            h           = modulate_2d(ctx, h, chunks[3], chunks[4]);
             h           = conv4->forward(ctx, h);
             h           = ggml_gelu(ctx->ggml_ctx, h);
             h           = conv5->forward(ctx, h);
@@ -179,17 +180,18 @@ namespace MageVAE {
         }
     };
 
-    inline ggml_tensor* replicate_pad_right_bottom(ggml_context* ctx,
+    inline ggml_tensor* replicate_pad_right_bottom(GGMLRunnerContext* runner_ctx,
                                                    ggml_tensor* x,
                                                    int pad_w,
                                                    int pad_h) {
+        auto ctx = runner_ctx->ggml_ctx;
         if (pad_w > 0) {
-            auto edge = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]);
+            auto edge = ggml_ext_slice(runner_ctx, x, 0, x->ne[0] - 1, x->ne[0]);
             edge      = ggml_repeat_4d(ctx, edge, pad_w, x->ne[1], x->ne[2], x->ne[3]);
             x         = ggml_concat(ctx, x, edge, 0);
         }
         if (pad_h > 0) {
-            auto edge = ggml_ext_slice(ctx, x, 1, x->ne[1] - 1, x->ne[1]);
+            auto edge = ggml_ext_slice(runner_ctx, x, 1, x->ne[1] - 1, x->ne[1]);
             edge      = ggml_repeat_4d(ctx, edge, x->ne[0], pad_h, x->ne[2], x->ne[3]);
             x         = ggml_concat(ctx, x, edge, 1);
         }
@@ -209,25 +211,27 @@ namespace MageVAE {
             blocks["proj_out"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
         }
 
-        ggml_tensor* to_patches(ggml_context* ctx, ggml_tensor* x) {
-            x = DiT::patchify(ctx, x, patch_size, patch_size);
-            x = ggml_reshape_4d(ctx, x, patch_size * patch_size, channels, x->ne[1], x->ne[2]);
+        ggml_tensor* to_patches(GGMLRunnerContext* runner_ctx, ggml_tensor* x) {
+            auto ctx = runner_ctx->ggml_ctx;
+            x        = DiT::patchify(runner_ctx, x, patch_size, patch_size);
+            x        = ggml_reshape_4d(ctx, x, patch_size * patch_size, channels, x->ne[1], x->ne[2]);
             // [N, np, C, P] -> [N, np, P, C] for attention over P pixels.
-            x = ggml_ext_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
+            x = ggml_ext_cont(runner_ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
             return ggml_reshape_3d(ctx, x, channels, patch_size * patch_size, x->ne[2] * x->ne[3]);
         }
 
-        ggml_tensor* from_patches(ggml_context* ctx,
+        ggml_tensor* from_patches(GGMLRunnerContext* runner_ctx,
                                   ggml_tensor* x,
                                   int64_t patch_count,
                                   int64_t batch_size,
                                   int64_t h_patches,
                                   int64_t w_patches) {
-            x = ggml_reshape_4d(ctx, x, channels, patch_size * patch_size, patch_count, batch_size);
+            auto ctx = runner_ctx->ggml_ctx;
+            x        = ggml_reshape_4d(ctx, x, channels, patch_size * patch_size, patch_count, batch_size);
             // [N, np, P, C] -> [N, np, C, P] before spatial unpatchify.
-            x = ggml_ext_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
+            x = ggml_ext_cont(runner_ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
             x = ggml_reshape_3d(ctx, x, patch_size * patch_size * channels, patch_count, batch_size);
-            return DiT::unpatchify(ctx, x, h_patches, w_patches, patch_size, patch_size);
+            return DiT::unpatchify(runner_ctx, x, h_patches, w_patches, patch_size, patch_size);
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override {
@@ -247,19 +251,19 @@ namespace MageVAE {
             int64_t np     = wp * hp;
 
             auto h = norm->forward(ctx, input);
-            auto q = replicate_pad_right_bottom(ctx->ggml_ctx, q_proj->forward(ctx, h), pad_w, pad_h);
-            auto k = replicate_pad_right_bottom(ctx->ggml_ctx, k_proj->forward(ctx, h), pad_w, pad_h);
-            auto v = replicate_pad_right_bottom(ctx->ggml_ctx, v_proj->forward(ctx, h), pad_w, pad_h);
-            q      = to_patches(ctx->ggml_ctx, q);
-            k      = to_patches(ctx->ggml_ctx, k);
-            v      = to_patches(ctx->ggml_ctx, v);
+            auto q = replicate_pad_right_bottom(ctx, q_proj->forward(ctx, h), pad_w, pad_h);
+            auto k = replicate_pad_right_bottom(ctx, k_proj->forward(ctx, h), pad_w, pad_h);
+            auto v = replicate_pad_right_bottom(ctx, v_proj->forward(ctx, h), pad_w, pad_h);
+            q      = to_patches(ctx, q);
+            k      = to_patches(ctx, k);
+            v      = to_patches(ctx, v);
             h      = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
-            h      = from_patches(ctx->ggml_ctx, h, np, batch, hp, wp);
+            h      = from_patches(ctx, h, np, batch, hp, wp);
             if (pad_h > 0) {
-                h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, height);
+                h = ggml_ext_slice(ctx, h, 1, 0, height);
             }
             if (pad_w > 0) {
-                h = ggml_ext_slice(ctx->ggml_ctx, h, 0, 0, width);
+                h = ggml_ext_slice(ctx, h, 0, 0, width);
             }
             return ggml_add(ctx->ggml_ctx, input, proj_out->forward(ctx, h));
         }
@@ -311,11 +315,11 @@ namespace MageVAE {
                 cond = std::dynamic_pointer_cast<EncoderDiCoBlock>(blocks["head_blocks." + std::to_string(i)])->forward(ctx, cond);
             }
             cond   = std::dynamic_pointer_cast<Conv2d>(blocks["proj_down"])->forward(ctx, cond);
-            auto z = ggml_ext_zeros(ctx->ggml_ctx, cond->ne[0], cond->ne[1], LATENT_CHANNELS, cond->ne[3]);
+            auto z = ggml_ext_zeros(ctx, cond->ne[0], cond->ne[1], LATENT_CHANNELS, cond->ne[3]);
             z      = std::dynamic_pointer_cast<Conv2d>(blocks["z_proj"])->forward(ctx, z);
             z      = ggml_concat(ctx->ggml_ctx, cond, z, 2);
             z      = std::dynamic_pointer_cast<Conv2d>(blocks["fuse_proj"])->forward(ctx, z);
-            auto t = ggml_ext_zeros(ctx->ggml_ctx, image->ne[3], 1, 1, 1);
+            auto t = ggml_ext_zeros(ctx, image->ne[3], 1, 1, 1);
             t      = ggml_reshape_1d(ctx->ggml_ctx, t, image->ne[3]);
             auto c = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"])->forward(ctx, t);
             for (int i = 0; i < 21; ++i) {
@@ -336,9 +340,9 @@ namespace MageVAE {
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* condition) {
             auto params = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"])->forward(ctx, ggml_silu(ctx->ggml_ctx, condition));
-            auto chunks = ggml_ext_chunk(ctx->ggml_ctx, params, 3, 0);
+            auto chunks = ggml_ext_chunk(ctx, params, 3, 0);
             auto h      = std::dynamic_pointer_cast<LayerNorm>(blocks["in_ln"])->forward(ctx, x);
-            h           = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, h, ggml_add(ctx->ggml_ctx, chunks[1], ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1))), chunks[0]);
+            h           = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, h, ggml_add(ctx->ggml_ctx, chunks[1], ggml_ext_ones(ctx, 1, 1, 1, 1))), chunks[0]);
             h           = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"])->forward(ctx, h);
             h           = ggml_silu(ctx->ggml_ctx, h);
             h           = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"])->forward(ctx, h);
@@ -373,8 +377,8 @@ namespace MageVAE {
             int64_t n      = cond->ne[3];
             int64_t length = w * h;
 
-            auto image = ggml_ext_zeros(ctx->ggml_ctx, w * PATCH_SIZE, h * PATCH_SIZE, 3, n);
-            auto t     = ggml_ext_zeros(ctx->ggml_ctx, n, 1, 1, 1);
+            auto image = ggml_ext_zeros(ctx, w * PATCH_SIZE, h * PATCH_SIZE, 3, n);
+            auto t     = ggml_ext_zeros(ctx, n, 1, 1, 1);
             t          = ggml_reshape_1d(ctx->ggml_ctx, t, n);
             auto c     = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"])->forward(ctx, t);
 
@@ -385,15 +389,15 @@ namespace MageVAE {
                 s = std::dynamic_pointer_cast<DiCoBlock>(blocks["blocks." + std::to_string(i)])->forward(ctx, s, c);
             }
             // [N, C, H, W] -> [N*H*W, C].
-            s = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, s, 1, 2, 0, 3));
+            s = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, s, 1, 2, 0, 3));
             s = ggml_reshape_2d(ctx->ggml_ctx, s, HIDDEN_SIZE, length * n);
 
             auto y = std::dynamic_pointer_cast<Conv2d>(blocks["y_embedder_x"])->forward(ctx, cond);
             // Split 32*P channels as [32, P], then produce [N*L, P, 32].
             y          = ggml_reshape_4d(ctx->ggml_ctx, y, length, PATCH_SIZE * PATCH_SIZE, 32, n);
-            y          = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, y, 2, 1, 0, 3));
+            y          = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, y, 2, 1, 0, 3));
             y          = ggml_reshape_3d(ctx->ggml_ctx, y, 32, PATCH_SIZE * PATCH_SIZE, length * n);
-            auto zeros = ggml_ext_zeros(ctx->ggml_ctx, 3, PATCH_SIZE * PATCH_SIZE, length * n, 1);
+            auto zeros = ggml_ext_zeros(ctx, 3, PATCH_SIZE * PATCH_SIZE, length * n, 1);
             dct        = ggml_repeat_4d(ctx->ggml_ctx, dct, 64, PATCH_SIZE * PATCH_SIZE, length * n, 1);
             auto x     = ggml_concat(ctx->ggml_ctx, zeros, y, 0);
             x          = ggml_concat(ctx->ggml_ctx, x, dct, 0);
@@ -409,9 +413,9 @@ namespace MageVAE {
             x = std::dynamic_pointer_cast<Linear>(blocks["final_layer.linear"])->forward(ctx, x);
             // [N*L, P, 3] -> [N, L, 3*P] for fold/unpatchify.
             x = ggml_reshape_4d(ctx->ggml_ctx, x, 3, PATCH_SIZE * PATCH_SIZE, length, n);
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
             x = ggml_reshape_3d(ctx->ggml_ctx, x, 3 * PATCH_SIZE * PATCH_SIZE, length, n);
-            return DiT::unpatchify(ctx->ggml_ctx, x, h, w, PATCH_SIZE, PATCH_SIZE);
+            return DiT::unpatchify(ctx, x, h, w, PATCH_SIZE, PATCH_SIZE);
         }
     };
 

src/model/vae/minimax_h3_audio_vae.hpp

@@ -28,7 +28,7 @@ namespace MiniMaxH3 {
             auto alpha       = params["alpha"];
             auto oscillation = ggml_sin(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, x, alpha));
             oscillation      = ggml_mul(ctx->ggml_ctx, oscillation, oscillation);
-            auto eps         = ggml_ext_scale(ctx->ggml_ctx, ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1), 1e-9f);
+            auto eps         = ggml_ext_scale(ctx, ggml_ext_ones(ctx, 1, 1, 1, 1), 1e-9f);
             return ggml_add(ctx->ggml_ctx,
                             x,
                             ggml_div(ctx->ggml_ctx, oscillation, ggml_add(ctx->ggml_ctx, alpha, eps)));
@@ -59,7 +59,7 @@ namespace MiniMaxH3 {
             auto h     = conv2->forward(ctx, act2->forward(ctx, conv1->forward(ctx, act1->forward(ctx, x))));
             if (x->ne[0] != h->ne[0]) {
                 int64_t pad = (x->ne[0] - h->ne[0]) / 2;
-                x           = ggml_ext_slice(ctx->ggml_ctx, x, 0, pad, x->ne[0] - pad);
+                x           = ggml_ext_slice(ctx, x, 0, pad, x->ne[0] - pad);
             }
             return ggml_add(ctx->ggml_ctx, x, h);
         }
@@ -162,7 +162,7 @@ namespace MiniMaxH3 {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
             auto qkv_layer  = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
             auto proj       = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
-            auto qkv        = ggml_ext_chunk(ctx->ggml_ctx, qkv_layer->forward(ctx, x), 3, 0);
+            auto qkv        = ggml_ext_chunk(ctx, qkv_layer->forward(ctx, x), 3, 0);
             auto bias_shape = [&](ggml_tensor* bias) {
                 return ggml_reshape_4d(ctx->ggml_ctx, bias, bias->ne[0], 1, 1, 1);
             };
@@ -172,7 +172,7 @@ namespace MiniMaxH3 {
 
             int64_t sequence = x->ne[1];
             auto mask        = ggml_diag_mask_inf(ctx->ggml_ctx,
-                                                  ggml_ext_zeros(ctx->ggml_ctx, sequence, sequence, 1, 1),
+                                                  ggml_ext_zeros(ctx, sequence, sequence, 1, 1),
                                                   0);
             auto attn_out    = ggml_ext_attention_ext(ctx,
                                                       q,
@@ -323,7 +323,7 @@ namespace MiniMaxH3 {
                     auto value = block->forward(ctx, x);
                     sum        = sum == nullptr ? value : ggml_add(ctx->ggml_ctx, sum, value);
                 }
-                x = ggml_ext_scale(ctx->ggml_ctx, sum, 1.f / num_kernels);
+                x = ggml_ext_scale(ctx, sum, 1.f / num_kernels);
             }
             auto activation = std::dynamic_pointer_cast<LTXV::Activation1D>(blocks["activation_post"]);
             auto conv_post  = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["conv_post"]);
@@ -412,7 +412,7 @@ namespace MiniMaxH3 {
                                                     streams);
             ggml_tensor* waveform = nullptr;
             for (int64_t stream = 0; stream < streams; ++stream) {
-                auto stream_latent   = ggml_ext_slice(ctx->ggml_ctx, latent, 2, stream, stream + 1);
+                auto stream_latent   = ggml_ext_slice(ctx, latent, 2, stream, stream + 1);
                 auto stream_waveform = decoder->forward(ctx, dec_in->forward(ctx, stream_latent));
                 waveform             = waveform == nullptr
                                            ? stream_waveform

src/model/vae/minimax_h3_vae.hpp

@@ -66,8 +66,8 @@ namespace MiniMaxH3VAE {
             auto reflect_pad = [&](ggml_tensor* value, int dim, int amount) {
                 for (int i = 0; i < amount; ++i) {
                     GGML_ASSERT(value->ne[dim] > 1);
-                    auto left  = ggml_ext_slice(ctx->ggml_ctx, value, dim, 1, 2);
-                    auto right = ggml_ext_slice(ctx->ggml_ctx,
+                    auto left  = ggml_ext_slice(ctx, value, dim, 1, 2);
+                    auto right = ggml_ext_slice(ctx,
                                                 value,
                                                 dim,
                                                 value->ne[dim] - 2,
@@ -82,8 +82,7 @@ namespace MiniMaxH3VAE {
             x                = reflect_pad(x, 1, std::get<1>(temporal_padding));
             int temporal_pad = std::get<0>(temporal_padding) * 2;
             if (temporal_pad > 0) {
-                x = ggml_ext_pad_ext(ctx->ggml_ctx,
-                                     ctx->backend,
+                x = ggml_ext_pad_ext(ctx,
                                      x,
                                      0,
                                      0,
@@ -105,7 +104,7 @@ namespace MiniMaxH3VAE {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
             ggml_tensor* result = nullptr;
             for (int64_t t = 0; t < x->ne[2]; ++t) {
-                auto frame = ggml_ext_slice(ctx->ggml_ctx, x, 2, t, t + 1);
+                auto frame = ggml_ext_slice(ctx, x, 2, t, t + 1);
                 GGML_ASSERT(frame->ne[3] % num_channels == 0);
                 int64_t batch_size = frame->ne[3] / num_channels;
                 frame              = ggml_cont(ctx->ggml_ctx, frame);
@@ -146,9 +145,9 @@ namespace MiniMaxH3VAE {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
             if (spatial_stride == 2) {
                 GGML_ASSERT(x->ne[0] > 1 && x->ne[1] > 1);
-                auto right  = ggml_ext_slice(ctx->ggml_ctx, x, 0, x->ne[0] - 2, x->ne[0] - 1);
+                auto right  = ggml_ext_slice(ctx, x, 0, x->ne[0] - 2, x->ne[0] - 1);
                 x           = ggml_concat(ctx->ggml_ctx, x, right, 0);
-                auto bottom = ggml_ext_slice(ctx->ggml_ctx, x, 1, x->ne[1] - 2, x->ne[1] - 1);
+                auto bottom = ggml_ext_slice(ctx, x, 1, x->ne[1] - 2, x->ne[1] - 1);
                 x           = ggml_concat(ctx->ggml_ctx, x, bottom, 1);
             }
             return std::dynamic_pointer_cast<CausalConv3d>(blocks["conv"])->forward(ctx, x);
@@ -263,14 +262,14 @@ namespace MiniMaxH3VAE {
         auto ctx        = runner_ctx->ggml_ctx;
         int64_t rot_dim = pe->ne[2] * 2;
         auto rotated    = Rope::apply_rope(runner_ctx,
-                                           ggml_ext_slice(ctx, x, 0, 0, rot_dim),
+                                           ggml_ext_slice(runner_ctx, x, 0, 0, rot_dim),
                                            pe,
                                            false);
         if (rot_dim == x->ne[0]) {
             return rotated;
         }
         auto tail = attention_layout(ctx,
-                                     ggml_ext_slice(ctx, x, 0, rot_dim, x->ne[0]));
+                                     ggml_ext_slice(runner_ctx, x, 0, rot_dim, x->ne[0]));
         return ggml_concat(ctx, rotated, tail, 0);
     }
 
@@ -298,7 +297,7 @@ namespace MiniMaxH3VAE {
                                                   num_head,
                                                   sequence,
                                                   batch_size);
-            auto qkv            = ggml_ext_chunk(ctx->ggml_ctx, qkv_projection, 3, 0);
+            auto qkv            = ggml_ext_chunk(ctx, qkv_projection, 3, 0);
             auto q              = ggml_reshape_4d(ctx->ggml_ctx,
                                                   qkv[0],
                                                   head_dim,
@@ -434,7 +433,7 @@ namespace MiniMaxH3VAE {
             GGML_ASSERT(batch_size == 1);
 
             z                   = ggml_cont(ctx->ggml_ctx,
-                                            ggml_ext_torch_permute(ctx->ggml_ctx, z, 3, 0, 1, 2));
+                                            ggml_ext_torch_permute(ctx, z, 3, 0, 1, 2));
             z                   = ggml_reshape_3d(ctx->ggml_ctx,
                                                   z,
                                                   24,
@@ -444,8 +443,8 @@ namespace MiniMaxH3VAE {
             auto h              = x_embedder->forward(ctx, z);
             int64_t num_patches = h->ne[1];
             h                   = ggml_concat(ctx->ggml_ctx, h, params["register_tokens"], 1);
-            auto zero           = ggml_ext_scale(ctx->ggml_ctx,
-                                                 ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, 1),
+            auto zero           = ggml_ext_scale(ctx,
+                                                 ggml_ext_slice(ctx, h, 1, 0, 1),
                                                  0.f);
             h                   = ggml_concat(ctx->ggml_ctx, h, zero, 1);
 
@@ -461,8 +460,8 @@ namespace MiniMaxH3VAE {
             auto norm_out = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"]);
             auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
             h             = proj_out->forward(ctx, norm_out->forward(ctx, h));
-            h             = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, num_patches);
-            return DiT::unpatchify_3d(ctx->ggml_ctx,
+            h             = ggml_ext_slice(ctx, h, 1, 0, num_patches);
+            return DiT::unpatchify_3d(ctx,
                                       h,
                                       num_frames,
                                       height,
@@ -499,7 +498,7 @@ namespace MiniMaxH3VAE {
             auto encoder = std::dynamic_pointer_cast<Encoder>(blocks["encoder"]);
             auto quant   = std::dynamic_pointer_cast<Conv3d>(blocks["quant_conv"]);
             auto moments = quant->forward(ctx, encoder->forward(ctx, pixels));
-            return single_frame ? moments : ggml_ext_slice(ctx->ggml_ctx, moments, 3, 0, 24);
+            return single_frame ? moments : ggml_ext_slice(ctx, moments, 3, 0, 24);
         }
 
         ggml_tensor* decode(GGMLRunnerContext* ctx,

src/model/vae/seedvr2_vae.hpp

@@ -45,20 +45,20 @@ namespace SeedVR2 {
                 // equivalent spatial filter to F16, instead of copying and convolving the halo.
                 auto w       = ggml_cast(c, params["weight"], GGML_TYPE_F32);
                 w            = ggml_reshape_3d(c, w, kw * kh, kt, in_channels * out_channels);
-                auto spatial = ggml_ext_slice(c, w, 1, 0, 1, false);
+                auto spatial = ggml_ext_slice(ctx, w, 1, 0, 1, false);
                 for (int i = 1; i < kt; ++i)
-                    spatial = ggml_add(c, spatial, ggml_ext_slice(c, w, 1, i, i + 1, false));
+                    spatial = ggml_add(c, spatial, ggml_ext_slice(ctx, w, 1, i, i + 1, false));
                 w = ggml_cast(c, ggml_reshape_4d(c, spatial, kw, kh, 1, in_channels * out_channels), GGML_TYPE_F16);
                 if (ctx->conv2d_direct_enabled) {
                     w = ggml_reshape_4d(c, w, kw, kh, in_channels, out_channels);
-                    x = ggml_reshape_4d(c, ggml_ext_cont(c, x), x->ne[0], x->ne[1], in_channels, 1);
-                    x = ggml_ext_conv_2d(c, x, w, bias ? params["bias"] : nullptr,
+                    x = ggml_reshape_4d(c, ggml_ext_cont(ctx, x), x->ne[0], x->ne[1], in_channels, 1);
+                    x = ggml_ext_conv_2d(ctx, x, w, bias ? params["bias"] : nullptr,
                                          std::get<2>(stride), std::get<1>(stride),
                                          std::get<2>(padding), std::get<1>(padding),
                                          std::get<2>(dilation), std::get<1>(dilation), true);
                     return ggml_reshape_4d(c, x, x->ne[0], x->ne[1], 1, out_channels);
                 }
-                return ggml_ext_conv_3d(c, ctx->backend, x, w, bias ? params["bias"] : nullptr, in_channels,
+                return ggml_ext_conv_3d(ctx, x, w, bias ? params["bias"] : nullptr, in_channels,
                                         std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
                                         std::get<2>(padding), std::get<1>(padding), 0,
                                         std::get<2>(dilation), std::get<1>(dilation), 1, force_prec_f32, true);
@@ -66,8 +66,8 @@ namespace SeedVR2 {
             if (temporal_padding_) {
                 const int64_t w = x->ne[0], h = x->ne[1], t = x->ne[2], ch = x->ne[3];
                 // Merge spatial axes so concat does not launch separately for every channel.
-                x          = ggml_reshape_3d(c, ggml_ext_cont(c, x), w * h, t, ch);
-                auto head  = ggml_ext_slice(c, x, 1, 0, 1);
+                x          = ggml_reshape_3d(c, ggml_ext_cont(ctx, x), w * h, t, ch);
+                auto head  = ggml_ext_slice(ctx, x, 1, 0, 1);
                 auto shape = ggml_new_tensor_3d(c, x->type, w * h, temporal_padding_, ch);
                 x          = ggml_concat(c, ggml_repeat(c, head, shape), x, 1);
                 x          = ggml_reshape_4d(c, x, w, h, t + temporal_padding_, ch);
@@ -92,16 +92,16 @@ namespace SeedVR2 {
                 ~RestorePadding() { target = value; }
             } restore{padding, padding};
             padding     = {0, 0, pw};
-            auto output = ggml_ext_zeros(c, ow, oh, ot, out_channels);
+            auto output = ggml_ext_zeros(ctx, ow, oh, ot, out_channels);
             for (int64_t t = 0; t < ot; ++t) {
-                auto frame = ggml_ext_slice(c, x, 2, t * st, t * st + kt, false);
+                auto frame = ggml_ext_slice(ctx, x, 2, t * st, t * st + kt, false);
                 for (int64_t y = 0; y < oh; y += rows) {
                     const int64_t end = std::min(oh, y + rows);
                     const int64_t top = y * sh - ph, bottom = (end - 1) * sh - ph + kh;
-                    auto strip = ggml_ext_slice(c, frame, 1, std::max<int64_t>(0, top), std::min(x->ne[1], bottom), false);
-                    strip      = ggml_ext_pad_ext(c, ctx->backend, strip, 0, 0, int(std::max<int64_t>(0, -top)),
+                    auto strip = ggml_ext_slice(ctx, frame, 1, std::max<int64_t>(0, top), std::min(x->ne[1], bottom), false);
+                    strip      = ggml_ext_pad_ext(ctx, strip, 0, 0, int(std::max<int64_t>(0, -top)),
                                                   int(std::max<int64_t>(0, bottom - x->ne[1])), 0, 0, 0, 0);
-                    strip      = ggml_ext_cont(c, strip);
+                    strip      = ggml_ext_cont(ctx, strip);
                     strip      = Conv3d::forward(ctx, strip);
                     output     = ggml_set_inplace(c, output, strip, output->nb[1], output->nb[2], output->nb[3], y * output->nb[1] + t * output->nb[2]);
                 }
@@ -153,7 +153,7 @@ namespace SeedVR2 {
             h           = ggml_reshape_3d(c, h, width * height, frames, channels);
             h           = ggml_cont(c, ggml_permute(c, h, 1, 2, 0, 3));
             auto linear = [&](const char* name, ggml_tensor* value) { return std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, value); };
-            h           = ggml_ext_attention_ext(c, ctx->backend, linear("to_q", h), linear("to_k", h), linear("to_v", h), 1, nullptr, false, ctx->flash_attn_enabled);
+            h           = ggml_ext_attention_ext(ctx, linear("to_q", h), linear("to_k", h), linear("to_v", h), 1, nullptr, false, ctx->flash_attn_enabled);
             h           = linear("to_out.0", h);
             h           = ggml_cont(c, ggml_permute(c, h, 2, 0, 1, 3));
             return ggml_add(c, x, ggml_reshape_4d(c, h, width, height, frames, channels));
@@ -184,8 +184,8 @@ namespace SeedVR2 {
             x = ggml_cont(c, ggml_permute(c, x, 0, 2, 3, 1));
             x = ggml_reshape_4d(c, x, w * 2, h * 2, t * temporal_, ch);
             if (temporal_ == 2) {
-                auto first = ggml_ext_slice(c, x, 2, 0, 1);
-                x          = t == 1 ? first : ggml_concat(c, first, ggml_ext_slice(c, x, 2, 2, t * 2), 2);
+                auto first = ggml_ext_slice(ctx, x, 2, 0, 1);
+                x          = t == 1 ? first : ggml_concat(c, first, ggml_ext_slice(ctx, x, 2, 2, t * 2), 2);
             }
             return std::dynamic_pointer_cast<UnaryBlock>(blocks["conv"])->forward(ctx, x);
         }
@@ -255,7 +255,7 @@ namespace SeedVR2 {
                 mid();
             x = ggml_silu(c, causal_norm(ctx, x, blocks[root + "conv_norm_out"]));
             apply("conv_out");
-            return decode ? x : ggml_ext_slice(c, x, 3, 0, 16);
+            return decode ? x : ggml_ext_slice(ctx, x, 3, 0, 16);
         }
     };
 

src/model/vae/tae.hpp

@@ -179,9 +179,9 @@ public:
         // z: [n, z_channels, h, w]
         // return: [n, out_channels, h*8, w*8]
 
-        auto h = ggml_ext_scale(ctx->ggml_ctx, z, 1.0f / 3.0f);
+        auto h = ggml_ext_scale(ctx, z, 1.0f / 3.0f);
         h      = ggml_tanh_inplace(ctx->ggml_ctx, h);
-        h      = ggml_ext_scale(ctx->ggml_ctx, h, 3.0f);
+        h      = ggml_ext_scale(ctx, h, 3.0f);
 
         for (int i = 0; i < num_blocks * 3 + 10; i++) {
             if (blocks.find(std::to_string(i)) == blocks.end()) {
@@ -322,10 +322,8 @@ public:
 };
 
 ggml_tensor*
-patchify(ggml_context* ctx,
-         ggml_tensor* x,
-         int64_t patch_size,
-         int64_t b = 1) {
+patchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t patch_size, int64_t b = 1) {
+    auto ctx = runner_ctx->ggml_ctx;
     // x: [f, b*c, h*q, w*r]
     // return: [f, b*c*r*q, h, w]
     if (patch_size == 1) {
@@ -342,19 +340,17 @@ patchify(ggml_context* ctx,
     int64_t w = W / r;
     int64_t h = H / q;
 
-    x = ggml_reshape_4d(ctx, x, W, q, h, C * f);                         // [W, q, h, C*f]
-    x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));  // [W, h, q, C*f]
-    x = ggml_reshape_4d(ctx, x, r, w, h, q * C * f);                     // [r, w, h, q*C*f]
-    x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));  // [w, h, r, q*C*f]
-    x = ggml_reshape_4d(ctx, x, w, h, r * q * C, f);                     // [f, b*c*r*q, h, w]
+    x = ggml_reshape_4d(ctx, x, W, q, h, C * f);                                       // [W, q, h, C*f]
+    x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [W, h, q, C*f]
+    x = ggml_reshape_4d(ctx, x, r, w, h, q * C * f);                                   // [r, w, h, q*C*f]
+    x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 1, 2, 0, 3));  // [w, h, r, q*C*f]
+    x = ggml_reshape_4d(ctx, x, w, h, r * q * C, f);                                   // [f, b*c*r*q, h, w]
 
     return x;
 }
 
-ggml_tensor* unpatchify(ggml_context* ctx,
-                        ggml_tensor* x,
-                        int64_t patch_size,
-                        int64_t b = 1) {
+ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t patch_size, int64_t b = 1) {
+    auto ctx = runner_ctx->ggml_ctx;
     // x: [f, b*c*r*q, h, w]
     // return: [f, b*c, h*q, w*r]
     if (patch_size == 1) {
@@ -367,10 +363,10 @@ ggml_tensor* unpatchify(ggml_context* ctx,
     int64_t h = x->ne[1];
     int64_t w = x->ne[0];
 
-    x = ggml_reshape_4d(ctx, x, w, h, r, q * c * b * f);                 // [q*c*b*f, r, h, w]
-    x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));  // [r, w, h, q*c*b*f]
-    x = ggml_reshape_4d(ctx, x, r * w, h, q, c * b * f);                 // [c*b*f, q, h, r*w]
-    x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));  // [r*w, q, h, c*b*f]
+    x = ggml_reshape_4d(ctx, x, w, h, r, q * c * b * f);                               // [q*c*b*f, r, h, w]
+    x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 2, 0, 1, 3));  // [r, w, h, q*c*b*f]
+    x = ggml_reshape_4d(ctx, x, r * w, h, q, c * b * f);                               // [c*b*f, q, h, r*w]
+    x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [r*w, q, h, c*b*f]
     x = ggml_reshape_4d(ctx, x, r * w, q * h, c * b, f);
 
     return x;
@@ -412,7 +408,7 @@ public:
         auto first_conv = std::dynamic_pointer_cast<Conv2d>(blocks["0"]);
 
         if (patch_size > 1) {
-            z = patchify(ctx->ggml_ctx, z, patch_size, 1);
+            z = patchify(ctx, z, patch_size, 1);
         }
 
         auto h = first_conv->forward(ctx, z);
@@ -427,7 +423,7 @@ public:
             h = conv->forward(ctx, h);
             for (int j = 0; j < num_blocks; j++) {
                 auto block = std::dynamic_pointer_cast<MemBlock>(blocks[std::to_string(index++)]);
-                auto mem   = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, h, 0, 0, 0, 0, 0, 0, 1, 0);
+                auto mem   = ggml_ext_pad_ext(ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
                 mem        = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
                 h          = block->forward(ctx, h, mem);
             }
@@ -487,9 +483,9 @@ public:
         auto first_conv = std::dynamic_pointer_cast<Conv2d>(blocks["1"]);
 
         // Clamp()
-        auto h = ggml_ext_scale(ctx->ggml_ctx,
+        auto h = ggml_ext_scale(ctx,
                                 ggml_tanh_inplace(ctx->ggml_ctx,
-                                                  ggml_ext_scale(ctx->ggml_ctx, z, 1.0f / 3.0f)),
+                                                  ggml_ext_scale(ctx, z, 1.0f / 3.0f)),
                                 3.0f,
                                 true);
 
@@ -498,7 +494,7 @@ public:
         int index = 3;
         for (int i = 0; i < num_layers; i++) {
             for (int j = 0; j < num_blocks; j++) {
-                auto mem = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, h, 0, 0, 0, 0, 0, 0, 1, 0);
+                auto mem = ggml_ext_pad_ext(ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
                 mem      = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
                 if (is_wide) {
                     auto block = std::dynamic_pointer_cast<WideMemBlock>(blocks[std::to_string(index++)]);
@@ -521,7 +517,7 @@ public:
         auto last_conv = std::dynamic_pointer_cast<Conv2d>(blocks[std::to_string(++index)]);
         h              = last_conv->forward(ctx, h);
         if (patch_size > 1) {
-            h = unpatchify(ctx->ggml_ctx, h, patch_size, 1);
+            h = unpatchify(ctx, h, patch_size, 1);
         }
         // shape(W, H, 3, (t_upscale - 1) + T) => shape(W, H, 3, T)
         h = ggml_view_4d(ctx->ggml_ctx, h, h->ne[0], h->ne[1], h->ne[2], h->ne[3] - (t_upscale - 1), h->nb[1], h->nb[2], h->nb[3], (t_upscale - 1) * h->nb[3]);
@@ -579,7 +575,7 @@ public:
             int64_t chunk_frames = 5 * decoder->t_upscale;
             int64_t pad          = (chunk_frames - (num_frames % chunk_frames)) % chunk_frames;
 
-            result = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, result, 0, 0, 0, 0, 0, 0, 0, static_cast<int>(pad), false, false);
+            result = ggml_ext_pad_ext(ctx, result, 0, 0, 0, 0, 0, 0, 0, static_cast<int>(pad), false, false);
 
             int64_t num_chunks                  = (num_frames + pad) / chunk_frames;
             auto to_trim                        = decoder->t_upscale - 1;
@@ -591,7 +587,7 @@ public:
                                           i * chunk_frames * result->nb[3]);
                 to_concat.push_back(chunk);
             }
-            result = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3);
+            result = ggml_ext_vec_concat(ctx, to_concat, 3);
             result = ggml_view_4d(ctx->ggml_ctx, result,
                                   result->ne[0], result->ne[1], result->ne[2],
                                   result->ne[3] - decoder->t_upscale * 3,
@@ -645,7 +641,7 @@ public:
 
                 to_concat.push_back(chunk_padded);
             }
-            ggml_tensor* x_in = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3);
+            ggml_tensor* x_in = ggml_ext_vec_concat(ctx, to_concat, 3);
             out               = encoder->forward(ctx, x_in);
         } else {
             std::vector<ggml_tensor*> to_concat = {};
@@ -662,7 +658,7 @@ public:
 
                 to_concat.push_back(chunk_out);
             }
-            out = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3);
+            out = ggml_ext_vec_concat(ctx, to_concat, 3);
         }
 
         // Return x[:, :-3] - drop the last 3 elements in the T dimension
@@ -738,21 +734,21 @@ public:
     ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
         auto decoder = std::dynamic_pointer_cast<TinyDecoder>(blocks["decoder.layers"]);
         if (taef2) {
-            z = unpatchify(ctx->ggml_ctx, z, 2);
+            z = unpatchify(ctx, z, 2);
         }
         auto x = decoder->forward(ctx, z);
-        return f16 ? unpatchify(ctx->ggml_ctx, x, 2) : x;
+        return f16 ? unpatchify(ctx, x, 2) : x;
     }
 
     ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
         if (f16) {
             GGML_ASSERT(x->ne[0] % 2 == 0 && x->ne[1] % 2 == 0);
-            x = patchify(ctx->ggml_ctx, x, 2);
+            x = patchify(ctx, x, 2);
         }
         auto encoder = std::dynamic_pointer_cast<TinyEncoder>(blocks["encoder.layers"]);
         auto z       = encoder->forward(ctx, x);
         if (taef2) {
-            z = patchify(ctx->ggml_ctx, z, 2);
+            z = patchify(ctx, z, 2);
         }
         return z;
     }

src/model/vae/wan_vae.hpp

@@ -86,26 +86,12 @@ namespace WAN {
                 lp2 -= (int)cache_x->ne[2];
             }
 
-            x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
-            if (w->ne[2] == 1 && x->ne[2] == 1 && x->ne[3] == in_channels) {
-                // One frame through a one-frame-deep kernel is a 2D conv; backends without
-                // im2col_3d (Metal) otherwise fall back to a much slower direct conv_3d.
-                if (!ggml_is_contiguous(x)) {
-                    x = ggml_cont(ctx->ggml_ctx, x);
-                }
-                ggml_tensor* x2 = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], x->ne[1], in_channels, 1);
-                ggml_tensor* w2 = ggml_reshape_4d(ctx->ggml_ctx, w, w->ne[0], w->ne[1], in_channels, out_channels);
-                x2              = ggml_ext_conv_2d(ctx->ggml_ctx, x2, w2, b,
-                                                   std::get<2>(stride), std::get<1>(stride), 0, 0,
-                                                   std::get<2>(dilation), std::get<1>(dilation),
-                                                   ctx->conv2d_direct_enabled, false, false, scale);
-                return ggml_reshape_4d(ctx->ggml_ctx, x2, x2->ne[0], x2->ne[1], 1, out_channels);
-            }
-            return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
+            x = ggml_ext_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
+            return ggml_ext_conv_3d(ctx, x, w, b, in_channels,
                                     std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
                                     0, 0, 0,
                                     std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation),
-                                    false, ctx->conv3d_direct_enabled, scale);
+                                    false, ctx->conv3d_direct_enabled, scale, true);
         }
     };
 
@@ -133,10 +119,10 @@ namespace WAN {
 
             ggml_tensor* w = params["gamma"];
             w              = ggml_reshape_1d(ctx->ggml_ctx, w, ggml_nelements(w));
-            auto h         = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 3, 0, 1, 2));  // [ID, IH, IW, N*IC]
+            auto h         = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 3, 0, 1, 2));  // [ID, IH, IW, N*IC]
             h              = ggml_rms_norm(ctx->ggml_ctx, h, 1e-12f);
             h              = ggml_mul(ctx->ggml_ctx, h, w);
-            h              = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, h, 1, 2, 3, 0));
+            h              = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, h, 1, 2, 3, 0));
 
             return h;
         }
@@ -223,16 +209,16 @@ namespace WAN {
                     } else {
                         auto time_conv = std::dynamic_pointer_cast<CausalConv3d>(blocks["time_conv"]);
 
-                        auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]);
+                        auto cache_x = ggml_ext_slice(ctx, x, 2, -CACHE_T, x->ne[2]);
                         if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) {  // chunk_idx >= 2
                             // cache last frame of last two chunk
                             cache_x = ggml_concat(ctx->ggml_ctx,
-                                                  ggml_ext_slice(ctx->ggml_ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]),
+                                                  ggml_ext_slice(ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]),
                                                   cache_x,
                                                   2);
                         }
                         if (chunk_idx == 1 && cache_x->ne[2] < 2) {  // Rep
-                            cache_x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, cache_x, 0, 0, 0, 0, (int)cache_x->ne[2], 0, 0, 0);
+                            cache_x = ggml_ext_pad_ext(ctx, cache_x, 0, 0, 0, 0, (int)cache_x->ne[2], 0, 0, 0);
                             // aka cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device),cache_x],dim=2)
                         }
                         if (chunk_idx == 1) {
@@ -241,9 +227,9 @@ namespace WAN {
                             x = time_conv->forward(ctx, x, feat_cache[idx]);
                         }
                         feat_cache[idx] = cache_x;
-                        x               = ggml_reshape_4d(ctx->ggml_ctx, x, w * h, t, c, 2);                                   // (2, c, t, h*w)
-                        x               = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 3, 1, 2));  // (c, t, 2, h*w)
-                        x               = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, 2 * t, c);                                   // (c, t*2, h, w)
+                        x               = ggml_reshape_4d(ctx->ggml_ctx, x, w * h, t, c, 2);               // (2, c, t, h*w)
+                        x               = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 3, 1, 2));  // (c, t, 2, h*w)
+                        x               = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, 2 * t, c);               // (c, t*2, h, w)
                     }
                 }
             }
@@ -252,18 +238,18 @@ namespace WAN {
             if (mode != "none") {
                 auto resample_1 = std::dynamic_pointer_cast<Conv2d>(blocks["resample.1"]);
 
-                x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // (t, c, h, w)
+                x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));  // (t, c, h, w)
                 if (mode == "upsample2d") {
                     x = ggml_upscale(ctx->ggml_ctx, x, 2, GGML_SCALE_MODE_NEAREST);
                 } else if (mode == "upsample3d") {
                     x = ggml_upscale(ctx->ggml_ctx, x, 2, GGML_SCALE_MODE_NEAREST);
                 } else if (mode == "downsample2d") {
-                    x = ggml_ext_pad(ctx->ggml_ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
+                    x = ggml_ext_pad(ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
                 } else if (mode == "downsample3d") {
-                    x = ggml_ext_pad(ctx->ggml_ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
+                    x = ggml_ext_pad(ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
                 }
                 x = resample_1->forward(ctx, x);
-                x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // (c, t, h, w)
+                x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));  // (c, t, h, w)
             }
 
             if (mode == "downsample3d") {
@@ -275,9 +261,9 @@ namespace WAN {
                     } else {
                         auto time_conv = std::dynamic_pointer_cast<CausalConv3d>(blocks["time_conv"]);
 
-                        auto cache_x    = ggml_ext_slice(ctx->ggml_ctx, x, 2, -1, x->ne[2]);
+                        auto cache_x    = ggml_ext_slice(ctx, x, 2, -1, x->ne[2]);
                         x               = ggml_concat(ctx->ggml_ctx,
-                                                      ggml_ext_slice(ctx->ggml_ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]),
+                                                      ggml_ext_slice(ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]),
                                                       x,
                                                       2);
                         x               = time_conv->forward(ctx, x);
@@ -320,19 +306,19 @@ namespace WAN {
 
             int pad_t = (factor_t - T % factor_t) % factor_t;
 
-            x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, 0, 0, 0, 0, pad_t, 0, 0, 0);
+            x = ggml_ext_pad_ext(ctx, x, 0, 0, 0, 0, pad_t, 0, 0, 0);
             T = x->ne[2];
 
             x = ggml_reshape_4d(ctx->ggml_ctx, x, W * H, factor_t, T / factor_t, C);                                                  // [C, T/factor_t, factor_t, H*W]
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));                                       // [C, factor_t, T/factor_t, H*W]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));                                                 // [C, factor_t, T/factor_t, H*W]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, W, factor_s, (H / factor_s) * (T / factor_t), factor_t * C);                        // [C*factor_t, T/factor_t*H/factor_s, factor_s, W]
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));                                       // [C*factor_t, factor_s, T/factor_t*H/factor_s, W]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));                                                 // [C*factor_t, factor_s, T/factor_t*H/factor_s, W]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s, W / factor_s, (H / factor_s) * (T / factor_t), factor_s * factor_t * C);  // [C*factor_t*factor_s, T/factor_t*H/factor_s, W/factor_s, factor_s]
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));                                       // [C*factor_t*factor_s, factor_s, T/factor_t*H/factor_s, W/factor_s]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));                                                 // [C*factor_t*factor_s, factor_s, T/factor_t*H/factor_s, W/factor_s]
             x = ggml_reshape_3d(ctx->ggml_ctx, x, (W / factor_s) * (H / factor_s) * (T / factor_t), group_size, out_channels);        // [out_channels, group_size, T/factor_t*H/factor_s*W/factor_s]
 
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));  // [out_channels, T/factor_t*H/factor_s*W/factor_s, group_size]
-            x = ggml_mean(ctx->ggml_ctx, x);                                                     // [out_channels, T/factor_t*H/factor_s*W/factor_s, 1]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3));  // [out_channels, T/factor_t*H/factor_s*W/factor_s, group_size]
+            x = ggml_mean(ctx->ggml_ctx, x);                                           // [out_channels, T/factor_t*H/factor_s*W/factor_s, 1]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, W / factor_s, H / factor_s, T / factor_t, out_channels);
             return x;
         }
@@ -374,15 +360,15 @@ namespace WAN {
             C = out_channels;
 
             x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, factor_s, factor_s * factor_t * C);  // [C*factor_t*factor_s, factor_s, T*H, W]
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));  // [C*factor_t*factor_s, T*H, W, factor_s]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));            // [C*factor_t*factor_s, T*H, W, factor_s]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, H * T, factor_s, factor_t * C);  // [C*factor_t, factor_s, T*H, W*factor_s]
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));  // [C*factor_t, T*H, factor_s, W*factor_s]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));            // [C*factor_t, T*H, factor_s, W*factor_s]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W * factor_s * H, T, factor_t, C);  // [C, factor_t, T, H*factor_s*W*factor_s]
-            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));  // [C, T, factor_t, H*factor_s*W*factor_s]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));            // [C, T, factor_t, H*factor_s*W*factor_s]
             x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, factor_s * H, factor_t * T, C);  // [C, T*factor_t, H*factor_s, W*factor_s]
 
             if (first_chunk) {
-                x = ggml_ext_slice(ctx->ggml_ctx, x, 2, factor_t - 1, x->ne[2]);
+                x = ggml_ext_slice(ctx, x, 2, factor_t - 1, x->ne[2]);
             }
 
             return x;
@@ -451,11 +437,11 @@ namespace WAN {
 
                     if (feat_cache.size() > 0) {
                         int idx      = feat_idx;
-                        auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]);
+                        auto cache_x = ggml_ext_slice(ctx, x, 2, -CACHE_T, x->ne[2]);
                         if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) {
                             // cache last frame of last two chunk
                             cache_x = ggml_concat(ctx->ggml_ctx,
-                                                  ggml_ext_slice(ctx->ggml_ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]),
+                                                  ggml_ext_slice(ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]),
                                                   cache_x,
                                                   2);
                         }
@@ -629,7 +615,7 @@ namespace WAN {
 
             x = norm->forward(ctx, x);
 
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // (t, c, h, w)
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));  // (t, c, h, w)
 
             const int64_t n = x->ne[3];
             const int64_t c = x->ne[2];
@@ -637,28 +623,28 @@ namespace WAN {
             const int64_t w = x->ne[0];
 
             auto qkv     = to_qkv->forward(ctx, x);
-            auto qkv_vec = split_image_qkv(ctx->ggml_ctx, qkv);
+            auto qkv_vec = ggml_ext_split_image_qkv(ctx, qkv);
 
             auto q = qkv_vec[0];
-            q      = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, q, 2, 0, 1, 3));  // [t, h, w, c]
-            q      = ggml_reshape_3d(ctx->ggml_ctx, q, c, h * w, n);                                      // [t, h * w, c]
+            q      = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, q, 2, 0, 1, 3));  // [t, h, w, c]
+            q      = ggml_reshape_3d(ctx->ggml_ctx, q, c, h * w, n);                  // [t, h * w, c]
 
             auto k = qkv_vec[1];
-            k      = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, k, 2, 0, 1, 3));  // [t, h, w, c]
-            k      = ggml_reshape_3d(ctx->ggml_ctx, k, c, h * w, n);                                      // [t, h * w, c]
+            k      = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, k, 2, 0, 1, 3));  // [t, h, w, c]
+            k      = ggml_reshape_3d(ctx->ggml_ctx, k, c, h * w, n);                  // [t, h * w, c]
 
             auto v = qkv_vec[2];
             v      = ggml_reshape_3d(ctx->ggml_ctx, v, h * w, c, n);  // [t, c, h * w]
 
-            v = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3));    // [t, h * w, c]
+            v = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, v, 1, 0, 2, 3));              // [t, h * w, c]
             x = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);  // [t, h * w, c]
 
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));  // [t, c, h * w]
-            x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, c, n);                             // [t, c, h, w]
+            x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));  // [t, c, h * w]
+            x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, c, n);                   // [t, c, h, w]
 
             x = proj->forward(ctx, x);
 
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // (c, t, h, w)
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2));  // (c, t, h, w)
 
             x = ggml_add(ctx->ggml_ctx, x, identity);
             return x;
@@ -770,11 +756,11 @@ namespace WAN {
                 x          = conv1->forward(ctx, x);
             } else if (feat_cache.size() > 0) {
                 int idx      = feat_idx;
-                auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]);
+                auto cache_x = ggml_ext_slice(ctx, x, 2, -CACHE_T, x->ne[2]);
                 if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) {
                     // cache last frame of last two chunk
                     cache_x = ggml_concat(ctx->ggml_ctx,
-                                          ggml_ext_slice(ctx->ggml_ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]),
+                                          ggml_ext_slice(ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]),
                                           cache_x,
                                           2);
                 }
@@ -828,11 +814,11 @@ namespace WAN {
                 x           = head_2->forward(ctx, x);
             } else if (feat_cache.size() > 0) {
                 int idx      = feat_idx;
-                auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]);
+                auto cache_x = ggml_ext_slice(ctx, x, 2, -CACHE_T, x->ne[2]);
                 if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) {
                     // cache last frame of last two chunk
                     cache_x = ggml_concat(ctx->ggml_ctx,
-                                          ggml_ext_slice(ctx->ggml_ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]),
+                                          ggml_ext_slice(ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]),
                                           cache_x,
                                           2);
                 }
@@ -962,11 +948,11 @@ namespace WAN {
                 x          = conv1->forward(ctx, x);
             } else if (feat_cache.size() > 0) {
                 int idx      = feat_idx;
-                auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]);
+                auto cache_x = ggml_ext_slice(ctx, x, 2, -CACHE_T, x->ne[2]);
                 if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) {
                     // cache last frame of last two chunk
                     cache_x = ggml_concat(ctx->ggml_ctx,
-                                          ggml_ext_slice(ctx->ggml_ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]),
+                                          ggml_ext_slice(ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]),
                                           cache_x,
                                           2);
                 }
@@ -1020,11 +1006,11 @@ namespace WAN {
                 x           = head_2->forward(ctx, x);
             } else if (feat_cache.size() > 0) {
                 int idx      = feat_idx;
-                auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]);
+                auto cache_x = ggml_ext_slice(ctx, x, 2, -CACHE_T, x->ne[2]);
                 if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) {
                     // cache last frame of last two chunk
                     cache_x = ggml_concat(ctx->ggml_ctx,
-                                          ggml_ext_slice(ctx->ggml_ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]),
+                                          ggml_ext_slice(ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]),
                                           cache_x,
                                           2);
                 }
@@ -1137,10 +1123,8 @@ namespace WAN {
             }
         }
 
-        static ggml_tensor* patchify(ggml_context* ctx,
-                                     ggml_tensor* x,
-                                     int64_t patch_size,
-                                     int64_t b = 1) {
+        static ggml_tensor* patchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t patch_size, int64_t b = 1) {
+            auto ctx = runner_ctx->ggml_ctx;
             // x: [b*c, f, h*q, w*r]
             // return: [b*c*r*q, f, h, w]
             if (patch_size == 1) {
@@ -1153,21 +1137,22 @@ namespace WAN {
             int64_t h = x->ne[1] / q;
             int64_t w = x->ne[0] / r;
 
-            x = ggml_reshape_4d(ctx, x, r * w, q, h, f * c * b);                 // [b*c*f, h, q, w*r]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));  // [b*c*f, q, h, w*r]
-            x = ggml_reshape_4d(ctx, x, r, w, h * q, f * c * b);                 // [b*c*f, q*h, w, r]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));  // [b*c*f, r, q*h, w]
-            x = ggml_reshape_4d(ctx, x, w * h, q * r, f, c * b);                 // [b*c, f, r*q, h*w]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));  // [b*c, r*q, f, h*w]
-            x = ggml_reshape_4d(ctx, x, w, h, f, q * r * c * b);                 // [b*c*r*q, f, h, w]
+            x = ggml_reshape_4d(ctx, x, r * w, q, h, f * c * b);                               // [b*c*f, h, q, w*r]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [b*c*f, q, h, w*r]
+            x = ggml_reshape_4d(ctx, x, r, w, h * q, f * c * b);                               // [b*c*f, q*h, w, r]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 1, 2, 0, 3));  // [b*c*f, r, q*h, w]
+            x = ggml_reshape_4d(ctx, x, w * h, q * r, f, c * b);                               // [b*c, f, r*q, h*w]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [b*c, r*q, f, h*w]
+            x = ggml_reshape_4d(ctx, x, w, h, f, q * r * c * b);                               // [b*c*r*q, f, h, w]
 
             return x;
         }
 
-        static ggml_tensor* unpatchify(ggml_context* ctx,
+        static ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx,
                                        ggml_tensor* x,
                                        int64_t patch_size,
                                        int64_t b = 1) {
+            auto ctx = runner_ctx->ggml_ctx;
             // x: [b*c*r*q, f, h, w]
             // return: [b*c, f, h*q, w*r]
             if (patch_size == 1) {
@@ -1180,13 +1165,13 @@ namespace WAN {
             int64_t h = x->ne[1];
             int64_t w = x->ne[0];
 
-            x = ggml_reshape_4d(ctx, x, w * h, f, q * r, c * b);                 // [b*c, r*q, f, h*w]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));  // [b*c, f, r*q, h*w]
-            x = ggml_reshape_4d(ctx, x, w, h * q, r, f * c * b);                 // [b*c*f, r, q*h, w]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3));  // [b*c*f, q*h, w, r]
-            x = ggml_reshape_4d(ctx, x, r * w, h, q, f * c * b);                 // [b*c*f, q, h, w*r]
-            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));  // [b*c*f, h, q, w*r]
-            x = ggml_reshape_4d(ctx, x, r * w, q * h, f, c * b);                 // [b*c, f, h*q, w*r]
+            x = ggml_reshape_4d(ctx, x, w * h, f, q * r, c * b);                               // [b*c, r*q, f, h*w]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [b*c, f, r*q, h*w]
+            x = ggml_reshape_4d(ctx, x, w, h * q, r, f * c * b);                               // [b*c*f, r, q*h, w]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 2, 0, 1, 3));  // [b*c*f, q*h, w, r]
+            x = ggml_reshape_4d(ctx, x, r * w, h, q, f * c * b);                               // [b*c*f, q, h, w*r]
+            x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3));  // [b*c*f, h, q, w*r]
+            x = ggml_reshape_4d(ctx, x, r * w, q * h, f, c * b);                               // [b*c, f, h*q, w*r]
             return x;
         }
 
@@ -1203,7 +1188,7 @@ namespace WAN {
 
             clear_cache();
 
-            x = patchify(ctx->ggml_ctx, x, patch_size, b);
+            x = patchify(ctx, x, patch_size, b);
             // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encode.prelude", "x");
 
             auto encoder = std::dynamic_pointer_cast<Encoder3d>(blocks["encoder"]);
@@ -1215,11 +1200,11 @@ namespace WAN {
             for (int i = 0; i < iter_; i++) {
                 _enc_conv_idx = 0;
                 if (i == 0) {
-                    auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);  // [b*c, 1, h, w]
+                    auto in = ggml_ext_slice(ctx, x, 2, 0, 1);  // [b*c, 1, h, w]
                     out     = encoder->forward(ctx, in, b, _enc_feat_map, _enc_conv_idx, i);
                 } else {
                     // if is_2D, drop 3 out of 4 frames
-                    auto in   = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1 + 4 * (i - 1), (is_2D ? 1 - 3 : 1) + 4 * i);  // [b*c, 4, h, w]
+                    auto in   = ggml_ext_slice(ctx, x, 2, 1 + 4 * (i - 1), (is_2D ? 1 - 3 : 1) + 4 * i);  // [b*c, 4, h, w]
                     auto out_ = encoder->forward(ctx, in, b, _enc_feat_map, _enc_conv_idx, i);
                     out       = ggml_concat(ctx->ggml_ctx, out, out_, 2);
                 }
@@ -1230,7 +1215,7 @@ namespace WAN {
             } else {
                 out = conv1->forward(ctx, out);
             }
-            auto mu = ggml_ext_chunk(ctx->ggml_ctx, out, 2, 3)[0];
+            auto mu = ggml_ext_chunk(ctx, out, 2, 3)[0];
             // sd::ggml_graph_cut::mark_graph_cut(mu, "wan_vae.encode.final", "mu");
             clear_cache();
             return mu;
@@ -1264,10 +1249,10 @@ namespace WAN {
             for (int i = 0; i < iter_; i++) {
                 _conv_idx = 0;
                 if (i == 0) {
-                    auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1);  // [b*c, 1, h, w]
+                    auto in = ggml_ext_slice(ctx, x, 2, i, i + 1);  // [b*c, 1, h, w]
                     out     = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i);
                 } else {
-                    auto in   = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1);  // [b*c, 1, h, w]
+                    auto in   = ggml_ext_slice(ctx, x, 2, i, i + 1);  // [b*c, 1, h, w]
                     auto out_ = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i);
                     out       = ggml_concat(ctx->ggml_ctx, out, out_, 2);
                     if (is_2D) {
@@ -1278,7 +1263,7 @@ namespace WAN {
                     }
                 }
             }
-            out = unpatchify(ctx->ggml_ctx, out, patch_size, b);
+            out = unpatchify(ctx, out, patch_size, b);
             // sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode.final", "out");
             clear_cache();
             return out;
@@ -1309,7 +1294,7 @@ namespace WAN {
                 const int global_frame = chunk_idx + static_cast<int>(frame);
                 // The first latent frame bypasses temporal upsampling and primes the causal caches.
                 const int64_t count = !is_2D && global_frame > 0 ? std::min<int64_t>(decode_frames, x->ne[2] - frame) : 1;
-                auto in             = ggml_ext_slice(ctx->ggml_ctx, x, 2, frame, frame + count);
+                auto in             = ggml_ext_slice(ctx, x, 2, frame, frame + count);
                 _conv_idx           = 0;
                 auto out_frame      = decoder->forward(ctx, in, b, _feat_map, _conv_idx, global_frame);
                 if (is_2D && global_frame > 0) {
@@ -1322,7 +1307,7 @@ namespace WAN {
                 out = out == nullptr ? out_frame : ggml_concat(ctx->ggml_ctx, out, out_frame, 2);
                 frame += count;
             }
-            return unpatchify(ctx->ggml_ctx, out, patch_size, b);
+            return unpatchify(ctx, out, patch_size, b);
         }
     };
 

src/pipeline/diffusion_engine.cpp

@@ -856,11 +856,13 @@ bool StableDiffusionGGML::set_sage_attention_enabled(bool enabled) {
         if (ctx == nullptr) {
             return false;
         }
-        auto* q        = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1);
-        auto* k        = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1);
-        auto* v        = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 128, 128, 1, 1);
-        auto* op       = ggml_sage_attn(ctx, q, k, v, 1.f / sqrtf(128.f), GGML_SAGE_ATTN_AUTO);
-        bool supported = true;
+        auto* q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1);
+        auto* k = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1);
+        auto* v = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 128, 128, 1, 1);
+        GGMLRunnerContext probe_ctx;
+        probe_ctx.ggml_ctx = ctx;
+        auto* op           = ggml_ext_sage_attn(&probe_ctx, q, k, v, 1.f / sqrtf(128.f), GGML_SAGE_ATTN_AUTO);
+        bool supported     = true;
         for (auto backend : backend_manager.runtime_backends(SDBackendModule::DIFFUSION)) {
             if (!ggml_backend_supports_op(backend, op)) {
                 LOG_ERROR("SageAttention is unavailable on %s; it requires patched GGML, CUDA Toolkit 12.0 or newer, and SM80 or newer kernels",
@@ -2442,9 +2444,9 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
 
 #ifdef SD_USE_UPSTREAM_GGML
     // workaround for sd.cpp #2149
-    const auto diffusion_backends = backend_manager.runtime_backends(SDBackendModule::DIFFUSION);
+    const auto diffusion_backends   = backend_manager.runtime_backends(SDBackendModule::DIFFUSION);
     const bool batch_cfg_backend_ok = std::none_of(diffusion_backends.begin(), diffusion_backends.end(),
-            [](ggml_backend_t backend) { return sd_backend_is(backend, "MTL"); });
+                                                   [](ggml_backend_t backend) { return sd_backend_is(backend, "MTL"); });
 #else
     const bool batch_cfg_backend_ok = true;
 #endif