cuda: fuse MoE expert weighting and reduction (+3% to +9% prefill, MoE models only)

ggml/src/ggml-cuda/binbcast.cu

@@ -542,6 +542,136 @@ void ggml_cuda_op_fused_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst,
     }
 }
 
+template <int n_experts, bool has_scale>
+static __global__ void k_moe_weighted_sum(
+        const float * __restrict__ experts,
+        const float * __restrict__ scale,
+        const float * __restrict__ weights,
+        float * __restrict__ dst,
+        const int n_embd) {
+    ggml_cuda_pdl_lc();
+
+    const int i = blockIdx.x * blockDim.x + threadIdx.x;
+    const int t = blockIdx.y;
+
+    if (i >= n_embd) {
+        return;
+    }
+
+    const size_t expert_base = size_t(t) * n_experts * n_embd + i;
+    const size_t factor_base = size_t(t) * n_experts;
+
+    ggml_cuda_pdl_sync();
+
+    float value = experts[expert_base];
+    if constexpr (has_scale) {
+        value = __fmul_rn(value, scale[factor_base]);
+    }
+    float sum = __fmul_rn(value, weights[factor_base]);
+
+#pragma unroll
+    for (int e = 1; e < n_experts; ++e) {
+        value = experts[expert_base + size_t(e) * n_embd];
+        if constexpr (has_scale) {
+            value = __fmul_rn(value, scale[factor_base + e]);
+        }
+        value = __fmul_rn(value, weights[factor_base + e]);
+        sum = __fadd_rn(sum, value);
+    }
+
+    dst[size_t(t) * n_embd + i] = sum;
+}
+
+template <bool has_scale>
+static void launch_moe_weighted_sum(
+        const float * experts,
+        const float * scale,
+        const float * weights,
+        float * dst,
+        const int n_embd,
+        const int n_tokens,
+        const int n_experts,
+        cudaStream_t stream) {
+    const dim3 block_dims(256, 1, 1);
+    const dim3 block_nums((n_embd + block_dims.x - 1) / block_dims.x, n_tokens, 1);
+    const ggml_cuda_kernel_launch_params launch_params(block_nums, block_dims, 0, stream);
+
+#define GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(N) \
+    ggml_cuda_kernel_launch(k_moe_weighted_sum<N, has_scale>, launch_params, experts, scale, weights, dst, n_embd)
+
+    switch (n_experts) {
+        case 2:  GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(2);  break;
+        case 3:  GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(3);  break;
+        case 4:  GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(4);  break;
+        case 5:  GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(5);  break;
+        case 6:  GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(6);  break;
+        case 7:  GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(7);  break;
+        case 8:  GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(8);  break;
+        case 9:  GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(9);  break;
+        case 10: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(10); break;
+        case 11: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(11); break;
+        case 12: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(12); break;
+        case 13: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(13); break;
+        case 14: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(14); break;
+        case 15: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(15); break;
+        default: GGML_ABORT("unsupported number of experts");
+    }
+
+#undef GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM
+}
+
+void ggml_cuda_op_moe_weighted_sum(
+        ggml_backend_cuda_context & ctx,
+        const ggml_tensor * experts,
+        const ggml_tensor * scale,
+        const ggml_tensor * weights,
+        ggml_tensor * dst) {
+    GGML_ASSERT(experts->type == GGML_TYPE_F32);
+    GGML_ASSERT(weights->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type == GGML_TYPE_F32);
+    GGML_ASSERT(scale == nullptr || scale->type == GGML_TYPE_F32);
+    GGML_ASSERT(ggml_is_contiguous(experts));
+    GGML_ASSERT(ggml_is_contiguous(weights));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+    GGML_ASSERT(scale == nullptr || ggml_is_contiguous(scale));
+
+    const int n_embd    = experts->ne[0];
+    const int n_experts = experts->ne[1];
+    const int n_tokens  = experts->ne[2];
+
+    const auto overlaps = [](const ggml_tensor * a, const ggml_tensor * b) {
+        const uintptr_t a_begin = (uintptr_t) a->data;
+        const uintptr_t a_end   = a_begin + ggml_nbytes(a);
+        const uintptr_t b_begin = (uintptr_t) b->data;
+        const uintptr_t b_end   = b_begin + ggml_nbytes(b);
+        return a_begin < b_end && b_begin < a_end;
+    };
+
+    const bool needs_tmp =
+            overlaps(dst, experts) || overlaps(dst, weights) || (scale && overlaps(dst, scale));
+    ggml_cuda_pool_alloc<float> tmp(ctx.pool());
+    float * dst_ptr = (float *) dst->data;
+    if (needs_tmp) {
+        dst_ptr = tmp.alloc(ggml_nelements(dst));
+    }
+
+    if (scale) {
+        launch_moe_weighted_sum<true>(
+                (const float *) experts->data, (const float *) scale->data, (const float *) weights->data,
+                dst_ptr,
+                n_embd, n_tokens, n_experts, ctx.stream());
+    } else {
+        launch_moe_weighted_sum<false>(
+                (const float *) experts->data, nullptr, (const float *) weights->data,
+                dst_ptr,
+                n_embd, n_tokens, n_experts, ctx.stream());
+    }
+
+    if (needs_tmp) {
+        CUDA_CHECK(cudaMemcpyAsync(dst->data, dst_ptr, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, ctx.stream()));
+    }
+}
+
 void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const ggml_tensor * src0 = dst->src[0];
 

ggml/src/ggml-cuda/binbcast.cuh

@@ -10,3 +10,10 @@ void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst
 
 void ggml_cuda_op_fused_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse);
 void ggml_cuda_op_fused_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse);
+
+void ggml_cuda_op_moe_weighted_sum(
+        ggml_backend_cuda_context & ctx,
+        const ggml_tensor * experts,
+        const ggml_tensor * scale,
+        const ggml_tensor * weights,
+        ggml_tensor * dst);

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -2701,6 +2701,86 @@ static int ggml_cuda_try_gdn_cache_fusion(
     return skip;
 }
 
+static int ggml_cuda_try_moe_weighted_sum_fusion(
+        ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int node_idx) {
+    ggml_tensor * first_mul = cgraph->nodes[node_idx];
+    if (first_mul->op != GGML_OP_MUL || first_mul->type != GGML_TYPE_F32 ||
+            first_mul->ne[1] < 2 || first_mul->ne[1] > 15 ||
+            first_mul->ne[2] < 1 || first_mul->ne[3] != 1) {
+        return 0;
+    }
+
+    int n_muls = 1;
+    ggml_tensor * weighted = first_mul;
+    if (node_idx + 1 < cgraph->n_nodes &&
+            cgraph->nodes[node_idx + 1]->op == GGML_OP_MUL &&
+            cgraph->nodes[node_idx + 1]->src[0] == first_mul &&
+            ggml_are_same_shape(cgraph->nodes[node_idx + 1], first_mul)) {
+        weighted = cgraph->nodes[node_idx + 1];
+        n_muls = 2;
+    }
+
+    const int n_experts = weighted->ne[1];
+    const int view_idx = node_idx + n_muls;
+    const int add_idx = view_idx + n_experts;
+    const int n_nodes = n_muls + n_experts + n_experts - 1;
+    if (node_idx + n_nodes > cgraph->n_nodes) {
+        return 0;
+    }
+
+    for (int e = 0; e < n_experts; ++e) {
+        const ggml_tensor * view = cgraph->nodes[view_idx + e];
+        if (view->op != GGML_OP_VIEW || view->view_src != weighted ||
+                view->view_offs != size_t(e) * weighted->nb[1] ||
+                view->ne[0] != weighted->ne[0] || view->ne[1] != weighted->ne[2] ||
+                view->ne[2] != 1 || view->ne[3] != 1) {
+            return 0;
+        }
+    }
+
+    for (int e = 1; e < n_experts; ++e) {
+        const ggml_tensor * add = cgraph->nodes[add_idx + e - 1];
+        const ggml_tensor * lhs = e == 1 ? cgraph->nodes[view_idx] : cgraph->nodes[add_idx + e - 2];
+        const ggml_tensor * rhs = cgraph->nodes[view_idx + e];
+        if (add->op != GGML_OP_ADD || add->src[0] != lhs || add->src[1] != rhs) {
+            return 0;
+        }
+    }
+
+    const ggml_tensor * experts = first_mul->src[0];
+    const ggml_tensor * scale   = n_muls == 2 ? first_mul->src[1] : nullptr;
+    const ggml_tensor * weights = n_muls == 2 ? weighted->src[1] : first_mul->src[1];
+    ggml_tensor * output = cgraph->nodes[node_idx + n_nodes - 1];
+
+    if (experts->type != GGML_TYPE_F32 || weights->type != GGML_TYPE_F32 ||
+            (scale && scale->type != GGML_TYPE_F32) ||
+            !ggml_is_contiguous(experts) || !ggml_is_contiguous(weights) ||
+            (scale && !ggml_is_contiguous(scale)) || !ggml_is_contiguous(output) ||
+            weights->ne[0] != 1 || weights->ne[1] != n_experts || weights->ne[2] != experts->ne[2] ||
+            (scale && (scale->ne[0] != 1 || scale->ne[1] != n_experts || scale->ne[2] != experts->ne[2]))) {
+        return 0;
+    }
+
+    std::vector<ggml_op> ops(n_nodes);
+    for (int j = 0; j < n_muls; ++j) {
+        ops[j] = GGML_OP_MUL;
+    }
+    for (int j = 0; j < n_experts; ++j) {
+        ops[n_muls + j] = GGML_OP_VIEW;
+    }
+    for (int j = 0; j < n_experts - 1; ++j) {
+        ops[n_muls + n_experts + j] = GGML_OP_ADD;
+    }
+
+    const int out_node = node_idx + n_nodes - 1;
+    if (!ggml_can_fuse_subgraph(cgraph, node_idx, n_nodes, ops.data(), &out_node, 1)) {
+        return 0;
+    }
+
+    ggml_cuda_op_moe_weighted_sum(*cuda_ctx, experts, scale, weights, output);
+    return n_nodes - 1;
+}
+
 static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int node_idx, ggml_cuda_topk_moe_args & args) {
     args.sigmoid         = false;
     args.sqrt_softplus   = false;
@@ -3150,6 +3230,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
 
     ggml_tensor * node = cgraph->nodes[i];
 
+    if (node->op == GGML_OP_MUL) {
+        const int nodes_to_skip = ggml_cuda_try_moe_weighted_sum_fusion(cuda_ctx, cgraph, i);
+        if (nodes_to_skip > 0) {
+            return nodes_to_skip;
+        }
+    }
+
     // gated_delta_net -> cpy: scatter recurrent-state snapshots into the cache
     if (node->op == GGML_OP_GATED_DELTA_NET) {
         ggml_cuda_gated_delta_net_fused_cache fused_state_cpy;