metal : add MoE weighted reduction fusion

Fuses MUL(experts, weights) plus the expert VIEW/ADD chain into one kernel
that computes the weighted sum directly. The graph_optimize hook keeps the
expert and weight buffers alive until the fused output so the allocator cannot
reuse them while the kernel is still reading them.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -1514,6 +1514,21 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_topk_moe(ggml_me
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_moe_weighted_reduction(ggml_metal_library_t lib) {
+    char base[256];
+    char name[256];
+
+    snprintf(base, 256, "kernel_moe_weighted_reduction_f32");
+    snprintf(name, 256, "%s", base);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+    }
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_pad(
         ggml_metal_library_t lib,
         const struct ggml_tensor * op,

ggml/src/ggml-metal/ggml-metal-device.h

@@ -148,6 +148,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix       (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge       (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_topk_moe          (ggml_metal_library_t lib);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_moe_weighted_reduction (ggml_metal_library_t lib);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin               (ggml_metal_library_t lib, const struct ggml_tensor * op, int32_t n_fuse );
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin_one           (ggml_metal_library_t lib, enum ggml_op op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_l2_norm           (ggml_metal_library_t lib, const struct ggml_tensor * op);

ggml/src/ggml-metal/ggml-metal-fusion.cpp

@@ -392,6 +392,136 @@ static bool ggml_metal_fusion_check_topk_moe(
     return true;
 }
 
+#define GGML_METAL_MOE_WEIGHTED_REDUCTION_MAX_EXPERTS 8
+
+bool ggml_metal_fusion_match_moe_weighted_reduction(
+        const ggml_cgraph * gf, int node_idx, ggml_metal_moe_weighted_reduction_match * match) {
+    if (match == nullptr || node_idx < 0 || node_idx + 3 > gf->n_nodes) {
+        return false;
+    }
+
+    const ggml_tensor * mul = gf->nodes[node_idx];
+    if (mul->op != GGML_OP_MUL || mul->type != GGML_TYPE_F32) {
+        return false;
+    }
+
+    // MUL, then one VIEW per expert, then one ADD per additional expert
+    int n_views = 0;
+    while (node_idx + 1 + n_views < gf->n_nodes &&
+           gf->nodes[node_idx + 1 + n_views]->op == GGML_OP_VIEW) {
+        n_views++;
+    }
+
+    if (n_views < 2 || n_views > GGML_METAL_MOE_WEIGHTED_REDUCTION_MAX_EXPERTS) {
+        return false;
+    }
+
+    const int n_expert_used = n_views;
+    const int raw_count     = 1 + n_expert_used + (n_expert_used - 1);
+
+    if (node_idx + raw_count > gf->n_nodes) {
+        return false;
+    }
+
+    for (int i = n_expert_used + 1; i < raw_count; ++i) {
+        if (gf->nodes[node_idx + i]->op != GGML_OP_ADD) {
+            return false;
+        }
+    }
+
+    std::vector<ggml_op> raw_ops(raw_count);
+    std::vector<int>     raw_idxs(raw_count);
+    for (int i = 0; i < raw_count; ++i) {
+        raw_idxs[i] = node_idx + i;
+        raw_ops[i]  = gf->nodes[node_idx + i]->op;
+    }
+
+    const ggml_tensor * experts = mul->src[0];
+    const ggml_tensor * weights = mul->src[1];
+    const ggml_tensor * dst     = gf->nodes[node_idx + raw_count - 1];
+
+    if (experts->type != GGML_TYPE_F32 || weights->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+        return false;
+    }
+
+    const int64_t n_embd   = experts->ne[0];
+    const int64_t n_tokens = experts->ne[2];
+
+    if (n_embd <= 0 || n_tokens <= 0 || experts->ne[1] != n_expert_used || experts->ne[3] != 1 ||
+        weights->ne[0] != 1 || weights->ne[1] != n_expert_used || weights->ne[2] != n_tokens || weights->ne[3] != 1 ||
+        dst->ne[0] != n_embd || dst->ne[1] != n_tokens || dst->ne[2] != 1 || dst->ne[3] != 1) {
+        return false;
+    }
+
+    if (!ggml_is_contiguous(experts) || !ggml_is_contiguous(weights) || !ggml_is_contiguous(dst)) {
+        return false;
+    }
+
+    for (int i = 1; i <= n_expert_used; ++i) {
+        const ggml_tensor * view = gf->nodes[node_idx + i];
+        if (view->view_src != mul || view->src[0] != mul ||
+            view->view_offs != (size_t) (i - 1) * mul->nb[1] ||
+            view->ne[0] != n_embd || view->ne[1] != n_tokens ||
+            view->nb[1] != mul->nb[2]) {
+            return false;
+        }
+    }
+
+    const ggml_tensor * prev_add = nullptr;
+    for (int j = 1; j < n_expert_used; ++j) {
+        const ggml_tensor * add = gf->nodes[node_idx + n_expert_used + j];
+        const ggml_tensor * rhs = gf->nodes[node_idx + j + 1];
+        const ggml_tensor * lhs = j == 1 ? gf->nodes[node_idx + 1] : prev_add;
+        if (add->src[0] != lhs || add->src[1] != rhs) {
+            return false;
+        }
+        prev_add = add;
+    }
+
+    const int outputs[1] = { node_idx + raw_count - 1 };
+    if (!ggml_can_fuse_subgraph_ext(gf, raw_idxs.data(), raw_count, raw_ops.data(), outputs, 1)) {
+        return false;
+    }
+
+    match->experts    = experts;
+    match->weights    = weights;
+    match->dst        = dst;
+    match->node_count = raw_count;
+    return true;
+}
+
+static bool ggml_metal_fusion_check_moe_weighted_reduction(
+        const ggml_metal_fusion      * fusion,
+        const ggml_tensor * const    * nodes,
+        const ggml_cgraph            * gf,
+        const int                    * node_idxs,
+              int                      idx,
+              ggml_metal_fusion_mode   mode) {
+    GGML_UNUSED(nodes);
+
+    ggml_metal_moe_weighted_reduction_match match;
+    if (!ggml_metal_fusion_match_moe_weighted_reduction(gf, node_idxs[idx], &match)) {
+        return false;
+    }
+
+    if (fusion->n_ops != match.experts->ne[1]) {
+        return false;
+    }
+
+    const int raw_end = node_idxs[idx] + match.node_count - 1;
+    if (node_idxs[idx + fusion->n_ops - 1] != raw_end) {
+        return false;
+    }
+
+    if (mode == GGML_METAL_FUSION_FULL) {
+        if (!match.experts->data || !match.weights->data || !match.dst->data) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
 // ---- patterns ------------------------------------------------------------
 
 static const ggml_op ops_norm_mul[]         = { GGML_OP_NORM, GGML_OP_MUL };
@@ -424,6 +554,14 @@ static const ggml_op ops_topk_moe_norm_scale[] = {
     GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_DIV, GGML_OP_SCALE
 };
 
+static const ggml_op ops_moe_weighted_reduction_2[] = { GGML_OP_MUL, GGML_OP_ADD };
+static const ggml_op ops_moe_weighted_reduction_3[] = { GGML_OP_MUL, GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_moe_weighted_reduction_4[] = { GGML_OP_MUL, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_moe_weighted_reduction_5[] = { GGML_OP_MUL, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_moe_weighted_reduction_6[] = { GGML_OP_MUL, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_moe_weighted_reduction_7[] = { GGML_OP_MUL, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_moe_weighted_reduction_8[] = { GGML_OP_MUL, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD };
+
 static const ggml_metal_fusion ggml_metal_fusions[] = {
     { GGML_METAL_FUSION_NORM_MUL,     ops_norm_mul,         2, false, ggml_metal_fusion_check_norm },
     { GGML_METAL_FUSION_NORM_MUL_ADD, ops_norm_mul_add,     3, false, ggml_metal_fusion_check_norm },
@@ -441,6 +579,13 @@ static const ggml_metal_fusion ggml_metal_fusions[] = {
     { GGML_METAL_FUSION_TOPK_MOE,     ops_topk_moe_scale,    4, true,  ggml_metal_fusion_check_topk_moe },
     { GGML_METAL_FUSION_TOPK_MOE,     ops_topk_moe_norm,     6, true,  ggml_metal_fusion_check_topk_moe },
     { GGML_METAL_FUSION_TOPK_MOE,     ops_topk_moe_norm_scale, 7, true, ggml_metal_fusion_check_topk_moe },
+    { GGML_METAL_FUSION_MOE_WEIGHTED_REDUCTION, ops_moe_weighted_reduction_2, 2, true, ggml_metal_fusion_check_moe_weighted_reduction },
+    { GGML_METAL_FUSION_MOE_WEIGHTED_REDUCTION, ops_moe_weighted_reduction_3, 3, true, ggml_metal_fusion_check_moe_weighted_reduction },
+    { GGML_METAL_FUSION_MOE_WEIGHTED_REDUCTION, ops_moe_weighted_reduction_4, 4, true, ggml_metal_fusion_check_moe_weighted_reduction },
+    { GGML_METAL_FUSION_MOE_WEIGHTED_REDUCTION, ops_moe_weighted_reduction_5, 5, true, ggml_metal_fusion_check_moe_weighted_reduction },
+    { GGML_METAL_FUSION_MOE_WEIGHTED_REDUCTION, ops_moe_weighted_reduction_6, 6, true, ggml_metal_fusion_check_moe_weighted_reduction },
+    { GGML_METAL_FUSION_MOE_WEIGHTED_REDUCTION, ops_moe_weighted_reduction_7, 7, true, ggml_metal_fusion_check_moe_weighted_reduction },
+    { GGML_METAL_FUSION_MOE_WEIGHTED_REDUCTION, ops_moe_weighted_reduction_8, 8, true, ggml_metal_fusion_check_moe_weighted_reduction },
 };
 
 const ggml_metal_fusion * ggml_metal_fusion_all(int * n) {

ggml/src/ggml-metal/ggml-metal-fusion.h

@@ -36,6 +36,7 @@ typedef enum ggml_metal_fusion_id {
     GGML_METAL_FUSION_SNAKE,        // MUL + SIN + SQR + MUL + ADD
     GGML_METAL_FUSION_GDN_CACHE,    // GATED_DELTA_NET + CPY (write snapshots into the recurrent cache)
     GGML_METAL_FUSION_TOPK_MOE,     // SOFT_MAX + ARGSORT + GET_ROWS + norm/scale (MoE routing)
+    GGML_METAL_FUSION_MOE_WEIGHTED_REDUCTION, // MUL + expert VIEWs + ADD chain (MoE output reduction)
 } ggml_metal_fusion_id;
 
 struct ggml_metal_fusion {
@@ -61,6 +62,18 @@ struct ggml_metal_fusion {
 
 typedef struct ggml_metal_fusion ggml_metal_fusion;
 
+struct ggml_metal_moe_weighted_reduction_match {
+    const struct ggml_tensor * experts;
+    const struct ggml_tensor * weights;
+    const struct ggml_tensor * dst;
+    int node_count;
+};
+
+// match MUL(experts, weights) + expert VIEWs + ADD chain; used by both the fusion
+// validator and the graph-optimize alloc-dependency hook
+bool ggml_metal_fusion_match_moe_weighted_reduction(
+        const struct ggml_cgraph * gf, int node_idx, struct ggml_metal_moe_weighted_reduction_match * match);
+
 // the single table of all fusions supported by the Metal backend
 const ggml_metal_fusion * ggml_metal_fusion_all(int * n);
 

ggml/src/ggml-metal/ggml-metal-impl.h

@@ -1232,6 +1232,12 @@ typedef struct {
     float    scale_val;
 } ggml_metal_kargs_topk_moe;
 
+typedef struct {
+    int32_t ne00; // n_embd
+    int32_t ne01; // n_expert_used
+    int32_t ne02; // n_tokens
+} ggml_metal_kargs_moe_weighted_reduction;
+
 typedef struct {
     int32_t nrows;
 } ggml_metal_kargs_fwht;

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -3777,6 +3777,12 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) {
             ctx->count_fusions(fusion);
             return ggml_metal_op_snake_fused(ctx, idx);
         }
+
+        // MoE output reduction: experts * weights -> weighted sum
+        if (fusion && fusion->id == GGML_METAL_FUSION_MOE_WEIGHTED_REDUCTION) {
+            ctx->count_fusions(fusion);
+            return ggml_metal_op_moe_weighted_reduction(ctx, idx);
+        }
     }
 
     ggml_tensor * op = ctx->node(idx);
@@ -5420,6 +5426,49 @@ int ggml_metal_op_topk_moe(ggml_metal_op_t ctx, int idx) {
     return n_fuse;
 }
 
+int ggml_metal_op_moe_weighted_reduction(ggml_metal_op_t ctx, int idx) {
+    ggml_metal_library_t lib = ctx->lib;
+    ggml_metal_encoder_t enc = ctx->enc;
+
+    int n_fuse = 1;
+    const ggml_metal_fusion * fusion = ctx->can_fuse(idx, GGML_METAL_FUSION_FULL, &n_fuse);
+    if (!fusion || fusion->id != GGML_METAL_FUSION_MOE_WEIGHTED_REDUCTION) {
+        return 1;
+    }
+
+    ggml_tensor * mul     = ctx->node(idx);
+    ggml_tensor * experts = mul->src[0];
+    ggml_tensor * weights = mul->src[1];
+    ggml_tensor * dst     = ctx->node(idx + n_fuse - 1);
+
+    ggml_metal_kargs_moe_weighted_reduction args = {
+        /*.ne00 =*/ (int32_t) experts->ne[0],
+        /*.ne01 =*/ (int32_t) experts->ne[1],
+        /*.ne02 =*/ (int32_t) experts->ne[2],
+    };
+
+    auto pipeline = ggml_metal_library_get_pipeline_moe_weighted_reduction(lib);
+
+    const int nth = std::min(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
+    const int n_col_tiles = (args.ne00 + nth - 1) / nth;
+
+    ggml_metal_encoder_set_pipeline(enc, pipeline);
+    ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(experts), 1);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(weights), 2);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(dst),     3);
+
+    ggml_metal_encoder_dispatch_threadgroups(enc, (uint32_t) args.ne02, (uint32_t) n_col_tiles, 1, nth, 1, 1);
+
+    ctx->count_fusions(fusion);
+
+    if (ggml_metal_fusion_info_debug(ctx->finfo) > 1) {
+        GGML_LOG_DEBUG("%s: fuse: MOE_WEIGHTED_REDUCTION\n", __func__);
+    }
+
+    return n_fuse;
+}
+
 int ggml_metal_op_top_k(ggml_metal_op_t ctx, int idx) {
     ggml_tensor * op = ctx->node(idx);
 

ggml/src/ggml-metal/ggml-metal-ops.h

@@ -98,6 +98,7 @@ int ggml_metal_op_argmax            (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_argsort           (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_top_k             (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_topk_moe          (ggml_metal_op_t ctx, int idx);
+int ggml_metal_op_moe_weighted_reduction (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_tri               (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_opt_step_adamw    (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_opt_step_sgd      (ggml_metal_op_t ctx, int idx);

ggml/src/ggml-metal/ggml-metal.cpp

@@ -561,7 +561,25 @@ static void ggml_backend_metal_event_wait(ggml_backend_t backend, ggml_backend_e
 }
 
 static void ggml_backend_metal_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph, ggml_backend_graph_optimize_params * params) {
-    GGML_UNUSED(params);
+    GGML_ASSERT(params && params->add_alloc_dep);
+
+    // keep the MoE weighted-reduction inputs alive until the fused output so the
+    // allocator cannot reuse them while the fused kernel is still reading them
+    for (int i = 0; i < cgraph->n_nodes; ++i) {
+        if (cgraph->nodes[i]->op != GGML_OP_MUL) {
+            continue;
+        }
+
+        ggml_metal_moe_weighted_reduction_match match;
+        if (!ggml_metal_fusion_match_moe_weighted_reduction(cgraph, i, &match)) {
+            continue;
+        }
+
+        params->add_alloc_dep(params->user_data, (ggml_tensor *) match.experts, (ggml_tensor *) match.dst);
+        params->add_alloc_dep(params->user_data, (ggml_tensor *) match.weights, (ggml_tensor *) match.dst);
+
+        i += match.node_count - 1;
+    }
 
     ggml_metal_t ctx = (ggml_metal_t)backend->context;
 

ggml/src/ggml-metal/kernels/argsort.metal

@@ -445,3 +445,27 @@ kernel void kernel_topk_moe_f32(
         }
     }
 }
+
+// fused MoE expert weighting + reduction: weighted = sum(experts[e] * weights[e]).
+// The host guarantees all tensors are contiguous F32.
+kernel void kernel_moe_weighted_reduction_f32(
+        constant   ggml_metal_kargs_moe_weighted_reduction & args,
+        device const float * experts,
+        device const float * weights,
+        device       float * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    const int64_t token = tgpig.x;
+    const int64_t col   = (int64_t) tgpig.y * ntg.x + tpitg.x;
+    if (token >= args.ne02 || col >= args.ne00) {
+        return;
+    }
+
+    const int64_t base = token * (int64_t) args.ne01 * args.ne00 + col;
+    float sum = 0.0f;
+    for (int e = 0; e < args.ne01; ++e) {
+        sum += experts[base + e * args.ne00] * weights[token * args.ne01 + e];
+    }
+    dst[token * args.ne00 + col] = sum;
+}