refactor: replace QKV split helpers with ggml_ext_chunk (#2322)

src/core/ggml_extend.cpp

@@ -541,36 +541,6 @@ ggml_tensor* ggml_ext_conv_3d(GGMLRunnerContext* runner_ctx,
     return x;
 }
 
-std::vector<ggml_tensor*> ggml_ext_split_qkv(GGMLRunnerContext* runner_ctx, ggml_tensor* qkv) {
-    auto ctx = runner_ctx->ggml_ctx;
-    qkv      = ggml_reshape_4d(ctx, qkv, qkv->ne[0] / 3, 3, qkv->ne[1], qkv->ne[2]);  // [N, L, 3, C]
-    qkv      = ggml_cont(ctx, ggml_permute(ctx, qkv, 0, 3, 1, 2));                    // [3, N, L, C]
-
-    int64_t offset = qkv->nb[2] * qkv->ne[2];
-    auto q         = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 0);  // [N, L, C]
-    auto k         = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 1);  // [N, L, C]
-    auto v         = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 2);  // [N, L, C]
-    return {q, k, v};
-}
-
-std::vector<ggml_tensor*> ggml_ext_split_image_qkv(GGMLRunnerContext* runner_ctx, ggml_tensor* qkv) {
-    auto ctx    = runner_ctx->ggml_ctx;
-    int64_t W   = qkv->ne[0];
-    int64_t H   = qkv->ne[1];
-    int64_t C   = qkv->ne[2] / 3;
-    int64_t N   = qkv->ne[3];
-    int64_t nb1 = qkv->nb[1];
-    int64_t nb2 = qkv->nb[2];
-    qkv         = ggml_reshape_4d(ctx, qkv, W * H, C, 3, N);                            // [N, 3, C, H*W]
-    qkv         = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, qkv, 0, 1, 3, 2));  // [3, N, C, H*W]
-
-    int64_t offset = qkv->nb[2] * qkv->ne[2];
-    auto q         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[2], offset * 0);  // [N, C, H, W]
-    auto k         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[2], offset * 1);  // [N, C, H, W]
-    auto v         = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[2], offset * 2);  // [N, C, H, W]
-    return {q, k, v};
-}
-
 ggml_tensor* ggml_ext_full(GGMLRunnerContext* runner_ctx,
                            float value,
                            int64_t ne0,

src/core/ggml_extend.h

@@ -150,14 +150,6 @@ ggml_tensor* ggml_ext_conv_3d(GGMLRunnerContext* runner_ctx,
                               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*> 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*> 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(GGMLRunnerContext* runner_ctx,
                            float value,

src/model/common/ggml_block.hpp

@@ -910,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 = ggml_ext_split_qkv(ctx, qkv);
+            auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0);
             q            = qkv_vec[0];
             k            = qkv_vec[1];
             v            = qkv_vec[2];

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 = ggml_ext_split_qkv(ctx, qkv);
+            auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0);
             auto q       = qkv_vec[0];
             auto k       = qkv_vec[1];
             auto v       = qkv_vec[2];

src/model/diffusion/ideogram4.hpp

@@ -222,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 = ggml_ext_split_qkv(ctx, qkv);
+            auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0);
             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);

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 = ggml_ext_split_qkv(ctx, std::dynamic_pointer_cast<Linear>(blocks["qkv"])->forward(ctx, x));
+            auto qkv = ggml_ext_chunk(ctx, std::dynamic_pointer_cast<Linear>(blocks["qkv"])->forward(ctx, x), 3, 0);
             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]);

src/model/diffusion/lens.hpp

@@ -151,8 +151,8 @@ namespace Lens {
             int64_t N       = img->ne[2];
             int64_t inner   = dim_head * num_heads;
 
-            auto img_qkv_vec = ggml_ext_split_qkv(ctx, img_qkv->forward(ctx, img));
+            auto img_qkv_vec = ggml_ext_chunk(ctx, img_qkv->forward(ctx, img), 3, 0);
-            auto txt_qkv_vec = ggml_ext_split_qkv(ctx, txt_qkv->forward(ctx, txt));
+            auto txt_qkv_vec = ggml_ext_chunk(ctx, txt_qkv->forward(ctx, txt), 3, 0);
 
             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);

src/model/diffusion/mmdit.hpp

@@ -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     = ggml_ext_split_qkv(ctx, qkv);
+        auto qkv_vec     = ggml_ext_chunk(ctx, qkv, 3, 0);
         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]

src/model/diffusion/pid.hpp

@@ -293,7 +293,7 @@ namespace Pid {
             auto proj     = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
 
             auto qkv         = qkv_proj->forward(ctx, x);
-            auto qkv_vec     = ggml_ext_split_qkv(ctx, qkv);
+            auto qkv_vec     = ggml_ext_chunk(ctx, qkv, 3, 0);
             int64_t L        = x->ne[1];
             int64_t N        = x->ne[2];
             int64_t head_dim = dim / num_heads;
@@ -343,14 +343,14 @@ namespace Pid {
             int64_t N        = x->ne[2];
             int64_t head_dim = dim / num_heads;
 
-            auto qkv_x = ggml_ext_split_qkv(ctx, qkv_x_proj->forward(ctx, x));
+            auto qkv_x = ggml_ext_chunk(ctx, qkv_x_proj->forward(ctx, x), 3, 0);
             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 = ggml_ext_split_qkv(ctx, qkv_y_proj->forward(ctx, y));
+            auto qkv_y = ggml_ext_chunk(ctx, qkv_y_proj->forward(ctx, y), 3, 0);
             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);

src/model/te/llm.hpp

@@ -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       = ggml_ext_split_qkv(ctx, qkv);
+                qkv_vec       = ggml_ext_chunk(ctx, qkv, 3, 0);
             }
 
             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]

src/model/vae/wan_vae.hpp

@@ -623,7 +623,7 @@ namespace WAN {
             const int64_t w = x->ne[0];
 
             auto qkv     = to_qkv->forward(ctx, x);
-            auto qkv_vec = ggml_ext_split_image_qkv(ctx, qkv);
+            auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 2);
 
             auto q = qkv_vec[0];
             q      = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, q, 2, 0, 1, 3));  // [t, h, w, c]