diff --git a/ggml/src/ggml-cuda/binbcast.cu b/ggml/src/ggml-cuda/binbcast.cu index 2e38077bf67..010b94f38d4 100644 --- a/ggml/src/ggml-cuda/binbcast.cu +++ b/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 +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 +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, 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 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( + (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( + (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]; diff --git a/ggml/src/ggml-cuda/binbcast.cuh b/ggml/src/ggml-cuda/binbcast.cuh index 12624785b44..6e057c04aab 100644 --- a/ggml/src/ggml-cuda/binbcast.cuh +++ b/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); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index e73a7b8906c..9e610e7eb30 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/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 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;