Various fused ops around expert selection (ikawrakow#840)

* Fuse sigmoid+add+grouped_topk+get_rows (CPU)

* Fix CPU + CUDA

but CUDA is somehow not 100% correct as I get a slightly different
PPL (lower!)

* Minor

* Fuse sigmoid+add+topk+get_rows (CUDA)

* Fuse sigmoid+add+topk+get_rows (CPU)

* Fuse topk+view+get_rows+reshape+softmax (CPU)

* Fuse topk+view+get_rows+reshape+softmax (CUDA)

* cpu: turn off the openai topk fusing for now

Something is not right and I don't see the bug.
On the CPU one doesn't gain much if anything, so not a big loss.

* Also fuse sum_rows and div

---------

Co-authored-by: Iwan Kawrakow <[email protected]>

Author: ikawrakow

ggml/src/ggml-cuda.cu

@@ -3173,7 +3173,25 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
                     ggml_cuda_op_relu(ctx, dst);
                     break;
                 case GGML_UNARY_OP_SIGMOID:
-                    ggml_cuda_op_sigmoid(ctx, dst);
+                    if (i + 5 < cgraph->n_nodes &&
+                        cgraph->nodes[i+1]->op == GGML_OP_RESHAPE &&
+                        cgraph->nodes[i+2]->op == GGML_OP_ADD &&
+                        cgraph->nodes[i+3]->op == GGML_OP_ARGSORT &&
+                        cgraph->nodes[i+4]->op == GGML_OP_VIEW &&
+                        cgraph->nodes[i+5]->op == GGML_OP_GET_ROWS) {
+                        cuda_glm45moe_experts(ctx, cgraph->nodes[i+5], cgraph->nodes[i+4]);
+                        i += 5;
+                    }
+                    else if (i + 4 < cgraph->n_nodes &&
+                        cgraph->nodes[i+1]->op == GGML_OP_RESHAPE &&
+                        cgraph->nodes[i+2]->op == GGML_OP_ADD &&
+                        cgraph->nodes[i+3]->op == GGML_OP_GROUPED_TOPK &&
+                        cgraph->nodes[i+4]->op == GGML_OP_GET_ROWS) {
+                        cuda_bailingmoev2_experts(ctx, cgraph->nodes[i+4], cgraph->nodes[i+4]);
+                        i += 4;
+                    } else {
+                        ggml_cuda_op_sigmoid(ctx, dst);
+                    }
                     break;
                 case GGML_UNARY_OP_HARDSIGMOID:
                     ggml_cuda_op_hardsigmoid(ctx, dst);
@@ -3315,10 +3333,28 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
             ggml_cuda_op_pool2d(ctx, dst);
             break;
         case GGML_OP_SUM_ROWS:
-            ggml_cuda_op_sum_rows(ctx, dst);
+            if (i + 1 < cgraph->n_nodes &&
+                cgraph->nodes[i+1]->op == GGML_OP_DIV &&
+                cgraph->nodes[i+1]->src[1] == dst &&
+                cgraph->nodes[i+1]->src[0] == dst->src[0]) {
+                ggml_cuda_op_sum_rows_div(ctx, cgraph->nodes[i+1]);
+                ++i;
+            } else {
+                ggml_cuda_op_sum_rows(ctx, dst);
+            }
             break;
         case GGML_OP_ARGSORT:
-            ggml_cuda_op_argsort(ctx, dst);
+            if (i + 5 < cgraph->n_nodes &&
+                cgraph->nodes[i+1]->op == GGML_OP_VIEW &&
+                cgraph->nodes[i+2]->op == GGML_OP_GET_ROWS &&
+                cgraph->nodes[i+3]->op == GGML_OP_RESHAPE &&
+                cgraph->nodes[i+4]->op == GGML_OP_SOFT_MAX &&
+                cgraph->nodes[i+5]->op == GGML_OP_RESHAPE) {
+                cuda_openai_experts(ctx, dst, cgraph->nodes[i+4]);
+                i += 5;
+            } else {
+                ggml_cuda_op_argsort(ctx, dst);
+            }
             break;
         case GGML_OP_ARGSORT_THRESH:
             ggml_cuda_op_argsort_thresh(ctx, dst);